use klieo_core::{
Agent, AgentContext, BusHandles, EpisodicMemory, Error, LlmClient, LongTermMemory,
MemoryHandles, RunId, ShortTermMemory, ToolDef, ToolInvoker,
};
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
#[cfg(feature = "tools")]
use klieo_core::Tool;
#[cfg(any(
feature = "llm-openai",
feature = "llm-anthropic",
feature = "llm-gemini"
))]
use secrecy::SecretString;
#[derive(Clone)]
pub struct App {
llm: Arc<dyn LlmClient>,
memory: MemoryHandles,
bus: BusHandles,
tools: Arc<dyn ToolInvoker>,
catalogue: Vec<ToolDef>,
parent_cancel: CancellationToken,
}
impl std::fmt::Debug for App {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("App")
.field("llm", &self.llm.name())
.field("tools", &self.catalogue.len())
.field("cancelled", &self.parent_cancel.is_cancelled())
.finish_non_exhaustive()
}
}
impl App {
#[cfg(all(
feature = "llm-ollama",
feature = "memory-sqlite",
feature = "bus-memory",
feature = "tools",
))]
pub fn local() -> AppBuilder {
AppBuilder::new()
.ollama("http://localhost:11434", "qwen2.5:14b")
.sqlite(":memory:")
.memory_bus()
}
#[cfg(all(
feature = "llm-ollama",
feature = "memory-sqlite",
feature = "bus-memory",
feature = "tools",
))]
pub fn local_at(db_path: impl Into<std::path::PathBuf>) -> AppBuilder {
AppBuilder::new()
.ollama("http://localhost:11434", "qwen2.5:14b")
.sqlite(db_path)
.memory_bus()
}
pub fn builder() -> AppBuilder {
AppBuilder::new()
}
pub fn context(&self, agent_name: impl Into<String>) -> AgentContext {
AgentContext::new(
self.llm.clone(),
self.memory.short_term.clone(),
self.memory.long_term.clone(),
self.memory.episodic.clone(),
self.bus.pubsub.clone(),
self.bus.kv.clone(),
self.bus.request_reply.clone(),
self.bus.jobs.clone(),
self.tools.clone(),
RunId::new(),
self.parent_cancel.child_token(),
agent_name,
)
}
pub fn tools_catalogue(&self) -> &[ToolDef] {
&self.catalogue
}
pub fn cancel_token(&self) -> CancellationToken {
self.parent_cancel.clone()
}
pub async fn run<A: Agent>(&self, agent: &A, input: A::Input) -> Result<A::Output, A::Error> {
let ctx = self.context(agent.name());
agent.run(ctx, input).await
}
}
enum LlmShortcut {
#[cfg(feature = "llm-ollama")]
Ollama { base_url: String, model: String },
#[cfg(feature = "llm-openai")]
OpenAi {
api_key: SecretString,
model: String,
},
#[cfg(feature = "llm-anthropic")]
Anthropic {
api_key: SecretString,
model: String,
},
#[cfg(feature = "llm-gemini")]
Gemini {
api_key: SecretString,
model: String,
},
}
enum BusShortcut {
#[cfg(feature = "bus-memory")]
Memory,
#[cfg(feature = "bus-nats")]
Nats(Box<klieo_bus_nats::NatsBusConfig>),
}
#[derive(Default)]
pub struct AppBuilder {
llm: Option<Arc<dyn LlmClient>>,
memory: Option<MemoryHandles>,
bus: Option<BusHandles>,
tools_invoker: Option<Arc<dyn ToolInvoker>>,
#[cfg(feature = "tools")]
pending_tools: Vec<Arc<dyn Tool>>,
parent_cancel: Option<CancellationToken>,
llm_shortcut: Option<LlmShortcut>,
bus_shortcut: Option<BusShortcut>,
short_term: Option<Arc<dyn ShortTermMemory>>,
long_term: Option<Arc<dyn LongTermMemory>>,
episodic: Option<Arc<dyn EpisodicMemory>>,
#[cfg(feature = "memory-sqlite")]
sqlite_path: Option<std::path::PathBuf>,
}
impl AppBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn llm(mut self, llm: Arc<dyn LlmClient>) -> Self {
self.llm = Some(llm);
self
}
pub fn memory(mut self, memory: impl Into<MemoryHandles>) -> Self {
self.memory = Some(memory.into());
self
}
pub fn bus(mut self, bus: impl Into<BusHandles>) -> Self {
self.bus = Some(bus.into());
self
}
pub fn tools_invoker(mut self, tools: Arc<dyn ToolInvoker>) -> Self {
self.tools_invoker = Some(tools);
self
}
#[cfg(feature = "tools")]
pub fn tool<T: Tool + 'static>(mut self, tool: T) -> Self {
self.pending_tools.push(Arc::new(tool));
self
}
pub fn cancel_token(mut self, token: CancellationToken) -> Self {
self.parent_cancel = Some(token);
self
}
pub fn short_term(mut self, short_term: Arc<dyn ShortTermMemory>) -> Self {
self.short_term = Some(short_term);
self
}
pub fn long_term(mut self, long_term: Arc<dyn LongTermMemory>) -> Self {
self.long_term = Some(long_term);
self
}
pub fn episodic(mut self, episodic: Arc<dyn EpisodicMemory>) -> Self {
self.episodic = Some(episodic);
self
}
#[cfg(feature = "llm-ollama")]
pub fn ollama(mut self, base_url: impl Into<String>, model: impl Into<String>) -> Self {
self.llm_shortcut = Some(LlmShortcut::Ollama {
base_url: base_url.into(),
model: model.into(),
});
self
}
#[cfg(feature = "llm-ollama")]
pub fn model(mut self, model: impl Into<String>) -> Self {
let base_url = match self.llm_shortcut.take() {
Some(LlmShortcut::Ollama { base_url, .. }) => base_url,
_ => "http://localhost:11434".to_string(),
};
self.llm_shortcut = Some(LlmShortcut::Ollama {
base_url,
model: model.into(),
});
self
}
#[cfg(feature = "llm-openai")]
pub fn openai(mut self, api_key: impl Into<SecretString>, model: impl Into<String>) -> Self {
self.llm_shortcut = Some(LlmShortcut::OpenAi {
api_key: api_key.into(),
model: model.into(),
});
self
}
#[cfg(feature = "llm-anthropic")]
pub fn anthropic(mut self, api_key: impl Into<SecretString>, model: impl Into<String>) -> Self {
self.llm_shortcut = Some(LlmShortcut::Anthropic {
api_key: api_key.into(),
model: model.into(),
});
self
}
#[cfg(feature = "llm-gemini")]
pub fn gemini(mut self, api_key: impl Into<SecretString>, model: impl Into<String>) -> Self {
self.llm_shortcut = Some(LlmShortcut::Gemini {
api_key: api_key.into(),
model: model.into(),
});
self
}
#[cfg(feature = "memory-sqlite")]
pub fn sqlite(mut self, path: impl Into<std::path::PathBuf>) -> Self {
self.sqlite_path = Some(path.into());
self
}
#[cfg(feature = "bus-memory")]
pub fn memory_bus(mut self) -> Self {
self.bus_shortcut = Some(BusShortcut::Memory);
self
}
#[cfg(feature = "bus-nats")]
pub fn nats(mut self, config: klieo_bus_nats::NatsBusConfig) -> Self {
self.bus_shortcut = Some(BusShortcut::Nats(Box::new(config)));
self
}
pub async fn build(mut self) -> Result<App, Error> {
self.resolve_shortcuts().await?;
let llm = self.llm.ok_or(Error::AppBuildError { missing: "llm" })?;
let memory = self
.memory
.ok_or(Error::AppBuildError { missing: "memory" })?;
let bus = self.bus.ok_or(Error::AppBuildError { missing: "bus" })?;
let (tools, catalogue) = resolve_tools(
self.tools_invoker,
#[cfg(feature = "tools")]
self.pending_tools,
)?;
Ok(App {
llm,
memory,
bus,
tools,
catalogue,
parent_cancel: self.parent_cancel.unwrap_or_default(),
})
}
async fn resolve_shortcuts(&mut self) -> Result<(), Error> {
self.resolve_llm_shortcut();
self.resolve_sqlite().await?;
self.resolve_memory_partials()?;
self.resolve_bus_shortcut().await?;
Ok(())
}
fn resolve_llm_shortcut(&mut self) {
if self.llm.is_some() {
return;
}
match self.llm_shortcut.take() {
None => {}
#[cfg(feature = "llm-ollama")]
Some(LlmShortcut::Ollama { base_url, model }) => {
self.llm = Some(Arc::new(klieo_llm_ollama::OllamaClient::new(
base_url, model,
)));
}
#[cfg(feature = "llm-openai")]
Some(LlmShortcut::OpenAi { api_key, model }) => {
self.llm = Some(Arc::new(klieo_llm_openai::OpenAiClient::new(
api_key, model,
)));
}
#[cfg(feature = "llm-anthropic")]
Some(LlmShortcut::Anthropic { api_key, model }) => {
self.llm = Some(Arc::new(klieo_llm_anthropic::AnthropicClient::new(
api_key, model,
)));
}
#[cfg(feature = "llm-gemini")]
Some(LlmShortcut::Gemini { api_key, model }) => {
self.llm = Some(Arc::new(klieo_llm_gemini::GeminiClient::new(
api_key, model,
)));
}
}
}
#[cfg(feature = "memory-sqlite")]
async fn resolve_sqlite(&mut self) -> Result<(), Error> {
if self.memory.is_some() {
return Ok(());
}
if let Some(path) = self.sqlite_path.take() {
let mem = klieo_memory_sqlite::MemorySqlite::new(
path,
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await?;
self.memory = Some(mem.into());
}
Ok(())
}
#[cfg(not(feature = "memory-sqlite"))]
#[allow(clippy::unused_async)]
async fn resolve_sqlite(&mut self) -> Result<(), Error> {
Ok(())
}
fn resolve_memory_partials(&mut self) -> Result<(), Error> {
if self.memory.is_some() {
return Ok(());
}
let any_partial =
self.short_term.is_some() || self.long_term.is_some() || self.episodic.is_some();
if !any_partial {
return Ok(());
}
let short_term = self.short_term.take().ok_or(Error::AppBuildError {
missing: "memory.short_term",
})?;
let long_term = self.long_term.take().ok_or(Error::AppBuildError {
missing: "memory.long_term",
})?;
let episodic = self.episodic.take().ok_or(Error::AppBuildError {
missing: "memory.episodic",
})?;
self.memory = Some(MemoryHandles::new(short_term, long_term, episodic));
Ok(())
}
async fn resolve_bus_shortcut(&mut self) -> Result<(), Error> {
if self.bus.is_some() {
return Ok(());
}
match self.bus_shortcut.take() {
None => {}
#[cfg(feature = "bus-memory")]
Some(BusShortcut::Memory) => {
self.bus = Some(klieo_bus_memory::MemoryBus::new().into());
}
#[cfg(feature = "bus-nats")]
Some(BusShortcut::Nats(config)) => {
let bus = klieo_bus_nats::NatsBus::connect(*config).await?;
self.bus = Some(bus.into());
}
}
Ok(())
}
}
fn resolve_tools(
invoker: Option<Arc<dyn ToolInvoker>>,
#[cfg(feature = "tools")] pending: Vec<Arc<dyn Tool>>,
) -> Result<(Arc<dyn ToolInvoker>, Vec<ToolDef>), Error> {
#[cfg(feature = "tools")]
{
let has_pending = !pending.is_empty();
match (invoker, has_pending) {
(Some(_), true) => Err(Error::AppBuildError {
missing: "tools (cannot combine .tool() with .tools_invoker())",
}),
(Some(inv), false) => {
let cat = inv.catalogue();
Ok((inv, cat))
}
(None, _) => {
let mut chained = klieo_tools::ChainedInvoker::new();
for t in pending {
chained
.add_tool(t)
.map_err(|e| Error::wrap("tool registration failed", e))?;
}
let cat = chained.catalogue();
Ok((Arc::new(chained), cat))
}
}
}
#[cfg(not(feature = "tools"))]
{
let inv = invoker.ok_or(Error::AppBuildError { missing: "tools" })?;
let cat = inv.catalogue();
Ok((inv, cat))
}
}
#[cfg(test)]
mod tests {
use super::*;
use klieo_core::test_utils::{FakeLlmClient, FakeLlmStep, FakeToolInvoker};
fn fake_llm() -> Arc<dyn LlmClient> {
Arc::new(FakeLlmClient::new("fake"))
}
fn fake_tools() -> Arc<dyn ToolInvoker> {
Arc::new(FakeToolInvoker::new())
}
#[tokio::test]
async fn build_missing_llm_errors() {
let err = App::builder()
.memory(
klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap(),
)
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(err, Error::AppBuildError { missing: "llm" }));
}
#[tokio::test]
async fn build_missing_memory_errors() {
let err = App::builder()
.llm(fake_llm())
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(err, Error::AppBuildError { missing: "memory" }));
}
#[tokio::test]
async fn build_missing_bus_errors() {
let err = App::builder()
.llm(fake_llm())
.memory(
klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap(),
)
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(err, Error::AppBuildError { missing: "bus" }));
}
#[tokio::test]
async fn build_local_with_only_model_succeeds() {
let app = App::local().model("dummy-model").build().await.unwrap();
assert!(app.tools_catalogue().is_empty());
let ctx = app.context("test-agent");
assert_eq!(ctx.agent_name, "test-agent");
}
#[tokio::test]
async fn context_mints_fresh_run_id_per_call() {
let app = App::local().model("dummy").build().await.unwrap();
let a = app.context("agent");
let b = app.context("agent");
assert_ne!(a.run_id, b.run_id);
}
#[tokio::test]
async fn cancel_token_propagates_to_minted_context() {
let app = App::local().model("dummy").build().await.unwrap();
let ctx = app.context("agent");
assert!(!ctx.cancel.is_cancelled());
app.cancel_token().cancel();
assert!(ctx.cancel.is_cancelled());
}
#[tokio::test]
async fn cannot_combine_tool_with_tools_invoker() {
let mut builder = App::local().model("dummy");
builder.pending_tools.push(Arc::new(EchoTool));
let err = builder
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(
err,
Error::AppBuildError {
missing: "tools (cannot combine .tool() with .tools_invoker())"
}
));
}
struct EchoTool;
#[async_trait::async_trait]
impl Tool for EchoTool {
fn name(&self) -> &str {
"echo"
}
fn description(&self) -> &str {
"echo"
}
fn json_schema(&self) -> &serde_json::Value {
static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
SCHEMA.get_or_init(|| serde_json::json!({"type": "object"}))
}
async fn invoke(
&self,
_args: serde_json::Value,
_ctx: klieo_core::ToolCtx,
) -> Result<serde_json::Value, klieo_core::ToolError> {
Ok(serde_json::Value::Null)
}
}
#[tokio::test]
async fn partial_memory_composes_into_handles() {
let mem = klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap();
let app = App::builder()
.llm(fake_llm())
.short_term(mem.short_term.clone())
.long_term(mem.long_term.clone())
.episodic(mem.episodic.clone())
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap();
let ctx = app.context("partial");
assert!(Arc::ptr_eq(&ctx.short_term, &mem.short_term));
assert!(Arc::ptr_eq(&ctx.long_term, &mem.long_term));
assert!(Arc::ptr_eq(&ctx.episodic, &mem.episodic));
}
#[tokio::test]
async fn partial_memory_missing_long_term_errors() {
let mem = klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap();
let err = App::builder()
.llm(fake_llm())
.short_term(mem.short_term.clone())
.episodic(mem.episodic.clone())
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(
err,
Error::AppBuildError {
missing: "memory.long_term"
}
));
}
#[tokio::test]
async fn explicit_memory_overrides_partials() {
let primary = klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap();
let secondary = klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap();
let primary_st = primary.short_term.clone();
let app = App::builder()
.llm(fake_llm())
.memory(primary)
.short_term(secondary.short_term)
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap();
let ctx = app.context("override");
assert!(Arc::ptr_eq(&ctx.short_term, &primary_st));
}
#[cfg(feature = "llm-openai")]
#[tokio::test]
async fn openai_shortcut_sets_llm() {
let app = App::local()
.openai("test-key".to_string(), "gpt-4o-mini")
.build()
.await
.unwrap();
let ctx = app.context("a");
assert!(
ctx.llm.name().starts_with("openai"),
"got {}",
ctx.llm.name()
);
}
#[cfg(feature = "llm-anthropic")]
#[tokio::test]
async fn anthropic_shortcut_sets_llm() {
let app = App::local()
.anthropic("test-key".to_string(), "claude-3-5-haiku-latest")
.build()
.await
.unwrap();
let ctx = app.context("a");
assert!(
ctx.llm.name().starts_with("anthropic"),
"got {}",
ctx.llm.name()
);
}
#[cfg(feature = "llm-gemini")]
#[tokio::test]
async fn gemini_shortcut_sets_llm() {
let app = App::local()
.gemini("test-key".to_string(), "gemini-2.0-flash")
.build()
.await
.unwrap();
let ctx = app.context("a");
assert!(ctx.llm.name().contains("gemini"), "got {}", ctx.llm.name());
}
#[tokio::test]
async fn last_llm_shortcut_wins() {
let app = App::local()
.ollama("http://localhost:11434", "qwen")
.ollama("http://other:11434", "llama")
.build()
.await
.unwrap();
let ctx = app.context("a");
assert!(
ctx.llm.name().starts_with("ollama"),
"got {}",
ctx.llm.name()
);
}
#[cfg(feature = "bus-nats")]
#[tokio::test]
async fn nats_shortcut_errors_on_unreachable_server() {
let cfg = klieo_bus_nats::NatsBusConfig {
url: "nats://127.0.0.1:14223".to_string(),
..klieo_bus_nats::NatsBusConfig::default()
};
let err = App::builder()
.llm(fake_llm())
.memory(
klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap(),
)
.nats(cfg)
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(err, Error::Bus(_)), "got {err:?}");
}
#[tokio::test]
async fn run_delegates_to_agent_with_fresh_context() {
let app = App::builder()
.llm(Arc::new(
FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]),
))
.memory(
klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap(),
)
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap();
let agent =
klieo_core::SimpleAgent::new("greeter", "be brief", app.tools_catalogue().to_vec());
let out = app.run(&agent, "hi".into()).await.unwrap();
assert_eq!(out, "done");
}
#[cfg(all(
feature = "llm-ollama",
feature = "memory-sqlite",
feature = "bus-memory",
feature = "tools",
))]
#[tokio::test]
async fn build_local_at_file_path_succeeds() {
let db_path = std::env::temp_dir().join("klieo_test_local_at.db");
let _ = std::fs::remove_file(&db_path);
let app = App::local_at(&db_path)
.model("dummy-model")
.build()
.await
.unwrap();
assert!(app.tools_catalogue().is_empty());
let ctx = app.context("test-agent");
assert_eq!(ctx.agent_name, "test-agent");
assert!(db_path.exists(), "sqlite file should have been created at {db_path:?}");
let _ = std::fs::remove_file(&db_path);
}
}