158 lines · 4.8 KB · Rust
RawDownload
use zed::{    http_client::HttpMethod,    http_client::HttpRequest,    serde_json::{self, json},};use zed_extension_api::{self as zed, http_client::RedirectPolicy, Result}; struct Perplexity; impl zed::Extension for Perplexity {    fn new() -> Self {        Self    }     fn run_slash_command(        &self,        command: zed::SlashCommand,        argument: Vec<String>,        worktree: Option<&zed::Worktree>,    ) -> zed::Result<zed::SlashCommandOutput> {        // Check if the command is 'perplexity'        if command.name != "perplexity" {            return Err("Invalid command. Expected 'perplexity'.".into());        }         let worktree = worktree.ok_or("Worktree is required")?;        // Join arguments with space as the query        let query = argument.join(" ");        if query.is_empty() {            return Ok(zed::SlashCommandOutput {                text: "Error: Query not provided. Please enter a question or topic.".to_string(),                sections: vec![],            });        }         // Get the API key from the environment        let env_vars = worktree.shell_env();        let api_key = env_vars            .iter()            .find(|(key, _)| key == "PERPLEXITY_API_KEY")            .map(|(_, value)| value.clone())            .ok_or("PERPLEXITY_API_KEY not found in environment")?;         // Prepare the request        let request = HttpRequest {            method: HttpMethod::Post,            url: "https://api.perplexity.ai/chat/completions".to_string(),            headers: vec![                ("Authorization".to_string(), format!("Bearer {}", api_key)),                ("Content-Type".to_string(), "application/json".to_string()),            ],            body: Some(                serde_json::to_vec(&json!({                    "model": "llama-3.1-sonar-small-128k-online",                    "messages": [{"role": "user", "content": query}],                    "stream": true,                }))                .unwrap(),            ),            redirect_policy: RedirectPolicy::FollowAll,        };         // Make the HTTP request        match zed::http_client::fetch_stream(&request) {            Ok(stream) => {                let mut full_content = String::new();                let mut buffer = String::new();                while let Ok(Some(chunk)) = stream.next_chunk() {                    buffer.push_str(&String::from_utf8_lossy(&chunk));                    for line in buffer.lines() {                        if let Some(json) = line.strip_prefix("data: ") {                            if let Ok(event) = serde_json::from_str::<StreamEvent>(json) {                                if let Some(choice) = event.choices.first() {                                    full_content.push_str(&choice.delta.content);                                }                            }                        }                    }                    buffer.clear();                }                Ok(zed::SlashCommandOutput {                    text: full_content,                    sections: vec![],                })            }            Err(e) => Ok(zed::SlashCommandOutput {                text: format!("API request failed. Error: {}. API Key: {}", e, api_key),                sections: vec![],            }),        }    }     fn complete_slash_command_argument(        &self,        _command: zed::SlashCommand,        query: Vec<String>,    ) -> zed::Result<Vec<zed::SlashCommandArgumentCompletion>> {        let suggestions = vec!["How do I develop a Zed extension?"];        let query = query.join(" ").to_lowercase();         Ok(suggestions            .into_iter()            .filter(|suggestion| suggestion.to_lowercase().contains(&query))            .map(|suggestion| zed::SlashCommandArgumentCompletion {                label: suggestion.to_string(),                new_text: suggestion.to_string(),                run_command: true,            })            .collect())    }     fn language_server_command(        &mut self,        _language_server_id: &zed_extension_api::LanguageServerId,        _worktree: &zed_extension_api::Worktree,    ) -> Result<zed_extension_api::Command> {        Err("Not implemented".into())    }} #[derive(serde::Deserialize)]struct StreamEvent {    id: String,    model: String,    created: u64,    usage: Usage,    object: String,    choices: Vec<Choice>,} #[derive(serde::Deserialize)]struct Usage {    prompt_tokens: u32,    completion_tokens: u32,    total_tokens: u32,} #[derive(serde::Deserialize)]struct Choice {    index: u32,    finish_reason: Option<String>,    message: Message,    delta: Delta,} #[derive(serde::Deserialize)]struct Message {    role: String,    content: String,} #[derive(serde::Deserialize)]struct Delta {    role: String,    content: String,} zed::register_extension!(Perplexity);