ai_lib/provider/
perplexity.rs

1use crate::api::ChatProvider;
2use crate::types::AiLibError;
3use crate::types::{ChatCompletionRequest, ChatCompletionResponse, Message, Role};
4use async_trait::async_trait;
5use futures::Stream;
6use reqwest::Client;
7use serde::{Deserialize, Serialize};
8
9/// Perplexity API adapter
10///
11/// Perplexity provides search-enhanced AI with a custom API format.
12/// Documentation: <https://docs.perplexity.ai/getting-started/overview>
13pub struct PerplexityAdapter {
14    client: Client,
15    api_key: String,
16}
17
18impl PerplexityAdapter {
19    pub fn new() -> Result<Self, AiLibError> {
20        let api_key = std::env::var("PERPLEXITY_API_KEY").map_err(|_| {
21            AiLibError::ConfigurationError(
22                "PERPLEXITY_API_KEY environment variable not set".to_string(),
23            )
24        })?;
25
26        Ok(Self {
27            client: Client::new(),
28            api_key,
29        })
30    }
31
32    pub async fn chat_completion(
33        &self,
34        request: ChatCompletionRequest,
35    ) -> Result<ChatCompletionResponse, AiLibError> {
36        let perplexity_request = self.convert_request(&request)?;
37
38        let response = self
39            .client
40            .post("https://api.perplexity.ai/chat/completions")
41            .header("Authorization", format!("Bearer {}", self.api_key))
42            .header("Content-Type", "application/json")
43            .json(&perplexity_request)
44            .send()
45            .await
46            .map_err(|e| {
47                AiLibError::NetworkError(format!("Perplexity API request failed: {}", e))
48            })?;
49
50        if !response.status().is_success() {
51            let status = response.status();
52            let error_text = response
53                .text()
54                .await
55                .unwrap_or_else(|_| "Unknown error".to_string());
56            return Err(AiLibError::ProviderError(format!(
57                "Perplexity API error {}: {}",
58                status, error_text
59            )));
60        }
61
62        let perplexity_response: PerplexityResponse = response.json().await.map_err(|e| {
63            AiLibError::DeserializationError(format!("Failed to parse Perplexity response: {}", e))
64        })?;
65
66        self.convert_response(perplexity_response)
67    }
68
69    pub async fn chat_completion_stream(
70        &self,
71        request: ChatCompletionRequest,
72    ) -> Result<
73        Box<
74            dyn futures::Stream<Item = Result<crate::api::ChatCompletionChunk, AiLibError>>
75                + Send
76                + Unpin,
77        >,
78        AiLibError,
79    > {
80        // For now, convert streaming request to non-streaming and return a single chunk
81        let response = self.chat_completion(request.clone()).await?;
82
83        // Create a single chunk from the response
84        let chunk = crate::api::ChatCompletionChunk {
85            id: response.id.clone(),
86            object: "chat.completion.chunk".to_string(),
87            created: response.created,
88            model: response.model.clone(),
89            choices: response
90                .choices
91                .into_iter()
92                .map(|choice| crate::api::ChoiceDelta {
93                    index: choice.index,
94                    delta: crate::api::MessageDelta {
95                        role: Some(choice.message.role),
96                        content: Some(match &choice.message.content {
97                            crate::Content::Text(text) => text.clone(),
98                            _ => "".to_string(),
99                        }),
100                    },
101                    finish_reason: choice.finish_reason,
102                })
103                .collect(),
104        };
105
106        let stream = futures::stream::once(async move { Ok(chunk) });
107        Ok(Box::new(Box::pin(stream)))
108    }
109
110    fn convert_request(
111        &self,
112        request: &ChatCompletionRequest,
113    ) -> Result<PerplexityRequest, AiLibError> {
114        // Convert messages to Perplexity format
115        let messages = request
116            .messages
117            .iter()
118            .map(|msg| PerplexityMessage {
119                role: match msg.role {
120                    Role::System => "system".to_string(),
121                    Role::User => "user".to_string(),
122                    Role::Assistant => "assistant".to_string(),
123                },
124                content: match &msg.content {
125                    crate::Content::Text(text) => text.clone(),
126                    _ => "Unsupported content type".to_string(),
127                },
128            })
129            .collect();
130
131        Ok(PerplexityRequest {
132            model: request.model.clone(),
133            messages,
134            max_tokens: request.max_tokens,
135            temperature: request.temperature,
136            top_p: request.top_p,
137            stream: Some(false),
138            extensions: request.extensions.clone(),
139        })
140    }
141
142    fn convert_response(
143        &self,
144        response: PerplexityResponse,
145    ) -> Result<ChatCompletionResponse, AiLibError> {
146        let choice = response.choices.first().ok_or_else(|| {
147            AiLibError::InvalidModelResponse("No choices in Perplexity response".to_string())
148        })?;
149
150        let message = Message {
151            role: match choice.message.role.as_str() {
152                "assistant" => Role::Assistant,
153                "user" => Role::User,
154                "system" => Role::System,
155                _ => Role::Assistant,
156            },
157            content: crate::Content::Text(choice.message.content.clone().unwrap_or_default()),
158            function_call: None,
159        };
160
161        Ok(ChatCompletionResponse {
162            id: response.id,
163            object: "chat.completion".to_string(),
164            created: response.created,
165            model: response.model,
166            choices: vec![crate::types::Choice {
167                index: 0,
168                message,
169                finish_reason: choice.finish_reason.clone(),
170            }],
171            usage: response
172                .usage
173                .map(|u| crate::types::Usage {
174                    prompt_tokens: u.prompt_tokens,
175                    completion_tokens: u.completion_tokens,
176                    total_tokens: u.total_tokens,
177                })
178                .unwrap_or_else(|| crate::types::Usage {
179                    prompt_tokens: 0,
180                    completion_tokens: 0,
181                    total_tokens: 0,
182                }),
183            usage_status: crate::types::response::UsageStatus::Finalized,
184        })
185    }
186}
187
188#[async_trait]
189impl ChatProvider for PerplexityAdapter {
190    fn name(&self) -> &str {
191        "Perplexity"
192    }
193
194    async fn chat(
195        &self,
196        request: ChatCompletionRequest,
197    ) -> Result<ChatCompletionResponse, AiLibError> {
198        self.chat_completion(request).await
199    }
200
201    async fn stream(
202        &self,
203        request: ChatCompletionRequest,
204    ) -> Result<
205        Box<dyn Stream<Item = Result<crate::api::ChatCompletionChunk, AiLibError>> + Send + Unpin>,
206        AiLibError,
207    > {
208        self.chat_completion_stream(request).await
209    }
210
211    async fn list_models(&self) -> Result<Vec<String>, AiLibError> {
212        // Return default models for Perplexity
213        Ok(vec![
214            "llama-3.1-sonar-small-128k-online".to_string(),
215            "llama-3.1-sonar-large-128k-online".to_string(),
216        ])
217    }
218
219    async fn get_model_info(&self, model_id: &str) -> Result<crate::api::ModelInfo, AiLibError> {
220        Ok(crate::api::ModelInfo {
221            id: model_id.to_string(),
222            object: "model".to_string(),
223            created: 0,
224            owned_by: "perplexity".to_string(),
225            permission: vec![],
226        })
227    }
228}
229
230#[derive(Serialize)]
231struct PerplexityRequest {
232    model: String,
233    messages: Vec<PerplexityMessage>,
234    max_tokens: Option<u32>,
235    temperature: Option<f32>,
236    top_p: Option<f32>,
237    stream: Option<bool>,
238    #[serde(flatten)]
239    extensions: Option<serde_json::Map<String, serde_json::Value>>,
240}
241
242#[derive(Serialize)]
243struct PerplexityMessage {
244    role: String,
245    content: String,
246}
247
248#[derive(Deserialize)]
249struct PerplexityResponse {
250    id: String,
251    #[allow(dead_code)]
252    object: String,
253    created: u64,
254    model: String,
255    choices: Vec<PerplexityChoice>,
256    usage: Option<PerplexityUsage>,
257}
258
259#[derive(Deserialize)]
260struct PerplexityChoice {
261    message: PerplexityMessageResponse,
262    finish_reason: Option<String>,
263}
264
265#[derive(Deserialize)]
266struct PerplexityMessageResponse {
267    role: String,
268    content: Option<String>,
269}
270
271#[derive(Deserialize)]
272struct PerplexityUsage {
273    prompt_tokens: u32,
274    completion_tokens: u32,
275    total_tokens: u32,
276}
ai_lib/provider/perplexity.rs

ai_lib/provider/
perplexity.rs
