use std::sync::Arc;
use agentyk_core::budget::BudgetChecker;
use agentyk_core::capability::Capability;
use agentyk_core::context::ContextAssembler;
use agentyk_core::driver::{ChatDriver, ModelSpec};
use agentyk_core::error::{Error, Result};
use agentyk_core::event::EventListener;
use agentyk_core::event_log::{EventLog, InMemoryEventLog};
use agentyk_core::hook::Hook;
use agentyk_core::id::SessionId;
use agentyk_core::message::Message;
use agentyk_core::middleware::TurnMiddleware;
use agentyk_core::profile::{ModelCatalog, ModelProfile};
use agentyk_core::tool::Tool;
use async_trait::async_trait;
use crate::config::AgentConfig;
use crate::host::TurnResult;
use crate::session::Session;
#[derive(Clone)]
pub struct AgentDefinition {
pub name: String,
pub instructions: String,
pub model: ModelSpec,
pub capabilities: Vec<Arc<dyn Capability>>,
pub max_iterations: usize,
pub middleware: Vec<Arc<dyn TurnMiddleware>>,
pub hooks: Vec<Arc<dyn Hook>>,
pub budget_checker: Option<Arc<dyn BudgetChecker>>,
pub context_assembler: Arc<dyn ContextAssembler>,
pub model_profile: Option<ModelProfile>,
pub context_token_limit: Option<usize>,
}
impl std::fmt::Debug for AgentDefinition {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AgentDefinition")
.field("name", &self.name)
.field("model", &self.model)
.field("capabilities", &self.capabilities.len())
.field("max_iterations", &self.max_iterations)
.finish_non_exhaustive()
}
}
#[derive(Clone)]
pub struct AgentEnvironment {
pub drivers: agentyk_core::driver::DriverRegistry,
pub listeners: Vec<Arc<dyn EventListener>>,
pub extensions: agentyk_core::extensions::Extensions,
}
impl std::fmt::Debug for AgentEnvironment {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AgentEnvironment")
.field("listeners", &self.listeners.len())
.finish_non_exhaustive()
}
}
#[derive(Clone)]
pub struct Agent {
definition: Arc<AgentDefinition>,
environment: Arc<AgentEnvironment>,
}
impl std::fmt::Debug for Agent {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Agent")
.field("definition", &self.definition)
.field("environment", &self.environment)
.finish_non_exhaustive()
}
}
impl Agent {
pub fn builder() -> AgentBuilder {
AgentBuilder::new()
}
pub fn definition(&self) -> &AgentDefinition {
&self.definition
}
pub fn environment(&self) -> &AgentEnvironment {
&self.environment
}
pub fn name(&self) -> &str {
&self.definition.name
}
pub fn model(&self) -> &ModelSpec {
&self.definition.model
}
pub fn model_profile(&self) -> Option<&ModelProfile> {
self.definition.model_profile.as_ref()
}
pub fn capabilities(&self) -> &[Arc<dyn Capability>] {
&self.definition.capabilities
}
pub fn driver_for_model(&self) -> Option<Arc<dyn ChatDriver>> {
self.environment.drivers.get(&self.definition.model.driver)
}
pub fn session(&self) -> Session {
self.session_with_log(Arc::new(InMemoryEventLog::new()))
}
pub fn session_with_log(&self, log: Arc<dyn EventLog>) -> Session {
Session::new(self.clone(), log)
}
pub async fn resume_session(
&self,
log: Arc<dyn EventLog>,
session_id: SessionId,
) -> Result<Session> {
Session::resume(self.clone(), log, session_id).await
}
pub async fn run(&self, input: impl Into<Message>) -> Result<TurnResult> {
self.session().run(input).await
}
}
struct AdHocTools {
tools: Vec<Arc<dyn Tool>>,
}
#[async_trait]
impl Capability for AdHocTools {
fn id(&self) -> &str {
"tools"
}
fn description(&self) -> &str {
"Tools attached directly to the agent."
}
async fn tools(&self) -> Result<Vec<Arc<dyn Tool>>> {
Ok(self.tools.clone())
}
}
pub struct AgentBuilder {
config: AgentConfig,
tools: Vec<Arc<dyn Tool>>,
}
impl AgentBuilder {
pub fn new() -> Self {
Self {
config: AgentConfig::default(),
tools: Vec::new(),
}
}
pub fn name(mut self, name: impl Into<String>) -> Self {
self.config.name = name.into();
self
}
pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
self.config.system_prompt = prompt.into();
self
}
pub fn model(mut self, model: ModelSpec) -> Self {
self.config.model = Some(model);
self
}
pub fn capability(mut self, capability: impl Capability + 'static) -> Self {
self.config.capabilities.push(Arc::new(capability));
self
}
pub fn capability_arc(mut self, capability: Arc<dyn Capability>) -> Self {
self.config.capabilities.push(capability);
self
}
pub fn tool(mut self, tool: impl Tool + 'static) -> Self {
self.tools.push(Arc::new(tool));
self
}
pub fn driver(mut self, driver: impl ChatDriver + 'static) -> Self {
self.config.drivers.register(driver);
self
}
pub fn listener(mut self, listener: impl EventListener + 'static) -> Self {
self.config.listeners.push(Arc::new(listener));
self
}
pub fn listener_arc(mut self, listener: Arc<dyn EventListener>) -> Self {
self.config.listeners.push(listener);
self
}
pub fn max_iterations(mut self, max_iterations: usize) -> Self {
self.config.max_iterations = max_iterations;
self
}
pub fn middleware(mut self, middleware: impl TurnMiddleware + 'static) -> Self {
self.config.middleware.push(Arc::new(middleware));
self
}
pub fn middleware_arc(mut self, middleware: Arc<dyn TurnMiddleware>) -> Self {
self.config.middleware.push(middleware);
self
}
pub fn hook(mut self, hook: impl Hook + 'static) -> Self {
self.config.hooks.push(Arc::new(hook));
self
}
pub fn hook_arc(mut self, hook: Arc<dyn Hook>) -> Self {
self.config.hooks.push(hook);
self
}
pub fn budget_checker(mut self, checker: impl BudgetChecker + 'static) -> Self {
self.config.budget_checker = Some(Arc::new(checker));
self
}
pub fn model_catalog(mut self, catalog: impl ModelCatalog + 'static) -> Self {
self.config.model_catalog = Some(Arc::new(catalog));
self
}
pub fn model_catalog_arc(mut self, catalog: Arc<dyn ModelCatalog>) -> Self {
self.config.model_catalog = Some(catalog);
self
}
pub fn context_assembler(mut self, assembler: impl ContextAssembler + 'static) -> Self {
self.config.context_assembler = Arc::new(assembler);
self
}
pub fn context_token_limit(mut self, token_limit: usize) -> Self {
self.config.context_token_limit = Some(token_limit);
self
}
pub fn extension<T: Send + Sync + 'static>(mut self, value: T) -> Self {
self.config.extensions.insert(value);
self
}
pub fn build(self) -> Result<Agent> {
let mut config = self.config;
if config.model.is_none() {
return Err(Error::InvalidAgent(
"a model is required — call .model(...)".into(),
));
}
if !self.tools.is_empty() {
config
.capabilities
.push(Arc::new(AdHocTools { tools: self.tools }));
}
if !config.drivers.contains(&config.model().driver) {
return Err(Error::UnknownDriver(config.model().driver.to_string()));
}
for hook in &config.hooks {
if let Some(matcher) = hook.matcher() {
if !hook.event().supports_matcher() {
return Err(Error::InvalidAgent(format!(
"hook `{}` has a matcher on non-tool event `{}`",
hook.id(),
hook.event().as_str()
)));
}
matcher.validate().map_err(|error| {
Error::InvalidAgent(format!("hook `{}`: {error}", hook.id()))
})?;
}
}
let model_profile = config
.model_catalog
.as_ref()
.and_then(|catalog| catalog.profile(&config.model().driver, &config.model().model));
if let Some(profile) = &model_profile
&& let Err(reason) = profile.validate(config.model())
{
return Err(Error::InvalidAgent(reason));
}
let context_token_limit = config.context_token_limit.or_else(|| {
let profile = model_profile.as_ref()?;
let window = usize::try_from(profile.context_window?).ok()?;
let output = usize::try_from(profile.max_output_tokens?).ok()?;
window.checked_sub(output)
});
let definition = AgentDefinition {
name: config.name.clone(),
instructions: config.system_prompt.clone(),
model: config.model().clone(),
capabilities: config.capabilities.clone(),
max_iterations: config.max_iterations,
middleware: config.middleware.clone(),
hooks: config.hooks.clone(),
budget_checker: config.budget_checker.clone(),
context_assembler: config.context_assembler.clone(),
model_profile,
context_token_limit,
};
let environment = AgentEnvironment {
drivers: config.drivers.clone(),
listeners: config.listeners.clone(),
extensions: config.extensions.clone(),
};
Ok(Agent {
definition: Arc::new(definition),
environment: Arc::new(environment),
})
}
}
impl Default for AgentBuilder {
fn default() -> Self {
Self::new()
}
}