agent module
This commit is contained in:
@@ -14,6 +14,7 @@ serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
shared-kernel = { workspace = true }
|
||||
spacetimedb = { workspace = true, optional = true }
|
||||
tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread"] }
|
||||
dotenvy = { workspace = true }
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,151 +0,0 @@
|
||||
//! Rig-like agent traits for Genarrative – generic, LLM-agnostic abstractions for
|
||||
//! building agents that can be backed by any LLM provider via [`platform_llm`].
|
||||
//!
|
||||
//! These traits follow the pattern popularized by the [`rig`](https://github.com/0xPlaygrounds/rig)
|
||||
//! library: define [`Agent`], [`Tool`], [`AgentMemory`], and [`AgentBuilder`] as generic
|
||||
//! building blocks, then plug in concrete implementations with `platform_llm::LlmClient`.
|
||||
|
||||
use std::future::Future;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Error
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Errors that can occur during agent execution.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PromptError {
|
||||
/// LLM completion API call failed.
|
||||
CompletionError(String),
|
||||
/// A tool execution failed.
|
||||
ToolError(String),
|
||||
/// Maximum conversation turns reached without producing a final response.
|
||||
MaxTurnsReached { max_turns: usize },
|
||||
}
|
||||
|
||||
impl std::fmt::Display for PromptError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::CompletionError(msg) => write!(f, "completion error: {msg}"),
|
||||
Self::ToolError(msg) => write!(f, "tool error: {msg}"),
|
||||
Self::MaxTurnsReached { max_turns } => {
|
||||
write!(f, "max turns reached: {max_turns}")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for PromptError {}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Core data types
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A tool call request emitted by the LLM.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCall {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub args: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Output of a single prompt cycle.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PromptOutput {
|
||||
pub text: String,
|
||||
pub tool_calls: Vec<ToolCall>,
|
||||
}
|
||||
|
||||
/// Result of executing a single tool call — holds the tool name and the
|
||||
/// serialized output so the agent loop can feed it back into the conversation.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCallResult {
|
||||
pub tool_name: String,
|
||||
pub output: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Flow control for tool execution within the agent loop.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Flow {
|
||||
/// Proceed with the tool call.
|
||||
Continue,
|
||||
/// Skip this tool call but continue processing other responses.
|
||||
Skip,
|
||||
/// Stop the agent execution entirely.
|
||||
Stop,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Core traits
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Core agent trait: processes a message and returns text output with optional
|
||||
/// tool calls.
|
||||
///
|
||||
/// The generic `Message` type represents the format of conversation messages
|
||||
/// (e.g., `Vec<LlmMessage>` when integrated with [`platform_llm`]).
|
||||
pub trait Agent<Message> {
|
||||
/// Process a prompt and return the output.
|
||||
fn prompt(
|
||||
&self,
|
||||
prompt: impl Into<Message> + Send,
|
||||
) -> impl Future<Output = Result<PromptOutput, PromptError>> + Send;
|
||||
|
||||
/// Called before executing a tool call to decide whether to continue, skip,
|
||||
/// or stop.
|
||||
fn before_tool_call(&self, tool_call: &ToolCall) -> Flow {
|
||||
let _ = tool_call;
|
||||
Flow::Continue
|
||||
}
|
||||
}
|
||||
|
||||
/// Conversation memory storing message history.
|
||||
pub trait AgentMemory<Message> {
|
||||
/// Returns a reference to the stored messages.
|
||||
fn get_memory(&self) -> &[Message];
|
||||
/// Appends a single message to the end of the memory.
|
||||
fn append_message(&mut self, message: Message);
|
||||
}
|
||||
|
||||
/// Builder pattern for constructing an agent with a system prompt, tools, and
|
||||
/// memory.
|
||||
pub trait AgentBuilder<Message, T: Agent<Message>> {
|
||||
/// Creates a new builder.
|
||||
fn new() -> Self;
|
||||
/// Sets the system prompt.
|
||||
fn system_prompt(self, system_prompt: impl Into<String>) -> Self;
|
||||
/// Registers a tool.
|
||||
fn tool(self, tool: impl Tool) -> Self;
|
||||
/// Sets the conversation memory.
|
||||
fn memory(self, memory: impl AgentMemory<Message>) -> Self;
|
||||
/// Builds the final agent.
|
||||
fn build(self) -> T;
|
||||
}
|
||||
|
||||
/// A tool that can be invoked by an agent.
|
||||
pub trait Tool: Sized {
|
||||
/// Unique name identifier for the tool.
|
||||
const NAME: &'static str;
|
||||
/// Error type returned by [`Tool::call`].
|
||||
type Error: std::error::Error + 'static;
|
||||
/// Arguments type, deserialized from the LLM's tool call JSON.
|
||||
type Args: for<'a> Deserialize<'a>;
|
||||
/// Output type from tool execution, serializable for the agent response.
|
||||
type Output: Serialize;
|
||||
|
||||
/// Returns the tool's name at runtime (defaults to [`Tool::NAME`]).
|
||||
fn tool_name(&self) -> &'static str {
|
||||
Self::NAME
|
||||
}
|
||||
|
||||
/// Human-readable description of what the tool does.
|
||||
fn description(&self) -> String;
|
||||
/// JSON Schema describing the tool's parameters.
|
||||
fn parameters(&self) -> serde_json::Value;
|
||||
/// Execute the tool with the given arguments.
|
||||
fn call(
|
||||
&self,
|
||||
args: Self::Args,
|
||||
) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send;
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
use std::pin::Pin;
|
||||
use crate::agent::{Tool, ToolDyn};
|
||||
use crate::agent::error::PromptError;
|
||||
use crate::agent::hook::Hook;
|
||||
use crate::agent::memory::AgentMemory;
|
||||
use crate::agent::run::{Flow, PromptRequest};
|
||||
|
||||
pub struct Agent<M: LlmApiAdaptor<Message>, Message> {
|
||||
pub model: M,
|
||||
pub tools: Vec<Box<dyn ToolDyn>>,
|
||||
pub hooks: Vec<Box<dyn Hook>>,
|
||||
pub default_max_turns: usize,
|
||||
pub system_prompt: Option<Message>,
|
||||
pub memory: Option<Box<dyn AgentMemory<Message> + Sync>>,
|
||||
pub context: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
impl<M, Message> Agent<M, Message>
|
||||
where
|
||||
M: LlmApiAdaptor<Message> + 'static,
|
||||
Message: Send + 'static,
|
||||
{
|
||||
pub fn new(model: M) -> Self {
|
||||
Self {
|
||||
model,
|
||||
tools: Vec::new(),
|
||||
hooks: Vec::new(),
|
||||
default_max_turns: 10,
|
||||
system_prompt: None,
|
||||
memory: None,
|
||||
context: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn tool(mut self, tool: impl Tool + Send + Sync + 'static) -> Self {
|
||||
self.tools.push(Box::new(tool));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn system_prompt(mut self, msg: Message) -> Self {
|
||||
self.system_prompt = Some(msg);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn memory(mut self, mem: impl AgentMemory<Message> + Sync + 'static) -> Self {
|
||||
self.memory = Some(Box::new(mem));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn hook(mut self, hook: impl Hook + 'static) -> Self {
|
||||
self.hooks.push(Box::new(hook));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn max_turns(mut self, n: usize) -> Self {
|
||||
self.default_max_turns = n;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn tools(&self) -> &[Box<dyn ToolDyn>] {
|
||||
&self.tools
|
||||
}
|
||||
|
||||
pub fn call_tool<'s>(
|
||||
&'s self,
|
||||
name: &str,
|
||||
args: serde_json::Value,
|
||||
) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send + 's>> {
|
||||
let tools: Vec<&Box<dyn ToolDyn>> = self.tools.iter().collect();
|
||||
let hooks: Vec<&Box<dyn Hook>> = self.hooks.iter().collect();
|
||||
let name = name.to_string();
|
||||
let context = self.context.clone().unwrap_or_default();
|
||||
Box::pin(async move {
|
||||
let mut json_output = {
|
||||
let mut found: Option<serde_json::Value> = None;
|
||||
for tool in tools {
|
||||
if tool.tool_name() == name {
|
||||
found = Some(tool.call_with_context(args, context).await?);
|
||||
break;
|
||||
}
|
||||
}
|
||||
found.ok_or_else(|| format!("unknown tool: {name}"))?
|
||||
};
|
||||
// Run after_tool_call hooks
|
||||
for hook in &hooks {
|
||||
match hook.after_tool_call(&name, &mut json_output) {
|
||||
Flow::Stop => return Err("tool call output rejected by hook".to_string()),
|
||||
Flow::Skip => {
|
||||
json_output = serde_json::Value::Null;
|
||||
break;
|
||||
}
|
||||
Flow::Continue => {}
|
||||
}
|
||||
}
|
||||
serde_json::to_string(&json_output).map_err(|e| e.to_string())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn prompt(&self, message: impl Into<Message> + Send) -> PromptRequest<'_, M, Message>
|
||||
where
|
||||
Message: 'static,
|
||||
{
|
||||
PromptRequest::new(self, message.into())
|
||||
}
|
||||
}
|
||||
|
||||
pub trait LlmApiAdaptor<Message>: Send + Sync {
|
||||
fn complete(
|
||||
&self,
|
||||
messages: &[Message],
|
||||
) -> impl Future<Output = Result<String, PromptError>> + Send;
|
||||
fn tool_result_message(&self, tool_name: &str, output: &str) -> Message;
|
||||
|
||||
fn build_assistant_message(&self, text: &str) -> Message;
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
use crate::agent::hook::Hook;
|
||||
use crate::agent::memory::AgentMemory;
|
||||
use crate::agent::agent::{Agent, LlmApiAdaptor};
|
||||
use crate::agent::tool::Tool;
|
||||
|
||||
pub trait AgentBuilder<Message, M: LlmApiAdaptor<Message>> {
|
||||
type Client;
|
||||
|
||||
fn new() -> Self;
|
||||
fn with_client(self, client: Self::Client) -> Self;
|
||||
fn system_prompt(self, system_prompt: impl Into<String>) -> Self;
|
||||
fn tool(self, tool: impl Tool + Send + Sync + 'static) -> Self;
|
||||
fn add_hook(self, hook: impl Hook + 'static) -> Self;
|
||||
fn max_turns(self, n: usize) -> Self;
|
||||
fn memory(self, memory: impl AgentMemory<Message> + Sync + 'static) -> Self;
|
||||
|
||||
fn context(self, context: serde_json::Value) -> Self;
|
||||
fn build(self) -> Agent<M, Message>;
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PromptError {
|
||||
CompletionError(String),
|
||||
ToolError(String),
|
||||
MaxTurnsReached { max_turns: usize },
|
||||
}
|
||||
|
||||
impl std::fmt::Display for PromptError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::CompletionError(msg) => write!(f, "completion error: {msg}"),
|
||||
Self::ToolError(msg) => write!(f, "tool error: {msg}"),
|
||||
Self::MaxTurnsReached { max_turns } => {
|
||||
write!(f, "max turns reached: {max_turns}")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for PromptError {}
|
||||
@@ -0,0 +1,24 @@
|
||||
use crate::agent::run::Flow;
|
||||
use crate::agent::tool::ToolCall;
|
||||
|
||||
pub trait Hook: Send + Sync {
|
||||
fn before_tool_call(&self, tool_call: &ToolCall) -> Flow;
|
||||
|
||||
/// Called after a tool call completes, before the JSON output is serialized to a string.
|
||||
/// The `output` value can be modified in place.
|
||||
/// Return `Flow::Stop` to abort the agent loop, `Flow::Skip` to discard this result,
|
||||
/// or `Flow::Continue` to proceed normally.
|
||||
fn after_tool_call(
|
||||
&self,
|
||||
_tool_name: &str,
|
||||
_output: &mut serde_json::Value,
|
||||
) -> Flow {
|
||||
Flow::Continue
|
||||
}
|
||||
}
|
||||
|
||||
impl Hook for () {
|
||||
fn before_tool_call(&self, _tool_call: &ToolCall) -> Flow {
|
||||
Flow::Continue
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
pub trait AgentMemory<Message> {
|
||||
fn get_memory(&self) -> &[Message];
|
||||
fn append_message(&mut self, message: Message);
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
use tool::{Tool, ToolDyn};
|
||||
|
||||
pub mod tool;
|
||||
pub mod agent;
|
||||
pub mod agent_builder;
|
||||
pub mod run;
|
||||
pub mod memory;
|
||||
pub mod error;
|
||||
pub mod hook;
|
||||
@@ -0,0 +1,273 @@
|
||||
use crate::agent::agent::Agent;
|
||||
use crate::agent::agent::LlmApiAdaptor;
|
||||
use crate::agent::error::PromptError;
|
||||
use crate::agent::hook::Hook;
|
||||
use crate::agent::tool::{ToolCall, ToolDyn};
|
||||
use serde::Deserialize;
|
||||
use std::pin::Pin;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PromptOutput {
|
||||
pub text: String,
|
||||
pub tool_calls: Vec<ToolCall>,
|
||||
}
|
||||
|
||||
/// Flow control for tool execution within the agent loop.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Flow {
|
||||
Continue,
|
||||
/// Skip this tool call but continue processing other responses.
|
||||
Skip,
|
||||
/// Stop the agent execution entirely.
|
||||
Stop,
|
||||
}
|
||||
|
||||
pub struct PromptRequest<'a, M: LlmApiAdaptor<Message> + 'a, Message: 'a> {
|
||||
agent: &'a Agent<M, Message>,
|
||||
message: Message,
|
||||
system_prompt: Option<Message>,
|
||||
hooks: Vec<Box<dyn Hook + 'static>>,
|
||||
max_turns: usize,
|
||||
context: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
impl<'a, M, Message> PromptRequest<'a, M, Message>
|
||||
where
|
||||
M: LlmApiAdaptor<Message> + 'a,
|
||||
Message: 'a,
|
||||
{
|
||||
pub fn new(agent: &'a Agent<M, Message>, message: Message) -> Self {
|
||||
Self {
|
||||
agent,
|
||||
message,
|
||||
system_prompt: None,
|
||||
hooks: Vec::new(),
|
||||
max_turns: agent.default_max_turns,
|
||||
context: agent.context.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn system_prompt(mut self, msg: Message) -> Self {
|
||||
self.system_prompt = Some(msg);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn add_hook(mut self, hook: impl Hook + 'static) -> Self {
|
||||
self.hooks.push(Box::new(hook));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn max_turns(mut self, n: usize) -> Self {
|
||||
self.max_turns = n;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, M, Message> IntoFuture for PromptRequest<'a, M, Message>
|
||||
where
|
||||
M: LlmApiAdaptor<Message> + Send + Sync + 'a,
|
||||
Message: Send + Sync + Clone + 'a,
|
||||
{
|
||||
type Output = Result<PromptOutput, PromptError>;
|
||||
type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;
|
||||
|
||||
fn into_future(self) -> Self::IntoFuture {
|
||||
let agent = self.agent;
|
||||
let message = self.message;
|
||||
let per_prompt_system_prompt = self.system_prompt;
|
||||
let extra_hooks = self.hooks;
|
||||
let max_turns = self.max_turns;
|
||||
|
||||
Box::pin(async move {
|
||||
let mut memory: Vec<Message> = Vec::new();
|
||||
|
||||
// Pre-populate from agent's memory backend if available
|
||||
if let Some(ref mem_backend) = agent.memory {
|
||||
memory.extend(mem_backend.get_memory().iter().cloned());
|
||||
}
|
||||
|
||||
// Effective system prompt: per-prompt override > agent default
|
||||
match per_prompt_system_prompt {
|
||||
Some(sp) => memory.push(sp),
|
||||
None => {
|
||||
if let Some(ref sp) = agent.system_prompt {
|
||||
memory.push(sp.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
memory.push(message);
|
||||
|
||||
for turn in 0..max_turns {
|
||||
// Delegate completion to the model (rig's CompletionModel pattern)
|
||||
let text = agent.model.complete(&memory).await?;
|
||||
|
||||
// Try to parse the LLM reply as JSON (handle markdown fences)
|
||||
let cleaned = clean_json_response(&text);
|
||||
|
||||
match serde_json::from_str::<LlmJsonResponse>(&cleaned) {
|
||||
Ok(json_resp) => {
|
||||
// Check for <end/> marker — stop the turn immediately
|
||||
let reply_has_end = json_resp.reply_text.contains("<end/>");
|
||||
let clean_text = json_resp
|
||||
.reply_text
|
||||
.replace("<end/>", "")
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
memory.push(agent.model.build_assistant_message(&clean_text));
|
||||
|
||||
if reply_has_end {
|
||||
return Ok(PromptOutput {
|
||||
text: clean_text,
|
||||
tool_calls: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
let tool_calls: Vec<ToolCall> = json_resp
|
||||
.tool_calls
|
||||
.into_iter()
|
||||
.map(|tc| ToolCall {
|
||||
id: format!("call_{turn}"),
|
||||
name: tc.tool_name,
|
||||
args: tc.args,
|
||||
})
|
||||
.collect();
|
||||
|
||||
if tool_calls.is_empty() {
|
||||
return Ok(PromptOutput {
|
||||
text: clean_text,
|
||||
tool_calls: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
for tc in &tool_calls {
|
||||
match run_hooks(&agent.hooks, &extra_hooks, tc) {
|
||||
Flow::Stop => {
|
||||
return Err(PromptError::ToolError(
|
||||
"tool call rejected by hook".to_string(),
|
||||
));
|
||||
}
|
||||
Flow::Skip => {
|
||||
let msg = agent
|
||||
.model
|
||||
.tool_result_message(&tc.name, "(skipped by hook)");
|
||||
memory.push(msg);
|
||||
continue;
|
||||
}
|
||||
Flow::Continue => {}
|
||||
}
|
||||
|
||||
let result = {
|
||||
let tools: Vec<&Box<dyn ToolDyn>> = agent.tools.iter().collect();
|
||||
let name = tc.name.clone();
|
||||
let args = tc.args.clone();
|
||||
let context:Value = self.context.clone().into();
|
||||
let fut = async move {
|
||||
for tool in tools {
|
||||
if tool.tool_name() == name {
|
||||
return tool.call_with_context(args, context.clone()).await
|
||||
}
|
||||
}
|
||||
Err(format!("unknown tool: {name}"))
|
||||
};
|
||||
fut.await
|
||||
};
|
||||
|
||||
match result {
|
||||
Ok(mut json_output) => {
|
||||
// Run after_tool_call hooks to allow output modification
|
||||
for hook in agent.hooks.iter().chain(extra_hooks.iter()) {
|
||||
match hook.after_tool_call(&tc.name, &mut json_output) {
|
||||
Flow::Stop => {
|
||||
return Err(PromptError::ToolError(
|
||||
"tool call output rejected by hook".to_string(),
|
||||
));
|
||||
}
|
||||
Flow::Skip => {
|
||||
json_output = Value::Null;
|
||||
break;
|
||||
}
|
||||
Flow::Continue => {}
|
||||
}
|
||||
}
|
||||
let output_str = serde_json::to_string(&json_output)
|
||||
.map_err(|e| PromptError::ToolError(e.to_string()))?;
|
||||
let msg = agent.model.tool_result_message(&tc.name, &output_str);
|
||||
memory.push(msg);
|
||||
}
|
||||
Err(e) => {
|
||||
let msg = agent
|
||||
.model
|
||||
.tool_result_message(&tc.name, &format!("error: {e}"));
|
||||
memory.push(msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
// Not valid JSON — treat as plain-text final response
|
||||
memory.push(agent.model.build_assistant_message(&text));
|
||||
return Ok(PromptOutput {
|
||||
text,
|
||||
tool_calls: vec![],
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(PromptError::MaxTurnsReached { max_turns })
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn run_hooks(
|
||||
agent_hooks: &[Box<dyn Hook>],
|
||||
extra_hooks: &[Box<dyn Hook>],
|
||||
tool_call: &ToolCall,
|
||||
) -> Flow {
|
||||
for hook in agent_hooks.iter().chain(extra_hooks.iter()) {
|
||||
match hook.before_tool_call(tool_call) {
|
||||
Flow::Stop => return Flow::Stop,
|
||||
Flow::Skip => return Flow::Skip,
|
||||
Flow::Continue => {}
|
||||
}
|
||||
}
|
||||
Flow::Continue
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct LlmJsonResponse {
|
||||
reply_text: String,
|
||||
#[serde(default)]
|
||||
tool_calls: Vec<LlmToolCallRequest>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct LlmToolCallRequest {
|
||||
tool_name: String,
|
||||
#[serde(default)]
|
||||
args: serde_json::Value,
|
||||
}
|
||||
|
||||
pub(crate) fn clean_json_response(text: &str) -> String {
|
||||
let text = text.trim();
|
||||
if text.starts_with("```") {
|
||||
let lines: Vec<&str> = text.lines().collect();
|
||||
let mut cleaned = Vec::new();
|
||||
let mut in_code = false;
|
||||
for line in lines {
|
||||
if line.trim().starts_with("```") {
|
||||
in_code = !in_code;
|
||||
continue;
|
||||
}
|
||||
if in_code {
|
||||
cleaned.push(line);
|
||||
}
|
||||
}
|
||||
if !cleaned.is_empty() {
|
||||
return cleaned.join("\n").trim().to_string();
|
||||
}
|
||||
}
|
||||
text.to_string()
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::pin::Pin;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCall {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub args: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCallResult {
|
||||
pub tool_name: String,
|
||||
pub output: serde_json::Value,
|
||||
}
|
||||
|
||||
pub trait Tool: Sized {
|
||||
const NAME: &'static str;
|
||||
type Error: std::error::Error + 'static;
|
||||
type Args: for<'a> Deserialize<'a>;
|
||||
type Output: Serialize;
|
||||
|
||||
fn tool_name(&self) -> &'static str {
|
||||
Self::NAME
|
||||
}
|
||||
|
||||
/// Human-readable description of what the tool does.
|
||||
fn description(&self) -> String;
|
||||
/// JSON Schema describing the tool's parameters.
|
||||
fn parameters(&self) -> serde_json::Value;
|
||||
fn call(
|
||||
&self,
|
||||
args: Self::Args,
|
||||
) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send;
|
||||
|
||||
fn call_with_context(
|
||||
&self,
|
||||
args: Self::Args,
|
||||
_context: serde_json::Value,
|
||||
) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send {
|
||||
self.call(args)
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper trait to allow for dynamic dispatch of simple tools.
|
||||
pub trait ToolDyn: Send + Sync {
|
||||
fn tool_name(&self) -> &'static str;
|
||||
fn description(&self) -> String;
|
||||
fn parameters(&self) -> serde_json::Value;
|
||||
fn call_with_context(
|
||||
&self,
|
||||
args: serde_json::Value,
|
||||
_context: serde_json::Value,
|
||||
) -> Pin<Box<dyn Future<Output = Result<serde_json::Value, String>> + Send + '_>>;
|
||||
}
|
||||
|
||||
impl<T: Tool + Send + Sync> ToolDyn for T {
|
||||
fn tool_name(&self) -> &'static str {
|
||||
T::NAME
|
||||
}
|
||||
|
||||
fn description(&self) -> String {
|
||||
self.description()
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
self.parameters()
|
||||
}
|
||||
|
||||
fn call_with_context(
|
||||
&self,
|
||||
args: serde_json::Value,
|
||||
context: serde_json::Value,
|
||||
) -> Pin<Box<dyn Future<Output = Result<serde_json::Value, String>> + Send + '_>> {
|
||||
Box::pin(async move {
|
||||
let parsed: T::Args = serde_json::from_value(args)
|
||||
.map_err(|e| format!("bad args for {}: {e}", T::NAME))?;
|
||||
let output = self
|
||||
.call_with_context(parsed, context)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
serde_json::to_value(&output).map_err(|e| e.to_string())
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user