chat agent
This commit is contained in:
parent
b9c6c831c1
commit
b44e4ed6f9
11 changed files with 626 additions and 16 deletions
|
|
@ -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() {
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue