test

src/agent.rs

@@ -0,0 +1,141 @@
+use chrono::Utc;
+
+use crate::api::{ChatRequest, ChatResponse, Message, Tool};
+use crate::tools;
+
+pub struct Agent {
+    client: reqwest::Client,
+    url: String,
+    zen_api_key: Option<String>,
+    tavily_api_key: Option<String>,
+    messages: Vec<Message>,
+    tools: Option<Vec<Tool>>,
+    logs: String,
+    answer: Option<String>,
+}
+
+impl Agent {
+    pub fn new(
+        zen_api_key: Option<String>,
+        tavily_api_key: Option<String>,
+        initial_message: String,
+    ) -> Self {
+        let intro = format!(
+                    "You are an autonomous agent. You have access to tools that can help 
+                    you achieve your goals. Use them wisely. Do not ask for clarification and use the
+                     answer tool once you to give your final answer. current date is {}",
+                    Utc::now().to_rfc3339()
+                );
+        println!("initial_message: {}", intro);
+        let messages = vec![
+            Message {
+                role: "system".to_string(),
+                content: Some(intro),
+                tool_calls: None,
+                tool_call_id: None,
+            },
+            Message {
+                role: "user".to_string(),
+                content: Some(initial_message),
+                tool_calls: None,
+                tool_call_id: None,
+            },
+        ];
+
+        let tools = Some(tools::get_tools());
+
+        Self {
+            client: reqwest::Client::new(),
+            url: "https://opencode.ai/zen/v1/chat/completions".to_string(),
+            zen_api_key,
+            tavily_api_key,
+            messages,
+            tools,
+            logs: String::new(),
+            answer: None,
+        }
+    }
+
+    fn log(&mut self, message: &str) {
+        println!("{}", message);
+        self.logs.push_str(message);
+        self.logs.push('\n');
+    }
+
+    pub async fn run(&mut self) -> Result<(String, Option<String>), Box<dyn std::error::Error>> {
+        let mut file_written = false;
+
+        while !file_written {
+            let request = ChatRequest {
+                model: "kimi-k2.5".to_string(),
+                messages: self.messages.clone(),
+                tools: self.tools.clone(),
+            };
+
+            self.log(&format!(
+                "--- Sending request to Zen API (Role: {}) ---",
+                self.messages.last().unwrap().role
+            ));
+
+            let mut request_builder = self.client.post(&self.url).json(&request);
+
+            if let Some(key) = &self.zen_api_key {
+                request_builder =
+                    request_builder.header("Authorization", format!("Bearer {}", key));
+            }
+
+            let response = request_builder.send().await?;
+
+            if !response.status().is_success() {
+                let status = response.status();
+                let error_text = response.text().await?;
+                self.log(&format!("Error: API request failed with status {}", status));
+                self.log(&format!("Error details: {}", error_text));
+                return Err(format!("API request failed: {}", status).into());
+            }
+
+            let chat_response: ChatResponse = response.json().await?;
+            let assistant_message = chat_response.choices.get(0).unwrap().message.clone();
+
+            self.messages.push(assistant_message.clone());
+
+            if let Some(content) = &assistant_message.content {
+                if !content.is_empty() {
+                    self.log(&format!("\nAssistant response:\n{}\n", content));
+                }
+            }
+
+            if let Some(tool_calls) = assistant_message.tool_calls {
+                for tool_call in tool_calls {
+                    let (tool_message, written, tool_answer) =
+                        tools::handle_tool_call(&tool_call, &self.tavily_api_key).await?;
+
+                    if let Some(ans) = tool_answer {
+                        self.answer = Some(ans);
+                    }
+
+                    if let Some(content) = &tool_message.content {
+                        self.log(&format!(
+                            "Tool result ({}): {}",
+                            tool_call.function.name, content
+                        ));
+                    }
+
+                    self.messages.push(tool_message);
+                    if written {
+                        file_written = true;
+                    }
+                }
+                // Continue the loop to send tool results back
+                continue;
+            }
+
+            // No more tool calls from assistant, but we only exit if file was written
+            if !file_written {
+                self.log("--- Assistant didn't use write_file yet. Waiting for next turn... ---");
+            }
+        }
+
+        Ok((self.logs.clone(), self.answer.clone()))
+    }
+}