287 lines
10 KiB
Rust
287 lines
10 KiB
Rust
use super::api::{self, FunctionDefinition, Message, Tool, ToolCall};
|
|
use sea_orm::{
|
|
ColumnTrait, Condition, DatabaseConnection, EntityTrait, QueryFilter, QueryOrder, QuerySelect,
|
|
};
|
|
use serde::Deserialize;
|
|
use std::sync::Arc;
|
|
use uuid::Uuid;
|
|
|
|
#[derive(Deserialize)]
|
|
struct GoogleSearchArgs {
|
|
query: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
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 {
|
|
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: "finish".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 final detailed answer to the task"
|
|
}
|
|
},
|
|
"required": ["result"]
|
|
}),
|
|
},
|
|
},
|
|
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"
|
|
},
|
|
|
|
},
|
|
}),
|
|
},
|
|
},
|
|
Tool {
|
|
tool_type: "function".to_string(),
|
|
function: FunctionDefinition {
|
|
name: "calendar_list_events".to_string(),
|
|
description: "List calendar events".to_string(),
|
|
parameters: serde_json::json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"upcoming": {
|
|
"type": "boolean",
|
|
"description": "If true, only upcoming events will be listed"
|
|
}
|
|
}
|
|
}),
|
|
},
|
|
},
|
|
Tool {
|
|
tool_type: "function".to_string(),
|
|
function: FunctionDefinition {
|
|
name: "calendar_create_event".to_string(),
|
|
description: "Create a new calendar event".to_string(),
|
|
parameters: serde_json::json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {
|
|
"type": "string",
|
|
"description": "Name of the event"
|
|
},
|
|
"from": {
|
|
"type": "string",
|
|
"description": "Start time in ISO 8601 format (e.g., 2023-10-27T10:00:00Z)"
|
|
},
|
|
"to": {
|
|
"type": "string",
|
|
"description": "End time in ISO 8601 format (e.g., 2023-10-27T11:00:00Z)"
|
|
}
|
|
},
|
|
"required": ["name", "from", "to"]
|
|
}),
|
|
},
|
|
},
|
|
]
|
|
}
|
|
|
|
pub async fn handle_tool_call(
|
|
tool_call: &ToolCall,
|
|
tavily_api_key: &Option<String>,
|
|
db: &DatabaseConnection,
|
|
calendar: &Arc<crate::domain::calendar::CalendarClient>,
|
|
user_sub: Option<&str>,
|
|
) -> Result<(Message, bool, Option<String>), Box<dyn std::error::Error>> {
|
|
let mut answer = None;
|
|
let name = &tool_call.function.name;
|
|
|
|
let (content, written) = if name == "google_search" {
|
|
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 {
|
|
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 == "finish" {
|
|
let args: FinishArgs = serde_json::from_str(&tool_call.function.arguments)?;
|
|
let result = &args.result;
|
|
|
|
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 if name == "calendar_list_events" {
|
|
let args: serde_json::Value = serde_json::from_str(&tool_call.function.arguments)?;
|
|
let upcoming = args["upcoming"].as_bool();
|
|
match calendar
|
|
.list_events(user_sub.map(|s| s.to_string()), upcoming)
|
|
.await
|
|
{
|
|
Ok(events) => {
|
|
tracing::info!("{:#?}", events);
|
|
(serde_json::to_string(&events)?, false)
|
|
}
|
|
Err(e) => (format!("Error listing events: {}", e), false),
|
|
}
|
|
} else if name == "calendar_create_event" {
|
|
let args: serde_json::Value = serde_json::from_str(&tool_call.function.arguments)?;
|
|
let name_val = args["name"].as_str().unwrap_or_default();
|
|
let from_val = args["from"].as_str().unwrap_or_default();
|
|
let to_val = args["to"].as_str().unwrap_or_default();
|
|
match calendar
|
|
.create_event(user_sub.map(|s| s.to_string()), name_val, from_val, to_val)
|
|
.await
|
|
{
|
|
Ok(event) => (
|
|
format!("Event created: {}", serde_json::to_string(&event)?),
|
|
false,
|
|
),
|
|
Err(e) => (format!("Error creating event: {}", e), false),
|
|
}
|
|
} 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,
|
|
))
|
|
}
|