use async_trait::async_trait;
use klieo_bus_memory::MemoryBus;
use klieo_core::agent::AgentContext;
use klieo_core::llm::LlmClient;
use klieo_core::test_utils::{
FakeLlmClient, FakeToolInvoker, InMemoryEpisodic, InMemoryLongTerm, InMemoryShortTerm,
ScriptedLlm,
};
use klieo_flows::{Flow, FlowError};
use serde_json::Value;
use std::sync::Arc;
pub fn dummy_llm() -> Arc<dyn LlmClient> {
Arc::new(FakeLlmClient::new("dummy"))
}
pub fn stub_llm(reply: &str) -> Arc<dyn LlmClient> {
Arc::new(ScriptedLlm::new(vec![reply.to_string()]))
}
pub fn const_subflow(tag: &'static str) -> Arc<dyn Flow> {
struct Const(&'static str);
#[async_trait]
impl Flow for Const {
fn name(&self) -> &str {
self.0
}
async fn run(&self, _ctx: AgentContext, _input: Value) -> Result<Value, FlowError> {
Ok(serde_json::json!({ "tag": self.0 }))
}
}
Arc::new(Const(tag))
}
pub fn test_ctx() -> AgentContext {
let bus = MemoryBus::new();
AgentContext::builder()
.llm(dummy_llm())
.short_term(Arc::new(InMemoryShortTerm::default()))
.long_term(Arc::new(InMemoryLongTerm::default()))
.episodic(Arc::new(InMemoryEpisodic::default()))
.pubsub(bus.pubsub)
.kv(bus.kv)
.request_reply(bus.request_reply)
.jobs(bus.jobs)
.tools(Arc::new(FakeToolInvoker::new()))
.agent_name("workflow-unit-test")
.build()
.expect("all required AgentContext fields set")
}