use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::Result;
use crate::harness::context::RunContext;
use crate::harness::events::EventSink;
use crate::harness::ids::{RunId, ThreadId};
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum ToolFormat {
#[default]
Json,
Xml,
PType {
parameters: Vec<String>,
},
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ToolSchema {
pub name: String,
pub description: String,
pub parameters: Value,
#[serde(default, skip_serializing_if = "ToolFormat::is_json")]
pub format: ToolFormat,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ToolCall {
pub id: String,
pub name: String,
#[serde(default)]
pub arguments: Value,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ToolResult {
pub call_id: String,
pub name: String,
pub content: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub raw: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
#[serde(default)]
pub elapsed_ms: u64,
}
#[derive(Clone)]
pub struct ToolExecutionContext {
pub run_id: RunId,
pub thread_id: Option<ThreadId>,
pub depth: usize,
pub max_turn_output_tokens: Option<u32>,
pub events: EventSink,
}
impl ToolExecutionContext {
pub fn from_run_context<Ctx>(ctx: &RunContext<Ctx>) -> Self {
Self {
run_id: ctx.config.run_id.clone(),
thread_id: ctx.config.thread_id.clone(),
depth: ctx.config.depth,
max_turn_output_tokens: ctx.config.max_turn_output_tokens,
events: ctx.events.clone(),
}
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ToolDelta {
pub call_id: String,
pub content: String,
}
#[async_trait]
pub trait Tool<State: Send + Sync>: Send + Sync {
fn name(&self) -> &str;
fn description(&self) -> &str;
fn schema(&self) -> ToolSchema;
async fn call(&self, state: &State, call: ToolCall) -> Result<ToolResult>;
async fn call_with_context(
&self,
state: &State,
call: ToolCall,
context: ToolExecutionContext,
) -> Result<ToolResult> {
let _ = context;
self.call(state, call).await
}
}
pub struct ToolRegistry<State> {
pub(crate) tools: HashMap<String, Arc<dyn Tool<State>>>,
}