diff --git a/Cargo.lock b/Cargo.lock index d4a5cfa..1be1df2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,18 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "once_cell", + "version_check", + "zerocopy", +] + [[package]] name = "android_system_properties" version = "0.1.5" @@ -186,6 +198,24 @@ dependencies = [ "syn", ] +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -288,6 +318,30 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +dependencies = [ + "ahash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "hashlink" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ba4ff7128dee98c7dc9794b6a411377e1404dba1c97deb8d1a55297bd25d8af" +dependencies = [ + "hashbrown 0.14.5", +] + [[package]] name = "http" version = "1.4.2" @@ -520,6 +574,16 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + [[package]] name = "ipnet" version = "2.12.0" @@ -549,6 +613,17 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libsqlite3-sys" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e99fb7a497b1e3339bc746195567ed8d3e24945ecd636e3619d20b9de9e9149" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + [[package]] name = "litemap" version = "0.8.2" @@ -623,6 +698,12 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + [[package]] name = "potential_utf" version = "0.1.5" @@ -729,10 +810,12 @@ dependencies = [ "chrono", "futures-util", "reqwest", + "rusqlite", "serde", "serde_json", "tokio", "tokio-stream", + "toml", "tower", "tower-http", "ureq", @@ -822,6 +905,20 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rusqlite" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7753b721174eb8ff87a9a0e799e2d7bc3749323e773db92e0984debb00019d6e" +dependencies = [ + "bitflags", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -930,6 +1027,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "serde_spanned" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3" +dependencies = [ + "serde", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -1118,6 +1224,47 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362" +dependencies = [ + "serde", + "serde_spanned", + "toml_datetime", + "toml_edit", +] + +[[package]] +name = "toml_datetime" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22cddaf88f4fbc13c51aebbf5f8eceb5c7c5a9da2ac40a13519eb5b0a0e8f11c" +dependencies = [ + "serde", +] + +[[package]] +name = "toml_edit" +version = "0.22.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" +dependencies = [ + "indexmap", + "serde", + "serde_spanned", + "toml_datetime", + "toml_write", + "winnow", +] + +[[package]] +name = "toml_write" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" + [[package]] name = "tower" version = "0.5.3" @@ -1235,6 +1382,18 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + [[package]] name = "want" version = "0.3.1" @@ -1506,6 +1665,15 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winnow" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" +dependencies = [ + "memchr", +] + [[package]] name = "wit-bindgen" version = "0.57.1" diff --git a/Cargo.toml b/Cargo.toml index 52563c4..f396fd0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,10 +9,12 @@ axum = "0.7" chrono = { version = "0.4", default-features = false, features = ["std", "clock", "serde"] } futures-util = "0.3" reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"] } +rusqlite = { version = "0.32", features = ["bundled"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync"] } tokio-stream = "0.1" +toml = "0.8" tower = { version = "0.5", features = ["util"] } tower-http = { version = "0.6", features = ["cors"] } ureq = { version = "2", default-features = false, features = ["tls"] } diff --git a/examples/openai_compatible_demo.rs b/examples/openai_compatible_demo.rs index 4f3f8a9..35e6240 100644 --- a/examples/openai_compatible_demo.rs +++ b/examples/openai_compatible_demo.rs @@ -60,7 +60,9 @@ fn main() -> Result<()> { description: "A tool-using assistant.".to_string(), system_prompt: "Use tools when useful, then answer the user.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(GetTimeTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/src/agent/agent.rs b/src/agent/agent.rs index 7297969..c0b15c9 100644 --- a/src/agent/agent.rs +++ b/src/agent/agent.rs @@ -1,3 +1,5 @@ +use std::{collections::HashSet, sync::Arc}; + use crate::{Result, event::Event}; use super::{AsyncEventStream, EventStream, InvocationContext, event_stream_from_result}; @@ -7,7 +9,7 @@ pub trait Agent: Send + Sync { fn description(&self) -> &str; - fn sub_agents(&self) -> &[Box] { + fn sub_agents(&self) -> &[Arc] { &[] } @@ -36,3 +38,55 @@ pub fn find_agent<'a>(agent: &'a dyn Agent, name: &str) -> Option<&'a dyn Agent> None } + +pub trait IntoAgentArc { + fn into_agent_arc(self) -> Arc; +} + +impl IntoAgentArc for Arc { + fn into_agent_arc(self) -> Arc { + self + } +} + +impl IntoAgentArc for Arc +where + T: Agent + 'static, +{ + fn into_agent_arc(self) -> Arc { + self + } +} + +impl IntoAgentArc for Box { + fn into_agent_arc(self) -> Arc { + Arc::from(self) + } +} + +impl IntoAgentArc for Box +where + T: Agent + 'static, +{ + fn into_agent_arc(self) -> Arc { + Arc::new(*self) + } +} + +pub fn validate_agent_names(agent: &dyn Agent) -> Result<()> { + let mut seen = HashSet::new(); + validate_agent_names_inner(agent, &mut seen) +} + +fn validate_agent_names_inner(agent: &dyn Agent, seen: &mut HashSet) -> Result<()> { + let name = agent.name().to_string(); + if !seen.insert(name.clone()) { + return Err(crate::Error::msg(format!("duplicate agent name: {name}"))); + } + + for sub_agent in agent.sub_agents() { + validate_agent_names_inner(sub_agent.as_ref(), seen)?; + } + + Ok(()) +} diff --git a/src/agent/custom_agent.rs b/src/agent/custom_agent.rs index 277867c..f54f60e 100644 --- a/src/agent/custom_agent.rs +++ b/src/agent/custom_agent.rs @@ -1,12 +1,14 @@ +use std::sync::Arc; + use crate::{Result, event::Event}; -use super::{Agent, InvocationContext}; +use super::{Agent, IntoAgentArc, InvocationContext}; pub struct CustomAgent { name: String, description: String, handler: Box Result> + Send + Sync>, - sub_agents: Vec>, + sub_agents: Vec>, } impl CustomAgent { @@ -23,12 +25,12 @@ impl CustomAgent { } } - pub fn with_sub_agent(mut self, sub_agent: Box) -> Self { - self.sub_agents.push(sub_agent); + pub fn with_sub_agent(mut self, sub_agent: impl IntoAgentArc) -> Self { + self.sub_agents.push(sub_agent.into_agent_arc()); self } - pub fn with_sub_agents(mut self, sub_agents: Vec>) -> Self { + pub fn with_sub_agents(mut self, sub_agents: Vec>) -> Self { self.sub_agents = sub_agents; self } @@ -43,7 +45,7 @@ impl Agent for CustomAgent { &self.description } - fn sub_agents(&self) -> &[Box] { + fn sub_agents(&self) -> &[Arc] { &self.sub_agents } diff --git a/src/agent/loader.rs b/src/agent/loader.rs new file mode 100644 index 0000000..15c42c0 --- /dev/null +++ b/src/agent/loader.rs @@ -0,0 +1,88 @@ +use std::{collections::HashMap, sync::Arc}; + +use crate::{Error, Result}; + +use super::{Agent, validate_agent_names}; + +pub trait AgentLoader: Send + Sync { + fn list_agents(&self) -> Result>; + + fn load_agent(&self, name: &str) -> Result>; + + fn root_agent(&self) -> Result>; +} + +pub struct SingleAgentLoader { + root_agent: Arc, +} + +impl SingleAgentLoader { + pub fn new(root_agent: Arc) -> Self { + Self { root_agent } + } +} + +impl AgentLoader for SingleAgentLoader { + fn list_agents(&self) -> Result> { + Ok(vec![self.root_agent.name().to_string()]) + } + + fn load_agent(&self, name: &str) -> Result> { + if self.root_agent.name() == name { + Ok(self.root_agent.clone()) + } else { + Err(Error::msg(format!("agent not found: {name}"))) + } + } + + fn root_agent(&self) -> Result> { + Ok(self.root_agent.clone()) + } +} + +pub struct MultiAgentLoader { + root_agent: Arc, + agents: HashMap>, +} + +impl MultiAgentLoader { + pub fn new(root_agent: Arc, agents: Vec>) -> Result { + validate_agent_names(root_agent.as_ref())?; + + let mut by_name = HashMap::new(); + for agent in agents { + let name = agent.name().to_string(); + if by_name.insert(name.clone(), agent).is_some() { + return Err(Error::msg(format!("duplicate agent name: {name}"))); + } + } + + if !by_name.contains_key(root_agent.name()) { + by_name.insert(root_agent.name().to_string(), root_agent.clone()); + } + + Ok(Self { + root_agent, + agents: by_name, + }) + } +} + +impl AgentLoader for MultiAgentLoader { + fn list_agents(&self) -> Result> { + let mut names = self.agents.keys().cloned().collect::>(); + names.sort(); + Ok(names) + } + + fn load_agent(&self, name: &str) -> Result> { + self.agents + .get(name) + .cloned() + .ok_or_else(|| Error::msg(format!("agent not found: {name}"))) + } + + fn root_agent(&self) -> Result> { + Ok(self.root_agent.clone()) + } +} diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 0b97c13..377c6da 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -1,11 +1,15 @@ mod agent; mod context; mod custom_agent; +mod loader; mod module; mod stream; +mod toml_loader; -pub use agent::{Agent, find_agent}; +pub use agent::{Agent, IntoAgentArc, find_agent, validate_agent_names}; pub use context::InvocationContext; pub use custom_agent::CustomAgent; +pub use loader::{AgentLoader, MultiAgentLoader, SingleAgentLoader}; pub use module::MODULE; pub use stream::{AsyncEventStream, EventStream, event_stream_from_result}; +pub use toml_loader::TomlAgentLoader; diff --git a/src/agent/toml_loader.rs b/src/agent/toml_loader.rs new file mode 100644 index 0000000..f694214 --- /dev/null +++ b/src/agent/toml_loader.rs @@ -0,0 +1,228 @@ +use std::{collections::HashMap, fs, path::Path, sync::Arc}; + +use serde::Deserialize; + +use crate::{ + Error, Result, + llm_agent::{IncludeContents, LlmAgent, LlmAgentCallbacks, LlmAgentConfig}, + model::{DefaultModelFactory, GenerateContentConfig, ModelConfig, ModelFactory}, + tool::{StaticToolset, ToolRegistry, ToolsetRegistry}, +}; + +use super::{Agent, AgentLoader, MultiAgentLoader}; + +pub struct TomlAgentLoader { + inner: MultiAgentLoader, +} + +impl TomlAgentLoader { + pub fn from_path( + path: impl AsRef, + model_factory: Arc, + tool_registry: ToolRegistry, + ) -> Result { + let source = fs::read_to_string(path) + .map_err(|err| Error::msg(format!("failed to read agent config: {err}")))?; + Self::from_str( + &source, + model_factory, + tool_registry, + ToolsetRegistry::new(), + ) + } + + pub fn from_str( + source: &str, + model_factory: Arc, + tool_registry: ToolRegistry, + toolset_registry: ToolsetRegistry, + ) -> Result { + let config = toml::from_str::(source) + .map_err(|err| Error::msg(format!("invalid agent config TOML: {err}")))?; + let mut builder = TomlAgentBuilder { + config, + model_factory, + tool_registry, + toolset_registry, + built_agents: HashMap::new(), + }; + + let root = builder.build_agent(&builder.config.root_agent.clone())?; + let agents = builder.built_agents.values().cloned().collect::>(); + Ok(Self { + inner: MultiAgentLoader::new(root, agents)?, + }) + } + + pub fn from_path_with_defaults(path: impl AsRef) -> Result { + Self::from_path( + path, + Arc::new(DefaultModelFactory::new()), + ToolRegistry::new(), + ) + } +} + +impl AgentLoader for TomlAgentLoader { + fn list_agents(&self) -> Result> { + self.inner.list_agents() + } + + fn load_agent(&self, name: &str) -> Result> { + self.inner.load_agent(name) + } + + fn root_agent(&self) -> Result> { + self.inner.root_agent() + } +} + +struct TomlAgentBuilder { + config: AdkTomlConfig, + model_factory: Arc, + tool_registry: ToolRegistry, + toolset_registry: ToolsetRegistry, + built_agents: HashMap>, +} + +impl TomlAgentBuilder { + fn build_agent(&mut self, id: &str) -> Result> { + if let Some(agent) = self.built_agents.get(id) { + return Ok(agent.clone()); + } + + let config = self + .config + .agents + .get(id) + .cloned() + .ok_or_else(|| Error::msg(format!("agent config not found: {id}")))?; + + if config.agent_type != "llm" { + return Err(Error::msg(format!( + "unsupported agent type for {id}: {}", + config.agent_type + ))); + } + + let model_config = self + .config + .models + .get(&config.model) + .cloned() + .ok_or_else(|| Error::msg(format!("model config not found: {}", config.model)))?; + let model = self.model_factory.create(model_config.into())?; + let sub_agents = config + .sub_agents + .iter() + .map(|name| self.build_agent(name)) + .collect::>>()?; + let static_tools = config + .tools + .iter() + .map(|name| self.tool_registry.get(name)) + .collect::>>()?; + let mut toolsets = config + .toolsets + .iter() + .map(|name| self.toolset_registry.get(name)) + .collect::>>()?; + if !static_tools.is_empty() { + toolsets.push(Arc::new(StaticToolset::new( + format!("{id}_static_tools"), + static_tools, + ))); + } + + let agent = Arc::new(LlmAgent::new(LlmAgentConfig { + name: config.name.unwrap_or_else(|| id.to_string()), + description: config.description, + system_prompt: config.system_prompt.unwrap_or_default(), + model, + generate_content_config: config.generate_content_config.map(Into::into), + tools: Vec::new(), + toolsets, + sub_agents, + max_steps: config.max_steps.unwrap_or(8), + include_contents: parse_include_contents(config.include_contents.as_deref())?, + callbacks: LlmAgentCallbacks::default(), + })); + self.built_agents.insert(id.to_string(), agent.clone()); + Ok(agent) + } +} + +fn parse_include_contents(value: Option<&str>) -> Result { + match value.unwrap_or("default") { + "default" => Ok(IncludeContents::Default), + "none" => Ok(IncludeContents::None), + other => Err(Error::msg(format!("unsupported include_contents: {other}"))), + } +} + +#[derive(Debug, Clone, Deserialize)] +struct AdkTomlConfig { + root_agent: String, + models: HashMap, + agents: HashMap, +} + +#[derive(Debug, Clone, Deserialize)] +struct TomlModelConfig { + provider: String, + protocol: String, + model: String, + base_url: Option, + api_key_env: Option, +} + +impl From for ModelConfig { + fn from(config: TomlModelConfig) -> Self { + Self { + provider: config.provider, + protocol: config.protocol, + model: config.model, + base_url: config.base_url, + api_key_env: config.api_key_env, + } + } +} + +#[derive(Debug, Clone, Deserialize)] +struct TomlAgentConfig { + #[serde(rename = "type")] + agent_type: String, + name: Option, + description: String, + model: String, + system_prompt: Option, + include_contents: Option, + max_steps: Option, + #[serde(default)] + tools: Vec, + #[serde(default)] + toolsets: Vec, + #[serde(default)] + sub_agents: Vec, + generate_content_config: Option, +} + +#[derive(Debug, Clone, Deserialize)] +struct TomlGenerateContentConfig { + temperature: Option, + top_p: Option, + max_output_tokens: Option, + #[serde(default)] + stop_sequences: Vec, +} + +impl From for GenerateContentConfig { + fn from(config: TomlGenerateContentConfig) -> Self { + Self { + temperature: config.temperature, + top_p: config.top_p, + max_output_tokens: config.max_output_tokens, + stop_sequences: config.stop_sequences, + } + } +} diff --git a/src/content/content.rs b/src/content/content.rs index 640ea64..18c6826 100644 --- a/src/content/content.rs +++ b/src/content/content.rs @@ -1,6 +1,8 @@ +use serde::{Deserialize, Serialize}; + use super::{Part, Role}; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Content { pub role: Role, pub parts: Vec, diff --git a/src/content/part.rs b/src/content/part.rs index f26db20..0860c4c 100644 --- a/src/content/part.rs +++ b/src/content/part.rs @@ -1,6 +1,8 @@ +use serde::{Deserialize, Serialize}; + use super::{ToolCall, ToolResponse}; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub enum Part { Text(String), ToolCall(ToolCall), diff --git a/src/content/role.rs b/src/content/role.rs index 9b2557f..5345a6c 100644 --- a/src/content/role.rs +++ b/src/content/role.rs @@ -1,4 +1,6 @@ -#[derive(Debug, Clone, PartialEq, Eq)] +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub enum Role { User, Model, diff --git a/src/content/tool_call.rs b/src/content/tool_call.rs index 602e780..f8bba4e 100644 --- a/src/content/tool_call.rs +++ b/src/content/tool_call.rs @@ -1,6 +1,7 @@ +use serde::{Deserialize, Serialize}; use serde_json::Value; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ToolCall { pub id: String, pub name: String, diff --git a/src/content/tool_response.rs b/src/content/tool_response.rs index 4be2082..e2e193a 100644 --- a/src/content/tool_response.rs +++ b/src/content/tool_response.rs @@ -1,6 +1,7 @@ +use serde::{Deserialize, Serialize}; use serde_json::Value; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ToolResponse { pub call_id: String, pub name: String, diff --git a/src/event/actions.rs b/src/event/actions.rs index aef14c9..25ce4c8 100644 --- a/src/event/actions.rs +++ b/src/event/actions.rs @@ -1,10 +1,11 @@ use std::collections::HashMap; +use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use crate::tool::ToolConfirmation; -#[derive(Debug, Clone, PartialEq, Eq, Default)] +#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] pub struct EventActions { pub state_delta: Map, pub artifact_delta: HashMap, diff --git a/src/event/event.rs b/src/event/event.rs index 954c735..d22b9c5 100644 --- a/src/event/event.rs +++ b/src/event/event.rs @@ -1,10 +1,12 @@ use std::time::SystemTime; +use serde::{Deserialize, Serialize}; + use crate::content::{Content, Part, Role}; use super::{EventActions, EventKind}; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Event { pub id: String, pub invocation_id: String, diff --git a/src/llm_agent/agent.rs b/src/llm_agent/agent.rs index 0eb3637..832c1be 100644 --- a/src/llm_agent/agent.rs +++ b/src/llm_agent/agent.rs @@ -1,5 +1,6 @@ use std::{ collections::HashMap, + sync::Arc, sync::atomic::{AtomicU64, Ordering}, time::SystemTime, }; @@ -15,9 +16,10 @@ use crate::{ session::Session, telemetry::span_attr, tool::{ - REQUEST_CONFIRMATION_TOOL_NAME, TRANSFER_TO_AGENT_TOOL_NAME, Tool, ToolConfirmation, - ToolContext, ToolDeclaration, TransferToAgentTool, confirmation_request_args, - original_tool_call_from_confirmation_args, tool_confirmation_from_request_args, + IntoToolArc, REQUEST_CONFIRMATION_TOOL_NAME, TRANSFER_TO_AGENT_TOOL_NAME, Tool, + ToolConfirmation, ToolContext, ToolDeclaration, Toolset, ToolsetContext, + TransferToAgentTool, confirmation_request_args, original_tool_call_from_confirmation_args, + tool_confirmation_from_request_args, }, }; @@ -30,8 +32,10 @@ pub struct LlmAgent { description: String, system_prompt: String, model: Box, - tools: Vec>, - sub_agents: Vec>, + generate_content_config: Option, + tools: Vec>, + toolsets: Vec>, + sub_agents: Vec>, max_steps: usize, include_contents: IncludeContents, callbacks: LlmAgentCallbacks, @@ -44,7 +48,11 @@ struct ConfirmedToolCall { impl LlmAgent { pub fn new(config: LlmAgentConfig) -> Self { - let mut tools = config.tools; + let mut tools = config + .tools + .into_iter() + .map(IntoToolArc::into_tool_arc) + .collect::>(); let transfer_targets = config .sub_agents .iter() @@ -56,7 +64,7 @@ impl LlmAgent { .iter() .any(|tool| tool.name() == TRANSFER_TO_AGENT_TOOL_NAME) { - tools.push(Box::new(TransferToAgentTool::new(transfer_targets))); + tools.push(Arc::new(TransferToAgentTool::new(transfer_targets))); } Self { @@ -64,7 +72,9 @@ impl LlmAgent { description: config.description, system_prompt: config.system_prompt, model: config.model, + generate_content_config: config.generate_content_config, tools, + toolsets: config.toolsets, sub_agents: config.sub_agents, max_steps: config.max_steps, include_contents: config.include_contents, @@ -76,12 +86,14 @@ impl LlmAgent { &self, ctx: &InvocationContext, session_snapshot: &Session, + visible_tools: &[Arc], ) -> Result { Ok(LlmRequest { model: self.model.name().to_string(), system_prompt: self.render_system_prompt(session_snapshot)?, contents: self.build_contents(ctx, session_snapshot), - tools: self.tool_declarations(), + tools: self.tool_declarations(visible_tools), + generate_content_config: self.generate_content_config.clone(), }) } @@ -109,8 +121,30 @@ impl LlmAgent { )?)) } - fn tool_declarations(&self) -> Vec { - self.tools + fn resolve_tools( + &self, + ctx: &InvocationContext, + session_snapshot: &Session, + ) -> Result>> { + let toolset_ctx = ToolsetContext::new(ctx, session_snapshot); + let mut tools = self.tools.clone(); + for toolset in &self.toolsets { + tools.extend(toolset.tools(&toolset_ctx)?); + } + + let mut names = std::collections::HashSet::new(); + for tool in &tools { + let name = tool.name().to_string(); + if !names.insert(name.clone()) { + return Err(Error::msg(format!("duplicate tool name: {name}"))); + } + } + + Ok(tools) + } + + fn tool_declarations(&self, tools: &[Arc]) -> Vec { + tools .iter() .map(|tool| ToolDeclaration { name: tool.name().to_string(), @@ -121,8 +155,8 @@ impl LlmAgent { .collect() } - fn tool_for_name(&self, name: &str) -> Option<&dyn Tool> { - self.tools + fn tool_for_name<'a>(&self, tools: &'a [Arc], name: &str) -> Option<&'a dyn Tool> { + tools .iter() .find(|tool| tool.name() == name) .map(|tool| tool.as_ref()) @@ -433,7 +467,7 @@ impl Agent for LlmAgent { &self.description } - fn sub_agents(&self) -> &[Box] { + fn sub_agents(&self) -> &[Arc] { &self.sub_agents } @@ -459,9 +493,15 @@ impl Agent for LlmAgent { let tool_event = if confirmed_tool_call.confirmation.confirmed == Some(false) { self.rejected_tool_response_event(&ctx, &tool_call) } else { + let visible_tools = self.resolve_tools(&ctx, &session_snapshot)?; let tool = self - .tool_for_name(&tool_call.name) - .ok_or_else(|| Error::msg(format!("tool not found: {}", tool_call.name)))?; + .tool_for_name(&visible_tools, &tool_call.name) + .ok_or_else(|| { + Error::msg(format!( + "tool not visible in current invocation: {}", + tool_call.name + )) + })?; let (result, actions) = self.run_tool_with_callbacks( &ctx, &session_snapshot, @@ -492,7 +532,8 @@ impl Agent for LlmAgent { continue; } - let request = self.build_request(&ctx, &session_snapshot)?; + let visible_tools = self.resolve_tools(&ctx, &session_snapshot)?; + let request = self.build_request(&ctx, &session_snapshot, &visible_tools)?; let responses = self.generate_stream_with_callbacks(&ctx, request)?; for response in responses { @@ -522,8 +563,13 @@ impl Agent for LlmAgent { for tool_call in tool_calls { let tool = self - .tool_for_name(&tool_call.name) - .ok_or_else(|| Error::msg(format!("tool not found: {}", tool_call.name)))?; + .tool_for_name(&visible_tools, &tool_call.name) + .ok_or_else(|| { + Error::msg(format!( + "tool not visible in current invocation: {}", + tool_call.name + )) + })?; if let Some(confirmation) = tool.confirmation(&tool_call.args) { let confirmation_event = self.confirmation_request_event(&ctx, &tool_call, confirmation); @@ -589,9 +635,15 @@ impl Agent for LlmAgent { let tool_event = if confirmed_tool_call.confirmation.confirmed == Some(false) { self.rejected_tool_response_event(&ctx, &tool_call) } else { + let visible_tools = self.resolve_tools(&ctx, &session_snapshot)?; let tool = self - .tool_for_name(&tool_call.name) - .ok_or_else(|| Error::msg(format!("tool not found: {}", tool_call.name)))?; + .tool_for_name(&visible_tools, &tool_call.name) + .ok_or_else(|| { + Error::msg(format!( + "tool not visible in current invocation: {}", + tool_call.name + )) + })?; let (result, actions) = self.run_tool_with_callbacks( &ctx, &session_snapshot, @@ -623,7 +675,8 @@ impl Agent for LlmAgent { continue; } - let mut request = self.build_request(&ctx, &session_snapshot)?; + let visible_tools = self.resolve_tools(&ctx, &session_snapshot)?; + let mut request = self.build_request(&ctx, &session_snapshot, &visible_tools)?; let short_circuit_response = if let Some(response) = ctx.plugin_manager.run_before_model(&ctx, &mut request)? { Some(response) @@ -711,8 +764,13 @@ impl Agent for LlmAgent { for tool_call in tool_calls { let tool = self - .tool_for_name(&tool_call.name) - .ok_or_else(|| Error::msg(format!("tool not found: {}", tool_call.name)))?; + .tool_for_name(&visible_tools, &tool_call.name) + .ok_or_else(|| { + Error::msg(format!( + "tool not visible in current invocation: {}", + tool_call.name + )) + })?; if let Some(confirmation) = tool.confirmation(&tool_call.args) { let confirmation_event = self.confirmation_request_event(&ctx, &tool_call, confirmation); diff --git a/src/llm_agent/config.rs b/src/llm_agent/config.rs index 6a033fa..92d462f 100644 --- a/src/llm_agent/config.rs +++ b/src/llm_agent/config.rs @@ -1,4 +1,10 @@ -use crate::{agent::Agent, model::Model, tool::Tool}; +use std::sync::Arc; + +use crate::{ + agent::Agent, + model::{GenerateContentConfig, Model}, + tool::{Tool, Toolset}, +}; use super::{IncludeContents, LlmAgentCallbacks}; @@ -7,8 +13,10 @@ pub struct LlmAgentConfig { pub description: String, pub system_prompt: String, pub model: Box, + pub generate_content_config: Option, pub tools: Vec>, - pub sub_agents: Vec>, + pub toolsets: Vec>, + pub sub_agents: Vec>, pub max_steps: usize, pub include_contents: IncludeContents, pub callbacks: LlmAgentCallbacks, diff --git a/src/main.rs b/src/main.rs index 2ba4cec..dd1c3a6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -60,7 +60,9 @@ fn main() -> Result<()> { description: "A DeepSeek-backed tool-using assistant.".to_string(), system_prompt: "Use tools when useful, then answer the user.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(GetTimeTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/src/model/factory.rs b/src/model/factory.rs new file mode 100644 index 0000000..5a4f8a5 --- /dev/null +++ b/src/model/factory.rs @@ -0,0 +1,29 @@ +use crate::{Error, Result}; + +use super::{Model, ModelConfig, OpenAiCompatibleModel}; + +pub trait ModelFactory: Send + Sync { + fn create(&self, config: ModelConfig) -> Result>; +} + +#[derive(Debug, Default, Clone)] +pub struct DefaultModelFactory; + +impl DefaultModelFactory { + pub fn new() -> Self { + Self + } +} + +impl ModelFactory for DefaultModelFactory { + fn create(&self, config: ModelConfig) -> Result> { + match config.protocol.as_str() { + "openai-compatible" | "openai-chat-completions" => { + Ok(Box::new(OpenAiCompatibleModel::new(config))) + } + protocol => Err(Error::msg(format!( + "unsupported model protocol: {protocol}" + ))), + } + } +} diff --git a/src/model/generate_config.rs b/src/model/generate_config.rs new file mode 100644 index 0000000..7afef2a --- /dev/null +++ b/src/model/generate_config.rs @@ -0,0 +1,18 @@ +#[derive(Debug, Clone, PartialEq)] +pub struct GenerateContentConfig { + pub temperature: Option, + pub top_p: Option, + pub max_output_tokens: Option, + pub stop_sequences: Vec, +} + +impl Default for GenerateContentConfig { + fn default() -> Self { + Self { + temperature: None, + top_p: None, + max_output_tokens: None, + stop_sequences: Vec::new(), + } + } +} diff --git a/src/model/llm.rs b/src/model/llm.rs index e330bd6..055b8fe 100644 --- a/src/model/llm.rs +++ b/src/model/llm.rs @@ -4,16 +4,19 @@ use futures_util::Stream; use crate::{content::Content, tool::ToolDeclaration}; +use super::GenerateContentConfig; + pub type LlmResponseStream<'a> = Box> + 'a>; pub type AsyncLlmResponseStream<'a> = Pin> + Send + 'a>>; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub struct LlmRequest { pub model: String, pub system_prompt: Option, pub contents: Vec, pub tools: Vec, + pub generate_content_config: Option, } #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/src/model/mod.rs b/src/model/mod.rs index edc7af0..0041865 100644 --- a/src/model/mod.rs +++ b/src/model/mod.rs @@ -1,12 +1,16 @@ mod config; +mod factory; mod fake_model; +mod generate_config; mod llm; mod model; mod module; pub mod openai_compatible; pub use config::ModelConfig; +pub use factory::{DefaultModelFactory, ModelFactory}; pub use fake_model::FakeModel; +pub use generate_config::GenerateContentConfig; pub use llm::{AsyncLlmResponseStream, LlmRequest, LlmResponse, ModelUsage}; pub use llm::{LlmResponseStream, llm_response_stream_from_result}; pub use model::Model; diff --git a/src/model/openai_compatible/convert.rs b/src/model/openai_compatible/convert.rs index c3512fb..781714e 100644 --- a/src/model/openai_compatible/convert.rs +++ b/src/model/openai_compatible/convert.rs @@ -44,6 +44,21 @@ pub fn request_json(request: &LlmRequest) -> Value { ); } + if let Some(config) = &request.generate_content_config { + if let Some(temperature) = config.temperature { + body["temperature"] = json!(temperature); + } + if let Some(top_p) = config.top_p { + body["top_p"] = json!(top_p); + } + if let Some(max_output_tokens) = config.max_output_tokens { + body["max_tokens"] = json!(max_output_tokens); + } + if !config.stop_sequences.is_empty() { + body["stop"] = json!(config.stop_sequences); + } + } + body } diff --git a/src/runner/runner.rs b/src/runner/runner.rs index efe77d2..d86a6ac 100644 --- a/src/runner/runner.rs +++ b/src/runner/runner.rs @@ -9,8 +9,8 @@ use futures_util::StreamExt; use crate::{ Error, Result, agent::{ - Agent, AsyncEventStream, EventStream, InvocationContext, event_stream_from_result, - find_agent, + Agent, AsyncEventStream, EventStream, IntoAgentArc, InvocationContext, + event_stream_from_result, find_agent, }, artifact::{ArtifactFacade, ArtifactService, InMemoryArtifactService}, content::Content, @@ -27,7 +27,7 @@ static NEXT_ID: AtomicU64 = AtomicU64::new(1); pub struct Runner { pub app_name: String, - pub root_agent: Box, + pub root_agent: Arc, pub session_service: S, pub auto_create_session: bool, pub artifact_service: Arc, @@ -39,13 +39,13 @@ pub struct Runner { impl Runner { pub fn new( app_name: impl Into, - root_agent: Box, + root_agent: impl IntoAgentArc, session_service: S, auto_create_session: bool, ) -> Self { Self { app_name: app_name.into(), - root_agent, + root_agent: root_agent.into_agent_arc(), session_service, auto_create_session, artifact_service: Arc::new(InMemoryArtifactService::new()), @@ -57,14 +57,14 @@ impl Runner { pub fn with_artifact_service( app_name: impl Into, - root_agent: Box, + root_agent: impl IntoAgentArc, session_service: S, auto_create_session: bool, artifact_service: Arc, ) -> Self { Self { app_name: app_name.into(), - root_agent, + root_agent: root_agent.into_agent_arc(), session_service, auto_create_session, artifact_service, @@ -76,14 +76,14 @@ impl Runner { pub fn with_memory_service( app_name: impl Into, - root_agent: Box, + root_agent: impl IntoAgentArc, session_service: S, auto_create_session: bool, memory_service: Arc, ) -> Self { Self { app_name: app_name.into(), - root_agent, + root_agent: root_agent.into_agent_arc(), session_service, auto_create_session, artifact_service: Arc::new(InMemoryArtifactService::new()), @@ -95,14 +95,14 @@ impl Runner { pub fn with_plugin_manager( app_name: impl Into, - root_agent: Box, + root_agent: impl IntoAgentArc, session_service: S, auto_create_session: bool, plugin_manager: PluginManager, ) -> Self { Self { app_name: app_name.into(), - root_agent, + root_agent: root_agent.into_agent_arc(), session_service, auto_create_session, artifact_service: Arc::new(InMemoryArtifactService::new()), @@ -114,7 +114,7 @@ impl Runner { pub fn with_services( app_name: impl Into, - root_agent: Box, + root_agent: impl IntoAgentArc, session_service: S, auto_create_session: bool, artifact_service: Arc, @@ -122,7 +122,7 @@ impl Runner { ) -> Self { Self { app_name: app_name.into(), - root_agent, + root_agent: root_agent.into_agent_arc(), session_service, auto_create_session, artifact_service, diff --git a/src/session/database.rs b/src/session/database.rs new file mode 100644 index 0000000..5835a5a --- /dev/null +++ b/src/session/database.rs @@ -0,0 +1,60 @@ +use crate::{Result, event::Event}; + +use super::{Session, SessionService}; + +pub trait SessionStore: Send + Sync { + fn create(&self, app_name: &str, user_id: &str, session_id: &str) -> Result; + + fn get(&self, app_name: &str, user_id: &str, session_id: &str) -> Result; + + fn list(&self, app_name: &str, user_id: &str) -> Result>; + + fn delete(&self, app_name: &str, user_id: &str, session_id: &str) -> Result<()>; + + fn append_event( + &self, + app_name: &str, + user_id: &str, + session_id: &str, + event: Event, + ) -> Result<()>; +} + +pub struct DatabaseSessionService { + store: S, +} + +impl DatabaseSessionService { + pub fn new(store: S) -> Result { + Ok(Self { store }) + } +} + +impl SessionService for DatabaseSessionService { + fn create(&self, app_name: &str, user_id: &str, session_id: &str) -> Result { + self.store.create(app_name, user_id, session_id) + } + + fn get(&self, app_name: &str, user_id: &str, session_id: &str) -> Result { + self.store.get(app_name, user_id, session_id) + } + + fn list(&self, app_name: &str, user_id: &str) -> Result> { + self.store.list(app_name, user_id) + } + + fn delete(&self, app_name: &str, user_id: &str, session_id: &str) -> Result<()> { + self.store.delete(app_name, user_id, session_id) + } + + fn append_event( + &self, + app_name: &str, + user_id: &str, + session_id: &str, + event: Event, + ) -> Result<()> { + self.store + .append_event(app_name, user_id, session_id, event) + } +} diff --git a/src/session/in_memory.rs b/src/session/in_memory.rs index 7474f49..15aed81 100644 --- a/src/session/in_memory.rs +++ b/src/session/in_memory.rs @@ -3,13 +3,17 @@ use std::{ sync::{Arc, RwLock}, }; +use serde_json::Map; + use crate::{Error, Result, event::Event}; -use super::{Session, SessionService}; +use super::{Session, SessionService, State, StateScope, state_scope}; #[derive(Debug, Clone, Default)] pub struct InMemorySessionService { sessions: Arc>>, + app_states: Arc>>, + user_states: Arc>>, } impl InMemorySessionService { @@ -33,12 +37,32 @@ impl SessionService for InMemorySessionService { fn get(&self, app_name: &str, user_id: &str, session_id: &str) -> Result { let key = SessionKey::new(app_name, user_id, session_id); - self.sessions + let mut session = self + .sessions .read() .expect("session lock poisoned") .get(&key) .cloned() - .ok_or_else(|| Error::msg("session not found")) + .ok_or_else(|| Error::msg("session not found"))?; + + if let Some(app_state) = self + .app_states + .read() + .expect("app state lock poisoned") + .get(app_name) + { + session.state.extend(app_state.clone()); + } + if let Some(user_state) = self + .user_states + .read() + .expect("user state lock poisoned") + .get(&UserKey::new(app_name, user_id)) + { + session.state.extend(user_state.clone()); + } + + Ok(session) } fn list(&self, app_name: &str, user_id: &str) -> Result> { @@ -72,14 +96,56 @@ impl SessionService for InMemorySessionService { event: Event, ) -> Result<()> { let key = SessionKey::new(app_name, user_id, session_id); + let mut event = event; let mut sessions = self.sessions.write().expect("session lock poisoned"); let session = sessions .get_mut(&key) .ok_or_else(|| Error::msg("session not found"))?; let state_delta = event.actions.state_delta.clone(); + let mut session_delta = State::new(); + let mut app_delta = State::new(); + let mut user_delta = State::new(); + let mut persisted_state_delta = Map::new(); + + for (key, value) in state_delta { + match state_scope(&key) { + StateScope::App => { + app_delta.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::User => { + user_delta.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::Session => { + session_delta.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::Temp => {} + } + } + + if !app_delta.is_empty() { + self.app_states + .write() + .expect("app state lock poisoned") + .entry(app_name.to_string()) + .or_default() + .extend(app_delta); + } + if !user_delta.is_empty() { + self.user_states + .write() + .expect("user state lock poisoned") + .entry(UserKey::new(app_name, user_id)) + .or_default() + .extend(user_delta); + } + + event.actions.state_delta = persisted_state_delta; session.events.push(event); - session.state.extend(state_delta); + session.state.extend(session_delta); Ok(()) } @@ -92,6 +158,21 @@ struct SessionKey { session_id: String, } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct UserKey { + app_name: String, + user_id: String, +} + +impl UserKey { + fn new(app_name: &str, user_id: &str) -> Self { + Self { + app_name: app_name.to_string(), + user_id: user_id.to_string(), + } + } +} + impl SessionKey { fn new(app_name: &str, user_id: &str, session_id: &str) -> Self { Self { diff --git a/src/session/mod.rs b/src/session/mod.rs index 22d1c70..ed8c9d6 100644 --- a/src/session/mod.rs +++ b/src/session/mod.rs @@ -1,11 +1,18 @@ +mod database; mod in_memory; mod module; mod service; mod session; +mod sqlite; mod state; +pub use database::{DatabaseSessionService, SessionStore}; pub use in_memory::InMemorySessionService; pub use module::MODULE; pub use service::SessionService; pub use session::Session; -pub use state::State; +pub use sqlite::SqliteSessionStore; +pub use state::{ + APP_STATE_PREFIX, SESSION_STATE_PREFIX, State, StateScope, TEMP_STATE_PREFIX, + USER_STATE_PREFIX, state_scope, +}; diff --git a/src/session/sqlite.rs b/src/session/sqlite.rs new file mode 100644 index 0000000..9efdd2a --- /dev/null +++ b/src/session/sqlite.rs @@ -0,0 +1,400 @@ +use std::{ + path::Path, + sync::{Arc, Mutex}, + time::{SystemTime, UNIX_EPOCH}, +}; + +use rusqlite::{Connection, OptionalExtension, params}; +use serde_json::Map; + +use crate::{Error, Result, event::Event}; + +use super::{Session, SessionStore, State, StateScope, state_scope}; + +#[derive(Clone)] +pub struct SqliteSessionStore { + connection: Arc>, +} + +impl SqliteSessionStore { + pub fn open(path: impl AsRef) -> Result { + let connection = Connection::open(path) + .map_err(|err| Error::msg(format!("failed to open sqlite session store: {err}")))?; + let store = Self { + connection: Arc::new(Mutex::new(connection)), + }; + store.migrate()?; + Ok(store) + } + + fn migrate(&self) -> Result<()> { + let connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + connection + .execute_batch( + r#" + PRAGMA foreign_keys = ON; + + CREATE TABLE IF NOT EXISTS sessions ( + app_name TEXT NOT NULL, + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + state_json TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + PRIMARY KEY (app_name, user_id, session_id) + ); + + CREATE TABLE IF NOT EXISTS events ( + app_name TEXT NOT NULL, + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + event_index INTEGER NOT NULL, + event_id TEXT NOT NULL, + event_json TEXT NOT NULL, + PRIMARY KEY (app_name, user_id, session_id, event_index), + FOREIGN KEY (app_name, user_id, session_id) + REFERENCES sessions(app_name, user_id, session_id) + ON DELETE CASCADE + ); + + CREATE TABLE IF NOT EXISTS app_states ( + app_name TEXT PRIMARY KEY, + state_json TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS user_states ( + app_name TEXT NOT NULL, + user_id TEXT NOT NULL, + state_json TEXT NOT NULL, + PRIMARY KEY (app_name, user_id) + ); + "#, + ) + .map_err(|err| Error::msg(format!("failed to migrate sqlite session store: {err}")))?; + Ok(()) + } +} + +impl SessionStore for SqliteSessionStore { + fn create(&self, app_name: &str, user_id: &str, session_id: &str) -> Result { + let now = unix_millis(SystemTime::now())?; + let state = State::new(); + let state_json = state_to_json(&state)?; + let connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + let inserted = connection + .execute( + r#" + INSERT OR IGNORE INTO sessions + (app_name, user_id, session_id, state_json, created_at, updated_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?6) + "#, + params![app_name, user_id, session_id, state_json, now, now], + ) + .map_err(|err| Error::msg(format!("failed to create sqlite session: {err}")))?; + if inserted == 0 { + return Err(Error::msg("session already exists")); + } + + Ok(Session::new(app_name, user_id, session_id)) + } + + fn get(&self, app_name: &str, user_id: &str, session_id: &str) -> Result { + let connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + load_session(&connection, app_name, user_id, session_id) + } + + fn list(&self, app_name: &str, user_id: &str) -> Result> { + let connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + let mut statement = connection + .prepare( + r#" + SELECT session_id + FROM sessions + WHERE app_name = ?1 AND user_id = ?2 + ORDER BY session_id + "#, + ) + .map_err(|err| Error::msg(format!("failed to list sqlite sessions: {err}")))?; + let ids = statement + .query_map(params![app_name, user_id], |row| row.get::<_, String>(0)) + .map_err(|err| Error::msg(format!("failed to list sqlite sessions: {err}")))? + .collect::, _>>() + .map_err(|err| Error::msg(format!("failed to read sqlite sessions: {err}")))?; + + ids.into_iter() + .map(|session_id| load_session(&connection, app_name, user_id, &session_id)) + .collect() + } + + fn delete(&self, app_name: &str, user_id: &str, session_id: &str) -> Result<()> { + let connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + let deleted = connection + .execute( + "DELETE FROM sessions WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3", + params![app_name, user_id, session_id], + ) + .map_err(|err| Error::msg(format!("failed to delete sqlite session: {err}")))?; + if deleted == 0 { + return Err(Error::msg("session not found")); + } + Ok(()) + } + + fn append_event( + &self, + app_name: &str, + user_id: &str, + session_id: &str, + event: Event, + ) -> Result<()> { + let now = unix_millis(event.timestamp)?; + let mut event = event; + let mut connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + let tx = connection + .transaction() + .map_err(|err| Error::msg(format!("failed to start sqlite transaction: {err}")))?; + + let session_state_json = tx + .query_row( + "SELECT state_json FROM sessions WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3", + params![app_name, user_id, session_id], + |row| row.get::<_, String>(0), + ) + .optional() + .map_err(|err| Error::msg(format!("failed to load sqlite session: {err}")))? + .ok_or_else(|| Error::msg("session not found"))?; + let mut session_state = state_from_json(&session_state_json)?; + let mut app_state = load_scoped_state(&tx, "app_states", app_name, None)?; + let mut user_state = load_scoped_state(&tx, "user_states", app_name, Some(user_id))?; + let mut persisted_state_delta = Map::new(); + + for (key, value) in event.actions.state_delta.clone() { + match state_scope(&key) { + StateScope::App => { + app_state.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::User => { + user_state.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::Session => { + session_state.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::Temp => {} + } + } + event.actions.state_delta = persisted_state_delta; + + save_app_state(&tx, app_name, &app_state)?; + save_user_state(&tx, app_name, user_id, &user_state)?; + + let event_index = tx + .query_row( + r#" + SELECT COALESCE(MAX(event_index), -1) + 1 + FROM events + WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3 + "#, + params![app_name, user_id, session_id], + |row| row.get::<_, i64>(0), + ) + .map_err(|err| Error::msg(format!("failed to allocate event index: {err}")))?; + tx.execute( + r#" + INSERT INTO events + (app_name, user_id, session_id, event_index, event_id, event_json) + VALUES (?1, ?2, ?3, ?4, ?5, ?6) + "#, + params![ + app_name, + user_id, + session_id, + event_index, + event.id, + event_to_json(&event)? + ], + ) + .map_err(|err| Error::msg(format!("failed to insert sqlite event: {err}")))?; + tx.execute( + r#" + UPDATE sessions + SET state_json = ?4, updated_at = ?5 + WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3 + "#, + params![ + app_name, + user_id, + session_id, + state_to_json(&session_state)?, + now + ], + ) + .map_err(|err| Error::msg(format!("failed to update sqlite session state: {err}")))?; + tx.commit() + .map_err(|err| Error::msg(format!("failed to commit sqlite transaction: {err}")))?; + Ok(()) + } +} + +fn load_session( + connection: &Connection, + app_name: &str, + user_id: &str, + session_id: &str, +) -> Result { + let state_json = connection + .query_row( + "SELECT state_json FROM sessions WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3", + params![app_name, user_id, session_id], + |row| row.get::<_, String>(0), + ) + .optional() + .map_err(|err| Error::msg(format!("failed to load sqlite session: {err}")))? + .ok_or_else(|| Error::msg("session not found"))?; + let mut session = Session::new(app_name, user_id, session_id); + session.state = state_from_json(&state_json)?; + session + .state + .extend(load_scoped_state(connection, "app_states", app_name, None)?); + session.state.extend(load_scoped_state( + connection, + "user_states", + app_name, + Some(user_id), + )?); + session.events = load_events(connection, app_name, user_id, session_id)?; + Ok(session) +} + +fn load_events( + connection: &Connection, + app_name: &str, + user_id: &str, + session_id: &str, +) -> Result> { + let mut statement = connection + .prepare( + r#" + SELECT event_json + FROM events + WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3 + ORDER BY event_index + "#, + ) + .map_err(|err| Error::msg(format!("failed to load sqlite events: {err}")))?; + statement + .query_map(params![app_name, user_id, session_id], |row| { + row.get::<_, String>(0) + }) + .map_err(|err| Error::msg(format!("failed to load sqlite events: {err}")))? + .map(|row| { + row.map_err(|err| Error::msg(format!("failed to read sqlite event: {err}"))) + .and_then(|json| event_from_json(&json)) + }) + .collect() +} + +fn load_scoped_state( + connection: &Connection, + table: &str, + app_name: &str, + user_id: Option<&str>, +) -> Result { + let state_json = match user_id { + Some(user_id) => connection + .query_row( + &format!("SELECT state_json FROM {table} WHERE app_name = ?1 AND user_id = ?2"), + params![app_name, user_id], + |row| row.get::<_, String>(0), + ) + .optional(), + None => connection + .query_row( + &format!("SELECT state_json FROM {table} WHERE app_name = ?1"), + params![app_name], + |row| row.get::<_, String>(0), + ) + .optional(), + } + .map_err(|err| Error::msg(format!("failed to load sqlite scoped state: {err}")))?; + + state_json + .map(|json| state_from_json(&json)) + .unwrap_or_else(|| Ok(State::new())) +} + +fn save_app_state(connection: &Connection, app_name: &str, state: &State) -> Result<()> { + connection + .execute( + r#" + INSERT INTO app_states (app_name, state_json) + VALUES (?1, ?2) + ON CONFLICT(app_name) DO UPDATE SET state_json = excluded.state_json + "#, + params![app_name, state_to_json(state)?], + ) + .map_err(|err| Error::msg(format!("failed to save app state: {err}")))?; + Ok(()) +} + +fn save_user_state( + connection: &Connection, + app_name: &str, + user_id: &str, + state: &State, +) -> Result<()> { + connection + .execute( + r#" + INSERT INTO user_states (app_name, user_id, state_json) + VALUES (?1, ?2, ?3) + ON CONFLICT(app_name, user_id) DO UPDATE SET state_json = excluded.state_json + "#, + params![app_name, user_id, state_to_json(state)?], + ) + .map_err(|err| Error::msg(format!("failed to save user state: {err}")))?; + Ok(()) +} + +fn state_to_json(state: &State) -> Result { + serde_json::to_string(state).map_err(|err| Error::msg(format!("failed to encode state: {err}"))) +} + +fn state_from_json(json: &str) -> Result { + serde_json::from_str(json).map_err(|err| Error::msg(format!("failed to decode state: {err}"))) +} + +fn event_to_json(event: &Event) -> Result { + serde_json::to_string(event).map_err(|err| Error::msg(format!("failed to encode event: {err}"))) +} + +fn event_from_json(json: &str) -> Result { + serde_json::from_str(json).map_err(|err| Error::msg(format!("failed to decode event: {err}"))) +} + +fn unix_millis(time: SystemTime) -> Result { + let duration = time + .duration_since(UNIX_EPOCH) + .map_err(|err| Error::msg(format!("timestamp before unix epoch: {err}")))?; + Ok(duration.as_millis() as i64) +} diff --git a/src/session/state.rs b/src/session/state.rs index 9d5c9e3..3944b7c 100644 --- a/src/session/state.rs +++ b/src/session/state.rs @@ -3,3 +3,28 @@ use std::collections::HashMap; use serde_json::Value; pub type State = HashMap; + +pub const APP_STATE_PREFIX: &str = "app:"; +pub const USER_STATE_PREFIX: &str = "user:"; +pub const TEMP_STATE_PREFIX: &str = "temp:"; +pub const SESSION_STATE_PREFIX: &str = ""; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StateScope { + App, + User, + Temp, + Session, +} + +pub fn state_scope(key: &str) -> StateScope { + if key.starts_with(APP_STATE_PREFIX) { + StateScope::App + } else if key.starts_with(USER_STATE_PREFIX) { + StateScope::User + } else if key.starts_with(TEMP_STATE_PREFIX) { + StateScope::Temp + } else { + StateScope::Session + } +} diff --git a/src/tool/mod.rs b/src/tool/mod.rs index e9f6b20..1077f32 100644 --- a/src/tool/mod.rs +++ b/src/tool/mod.rs @@ -3,8 +3,10 @@ mod context; mod declaration; mod function_tool; mod module; +mod registry; mod schema; mod tool; +mod toolset; mod transfer_tool; pub use confirmation::{ @@ -18,5 +20,9 @@ pub use context::ToolContext; pub use declaration::ToolDeclaration; pub use function_tool::FunctionTool; pub use module::MODULE; -pub use tool::Tool; +pub use registry::{ToolRegistry, ToolsetRegistry}; +pub use tool::{IntoToolArc, Tool}; +pub use toolset::{ + StaticToolset, ToolPredicate, Toolset, ToolsetContext, allowed_tools_predicate, filter_toolset, +}; pub use transfer_tool::{TRANSFER_TO_AGENT_TOOL_NAME, TransferToAgentTool}; diff --git a/src/tool/registry.rs b/src/tool/registry.rs new file mode 100644 index 0000000..92a31df --- /dev/null +++ b/src/tool/registry.rs @@ -0,0 +1,57 @@ +use std::{collections::HashMap, sync::Arc}; + +use crate::{Error, Result}; + +use super::{Tool, Toolset}; + +#[derive(Default, Clone)] +pub struct ToolRegistry { + tools: HashMap>, +} + +impl ToolRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn register(&mut self, tool: Arc) -> Result<()> { + let name = tool.name().to_string(); + if self.tools.insert(name.clone(), tool).is_some() { + return Err(Error::msg(format!("duplicate tool name: {name}"))); + } + Ok(()) + } + + pub fn get(&self, name: &str) -> Result> { + self.tools + .get(name) + .cloned() + .ok_or_else(|| Error::msg(format!("unknown tool in config: {name}"))) + } +} + +#[derive(Default, Clone)] +pub struct ToolsetRegistry { + toolsets: HashMap>, +} + +impl ToolsetRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn register(&mut self, toolset: Arc) -> Result<()> { + let name = toolset.name().to_string(); + if self.toolsets.insert(name.clone(), toolset).is_some() { + return Err(Error::msg(format!("duplicate toolset name: {name}"))); + } + Ok(()) + } + + pub fn get(&self, name: &str) -> Result> { + self.toolsets + .get(name) + .cloned() + .ok_or_else(|| Error::msg(format!("unknown toolset in config: {name}"))) + } +} diff --git a/src/tool/tool.rs b/src/tool/tool.rs index 7f81e56..122d8a4 100644 --- a/src/tool/tool.rs +++ b/src/tool/tool.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use serde_json::Value; use crate::Result; @@ -25,3 +27,37 @@ pub trait Tool: Send + Sync { fn run(&self, ctx: &mut ToolContext, args: Value) -> Result; } + +pub trait IntoToolArc { + fn into_tool_arc(self) -> Arc; +} + +impl IntoToolArc for Arc { + fn into_tool_arc(self) -> Arc { + self + } +} + +impl IntoToolArc for Arc +where + T: Tool + 'static, +{ + fn into_tool_arc(self) -> Arc { + self + } +} + +impl IntoToolArc for Box { + fn into_tool_arc(self) -> Arc { + Arc::from(self) + } +} + +impl IntoToolArc for Box +where + T: Tool + 'static, +{ + fn into_tool_arc(self) -> Arc { + Arc::new(*self) + } +} diff --git a/src/tool/toolset.rs b/src/tool/toolset.rs new file mode 100644 index 0000000..258128d --- /dev/null +++ b/src/tool/toolset.rs @@ -0,0 +1,108 @@ +use std::{collections::HashSet, sync::Arc}; + +use crate::{Result, agent::InvocationContext, session::Session}; + +use super::Tool; + +#[derive(Debug, Clone)] +pub struct ToolsetContext { + pub app_name: String, + pub user_id: String, + pub session_id: String, + pub invocation_id: String, + pub agent_name: String, + pub session_snapshot: Session, +} + +impl ToolsetContext { + pub fn new(ctx: &InvocationContext, session_snapshot: &Session) -> Self { + Self { + app_name: ctx.app_name.clone(), + user_id: ctx.user_id.clone(), + session_id: ctx.session_id.clone(), + invocation_id: ctx.invocation_id.clone(), + agent_name: ctx.agent_name.clone(), + session_snapshot: session_snapshot.clone(), + } + } + + pub fn for_test() -> Self { + Self { + app_name: "test_app".to_string(), + user_id: "test_user".to_string(), + session_id: "test_session".to_string(), + invocation_id: "test_invocation".to_string(), + agent_name: "test_agent".to_string(), + session_snapshot: Session::new("test_app", "test_user", "test_session"), + } + } +} + +pub trait Toolset: Send + Sync { + fn name(&self) -> &str; + + fn tools(&self, ctx: &ToolsetContext) -> Result>>; +} + +pub type ToolPredicate = Arc bool + Send + Sync>; + +pub fn allowed_tools_predicate(allowed_tools: I) -> ToolPredicate +where + I: IntoIterator, + S: Into, +{ + let allowed = allowed_tools + .into_iter() + .map(Into::into) + .collect::>(); + + Arc::new(move |_ctx, tool| allowed.contains(tool.name())) +} + +pub fn filter_toolset(toolset: Arc, predicate: ToolPredicate) -> Arc { + Arc::new(FilterToolset { toolset, predicate }) +} + +pub struct StaticToolset { + name: String, + tools: Vec>, +} + +impl StaticToolset { + pub fn new(name: impl Into, tools: Vec>) -> Self { + Self { + name: name.into(), + tools, + } + } +} + +impl Toolset for StaticToolset { + fn name(&self) -> &str { + &self.name + } + + fn tools(&self, _ctx: &ToolsetContext) -> Result>> { + Ok(self.tools.clone()) + } +} + +struct FilterToolset { + toolset: Arc, + predicate: ToolPredicate, +} + +impl Toolset for FilterToolset { + fn name(&self) -> &str { + self.toolset.name() + } + + fn tools(&self, ctx: &ToolsetContext) -> Result>> { + Ok(self + .toolset + .tools(ctx)? + .into_iter() + .filter(|tool| (self.predicate)(ctx, tool.as_ref())) + .collect()) + } +} diff --git a/src/workflow/loop_agent.rs b/src/workflow/loop_agent.rs index ee8e37e..28dcebf 100644 --- a/src/workflow/loop_agent.rs +++ b/src/workflow/loop_agent.rs @@ -1,6 +1,8 @@ +use std::sync::Arc; + use crate::{ Result, - agent::{Agent, InvocationContext}, + agent::{Agent, IntoAgentArc, InvocationContext}, event::Event, }; @@ -10,21 +12,27 @@ pub struct LoopAgent { name: String, description: String, max_iterations: usize, - sub_agents: Vec>, + sub_agents: Vec>, } impl LoopAgent { - pub fn new( + pub fn new( name: impl Into, description: impl Into, max_iterations: usize, - sub_agents: Vec>, - ) -> Self { + sub_agents: Vec, + ) -> Self + where + T: IntoAgentArc, + { Self { name: name.into(), description: description.into(), max_iterations, - sub_agents, + sub_agents: sub_agents + .into_iter() + .map(IntoAgentArc::into_agent_arc) + .collect(), } } } @@ -38,7 +46,7 @@ impl Agent for LoopAgent { &self.description } - fn sub_agents(&self) -> &[Box] { + fn sub_agents(&self) -> &[Arc] { &self.sub_agents } diff --git a/src/workflow/sequential_agent.rs b/src/workflow/sequential_agent.rs index 5f22f5a..14e5ccf 100644 --- a/src/workflow/sequential_agent.rs +++ b/src/workflow/sequential_agent.rs @@ -1,6 +1,8 @@ +use std::sync::Arc; + use crate::{ Result, - agent::{Agent, InvocationContext}, + agent::{Agent, IntoAgentArc, InvocationContext}, event::Event, session::Session, }; @@ -8,19 +10,25 @@ use crate::{ pub struct SequentialAgent { name: String, description: String, - sub_agents: Vec>, + sub_agents: Vec>, } impl SequentialAgent { - pub fn new( + pub fn new( name: impl Into, description: impl Into, - sub_agents: Vec>, - ) -> Self { + sub_agents: Vec, + ) -> Self + where + T: IntoAgentArc, + { Self { name: name.into(), description: description.into(), - sub_agents, + sub_agents: sub_agents + .into_iter() + .map(IntoAgentArc::into_agent_arc) + .collect(), } } } @@ -34,7 +42,7 @@ impl Agent for SequentialAgent { &self.description } - fn sub_agents(&self) -> &[Box] { + fn sub_agents(&self) -> &[Arc] { &self.sub_agents } @@ -44,7 +52,7 @@ impl Agent for SequentialAgent { } pub(crate) fn run_sub_agents_once( - sub_agents: &[Box], + sub_agents: &[Arc], ctx: InvocationContext, ) -> Result> { let mut session_snapshot = ctx.session_snapshot.clone(); diff --git a/tests/agent_loader.rs b/tests/agent_loader.rs new file mode 100644 index 0000000..97f2d2f --- /dev/null +++ b/tests/agent_loader.rs @@ -0,0 +1,54 @@ +use std::sync::Arc; + +use r_agent::{ + agent::{AgentLoader, CustomAgent, MultiAgentLoader, SingleAgentLoader, validate_agent_names}, + runner::Runner, + session::InMemorySessionService, +}; + +#[test] +fn single_agent_loader_returns_root_agent() { + let root = Arc::new(CustomAgent::new("root", "Root agent.", |_| Ok(Vec::new()))); + let loader = SingleAgentLoader::new(root.clone()); + + assert_eq!(loader.list_agents().unwrap(), vec!["root".to_string()]); + assert_eq!(loader.root_agent().unwrap().name(), "root"); + assert_eq!(loader.load_agent("root").unwrap().name(), "root"); +} + +#[test] +fn multi_agent_loader_loads_agents_by_name() { + let root = Arc::new(CustomAgent::new("root", "Root agent.", |_| Ok(Vec::new()))); + let helper = Arc::new(CustomAgent::new("helper", "Helper agent.", |_| { + Ok(Vec::new()) + })); + let loader = MultiAgentLoader::new(root.clone(), vec![root, helper]).unwrap(); + + assert_eq!( + loader.list_agents().unwrap(), + vec!["helper".to_string(), "root".to_string()] + ); + assert_eq!(loader.root_agent().unwrap().name(), "root"); + assert_eq!(loader.load_agent("helper").unwrap().name(), "helper"); +} + +#[test] +fn duplicate_agent_names_are_rejected() { + let child_a = CustomAgent::new("child", "Child A.", |_| Ok(Vec::new())); + let child_b = CustomAgent::new("child", "Child B.", |_| Ok(Vec::new())); + let root = CustomAgent::new("root", "Root agent.", |_| Ok(Vec::new())) + .with_sub_agent(Box::new(child_a)) + .with_sub_agent(Box::new(child_b)); + + let err = validate_agent_names(&root).unwrap_err(); + + assert_eq!(err.to_string(), "duplicate agent name: child"); +} + +#[test] +fn runner_accepts_arc_agent() { + let root = Arc::new(CustomAgent::new("root", "Root agent.", |_| Ok(Vec::new()))); + let runner = Runner::new("app", root, InMemorySessionService::new(), true); + + assert_eq!(runner.root_agent.name(), "root"); +} diff --git a/tests/artifact.rs b/tests/artifact.rs index 74ef367..2f74afe 100644 --- a/tests/artifact.rs +++ b/tests/artifact.rs @@ -132,7 +132,9 @@ fn tool_context_save_artifact_records_event_delta() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(tool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/async_streaming.rs b/tests/async_streaming.rs index 2756be3..5605810 100644 --- a/tests/async_streaming.rs +++ b/tests/async_streaming.rs @@ -98,7 +98,9 @@ async fn llm_agent_async_stream_yields_model_partials_before_final() { description: "llm".to_string(), system_prompt: String::new(), model: Box::new(model), + generate_content_config: None, tools: vec![], + toolsets: Vec::new(), sub_agents: vec![], max_steps: 4, include_contents: IncludeContents::Default, @@ -181,7 +183,9 @@ async fn async_tool_loop_waits_for_final_tool_call() { description: "llm".to_string(), system_prompt: String::new(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(tool)], + toolsets: Vec::new(), sub_agents: vec![], max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/callbacks.rs b/tests/callbacks.rs index 01a68b5..7d3181c 100644 --- a/tests/callbacks.rs +++ b/tests/callbacks.rs @@ -128,7 +128,9 @@ fn before_agent_callback_can_return_content_and_skip_agent_flow() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -164,7 +166,9 @@ fn after_agent_callback_can_append_agent_event() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -208,7 +212,9 @@ fn before_model_callback_can_return_cached_response_and_skip_model() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -248,7 +254,9 @@ fn after_model_callback_can_replace_model_response() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -279,7 +287,9 @@ fn on_model_error_callback_can_recover_with_response() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(FailingModel), + generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -323,9 +333,11 @@ fn before_tool_callback_can_return_result_and_skip_tool() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(CountingWeatherTool { calls: tool_calls.clone(), })], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, @@ -380,9 +392,11 @@ fn after_tool_callback_can_replace_tool_result() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(CountingWeatherTool { calls: tool_calls.clone(), })], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, @@ -435,7 +449,9 @@ fn on_tool_error_callback_can_recover_with_result() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(FailingTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/content_event.rs b/tests/content_event.rs index fde1dd7..1380dc2 100644 --- a/tests/content_event.rs +++ b/tests/content_event.rs @@ -71,6 +71,7 @@ fn llm_request_and_response_use_content_as_the_model_boundary() { system_prompt: Some("You are helpful.".to_string()), contents: vec![user_content.clone()], tools: Vec::new(), + generate_content_config: None, }; let response = LlmResponse { content: Some(model_content.clone()), diff --git a/tests/database_session.rs b/tests/database_session.rs new file mode 100644 index 0000000..8510012 --- /dev/null +++ b/tests/database_session.rs @@ -0,0 +1,104 @@ +use std::{ + fs, + time::{SystemTime, UNIX_EPOCH}, +}; + +use serde_json::json; + +use r_agent::{ + event::{Event, EventActions}, + session::{DatabaseSessionService, SessionService, SqliteSessionStore}, +}; + +fn sqlite_path() -> std::path::PathBuf { + std::env::temp_dir().join(format!( + "r_agent_session_{}.sqlite", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + )) +} + +fn event_with_delta(id: &str, state_delta: Vec<(&str, serde_json::Value)>) -> Event { + let mut actions = EventActions::default(); + for (key, value) in state_delta { + actions.state_delta.insert(key.to_string(), value); + } + + Event { + id: id.to_string(), + invocation_id: "invocation-1".to_string(), + branch: None, + author: "agent".to_string(), + timestamp: SystemTime::UNIX_EPOCH, + content: None, + actions, + partial: false, + } +} + +#[test] +fn database_session_service_persists_events_and_state() { + let path = sqlite_path(); + let service = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()).unwrap(); + service.create("app", "user-1", "session-1").unwrap(); + service + .append_event( + "app", + "user-1", + "session-1", + event_with_delta("event-1", vec![("answer", json!(42))]), + ) + .unwrap(); + drop(service); + + let reopened = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()).unwrap(); + let session = reopened.get("app", "user-1", "session-1").unwrap(); + + assert_eq!(session.events.len(), 1); + assert_eq!(session.events[0].id, "event-1"); + assert_eq!(session.state.get("answer"), Some(&json!(42))); + + fs::remove_file(path).unwrap(); +} + +#[test] +fn database_session_service_applies_scoped_state_and_discards_temp_state() { + let path = sqlite_path(); + let service = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()).unwrap(); + service.create("app", "user-1", "s1").unwrap(); + service.create("app", "user-2", "s2").unwrap(); + service + .append_event( + "app", + "user-1", + "s1", + event_with_delta( + "event-1", + vec![ + ("app:timezone", json!("UTC")), + ("user:name", json!("Yayoi")), + ("temp:scratch", json!("discard")), + ], + ), + ) + .unwrap(); + + let same_user = service.get("app", "user-1", "s1").unwrap(); + let other_user = service.get("app", "user-2", "s2").unwrap(); + + assert_eq!(same_user.state.get("app:timezone"), Some(&json!("UTC"))); + assert_eq!(same_user.state.get("user:name"), Some(&json!("Yayoi"))); + assert_eq!(same_user.state.get("temp:scratch"), None); + assert_eq!(other_user.state.get("app:timezone"), Some(&json!("UTC"))); + assert_eq!(other_user.state.get("user:name"), None); + assert!( + !same_user.events[0] + .actions + .state_delta + .contains_key("temp:scratch") + ); + + fs::remove_file(path).unwrap(); +} diff --git a/tests/instruction_history.rs b/tests/instruction_history.rs index 17e4bf1..56e0b9c 100644 --- a/tests/instruction_history.rs +++ b/tests/instruction_history.rs @@ -83,7 +83,9 @@ fn agent_with_prompt( description: "A helpful assistant.".to_string(), system_prompt: system_prompt.to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents, diff --git a/tests/llm_agent.rs b/tests/llm_agent.rs index 1e606fa..41344fa 100644 --- a/tests/llm_agent.rs +++ b/tests/llm_agent.rs @@ -63,6 +63,7 @@ fn fake_model_returns_configured_response_and_records_request() { system_prompt: Some("You are helpful.".to_string()), contents: vec![text_content(Role::User, "hello")], tools: Vec::new(), + generate_content_config: None, }; let actual = model.generate(request.clone()).unwrap(); @@ -85,7 +86,9 @@ fn llm_agent_builds_request_from_session_snapshot_and_returns_final_event() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -139,7 +142,9 @@ fn runner_persists_llm_agent_response() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(FakeModel::new("fake-model", response.clone())), + generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, diff --git a/tests/long_running_tool.rs b/tests/long_running_tool.rs index 51adfd3..ae8ea40 100644 --- a/tests/long_running_tool.rs +++ b/tests/long_running_tool.rs @@ -79,7 +79,9 @@ fn long_running_agent(model: FakeModel) -> LlmAgent { description: "A helpful assistant.".to_string(), system_prompt: String::new(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(SandboxJobTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/memory.rs b/tests/memory.rs index f7998c7..fa9a138 100644 --- a/tests/memory.rs +++ b/tests/memory.rs @@ -121,7 +121,9 @@ fn load_memory_tool_reads_memory_through_runner_context() -> Result<()> { description: "A helpful assistant.".to_string(), system_prompt: "Use memory when needed.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(LoadMemoryTool::new())], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/model_factory.rs b/tests/model_factory.rs new file mode 100644 index 0000000..3bf8593 --- /dev/null +++ b/tests/model_factory.rs @@ -0,0 +1,36 @@ +use r_agent::model::{DefaultModelFactory, ModelConfig, ModelFactory}; + +#[test] +fn default_model_factory_creates_openai_compatible_model() { + let model = DefaultModelFactory::new() + .create(ModelConfig { + provider: "deepseek".to_string(), + protocol: "openai-compatible".to_string(), + model: "deepseek-chat".to_string(), + base_url: Some("https://api.deepseek.com/v1".to_string()), + api_key_env: Some("DEEPSEEK_API_KEY".to_string()), + }) + .unwrap(); + + assert_eq!(model.name(), "deepseek-chat"); +} + +#[test] +fn default_model_factory_rejects_unknown_protocol() { + let result = DefaultModelFactory::new().create(ModelConfig { + provider: "example".to_string(), + protocol: "unknown-protocol".to_string(), + model: "example-model".to_string(), + base_url: None, + api_key_env: None, + }); + let err = match result { + Ok(_) => panic!("expected unsupported protocol error"), + Err(err) => err, + }; + + assert_eq!( + err.to_string(), + "unsupported model protocol: unknown-protocol" + ); +} diff --git a/tests/multi_agent.rs b/tests/multi_agent.rs index 98cfae6..229a61c 100644 --- a/tests/multi_agent.rs +++ b/tests/multi_agent.rs @@ -1,4 +1,4 @@ -use std::time::SystemTime; +use std::{sync::Arc, time::SystemTime}; use serde_json::json; @@ -152,8 +152,10 @@ fn llm_agent_can_transfer_to_sub_agent_through_transfer_tool() { description: "Root agent.".to_string(), system_prompt: "Route to a sub agent when useful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), - sub_agents: vec![Box::new(child)], + toolsets: Vec::new(), + sub_agents: vec![Arc::new(child)], max_steps: 4, include_contents: IncludeContents::Default, callbacks: LlmAgentCallbacks::default(), diff --git a/tests/openai_compatible_model.rs b/tests/openai_compatible_model.rs index ac30eba..a810cf9 100644 --- a/tests/openai_compatible_model.rs +++ b/tests/openai_compatible_model.rs @@ -9,7 +9,7 @@ use serde_json::json; use r_agent::{ content::{Content, Part, Role, ToolCall, ToolResponse}, model::{ - LlmRequest, Model, ModelConfig, ModelUsage, OpenAiCompatibleModel, + GenerateContentConfig, LlmRequest, Model, ModelConfig, ModelUsage, OpenAiCompatibleModel, openai_compatible::{request_json, response_from_json}, }, tool::ToolDeclaration, @@ -58,6 +58,7 @@ fn converts_internal_request_to_openai_compatible_json() { })), output_schema: None, }], + generate_content_config: None, }; let actual = request_json(&request); @@ -110,6 +111,29 @@ fn converts_internal_request_to_openai_compatible_json() { ); } +#[test] +fn request_json_includes_generate_content_config() { + let request = LlmRequest { + model: "test-model".to_string(), + system_prompt: None, + contents: vec![text_content(Role::User, "hello")], + tools: Vec::new(), + generate_content_config: Some(GenerateContentConfig { + temperature: Some(0.2), + top_p: Some(0.9), + max_output_tokens: Some(1024), + stop_sequences: vec!["END".to_string()], + }), + }; + + let actual = request_json(&request); + + assert_eq!(actual["temperature"], json!(0.2)); + assert_eq!(actual["top_p"], json!(0.9)); + assert_eq!(actual["max_tokens"], json!(1024)); + assert_eq!(actual["stop"], json!(["END"])); +} + #[test] fn converts_openai_compatible_text_response_to_llm_response() { let response = response_from_json(&json!({ @@ -263,6 +287,7 @@ fn generate_posts_chat_completion_request_and_parses_response() { system_prompt: None, contents: vec![text_content(Role::User, "hello")], tools: Vec::new(), + generate_content_config: None, }) .unwrap(); @@ -289,6 +314,7 @@ fn generate_errors_when_api_key_env_is_missing() { system_prompt: None, contents: vec![text_content(Role::User, "hello")], tools: Vec::new(), + generate_content_config: None, }) .unwrap_err(); diff --git a/tests/openai_compatible_real.rs b/tests/openai_compatible_real.rs index 6de7a82..5e2b2da 100644 --- a/tests/openai_compatible_real.rs +++ b/tests/openai_compatible_real.rs @@ -39,6 +39,7 @@ fn real_openai_compatible_text_generation_when_enabled() { system_prompt: Some("Reply with exactly: pong".to_string()), contents: vec![text_content(Role::User, "ping")], tools: Vec::new(), + generate_content_config: None, }) .unwrap(); diff --git a/tests/scoped_state.rs b/tests/scoped_state.rs new file mode 100644 index 0000000..66b96cb --- /dev/null +++ b/tests/scoped_state.rs @@ -0,0 +1,112 @@ +use std::time::SystemTime; + +use serde_json::json; + +use r_agent::{ + event::{Event, EventActions}, + session::{ + APP_STATE_PREFIX, InMemorySessionService, SESSION_STATE_PREFIX, SessionService, + TEMP_STATE_PREFIX, USER_STATE_PREFIX, + }, +}; + +fn event_with_delta(id: &str, state_delta: Vec<(&str, serde_json::Value)>) -> Event { + let mut actions = EventActions::default(); + for (key, value) in state_delta { + actions.state_delta.insert(key.to_string(), value); + } + + Event { + id: id.to_string(), + invocation_id: "invocation-1".to_string(), + branch: None, + author: "agent".to_string(), + timestamp: SystemTime::now(), + content: None, + actions, + partial: false, + } +} + +#[test] +fn app_state_is_shared_across_users_and_sessions() { + let service = InMemorySessionService::new(); + service.create("app", "user-1", "s1").unwrap(); + service.create("app", "user-2", "s2").unwrap(); + + service + .append_event( + "app", + "user-1", + "s1", + event_with_delta("event-1", vec![("app:timezone", json!("UTC"))]), + ) + .unwrap(); + + let session = service.get("app", "user-2", "s2").unwrap(); + + assert_eq!(session.state.get("app:timezone"), Some(&json!("UTC"))); +} + +#[test] +fn user_state_is_shared_only_for_same_user() { + let service = InMemorySessionService::new(); + service.create("app", "user-1", "s1").unwrap(); + service.create("app", "user-1", "s2").unwrap(); + service.create("app", "user-2", "s3").unwrap(); + + service + .append_event( + "app", + "user-1", + "s1", + event_with_delta("event-1", vec![("user:name", json!("Yayoi"))]), + ) + .unwrap(); + + let same_user = service.get("app", "user-1", "s2").unwrap(); + let other_user = service.get("app", "user-2", "s3").unwrap(); + + assert_eq!(same_user.state.get("user:name"), Some(&json!("Yayoi"))); + assert_eq!(other_user.state.get("user:name"), None); +} + +#[test] +fn temp_state_is_not_persisted_to_state_or_event_delta() { + let service = InMemorySessionService::new(); + service.create("app", "user-1", "s1").unwrap(); + + service + .append_event( + "app", + "user-1", + "s1", + event_with_delta( + "event-1", + vec![ + ("temp:scratch", json!("discard")), + ("session_value", json!("keep")), + ], + ), + ) + .unwrap(); + + let session = service.get("app", "user-1", "s1").unwrap(); + + assert_eq!(session.state.get("temp:scratch"), None); + assert_eq!(session.state.get("session_value"), Some(&json!("keep"))); + assert!( + !session.events[0] + .actions + .state_delta + .contains_key("temp:scratch") + ); +} + +#[test] +fn state_prefix_constants_match_google_adk_style() { + assert_eq!(APP_STATE_PREFIX, "app:"); + assert_eq!(USER_STATE_PREFIX, "user:"); + assert_eq!(TEMP_STATE_PREFIX, "temp:"); + assert_eq!(SESSION_STATE_PREFIX, ""); +} diff --git a/tests/streaming.rs b/tests/streaming.rs index ead7ad5..dd1e868 100644 --- a/tests/streaming.rs +++ b/tests/streaming.rs @@ -111,7 +111,9 @@ fn llm_agent_turns_streaming_model_responses_into_partial_and_final_events() -> description: "A streaming assistant.".to_string(), system_prompt: "You stream.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, diff --git a/tests/telemetry.rs b/tests/telemetry.rs index 81e232f..1179b67 100644 --- a/tests/telemetry.rs +++ b/tests/telemetry.rs @@ -40,7 +40,9 @@ fn agent(model: FakeModel) -> LlmAgent { description: "A helpful assistant.".to_string(), system_prompt: String::new(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(tool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/toml_agent_loader.rs b/tests/toml_agent_loader.rs new file mode 100644 index 0000000..b9b114a --- /dev/null +++ b/tests/toml_agent_loader.rs @@ -0,0 +1,127 @@ +use std::{ + fs, + sync::Arc, + time::{SystemTime, UNIX_EPOCH}, +}; + +use serde_json::{Value, json}; + +use r_agent::{ + agent::{AgentLoader, TomlAgentLoader}, + model::DefaultModelFactory, + tool::{FunctionTool, ToolContext, ToolRegistry}, +}; + +fn write_config(contents: &str) -> std::path::PathBuf { + let path = std::env::temp_dir().join(format!( + "r_agent_config_{}.toml", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + fs::write(&path, contents).unwrap(); + path +} + +fn registry() -> ToolRegistry { + let mut registry = ToolRegistry::new(); + registry + .register(Arc::new(FunctionTool::new( + "get_weather", + "Gets weather.", + |_ctx: &mut ToolContext, _args: Value| Ok(json!({ "forecast": "sunny" })), + ))) + .unwrap(); + registry +} + +#[test] +fn toml_agent_loader_loads_llm_agent_tree() { + let path = write_config( + r#" +root_agent = "assistant" + +[models.fake_openai] +provider = "test" +protocol = "openai-compatible" +model = "test-model" +base_url = "http://localhost:11434/v1" +api_key_env = "TEST_API_KEY" + +[agents.assistant] +type = "llm" +description = "Main assistant." +model = "fake_openai" +system_prompt = "You are helpful." +include_contents = "default" +max_steps = 4 +tools = ["get_weather"] +sub_agents = ["researcher"] + +[agents.assistant.generate_content_config] +temperature = 0.2 +top_p = 0.9 +max_output_tokens = 128 +stop_sequences = ["END"] + +[agents.researcher] +type = "llm" +description = "Research assistant." +model = "fake_openai" +system_prompt = "Research carefully." +include_contents = "none" +max_steps = 2 +tools = [] +sub_agents = [] +"#, + ); + + let loader = + TomlAgentLoader::from_path(&path, Arc::new(DefaultModelFactory::new()), registry()) + .unwrap(); + + assert_eq!( + loader.list_agents().unwrap(), + vec!["assistant".to_string(), "researcher".to_string()] + ); + let root = loader.root_agent().unwrap(); + assert_eq!(root.name(), "assistant"); + assert_eq!(root.sub_agents()[0].name(), "researcher"); +} + +#[test] +fn toml_agent_loader_rejects_unknown_tool() { + let path = write_config( + r#" +root_agent = "assistant" + +[models.fake_openai] +provider = "test" +protocol = "openai-compatible" +model = "test-model" +base_url = "http://localhost:11434/v1" + +[agents.assistant] +type = "llm" +description = "Main assistant." +model = "fake_openai" +system_prompt = "You are helpful." +include_contents = "default" +max_steps = 4 +tools = ["missing_tool"] +sub_agents = [] +"#, + ); + + let err = match TomlAgentLoader::from_path( + &path, + Arc::new(DefaultModelFactory::new()), + ToolRegistry::new(), + ) { + Ok(_) => panic!("expected unknown tool error"), + Err(err) => err, + }; + + assert_eq!(err.to_string(), "unknown tool in config: missing_tool"); +} diff --git a/tests/tool_confirmation.rs b/tests/tool_confirmation.rs index ae0c66d..6a1ce3e 100644 --- a/tests/tool_confirmation.rs +++ b/tests/tool_confirmation.rs @@ -79,7 +79,9 @@ fn agent_with_confirmed_delete_tool(calls: Arc, model: FakeModel) - description: "A helpful assistant.".to_string(), system_prompt: String::new(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(delete_tool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/tool_loop.rs b/tests/tool_loop.rs index ab7bbc8..4ac7f34 100644 --- a/tests/tool_loop.rs +++ b/tests/tool_loop.rs @@ -75,7 +75,9 @@ fn llm_agent_runs_tool_loop_until_final_response() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(WeatherTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, @@ -155,7 +157,9 @@ fn llm_agent_errors_when_tool_loop_exceeds_max_steps() { usage: None, }, )), + generate_content_config: None, tools: vec![Box::new(WeatherTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, diff --git a/tests/toolset.rs b/tests/toolset.rs new file mode 100644 index 0000000..4895dd5 --- /dev/null +++ b/tests/toolset.rs @@ -0,0 +1,184 @@ +use std::sync::Arc; +use std::time::SystemTime; + +use serde_json::{Value, json}; + +use r_agent::{ + Result, + content::{Content, Part, Role, ToolCall}, + event::{Event, EventActions}, + llm_agent::{IncludeContents, LlmAgent, LlmAgentCallbacks, LlmAgentConfig}, + model::{FakeModel, LlmResponse}, + runner::Runner, + session::{InMemorySessionService, SessionService}, + tool::{ + FunctionTool, StaticToolset, Tool, ToolContext, Toolset, ToolsetContext, + allowed_tools_predicate, filter_toolset, + }, +}; + +fn text_content(role: Role, text: &str) -> Content { + Content { + role, + parts: vec![Part::Text(text.to_string())], + } +} + +fn tool_call_content(name: &str) -> Content { + Content { + role: Role::Model, + parts: vec![Part::ToolCall(ToolCall { + id: "call-1".to_string(), + name: name.to_string(), + args: json!({}), + })], + } +} + +fn named_tool(name: &str) -> Arc { + Arc::new(FunctionTool::new( + name, + format!("Runs {name}."), + |_ctx: &mut ToolContext, _args: Value| Ok(json!({ "ok": true })), + )) +} + +struct StateToolset { + required_key: String, + tool: Arc, +} + +impl Toolset for StateToolset { + fn name(&self) -> &str { + "state_toolset" + } + + fn tools(&self, ctx: &ToolsetContext) -> Result>> { + if ctx.session_snapshot.state.contains_key(&self.required_key) { + Ok(vec![self.tool.clone()]) + } else { + Ok(Vec::new()) + } + } +} + +#[test] +fn dynamic_toolset_exposes_tools_from_current_session_state() { + let model = FakeModel::new( + "fake-model", + LlmResponse { + content: Some(text_content(Role::Model, "done")), + partial: false, + usage: None, + }, + ); + let model_handle = model.clone(); + let agent = LlmAgent::new(LlmAgentConfig { + name: "assistant".to_string(), + description: "Assistant.".to_string(), + system_prompt: String::new(), + model: Box::new(model), + generate_content_config: None, + tools: Vec::new(), + toolsets: vec![Arc::new(StateToolset { + required_key: "project_id".to_string(), + tool: named_tool("project_search"), + })], + sub_agents: Vec::new(), + max_steps: 2, + include_contents: IncludeContents::Default, + callbacks: LlmAgentCallbacks::default(), + }); + let session_service = InMemorySessionService::new(); + session_service + .create("app", "user-1", "session-1") + .unwrap(); + let mut actions = EventActions::default(); + actions + .state_delta + .insert("project_id".to_string(), json!("p1")); + session_service + .append_event( + "app", + "user-1", + "session-1", + Event { + id: "state-event".to_string(), + invocation_id: "setup".to_string(), + branch: None, + author: "setup".to_string(), + timestamp: SystemTime::now(), + content: None, + actions, + partial: false, + }, + ) + .unwrap(); + let runner = Runner::new("app", Box::new(agent), session_service, true); + + runner + .run( + "user-1", + "session-1", + text_content(Role::User, "search project"), + ) + .unwrap(); + + let requests = model_handle.requests(); + assert_eq!(requests[0].tools.len(), 1); + assert_eq!(requests[0].tools[0].name, "project_search"); +} + +#[test] +fn filter_toolset_uses_predicate_to_limit_visible_tools() { + let source = Arc::new(StaticToolset::new( + "all_tools", + vec![named_tool("read_file"), named_tool("write_file")], + )); + let filtered = filter_toolset(source, allowed_tools_predicate(["read_file"])); + let ctx = ToolsetContext::for_test(); + + let tools = filtered.tools(&ctx).unwrap(); + + assert_eq!(tools.len(), 1); + assert_eq!(tools[0].name(), "read_file"); +} + +#[test] +fn llm_agent_rejects_tool_call_that_was_not_visible_to_model() { + let model = FakeModel::new( + "fake-model", + LlmResponse { + content: Some(tool_call_content("hidden_tool")), + partial: false, + usage: None, + }, + ); + let agent = LlmAgent::new(LlmAgentConfig { + name: "assistant".to_string(), + description: "Assistant.".to_string(), + system_prompt: String::new(), + model: Box::new(model), + generate_content_config: None, + tools: Vec::new(), + toolsets: Vec::new(), + sub_agents: Vec::new(), + max_steps: 2, + include_contents: IncludeContents::Default, + callbacks: LlmAgentCallbacks::default(), + }); + let runner = Runner::new("app", Box::new(agent), InMemorySessionService::new(), true); + + let err = runner + .run( + "user-1", + "session-1", + text_content(Role::User, "call hidden tool"), + ) + .unwrap_err(); + + assert_eq!( + err.to_string(), + "tool not visible in current invocation: hidden_tool" + ); +}