chat agent

This commit is contained in:
pavel 2026-02-11 16:43:05 +01:00
commit b44e4ed6f9
11 changed files with 626 additions and 16 deletions

View file

@ -3,7 +3,6 @@ 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>>,
@ -29,6 +28,7 @@ pub struct FunctionCall {
pub struct ChatRequest {
pub model: String,
pub messages: Vec<Message>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<Vec<Tool>>,
}
@ -72,6 +72,7 @@ pub async fn perform_search(
query: &str,
api_key: &str,
) -> Result<String, Box<dyn std::error::Error>> {
tracing::info!(query = %query, "Performing Tavily web search");
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(30))
.build()?;
@ -94,6 +95,7 @@ pub async fn perform_search(
}
let search_data: TavilyResponse = response.json().await?;
tracing::info!("Web search yielded {} results", search_data.results.len());
let mut results_text = String::new();
for (i, result) in search_data.results.iter().enumerate() {

View file

@ -2,11 +2,13 @@ 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>,
@ -21,6 +23,7 @@ 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,
@ -55,6 +58,7 @@ impl Agent {
.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,
@ -122,7 +126,7 @@ impl Agent {
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)
tools::handle_tool_call(&tool_call, &self.tavily_api_key, &self.db)
.await
.map_err(|e| {
AppError::Internal(format!("Tool execution failed: {}", e))

View file

@ -1,5 +1,9 @@
use super::api::{self, FunctionDefinition, Message, Tool, ToolCall};
use sea_orm::{
ColumnTrait, Condition, DatabaseConnection, EntityTrait, QueryFilter, QueryOrder, QuerySelect,
};
use serde::Deserialize;
use uuid::Uuid;
#[derive(Deserialize)]
struct GoogleSearchArgs {
@ -11,6 +15,13 @@ struct FinishArgs {
result: String,
}
#[derive(Deserialize)]
struct ListRunsArgs {
from: Option<String>,
to: Option<String>,
task_ids: Option<Vec<String>>,
}
pub fn get_tools() -> Vec<Tool> {
vec![
Tool {
@ -47,12 +58,51 @@ pub fn get_tools() -> Vec<Tool> {
}),
},
},
Tool {
tool_type: "function".to_string(),
function: FunctionDefinition {
name: "list_tasks".to_string(),
description: "List all existing tasks and their goals".to_string(),
parameters: serde_json::json!({
"type": "object",
"properties": {},
}),
},
},
Tool {
tool_type: "function".to_string(),
function: FunctionDefinition {
name: "list_runs".to_string(),
description: "List task runs with optional filters for date and task IDs"
.to_string(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"from": {
"type": "string",
"description": "ISO 8601 date string to filter runs from"
},
"to": {
"type": "string",
"description": "ISO 8601 date string to filter runs to"
},
"task_ids": {
"type": "array",
"items": { "type": "string", "format": "uuid" },
"description": "List of task IDs to filter runs for"
},
},
}),
},
},
]
}
pub async fn handle_tool_call(
tool_call: &ToolCall,
tavily_api_key: &Option<String>,
db: &DatabaseConnection,
) -> Result<(Message, bool, Option<String>), Box<dyn std::error::Error>> {
let mut answer = None;
let name = &tool_call.function.name;
@ -76,6 +126,78 @@ pub async fn handle_tool_call(
answer = Some(result.clone());
(result.clone(), true)
} else if name == "list_tasks" {
tracing::info!("Listing tasks from database");
use crate::entities::task;
let tasks = task::Entity::find()
.order_by_desc(task::Column::CreatedAt)
.all(db)
.await?;
let mut out = String::from("Tasks:\n");
for t in tasks {
out.push_str(&format!("- ID: {}, Goal: {}\n", t.id, t.goal));
}
(out, false)
} else if name == "list_runs" {
use crate::entities::task_run;
tracing::info!(
"Listing runs from database, args: {}",
tool_call.function.arguments
);
let args: ListRunsArgs = serde_json::from_str(&tool_call.function.arguments)?;
tracing::info!(
"Listing runs from database with filters: from={:?}, to={:?}, task_ids={:?}",
args.from,
args.to,
args.task_ids
);
let query = task_run::Entity::find();
let mut condition = Condition::all();
if let Some(from_str) = args.from {
if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(&from_str) {
condition = condition.add(task_run::Column::CreatedAt.gte(dt));
}
}
if let Some(to_str) = args.to {
if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(&to_str) {
condition = condition.add(task_run::Column::CreatedAt.lte(dt));
}
}
if let Some(id_strs) = args.task_ids {
let mut valid_ids = Vec::new();
for id_str in id_strs {
match id_str.parse::<Uuid>() {
Ok(uuid) => valid_ids.push(uuid),
Err(e) => tracing::warn!("Skipping invalid UUID '{}' from LLM: {}", id_str, e),
}
}
if !valid_ids.is_empty() {
condition = condition.add(task_run::Column::TaskId.is_in(valid_ids));
}
}
let runs = query
.filter(condition)
.order_by_desc(task_run::Column::CreatedAt)
.limit(20)
.all(db)
.await?;
let mut out = String::from("Recent Runs:\n");
for r in runs {
out.push_str(&format!(
"- ID: {}, Task ID: {}, Status: {}, Created At: {}, Answer: {:?}\n",
r.id, r.task_id, r.status, r.created_at, r.answer
));
}
(out, false)
} else {
(format!("Error: Unknown tool {}", name), false)
};

View file

@ -76,6 +76,7 @@ pub async fn execute_agent_run(
.map_err(crate::error::AppError::Database)?;
let mut agent = Agent::new(
db.clone(),
config.zen_api_key.clone(),
config.tavily_api_key.clone(),
goal.clone(),

View file

@ -107,6 +107,7 @@ impl Scheduler {
// Start agent in background
let mut agent = Agent::new(
db.clone(),
config.zen_api_key.clone(),
config.tavily_api_key.clone(),
task.goal.clone(),

130
src/server/chat.rs Normal file
View file

@ -0,0 +1,130 @@
use crate::domain::agent::api::{ChatRequest, ChatResponse, Message};
use crate::error::AppError;
use crate::server::AppState;
use axum::{Json, extract::State};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
#[derive(Debug, Deserialize)]
pub struct ChatPayload {
pub messages: Vec<Message>,
}
#[derive(Debug, Serialize)]
pub struct ChatResult {
pub message: Message,
}
pub async fn chat_handler(
State(state): State<Arc<AppState>>,
Json(payload): Json<ChatPayload>,
) -> Result<Json<ChatResult>, AppError> {
let msg_count = payload.messages.len();
tracing::info!("Received chat request with {} messages", msg_count);
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(60))
.build()
.map_err(|e| {
let err = format!("Failed to build HTTP client: {}", e);
tracing::error!("{}", err);
AppError::Internal(err)
})?;
let url = "https://opencode.ai/zen/v1/chat/completions";
let mut messages = payload.messages;
let tools = crate::domain::agent::tools::get_tools();
let max_turns = 10;
let mut turns = 0;
loop {
turns += 1;
if turns > max_turns {
return Err(AppError::Internal("Chat turn limit exceeded".into()));
}
let request = ChatRequest {
model: "big-pickle".to_string(),
messages: messages.clone(),
tools: Some(tools.clone()),
};
let mut request_builder = client.post(url).json(&request);
if let Some(key) = &state.config.zen_api_key {
request_builder = request_builder.header("Authorization", format!("Bearer {}", key));
}
let response = request_builder.send().await.map_err(|e| {
tracing::error!("Network error during chat completion: {}", e);
AppError::Network(e)
})?;
if !response.status().is_success() {
let status = response.status();
let error_text = response
.text()
.await
.unwrap_or_else(|_| "Unknown error".into());
let err = format!("API request failed: {} - {}", status, error_text);
tracing::error!("{}", err);
return Err(AppError::Internal(err));
}
let chat_response: ChatResponse = response.json().await.map_err(|e| {
let err = format!("Failed to parse LLM response: {}", e);
tracing::error!("{}", err);
AppError::Internal(err)
})?;
let assistant_message = chat_response
.choices
.get(0)
.ok_or_else(|| {
let err = "Missing assistant response choices";
tracing::error!("{}", err);
AppError::Internal(err.into())
})?
.message
.clone();
messages.push(assistant_message.clone());
if let Some(tool_calls) = &assistant_message.tool_calls {
for tool_call in tool_calls {
tracing::info!("Chat agent calling tool: {}", tool_call.function.name);
let (tool_message, is_final, answer) =
crate::domain::agent::tools::handle_tool_call(
tool_call,
&state.config.tavily_api_key,
&state.db,
)
.await
.map_err(|e| AppError::Internal(format!("Tool execution failed: {}", e)))?;
messages.push(tool_message.clone());
if is_final {
tracing::info!("Chat agent finished via tool");
return Ok(Json(ChatResult {
message: Message {
role: "assistant".to_string(),
content: answer.or(tool_message.content),
tool_calls: None,
tool_call_id: None,
},
}));
}
}
// After tool calls, we loop back to get another assistant response
continue;
}
// If no tool calls, it's a final response for this turn
tracing::info!("Chat completion successful");
return Ok(Json(ChatResult {
message: assistant_message,
}));
}
}

View file

@ -1,4 +1,5 @@
pub mod auth;
pub mod chat;
pub mod tasks;
use axum::{
@ -120,6 +121,7 @@ fn build_app(state: Arc<AppState>, config: &crate::config::Config) -> Router {
.route("/api/auth/callback", get(auth::auth_callback))
.route("/api/auth/refresh", post(auth::auth_refresh))
.route("/api/auth/logout", post(auth::auth_logout))
.route("/api/chat", post(chat::chat_handler))
.layer(cors)
.layer(tower_http::set_header::SetResponseHeaderLayer::overriding(
axum::http::header::CONTENT_SECURITY_POLICY,