Skip to main content

lsp_max/
agent.rs

1use lsp_max_protocol::AnalysisBundle;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::fs;
5use std::path::PathBuf;
6
7pub struct AgentExporter;
8
9impl AgentExporter {
10    pub fn export_bundle(bundle: &AnalysisBundle) -> String {
11        serde_json::to_string_pretty(bundle).unwrap()
12    }
13}
14
15pub struct AgentConfig {
16    pub api_key: Option<String>,
17    pub api_base: String,
18    pub model: String,
19}
20
21impl AgentConfig {
22    pub fn load() -> Self {
23        let api_key = std::env::var("LSP_MAX_API_KEY")
24            .ok()
25            .or_else(|| std::env::var("OPENAI_API_KEY").ok());
26        let api_base = std::env::var("LSP_MAX_API_BASE")
27            .ok()
28            .or_else(|| std::env::var("OPENAI_API_BASE").ok())
29            .unwrap_or_else(|| "https://api.openai.com/v1".to_string());
30        let model = std::env::var("LSP_MAX_MODEL")
31            .ok()
32            .or_else(|| std::env::var("OPENAI_MODEL").ok())
33            .unwrap_or_else(|| "gpt-4o".to_string());
34
35        let mut config_file_map = HashMap::new();
36        let config_path = if let Ok(path_str) = std::env::var("LSP_MAX_CONFIG") {
37            Some(PathBuf::from(path_str))
38        } else if let Ok(home) = std::env::var("HOME") {
39            Some(PathBuf::from(home).join(".lsp-max-config.json"))
40        } else {
41            Some(PathBuf::from(".lsp-max-config.json"))
42        };
43
44        if let Some(path) = config_path {
45            if path.exists() {
46                if let Ok(content) = fs::read_to_string(path) {
47                    if let Ok(map) = serde_json::from_str::<HashMap<String, String>>(&content) {
48                        config_file_map = map;
49                    }
50                }
51            }
52        }
53
54        let final_api_key = api_key.or_else(|| {
55            config_file_map
56                .get("api_key")
57                .or_else(|| config_file_map.get("openai_api_key"))
58                .cloned()
59        });
60
61        let final_api_base = if std::env::var("LSP_MAX_API_BASE").is_ok()
62            || std::env::var("OPENAI_API_BASE").is_ok()
63        {
64            api_base
65        } else {
66            config_file_map
67                .get("api_base")
68                .or_else(|| config_file_map.get("openai_api_base"))
69                .cloned()
70                .unwrap_or(api_base)
71        };
72
73        let final_model =
74            if std::env::var("LSP_MAX_MODEL").is_ok() || std::env::var("OPENAI_MODEL").is_ok() {
75                model
76            } else {
77                config_file_map
78                    .get("model")
79                    .or_else(|| config_file_map.get("openai_model"))
80                    .cloned()
81                    .unwrap_or(model)
82            };
83
84        Self {
85            api_key: final_api_key,
86            api_base: final_api_base,
87            model: final_model,
88        }
89    }
90}
91
92#[derive(Serialize)]
93struct ChatMessage {
94    role: String,
95    content: String,
96}
97
98#[derive(Serialize)]
99struct ChatRequest {
100    model: String,
101    messages: Vec<ChatMessage>,
102    #[serde(skip_serializing_if = "Option::is_none")]
103    temperature: Option<f32>,
104}
105
106#[derive(Deserialize)]
107struct ChatChoiceMessage {
108    content: String,
109}
110
111#[derive(Deserialize)]
112struct ChatChoice {
113    message: ChatChoiceMessage,
114}
115
116#[derive(Deserialize)]
117struct ChatResponse {
118    choices: Vec<ChatChoice>,
119}
120
121pub struct LspAgent;
122
123impl LspAgent {
124    pub fn invoke_task(task: &str) -> Result<String, String> {
125        let config = AgentConfig::load();
126        let system_prompt = "You are the lsp-max AI Agent. Your role is to perform analysis, code generation, and diagnostic checking for the Language Server Protocol implementation. Please process the user's task and return a precise, structured, and complete answer.";
127        Self::run_query(&config, system_prompt, task)
128    }
129
130    pub fn chat(message: &str) -> Result<String, String> {
131        let config = AgentConfig::load();
132        let system_prompt = "You are the lsp-max AI Agent. You are in a chat session with the user. Help them with their queries regarding Language Server Protocol, cargo workspaces, capabilities, and system configuration.";
133        Self::run_query(&config, system_prompt, message)
134    }
135
136    fn run_query(
137        config: &AgentConfig,
138        system_prompt: &str,
139        user_prompt: &str,
140    ) -> Result<String, String> {
141        let api_key = match &config.api_key {
142            Some(k) => k,
143            None => return Err("API key is not configured. Please set the LSP_MAX_API_KEY environment variable or run `lsp-max-cli config set api_key <your-key>`.".to_string()),
144        };
145
146        let url = format!("{}/chat/completions", config.api_base.trim_end_matches('/'));
147        let body = ChatRequest {
148            model: config.model.clone(),
149            messages: vec![
150                ChatMessage {
151                    role: "system".to_string(),
152                    content: system_prompt.to_string(),
153                },
154                ChatMessage {
155                    role: "user".to_string(),
156                    content: user_prompt.to_string(),
157                },
158            ],
159            temperature: Some(0.7),
160        };
161
162        let response = ureq::post(&url)
163            .header("Authorization", &format!("Bearer {}", api_key))
164            .header("Content-Type", "application/json")
165            .send_json(&body)
166            .map_err(|e| format!("HTTP request failed: {}", e))?;
167
168        let res_body: ChatResponse = response
169            .into_body()
170            .read_json::<ChatResponse>()
171            .map_err(|e| format!("Failed to parse response JSON: {}", e))?;
172
173        if let Some(choice) = res_body.choices.first() {
174            Ok(choice.message.content.clone())
175        } else {
176            Err("Received empty response from the AI provider".to_string())
177        }
178    }
179}