simple_llm_client/perplexity/

mod.rs

use dotenv::dotenv;
use futures_util::StreamExt;

use reqwest::Client;
use std::{env, fs, path::Path};
use tokio::fs::File;
use tokio::io::{stdout, AsyncWriteExt};

pub mod models;

use models::*;

// Function to format text as Markdown
#[allow(dead_code)]
fn format_as_markdown(text: &str, citations: &Option<Vec<Citation>>) -> String {
    // Just return the raw text as it already contains markdown formatting
    let mut formatted_text = text.to_string();

    // If citations exist, add them to the bottom of the document
    if let Some(citation_vec) = citations {
        if !citation_vec.is_empty() {
            formatted_text.push_str("\n\n## Sources\n\n");

            for (i, citation) in citation_vec.iter().enumerate() {
                if let Some(url) = citation.url.strip_prefix("http") {
                    formatted_text.push_str(&format!(
                        "{}. [{}](http{})\n",
                        i + 1,
                        citation.url,
                        url
                    ));
                } else {
                    formatted_text.push_str(&format!("{}. {}\n", i + 1, citation.url));
                }
            }
        }
    }

    formatted_text
}

#[allow(dead_code)]
pub async fn chat_completion(
    model: &str,
    messages: Vec<ChatMessage>,
) -> Result<(), Box<dyn std::error::Error>> {
    dotenv().ok();

    let api_key = env::var("PERPLEXITY_API_KEY").expect("PERPLEXITY_API_KEY not set in .env");

    let client = Client::new();
    let request = ChatCompletionRequest {
        model: model.to_string(),
        messages,
        stream: true,
    };

    let response = client
        .post("https://api.perplexity.ai/chat/completions")
        .header("Authorization", format!("Bearer {}", api_key))
        .header("Content-Type", "application/json")
        .json(&request)
        .send()
        .await?;

    if !response.status().is_success() {
        println!("Response status code: {}", response.status());
        let response_body = response.text().await?;
        println!("Response body: {}", response_body);
        return Err("API request failed".into());
    }

    let mut stream = response.bytes_stream();
    let mut stdout = stdout();

    while let Some(chunk) = stream.next().await {
        let chunk = chunk?;
        let chunk_str = String::from_utf8(chunk.to_vec())?;

        let events: Vec<&str> = chunk_str.split("\n\n").collect();

        for event in events {
            if event.starts_with("data:") {
                let data = event.trim_start_matches("data:").trim();

                if data == "[DONE]" {
                    break;
                }

                if let Ok(completion_response) =
                    serde_json::from_str::<ChatCompletionResponse>(data)
                {
                    if let Some(choice) = completion_response.choices.get(0) {
                        if let Some(message) = &choice.message {
                            stdout.write_all(message.content.as_bytes()).await?;
                            stdout.flush().await?;
                        }
                    }
                } else {
                    eprintln!("Failed to parse SSE data: {}", data);
                }
            }
        }
    }

    Ok(())
}

#[allow(dead_code)]
pub async fn chat_completion_markdown(
    model: &str,
    messages: Vec<ChatMessage>,
    path: Option<&Path>,
    filename: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    dotenv().ok();

    let api_key = env::var("PERPLEXITY_API_KEY").expect("PERPLEXITY_API_KEY not set in .env");

    // Use a non-streaming request to get the complete response at once
    let client = Client::new();
    let request = ChatCompletionRequest {
        model: model.to_string(),
        messages,
        stream: false,
    };

    let response = client
        .post("https://api.perplexity.ai/chat/completions")
        .header("Authorization", format!("Bearer {}", api_key))
        .header("Content-Type", "application/json")
        .json(&request)
        .send()
        .await?;

    if !response.status().is_success() {
        println!("Response status code: {}", response.status());
        let response_body = response.text().await?;
        println!("Response body: {}", response_body);
        return Err("API request failed".into());
    }

    // Get the complete JSON response
    let response_text = response.text().await?;

    // For debugging
    let _ = fs::write("raw_response.json", &response_text);

    // Parse the response
    let completion_response: ChatCompletionResponse = serde_json::from_str(&response_text)?;

    // Extract the content and citations from the response
    let (content, citations) = if let Some(choice) = completion_response.choices.get(0) {
        let content = choice
            .message
            .as_ref()
            .map_or(String::new(), |m| m.content.clone());
        (content, choice.citations.clone())
    } else {
        (String::new(), None)
    };

    // Format the content with proper citations
    let formatted_content = format_as_markdown(&content, &citations);

    // For debugging
    let _ = fs::write("extracted_content.txt", &content);
    let _ = fs::write("formatted_content.txt", &formatted_content);

    // Write to file
    let file_path = match path {
        Some(p) => p.join(filename),
        None => Path::new(filename).to_path_buf(),
    };

    let mut file = File::create(file_path).await?;
    file.write_all(formatted_content.as_bytes()).await?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::{env, path::Path};
    use tokio;

    #[tokio::test]
    async fn test_chat_completion() {
        dotenv::dotenv().ok();

        // This test requires the PERPLEXITY_API_KEY environment variable to be set.
        if env::var("PERPLEXITY_API_KEY").is_ok() {
            let system_message = "Be precise and concise.".to_string();
            let user_message = "How many stars are there in our galaxy?".to_string();

            let messages = vec![
                ChatMessage {
                    role: "system".to_string(),
                    content: system_message.clone(),
                },
                ChatMessage {
                    role: "user".to_string(),
                    content: user_message.clone(),
                },
            ];

            let result = super::chat_completion("sonar-pro", messages).await;
            assert!(result.is_ok());
        } else {
            println!("Skipping test because PERPLEXITY_API_KEY is not set");
        }
    }

    #[tokio::test]
    async fn test_chat_completion_markdown() {
        dotenv::dotenv().ok();

        // This test requires the PERPLEXITY_API_KEY environment variable to be set.
        if env::var("PERPLEXITY_API_KEY").is_ok() {
            let system_message =
                "Be precise and concise. Return the response as a markdown, not a JSON."
                    .to_string();
            let user_message = "How many stars are there in our galaxy?".to_string();

            let messages = vec![
                ChatMessage {
                    role: "system".to_string(),
                    content: system_message.clone(),
                },
                ChatMessage {
                    role: "user".to_string(),
                    content: user_message.clone(),
                },
            ];

            let path = Path::new(".");
            let filename = "test_output.md";

            let result =
                super::chat_completion_markdown("sonar-pro", messages, Some(path), filename).await;
            assert!(result.is_ok());
        } else {
            println!("Skipping test because PERPLEXITY_API_KEY is not set");
        }
    }
}