Skip to main content

atomr_agents_agent/
trait.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use atomr_agents_callable::Callable;
5use atomr_agents_core::{AgentId, CallCtx, Result, Value};
6
7use crate::inference::TurnResult;
8
9/// Public, type-erased handle to an agent. Implements `Callable`,
10/// so an agent can be passed wherever any executable unit is
11/// expected (workflow steps, team routing targets).
12pub struct AgentRef {
13    pub id: AgentId,
14    inner: Arc<dyn AgentDispatch>,
15}
16
17impl AgentRef {
18    pub fn new(id: AgentId, inner: Arc<dyn AgentDispatch>) -> Self {
19        Self { id, inner }
20    }
21
22    pub async fn turn(&self, user: String, ctx: CallCtx) -> Result<TurnResult> {
23        self.inner.dispatch(user, ctx).await
24    }
25}
26
27#[async_trait]
28pub trait AgentDispatch: Send + Sync + 'static {
29    async fn dispatch(&self, user: String, ctx: CallCtx) -> Result<TurnResult>;
30}
31
32#[async_trait]
33impl Callable for AgentRef {
34    async fn call(&self, input: Value, ctx: CallCtx) -> Result<Value> {
35        // Treat the input as either a plain string or `{"user": "..."}`.
36        let user = match input {
37            Value::String(s) => s,
38            Value::Object(ref m) => m
39                .get("user")
40                .and_then(|v| v.as_str())
41                .unwrap_or_default()
42                .to_string(),
43            _ => input.to_string(),
44        };
45        let r = self.turn(user, ctx).await?;
46        Ok(serde_json::json!({
47            "text": r.text,
48            "input_tokens": r.usage.input_tokens,
49            "output_tokens": r.usage.output_tokens,
50        }))
51    }
52
53    fn label(&self) -> &str {
54        self.id.as_str()
55    }
56}