bot/src/domain/agent/mod.rs

pub mod api;
pub mod tools;

use chrono::Utc;
use sea_orm::DatabaseConnection;
use std::time::{Duration, Instant};

use self::api::{ChatRequest, ChatResponse, Message, Tool};

pub struct Agent {
    db: DatabaseConnection,
    client: reqwest::Client,
    url: String,
    zen_api_key: Option<String>,
    tavily_api_key: Option<String>,
    pub messages: Vec<Message>,
    tools: Option<Vec<Tool>>,
    logs: String,
    answer: Option<String>,
}

use crate::error::{AppError, AppResult};

impl Agent {
    pub fn new(
        db: DatabaseConnection,
        zen_api_key: Option<String>,
        tavily_api_key: Option<String>,
        initial_message: String,
    ) -> AppResult<Self> {
        let intro = format!(
            "You are an autonomous agent. You have access to tools that can help 
                    you achieve your goals. Use them wisely. The user is unable to respond to you
                    so do not ask for clarification and use the
                    answer tool once you to give your final answer. current date is {}",
            Utc::now().format("%B %d, %Y %H:%M:%S UTC").to_string()
        );
        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,
            },
        ];

        Self::with_messages(db, zen_api_key, tavily_api_key, messages)
    }

    pub fn with_messages(
        db: DatabaseConnection,
        zen_api_key: Option<String>,
        tavily_api_key: Option<String>,
        messages: Vec<Message>,
    ) -> AppResult<Self> {
        let tools = Some(tools::get_tools());

        let client = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(120))
            .build()
            .map_err(|e| AppError::Internal(format!("Failed to build HTTP client: {}", e)))?;

        Ok(Self {
            db,
            client,
            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) {
        self.logs.push_str(message);
        self.logs.push('\n');
    }

    pub async fn run(
        &mut self,
        config: &crate::config::Config,
    ) -> AppResult<(String, Option<String>)> {
        let mut finished = false;
        let start_time = Instant::now();
        let max_duration = Duration::from_secs(config.agent_max_duration_secs);
        let max_turns = config.agent_max_turns;
        let mut turns = 0;

        while !finished {
            if start_time.elapsed() > max_duration {
                return Err(AppError::Internal("Agent run timed out".into()));
            }

            if turns >= max_turns {
                return Err(AppError::Internal("Agent run exceeded max turns".into()));
            }

            turns += 1;
            let current_role = self
                .messages
                .last()
                .map(|m| m.role.as_str())
                .unwrap_or("unknown");
            self.log(&format!(
                "\n[Turn {}] Sending request (Last role: {})",
                turns, current_role
            ));

            let assistant_message = self.execute_turn().await?;

            if let Some(tool_calls) = &assistant_message.tool_calls {
                if tool_calls.iter().any(|tc| tc.function.name == "answer") {
                    finished = true;
                }
            }

            if self.answer.is_some() {
                finished = true;
            }
        }

        self.log("\n--- Execution Finished ---");
        Ok((self.logs.clone(), self.answer.clone()))
    }

    pub async fn execute_turn(&mut self) -> AppResult<Message> {
        let max_sub_turns = 10;
        let mut sub_turns = 0;

        loop {
            sub_turns += 1;
            if sub_turns > max_sub_turns {
                return Err(AppError::Internal(
                    "Interaction cycle turn limit exceeded".into(),
                ));
            }

            let chat_response = self.call_llm().await?;
            let assistant_message = chat_response
                .choices
                .get(0)
                .ok_or_else(|| AppError::Internal("Missing assistant response".into()))?
                .message
                .clone();

            self.messages.push(assistant_message.clone());

            if let Some(content) = &assistant_message.content {
                if !content.is_empty() {
                    self.log(&format!("\nAssistant: {}", content));
                }
            }

            if let Some(tool_calls) = &assistant_message.tool_calls {
                let mut is_final_cycle = false;
                let mut final_answer = None;

                for tool_call in tool_calls {
                    self.log(&format!("Calling tool: {}", tool_call.function.name));

                    let (tool_message, is_final, tool_answer) =
                        tools::handle_tool_call(tool_call, &self.tavily_api_key, &self.db)
                            .await
                            .map_err(|e| {
                                AppError::Internal(format!("Tool execution failed: {}", e))
                            })?;

                    if let Some(ans) = tool_answer {
                        self.answer = Some(ans.clone());
                        final_answer = Some(ans);
                        self.log("Interaction marked as finished by tool.");
                    }

                    if let Some(content) = &tool_message.content {
                        self.log(&format!("Tool result: {}", content));
                    }

                    self.messages.push(tool_message);
                    if is_final {
                        is_final_cycle = true;
                    }
                }

                if is_final_cycle {
                    return Ok(Message {
                        role: "assistant".to_string(),
                        content: final_answer.or(assistant_message.content),
                        tool_calls: None,
                        tool_call_id: None,
                    });
                }

                continue;
            }

            return Ok(assistant_message);
        }
    }

    async fn call_llm(&self) -> AppResult<ChatResponse> {
        let request = ChatRequest {
            model: "kimi-k2.5".to_string(),
            messages: self.messages.clone(),
            tools: self.tools.clone(),
        };

        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 start = std::time::Instant::now();
        let response = request_builder.send().await.map_err(|e| {
            let duration = start.elapsed();
            let is_timeout = e.is_timeout();
            tracing::error!(
                "Network error after {:?} during LLM call (Timeout: {}): {:?}",
                duration,
                is_timeout,
                e
            );
            AppError::Network(e)
        })?;

        let duration = start.elapsed();
        tracing::info!("LLM request completed in {:?}", duration);

        if !response.status().is_success() {
            let status = response.status();
            let body_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown body".into());
            let err = format!("API request failed: {} - {}", status, body_text);
            tracing::error!("{}", err);
            return Err(AppError::Internal(err));
        }

        let response_text = response.text().await.map_err(|e| {
            let err = format!("Failed to read response text: {}", e);
            tracing::error!("{}", err);
            AppError::Internal(err)
        })?;

        serde_json::from_str(&response_text).map_err(|e| {
            let err = format!(
                "Failed to parse LLM response: {} | Raw Body: {}",
                e, response_text
            );
            tracing::error!("{}", err);
            AppError::Internal(err)
        })
    }
}