pavel/bot
1
0
Fork 0
Code Issues Pull requests Projects Releases 22 Packages Wiki Activity Actions

more refactoring

Browse source
This commit is contained in:
pavel 2026-02-11 01:41:07 +01:00
commit 7268d49b4a
12 changed files with 319 additions and 189 deletions
Download patch file Download diff file Expand all files Collapse all files

116
src/domain/agent/mod.rs
View file

@ -17,12 +17,14 @@ pub struct Agent {
answer: Option<String>,
}
use crate::error::{AppError, AppResult};
impl Agent {
pub fn new(
zen_api_key: Option<String>,
tavily_api_key: Option<String>,
initial_message: String,
) -> Result<Self, Box<dyn std::error::Error>> {
) -> 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
@ -49,7 +51,8 @@ impl Agent {
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(60))
.build()?;
.build()
.map_err(|e| AppError::Internal(format!("Failed to build HTTP client: {}", e)))?;
Ok(Self {
client,
@ -71,7 +74,7 @@ impl Agent {
pub async fn run(
&mut self,
config: &crate::config::Config,
) -> Result<(String, Option<String>), Box<dyn std::error::Error>> {
) -> AppResult<(String, Option<String>)> {
let mut finished = false;
let start_time = Instant::now();
let max_duration = Duration::from_secs(config.agent_max_duration_secs);
@ -80,47 +83,29 @@ impl Agent {
while !finished {
if start_time.elapsed() > max_duration {
return Err("Agent run timed out".into());
return Err(AppError::Internal("Agent run timed out".into()));
}
if turns >= max_turns {
return Err("Agent run exceeded max turns".into());
return Err(AppError::Internal("Agent run exceeded max turns".into()));
}
turns += 1;
let request = ChatRequest {
model: "kimi-k2.5".to_string(),
messages: self.messages.clone(),
tools: self.tools.clone(),
};
let current_role = self
.messages
.last()
.map(|m| m.role.as_str())
.unwrap_or("unknown");
self.log(&format!(
"--- Sending request to Zen API (Role: {}) ---",
self.messages.last().unwrap().role
"\n[Turn {}] Sending request (Last role: {})",
turns, current_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 chat_response = self.call_llm().await?;
let assistant_message = chat_response
.choices
.get(0)
.ok_or("Missing assistant response")?
.ok_or_else(|| AppError::Internal("Missing assistant response".into()))?
.message
.clone();
@ -128,41 +113,78 @@ impl Agent {
if let Some(content) = &assistant_message.content {
if !content.is_empty() {
self.log(&format!("\nAssistant response:\n{}\n", content));
self.log(&format!("\nAssistant: {}", 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?;
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)
.await
.map_err(|e| {
AppError::Internal(format!("Tool execution failed: {}", e))
})?;
if let Some(ans) = tool_answer {
self.answer = Some(ans);
self.log("Task marked as finished by tool.");
}
if let Some(content) = &tool_message.content {
self.log(&format!(
"Tool result ({}): {}",
tool_call.function.name, content
));
self.log(&format!("Tool result: {}", content));
}
self.messages.push(tool_message);
if written {
if is_final {
finished = true;
}
}
// Continue the loop to send tool results back
continue;
}
// No tool calls from assistant, but we only exit if the task was finished
if !finished {
self.log("--- Assistant didn't finish yet. Waiting for next turn... ---");
} else if assistant_message.content.is_some() {
// If assistant just talked without tools, we might be stuck or finished.
// But typically we expect a 'finish' tool call.
self.log("Assistant responded without tool calls.");
// For now we continue unless the assistant explicitly uses a tool to finish,
// or we could add heuristic here if needed.
}
}
self.log("\n--- Execution Finished ---");
Ok((self.logs.clone(), self.answer.clone()))
}
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 response = request_builder.send().await.map_err(AppError::Network)?;
if !response.status().is_success() {
let status = response.status();
let error_text = response
.text()
.await
.unwrap_or_else(|_| "Unknown error".into());
return Err(AppError::Internal(format!(
"API request failed: {} - {}",
status, error_text
)));
}
response
.json()
.await
.map_err(|e| AppError::Internal(format!("Failed to parse LLM response: {}", e)))
}
}

24
src/domain/agent/tools.rs
View file

@ -1,5 +1,15 @@
use super::api::{self, FunctionDefinition, Message, Tool, ToolCall};
use std::collections::HashMap;
use serde::Deserialize;
#[derive(Deserialize)]
struct GoogleSearchArgs {
query: String,
}
#[derive(Deserialize)]
struct FinishArgs {
result: String,
}
pub fn get_tools() -> Vec<Tool> {
vec![
@ -24,13 +34,13 @@ pub fn get_tools() -> Vec<Tool> {
tool_type: "function".to_string(),
function: FunctionDefinition {
name: "finish".to_string(),
description: "Finish the task".to_string(),
description: "Finish the task and provide a final answer".to_string(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"result": {
"type": "string",
"description": "The result of the task"
"description": "The final detailed answer to the task"
}
},
"required": ["result"]
@ -48,8 +58,8 @@ pub async fn handle_tool_call(
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")?;
let args: GoogleSearchArgs = serde_json::from_str(&tool_call.function.arguments)?;
let query = &args.query;
let search_result = if let Some(key) = tavily_api_key {
match api::perform_search(query, key).await {
@ -61,8 +71,8 @@ pub async fn handle_tool_call(
};
(search_result, false)
} 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")?;
let args: FinishArgs = serde_json::from_str(&tool_call.function.arguments)?;
let result = &args.result;
answer = Some(result.clone());
(result.clone(), true)

34
src/domain/tasks.rs
View file

@ -49,13 +49,15 @@ pub struct RecentRunResponse {
pub created_at: chrono::DateTime<chrono::FixedOffset>,
}
use crate::error::AppResult;
pub async fn execute_agent_run(
db: &DatabaseConnection,
_scheduler: &Arc<Scheduler>,
config: &Arc<Config>,
task_id: Uuid,
goal: String,
) -> Result<TaskResponse, Box<dyn std::error::Error>> {
) -> AppResult<TaskResponse> {
let run_id = Uuid::new_v4();
let new_run = task_run::ActiveModel {
@ -67,7 +69,10 @@ pub async fn execute_agent_run(
created_at: Set(Utc::now().into()),
};
new_run.insert(db).await?;
new_run
.insert(db)
.await
.map_err(crate::error::AppError::Database)?;
let mut agent = Agent::new(
config.zen_api_key.clone(),
@ -84,31 +89,36 @@ pub async fn execute_agent_run(
),
};
let mut run: task_run::ActiveModel = TaskRun::find_by_id(run_id)
let run: task_run::ActiveModel = TaskRun::find_by_id(run_id)
.one(db)
.await?
.ok_or("Run not found after insert")?
.await
.map_err(crate::error::AppError::Database)?
.ok_or_else(|| crate::error::AppError::NotFound("Run not found after insert".into()))?
.into();
let mut run = run;
run.logs = Set(logs.clone());
run.answer = Set(answer.clone());
run.status = Set(status);
run.update(db).await?;
run.update(db)
.await
.map_err(crate::error::AppError::Database)?;
get_task_inner(task_id, db).await
}
pub async fn get_task_inner(
id: Uuid,
db: &DatabaseConnection,
) -> Result<TaskResponse, Box<dyn std::error::Error>> {
pub async fn get_task_inner(id: Uuid, db: &DatabaseConnection) -> AppResult<TaskResponse> {
let results = Task::find_by_id(id)
.find_with_related(TaskRun)
.all(db)
.await?;
.await
.map_err(crate::error::AppError::Database)?;
let (t, mut runs) = results.into_iter().next().ok_or("Task not found")?;
let (t, mut runs) = results
.into_iter()
.next()
.ok_or_else(|| crate::error::AppError::NotFound("Task not found".into()))?;
runs.sort_by(|a, b| b.created_at.cmp(&a.created_at));