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()))
+    }
+}

src/api.rs

@@ -0,0 +1,112 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct Message {
+    pub role: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub content: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_calls: Option<Vec<ToolCall>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_call_id: Option<String>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct ToolCall {
+    pub id: String,
+    #[serde(rename = "type")]
+    pub call_type: String,
+    pub function: FunctionCall,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct FunctionCall {
+    pub name: String,
+    pub arguments: String,
+}
+
+#[derive(Serialize, Debug)]
+pub struct ChatRequest {
+    pub model: String,
+    pub messages: Vec<Message>,
+    pub tools: Option<Vec<Tool>>,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct ChatResponse {
+    pub choices: Vec<Choice>,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct Choice {
+    pub message: Message,
+}
+
+#[derive(Serialize, Debug, Clone)]
+pub struct Tool {
+    #[serde(rename = "type")]
+    pub tool_type: String,
+    pub function: FunctionDefinition,
+}
+
+#[derive(Serialize, Debug, Clone)]
+pub struct FunctionDefinition {
+    pub name: String,
+    pub description: String,
+    pub parameters: serde_json::Value,
+}
+
+#[derive(Deserialize, Debug)]
+struct TavilyResponse {
+    results: Vec<TavilyResult>,
+}
+
+#[derive(Deserialize, Debug)]
+struct TavilyResult {
+    title: String,
+    url: String,
+    content: String,
+}
+
+pub async fn perform_search(
+    query: &str,
+    api_key: &str,
+) -> Result<String, Box<dyn std::error::Error>> {
+    let client = reqwest::Client::new();
+    let response = client
+        .post("https://api.tavily.com/search")
+        .json(&serde_json::json!({
+            "api_key": api_key,
+            "query": query,
+            "search_depth": "basic",
+            "include_answer": false,
+            "include_images": false,
+            "include_raw_content": false,
+            "max_results": 5,
+        }))
+        .send()
+        .await?;
+
+    if !response.status().is_success() {
+        return Err(format!("Tavily API error: {}", response.status()).into());
+    }
+
+    let search_data: TavilyResponse = response.json().await?;
+    let mut results_text = String::new();
+
+    for (i, result) in search_data.results.iter().enumerate() {
+        results_text.push_str(&format!(
+            "{}. {} ({})\nSnippet: {}\n\n",
+            i + 1,
+            result.title,
+            result.url,
+            result.content
+        ));
+    }
+
+    if results_text.is_empty() {
+        Ok("No results found.".to_string())
+    } else {
+        Ok(results_text)
+    }
+}

src/entities/mod.rs

@@ -0,0 +1 @@
+pub mod task;

src/entities/task.rs

@@ -0,0 +1,21 @@
+use sea_orm::entity::prelude::*;
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)]
+#[sea_orm(table_name = "tasks")]
+pub struct Model {
+    #[sea_orm(primary_key, auto_increment = false)]
+    pub id: Uuid,
+    pub goal: String,
+    pub status: String,
+    #[sea_orm(column_type = "Text")]
+    pub logs: String,
+    #[sea_orm(column_type = "Text", nullable)]
+    pub answer: Option<String>,
+    pub created_at: DateTimeWithTimeZone,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}

src/main.rs

@@ -0,0 +1,14 @@
+mod agent;
+mod api;
+mod entities;
+mod server;
+mod tools;
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let db_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set");
+
+    server::start(&db_url).await?;
+
+    Ok(())
+}

src/server.rs

@@ -0,0 +1,179 @@
+use axum::{
+    Json, Router,
+    extract::{Path, State},
+    http::{Method, StatusCode},
+    routing::{get, post},
+};
+use chrono::Utc;
+use sea_orm::{ActiveModelTrait, Database, DatabaseConnection, EntityTrait, QueryOrder, Set};
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+use tower_http::cors::{Any, CorsLayer};
+use uuid::Uuid;
+
+use crate::agent::Agent;
+use crate::entities::task::{self, Entity as Task};
+use migration::{Migrator, MigratorTrait};
+
+#[derive(Clone)]
+pub struct AppState {
+    pub db: DatabaseConnection,
+    pub zen_api_key: Option<String>,
+    pub tavily_api_key: Option<String>,
+}
+
+#[derive(Deserialize)]
+pub struct CreateTaskRequest {
+    pub goal: String,
+}
+
+#[derive(Serialize)]
+pub struct TaskResponse {
+    pub id: Uuid,
+    pub goal: String,
+    pub status: String,
+    pub logs: String,
+    pub answer: Option<String>,
+    pub created_at: chrono::DateTime<chrono::FixedOffset>,
+}
+
+pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
+    let db = Database::connect(db_url).await?;
+    Migrator::up(&db, None).await?;
+
+    let zen_api_key = std::env::var("ZEN_API_KEY").ok();
+    let tavily_api_key = std::env::var("TAVILY_API_KEY").ok();
+
+    let state = Arc::new(AppState {
+        db,
+        zen_api_key,
+        tavily_api_key,
+    });
+
+    let cors = CorsLayer::new()
+        .allow_origin(Any)
+        .allow_methods([Method::GET, Method::POST])
+        .allow_headers(Any);
+
+    let app = Router::new()
+        .route("/tasks", post(create_task).get(list_tasks))
+        .route("/tasks/:id", get(get_task))
+        .layer(cors)
+        .with_state(state);
+
+    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await?;
+    println!("Server running on http://localhost:3000");
+    axum::serve(listener, app).await?;
+
+    Ok(())
+}
+
+async fn list_tasks(
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<Vec<TaskResponse>>, (StatusCode, String)> {
+    let tasks = Task::find()
+        .order_by_desc(task::Column::CreatedAt)
+        .all(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    let response = tasks
+        .into_iter()
+        .map(|t| TaskResponse {
+            id: t.id,
+            goal: t.goal,
+            status: t.status,
+            logs: t.logs,
+            answer: t.answer,
+            created_at: t.created_at,
+        })
+        .collect();
+
+    Ok(Json(response))
+}
+
+async fn create_task(
+    State(state): State<Arc<AppState>>,
+    Json(payload): Json<CreateTaskRequest>,
+) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+    let task_id = Uuid::new_v4();
+
+    // Initial save
+    let new_task = task::ActiveModel {
+        id: Set(task_id),
+        goal: Set(payload.goal.clone()),
+        status: Set("running".to_string()),
+        logs: Set(String::new()),
+        answer: Set(None),
+        created_at: Set(Utc::now().into()),
+    };
+
+    new_task
+        .insert(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    // Run agent in the background or synchronously for now as requested
+    // "it can be created and executed on the agent with an api. the resulting logs after execution will get returned by the endpoint"
+    // This implies we wait for it to finish.
+
+    let mut agent = Agent::new(
+        state.zen_api_key.clone(),
+        state.tavily_api_key.clone(),
+        payload.goal.clone(),
+    );
+
+    let (logs, answer) = match agent.run().await {
+        Ok((logs, answer)) => (logs, answer),
+        Err(e) => (format!("Execution failed: {}", e), None),
+    };
+
+    // Update with final logs and status
+    let mut task: task::ActiveModel = Task::find_by_id(task_id)
+        .one(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+        .ok_or((
+            StatusCode::NOT_FOUND,
+            "Task not found after insert".to_string(),
+        ))?
+        .into();
+
+    task.logs = Set(logs.clone());
+    task.answer = Set(answer.clone());
+    task.status = Set("completed".to_string());
+
+    let updated_task = task
+        .update(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    Ok(Json(TaskResponse {
+        id: updated_task.id,
+        goal: updated_task.goal,
+        status: updated_task.status,
+        logs: updated_task.logs,
+        answer: updated_task.answer,
+        created_at: updated_task.created_at,
+    }))
+}
+
+async fn get_task(
+    Path(id): Path<Uuid>,
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+    let task = Task::find_by_id(id)
+        .one(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+        .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))?;
+
+    Ok(Json(TaskResponse {
+        id: task.id,
+        goal: task.goal,
+        status: task.status,
+        logs: task.logs,
+        answer: task.answer,
+        created_at: task.created_at,
+    }))
+}

src/tools.rs

@@ -0,0 +1,129 @@
+use crate::api::{self, FunctionDefinition, Message, Tool, ToolCall};
+use std::collections::HashMap;
+
+pub fn get_tools() -> Vec<Tool> {
+    vec![
+        Tool {
+            tool_type: "function".to_string(),
+            function: FunctionDefinition {
+                name: "google_search".to_string(),
+                description: "Search the web for information".to_string(),
+                parameters: serde_json::json!({
+                    "type": "object",
+                    "properties": {
+                        "query": {
+                            "type": "string",
+                            "description": "The search query"
+                        }
+                    },
+                    "required": ["query"]
+                }),
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: FunctionDefinition {
+                name: "write_file".to_string(),
+                description: "Write content to a file. Ensure parent directories exist."
+                    .to_string(),
+                parameters: serde_json::json!({
+                    "type": "object",
+                    "properties": {
+                        "path": {
+                            "type": "string",
+                            "description": "The path to the file to write"
+                        },
+                        "content": {
+                            "type": "string",
+                            "description": "The content to write to the file"
+                        }
+                    },
+                    "required": ["path", "content"]
+                }),
+            },
+        },
+        Tool {
+            tool_type: "function".to_string(),
+            function: FunctionDefinition {
+                name: "finish".to_string(),
+                description: "Finish the task".to_string(),
+                parameters: serde_json::json!({
+                    "type": "object",
+                    "properties": {
+                        "result": {
+                            "type": "string",
+                            "description": "The result of the task"
+                        }
+                    },
+                    "required": ["result"]
+                }),
+            },
+        },
+    ]
+}
+
+pub async fn handle_tool_call(
+    tool_call: &ToolCall,
+    tavily_api_key: &Option<String>,
+) -> Result<(Message, bool, Option<String>), Box<dyn std::error::Error>> {
+    let mut file_written = false;
+    let mut answer = None;
+    let name = &tool_call.function.name;
+
+    let (content, written) = if name == "google_search" {
+        let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
+        let query = args.get("query").ok_or("Missing query argument")?;
+
+        println!(
+            "--- Executing tool: google_search(query: \"{}\") ---",
+            query
+        );
+
+        let search_result = if let Some(key) = tavily_api_key {
+            match api::perform_search(query, key).await {
+                Ok(results) => results,
+                Err(e) => format!("Search error: {}", e),
+            }
+        } else {
+            "Error: TAVILY_API_KEY is not set. Cannot perform real search.".to_string()
+        };
+        (search_result, false)
+    } else if name == "write_file" {
+        let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
+        let path = args.get("path").ok_or("Missing path argument")?;
+        let content = args.get("content").ok_or("Missing content argument")?;
+
+        println!("--- Executing tool: write_file(path: \"{}\") ---", path);
+
+        let write_result = match std::fs::write(path, content) {
+            Ok(_) => {
+                file_written = true;
+                format!("Successfully wrote content to {}", path)
+            }
+            Err(e) => format!("Error writing to {}: {}", path, e),
+        };
+        (write_result, file_written)
+    } else if name == "finish" {
+        let args: HashMap<String, String> = serde_json::from_str(&tool_call.function.arguments)?;
+        let result = args.get("result").ok_or("Missing result argument")?;
+
+        println!("--- Finishing task: {}", result);
+
+        file_written = true;
+        answer = Some(result.clone());
+        (result.clone(), file_written)
+    } else {
+        (format!("Error: Unknown tool {}", name), false)
+    };
+
+    Ok((
+        Message {
+            role: "tool".to_string(),
+            content: Some(content),
+            tool_calls: None,
+            tool_call_id: Some(tool_call.id.clone()),
+        },
+        written,
+        answer,
+    ))
+}