ai_lib/provider/
perplexity.rs

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

ai_lib/provider/
perplexity.rs
