rig/providers/
perplexity.rs

1//! Perplexity API client and Rig integration
2//!
3//! # Example
4//! ```
5//! use rig::providers::perplexity;
6//!
7//! let client = perplexity::Client::new("YOUR_API_KEY");
8//!
9//! let llama_3_1_sonar_small_online = client.completion_model(perplexity::LLAMA_3_1_SONAR_SMALL_ONLINE);
10//! ```
11use crate::{
12    OneOrMany,
13    client::{VerifyClient, VerifyError},
14    completion::{self, CompletionError, MessageError, message},
15    impl_conversion_traits, json_utils,
16};
17
18use crate::client::{ClientBuilderError, CompletionClient, ProviderClient};
19use crate::completion::CompletionRequest;
20use crate::json_utils::merge;
21use crate::providers::openai;
22use crate::providers::openai::send_compatible_streaming_request;
23use crate::streaming::StreamingCompletionResponse;
24use serde::{Deserialize, Serialize};
25use serde_json::{Value, json};
26
27// ================================================================
28// Main Cohere Client
29// ================================================================
30const PERPLEXITY_API_BASE_URL: &str = "https://api.perplexity.ai";
31
32pub struct ClientBuilder<'a> {
33    api_key: &'a str,
34    base_url: &'a str,
35    http_client: Option<reqwest::Client>,
36}
37
38impl<'a> ClientBuilder<'a> {
39    pub fn new(api_key: &'a str) -> Self {
40        Self {
41            api_key,
42            base_url: PERPLEXITY_API_BASE_URL,
43            http_client: None,
44        }
45    }
46
47    pub fn base_url(mut self, base_url: &'a str) -> Self {
48        self.base_url = base_url;
49        self
50    }
51
52    pub fn custom_client(mut self, client: reqwest::Client) -> Self {
53        self.http_client = Some(client);
54        self
55    }
56
57    pub fn build(self) -> Result<Client, ClientBuilderError> {
58        let http_client = if let Some(http_client) = self.http_client {
59            http_client
60        } else {
61            reqwest::Client::builder().build()?
62        };
63
64        Ok(Client {
65            base_url: self.base_url.to_string(),
66            api_key: self.api_key.to_string(),
67            http_client,
68        })
69    }
70}
71
72#[derive(Clone)]
73pub struct Client {
74    base_url: String,
75    api_key: String,
76    http_client: reqwest::Client,
77}
78
79impl std::fmt::Debug for Client {
80    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81        f.debug_struct("Client")
82            .field("base_url", &self.base_url)
83            .field("http_client", &self.http_client)
84            .field("api_key", &"<REDACTED>")
85            .finish()
86    }
87}
88
89impl Client {
90    /// Create a new Perplexity client builder.
91    ///
92    /// # Example
93    /// ```
94    /// use rig::providers::perplexity::{ClientBuilder, self};
95    ///
96    /// // Initialize the Perplexity client
97    /// let perplexity = Client::builder("your-perplexity-api-key")
98    ///    .build()
99    /// ```
100    pub fn builder(api_key: &str) -> ClientBuilder<'_> {
101        ClientBuilder::new(api_key)
102    }
103
104    /// Create a new Perplexity client. For more control, use the `builder` method.
105    ///
106    /// # Panics
107    /// - If the reqwest client cannot be built (if the TLS backend cannot be initialized).
108    pub fn new(api_key: &str) -> Self {
109        Self::builder(api_key)
110            .build()
111            .expect("Perplexity client should build")
112    }
113
114    pub(crate) fn post(&self, path: &str) -> reqwest::RequestBuilder {
115        let url = format!("{}/{}", self.base_url, path).replace("//", "/");
116        self.http_client.post(url).bearer_auth(&self.api_key)
117    }
118}
119
120impl ProviderClient for Client {
121    /// Create a new Perplexity client from the `PERPLEXITY_API_KEY` environment variable.
122    /// Panics if the environment variable is not set.
123    fn from_env() -> Self {
124        let api_key = std::env::var("PERPLEXITY_API_KEY").expect("PERPLEXITY_API_KEY not set");
125        Self::new(&api_key)
126    }
127
128    fn from_val(input: crate::client::ProviderValue) -> Self {
129        let crate::client::ProviderValue::Simple(api_key) = input else {
130            panic!("Incorrect provider value type")
131        };
132        Self::new(&api_key)
133    }
134}
135
136impl CompletionClient for Client {
137    type CompletionModel = CompletionModel;
138
139    fn completion_model(&self, model: &str) -> CompletionModel {
140        CompletionModel::new(self.clone(), model)
141    }
142}
143
144impl VerifyClient for Client {
145    #[cfg_attr(feature = "worker", worker::send)]
146    async fn verify(&self) -> Result<(), VerifyError> {
147        // No API endpoint to verify the API key
148        Ok(())
149    }
150}
151
152impl_conversion_traits!(
153    AsTranscription,
154    AsEmbeddings,
155    AsImageGeneration,
156    AsAudioGeneration for Client
157);
158
159#[derive(Debug, Deserialize)]
160struct ApiErrorResponse {
161    message: String,
162}
163
164#[derive(Debug, Deserialize)]
165#[serde(untagged)]
166enum ApiResponse<T> {
167    Ok(T),
168    Err(ApiErrorResponse),
169}
170
171// ================================================================
172// Perplexity Completion API
173// ================================================================
174/// `sonar-pro` completion model
175pub const SONAR_PRO: &str = "sonar-pro";
176/// `sonar` completion model
177pub const SONAR: &str = "sonar";
178
179#[derive(Debug, Deserialize, Serialize)]
180pub struct CompletionResponse {
181    pub id: String,
182    pub model: String,
183    pub object: String,
184    pub created: u64,
185    #[serde(default)]
186    pub choices: Vec<Choice>,
187    pub usage: Usage,
188}
189
190#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
191pub struct Message {
192    pub role: Role,
193    pub content: String,
194}
195
196#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
197#[serde(rename_all = "lowercase")]
198pub enum Role {
199    System,
200    User,
201    Assistant,
202}
203
204#[derive(Deserialize, Debug, Serialize)]
205pub struct Delta {
206    pub role: Role,
207    pub content: String,
208}
209
210#[derive(Deserialize, Debug, Serialize)]
211pub struct Choice {
212    pub index: usize,
213    pub finish_reason: String,
214    pub message: Message,
215    pub delta: Delta,
216}
217
218#[derive(Deserialize, Debug, Serialize)]
219pub struct Usage {
220    pub prompt_tokens: u32,
221    pub completion_tokens: u32,
222    pub total_tokens: u32,
223}
224
225impl std::fmt::Display for Usage {
226    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
227        write!(
228            f,
229            "Prompt tokens: {}\nCompletion tokens: {} Total tokens: {}",
230            self.prompt_tokens, self.completion_tokens, self.total_tokens
231        )
232    }
233}
234
235impl TryFrom<CompletionResponse> for completion::CompletionResponse<CompletionResponse> {
236    type Error = CompletionError;
237
238    fn try_from(response: CompletionResponse) -> Result<Self, Self::Error> {
239        let choice = response.choices.first().ok_or_else(|| {
240            CompletionError::ResponseError("Response contained no choices".to_owned())
241        })?;
242
243        match &choice.message {
244            Message {
245                role: Role::Assistant,
246                content,
247            } => Ok(completion::CompletionResponse {
248                choice: OneOrMany::one(content.clone().into()),
249                usage: completion::Usage {
250                    input_tokens: response.usage.prompt_tokens as u64,
251                    output_tokens: response.usage.completion_tokens as u64,
252                    total_tokens: response.usage.total_tokens as u64,
253                },
254                raw_response: response,
255            }),
256            _ => Err(CompletionError::ResponseError(
257                "Response contained no assistant message".to_owned(),
258            )),
259        }
260    }
261}
262
263#[derive(Clone)]
264pub struct CompletionModel {
265    client: Client,
266    pub model: String,
267}
268
269impl CompletionModel {
270    pub fn new(client: Client, model: &str) -> Self {
271        Self {
272            client,
273            model: model.to_string(),
274        }
275    }
276
277    fn create_completion_request(
278        &self,
279        completion_request: CompletionRequest,
280    ) -> Result<Value, CompletionError> {
281        // Build up the order of messages (context, chat_history, prompt)
282        let mut partial_history = vec![];
283        if let Some(docs) = completion_request.normalized_documents() {
284            partial_history.push(docs);
285        }
286        partial_history.extend(completion_request.chat_history);
287
288        // Initialize full history with preamble (or empty if non-existent)
289        let mut full_history: Vec<Message> =
290            completion_request
291                .preamble
292                .map_or_else(Vec::new, |preamble| {
293                    vec![Message {
294                        role: Role::System,
295                        content: preamble,
296                    }]
297                });
298
299        // Convert and extend the rest of the history
300        full_history.extend(
301            partial_history
302                .into_iter()
303                .map(message::Message::try_into)
304                .collect::<Result<Vec<Message>, _>>()?,
305        );
306
307        // Compose request
308        let request = json!({
309            "model": self.model,
310            "messages": full_history,
311            "temperature": completion_request.temperature,
312        });
313
314        let request = if let Some(ref params) = completion_request.additional_params {
315            json_utils::merge(request, params.clone())
316        } else {
317            request
318        };
319
320        Ok(request)
321    }
322}
323
324impl TryFrom<message::Message> for Message {
325    type Error = MessageError;
326
327    fn try_from(message: message::Message) -> Result<Self, Self::Error> {
328        Ok(match message {
329            message::Message::User { content } => {
330                let collapsed_content = content
331                    .into_iter()
332                    .map(|content| match content {
333                        message::UserContent::Text(message::Text { text }) => Ok(text),
334                        _ => Err(MessageError::ConversionError(
335                            "Only text content is supported by Perplexity".to_owned(),
336                        )),
337                    })
338                    .collect::<Result<Vec<_>, _>>()?
339                    .join("\n");
340
341                Message {
342                    role: Role::User,
343                    content: collapsed_content,
344                }
345            }
346
347            message::Message::Assistant { content, .. } => {
348                let collapsed_content = content
349                    .into_iter()
350                    .map(|content| {
351                        Ok(match content {
352                            message::AssistantContent::Text(message::Text { text }) => text,
353                            _ => return Err(MessageError::ConversionError(
354                                "Only text assistant message content is supported by Perplexity"
355                                    .to_owned(),
356                            )),
357                        })
358                    })
359                    .collect::<Result<Vec<_>, _>>()?
360                    .join("\n");
361
362                Message {
363                    role: Role::Assistant,
364                    content: collapsed_content,
365                }
366            }
367        })
368    }
369}
370
371impl From<Message> for message::Message {
372    fn from(message: Message) -> Self {
373        match message.role {
374            Role::User => message::Message::user(message.content),
375            Role::Assistant => message::Message::assistant(message.content),
376
377            // System messages get coerced into user messages for ease of error handling.
378            // They should be handled on the outside of `Message` conversions via the preamble.
379            Role::System => message::Message::user(message.content),
380        }
381    }
382}
383
384impl completion::CompletionModel for CompletionModel {
385    type Response = CompletionResponse;
386    type StreamingResponse = openai::StreamingCompletionResponse;
387
388    #[cfg_attr(feature = "worker", worker::send)]
389    async fn completion(
390        &self,
391        completion_request: completion::CompletionRequest,
392    ) -> Result<completion::CompletionResponse<CompletionResponse>, CompletionError> {
393        let request = self.create_completion_request(completion_request)?;
394
395        let response = self
396            .client
397            .post("/chat/completions")
398            .json(&request)
399            .send()
400            .await?;
401
402        if response.status().is_success() {
403            match response.json::<ApiResponse<CompletionResponse>>().await? {
404                ApiResponse::Ok(completion) => {
405                    tracing::info!(target: "rig",
406                        "Perplexity completion token usage: {}",
407                        completion.usage
408                    );
409                    Ok(completion.try_into()?)
410                }
411                ApiResponse::Err(error) => Err(CompletionError::ProviderError(error.message)),
412            }
413        } else {
414            Err(CompletionError::ProviderError(response.text().await?))
415        }
416    }
417
418    #[cfg_attr(feature = "worker", worker::send)]
419    async fn stream(
420        &self,
421        completion_request: completion::CompletionRequest,
422    ) -> Result<StreamingCompletionResponse<Self::StreamingResponse>, CompletionError> {
423        let mut request = self.create_completion_request(completion_request)?;
424
425        request = merge(request, json!({"stream": true}));
426
427        let builder = self.client.post("/chat/completions").json(&request);
428
429        send_compatible_streaming_request(builder).await
430    }
431}
432
433#[cfg(test)]
434mod tests {
435    use super::*;
436
437    #[test]
438    fn test_deserialize_message() {
439        let json_data = r#"
440        {
441            "role": "user",
442            "content": "Hello, how can I help you?"
443        }
444        "#;
445
446        let message: Message = serde_json::from_str(json_data).unwrap();
447        assert_eq!(message.role, Role::User);
448        assert_eq!(message.content, "Hello, how can I help you?");
449    }
450
451    #[test]
452    fn test_serialize_message() {
453        let message = Message {
454            role: Role::Assistant,
455            content: "I am here to assist you.".to_string(),
456        };
457
458        let json_data = serde_json::to_string(&message).unwrap();
459        let expected_json = r#"{"role":"assistant","content":"I am here to assist you."}"#;
460        assert_eq!(json_data, expected_json);
461    }
462
463    #[test]
464    fn test_message_to_message_conversion() {
465        let user_message = message::Message::user("User message");
466        let assistant_message = message::Message::assistant("Assistant message");
467
468        let converted_user_message: Message = user_message.clone().try_into().unwrap();
469        let converted_assistant_message: Message = assistant_message.clone().try_into().unwrap();
470
471        assert_eq!(converted_user_message.role, Role::User);
472        assert_eq!(converted_user_message.content, "User message");
473
474        assert_eq!(converted_assistant_message.role, Role::Assistant);
475        assert_eq!(converted_assistant_message.content, "Assistant message");
476
477        let back_to_user_message: message::Message = converted_user_message.into();
478        let back_to_assistant_message: message::Message = converted_assistant_message.into();
479
480        assert_eq!(user_message, back_to_user_message);
481        assert_eq!(assistant_message, back_to_assistant_message);
482    }
483}
rig/providers/perplexity.rs

rig/providers/
perplexity.rs
