atomr_agents_agent/
inference.rs

1//! Abstraction over a `ModelRunner` that produces a `TurnResult`
2//! (text, usage, finish reason, parsed tool calls) per `ExecuteBatch`.
3
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use atomr_agents_core::{AgentError, Result};
8use atomr_agents_tool::{Provider, ToolCallParser};
9use atomr_infer_core::batch::ExecuteBatch;
10use atomr_infer_core::runner::ModelRunner;
11use atomr_infer_core::tokens::{FinishReason, TokenUsage};
12use futures::stream::StreamExt;
13use tokio::sync::Mutex;
14
15use atomr_agents_tool::ParsedToolCall;
16
17#[derive(Debug, Default)]
18pub struct TurnResult {
19    pub text: String,
20    pub usage: TokenUsage,
21    pub finish_reason: Option<FinishReason>,
22    pub tool_calls: Vec<ParsedToolCall>,
23}
24
25/// Implemented by anything the agent can use to drive an inference
26/// request. The provider matters for tool-call delta parsing.
27#[async_trait]
28pub trait InferenceClient: Send + Sync + 'static {
29    fn provider(&self) -> Provider;
30    async fn run(&self, batch: ExecuteBatch) -> Result<TurnResult>;
31}
32
33/// Wrap any `ModelRunner` (including `MockRunner`) as an
34/// `InferenceClient`. Single-runner concurrency is bounded by the
35/// internal mutex; production setups should use the
36/// `EngineCoreActor` from `atomr-infer-runtime` instead.
37pub struct LocalRunnerClient<R: ModelRunner> {
38    runner: Arc<Mutex<R>>,
39    provider: Provider,
40}
41
42impl<R: ModelRunner + 'static> LocalRunnerClient<R> {
43    pub fn new(runner: R, provider: Provider) -> Self {
44        Self {
45            runner: Arc::new(Mutex::new(runner)),
46            provider,
47        }
48    }
49
50    pub fn from_arc(runner: Arc<Mutex<R>>, provider: Provider) -> Self {
51        Self { runner, provider }
52    }
53}
54
55#[async_trait]
56impl<R: ModelRunner + 'static> InferenceClient for LocalRunnerClient<R> {
57    fn provider(&self) -> Provider {
58        self.provider
59    }
60
61    async fn run(&self, batch: ExecuteBatch) -> Result<TurnResult> {
62        let mut g = self.runner.lock().await;
63        let handle = g
64            .execute(batch)
65            .await
66            .map_err(|e| AgentError::Inference(e.to_string()))?;
67        // Drop the mutex before consuming the stream — the stream is
68        // produced by `execute` and is independent of `&mut self`.
69        drop(g);
70        let mut text = String::new();
71        let mut usage = TokenUsage::default();
72        let mut finish: Option<FinishReason> = None;
73        let mut parser = ToolCallParser::new(self.provider);
74        let mut stream = handle.into_stream();
75        while let Some(item) = stream.next().await {
76            let chunk = item.map_err(|e| AgentError::Inference(e.to_string()))?;
77            text.push_str(&chunk.text_delta);
78            if let Some(d) = chunk.tool_call_delta.as_ref() {
79                parser.feed(d)?;
80            }
81            if let Some(u) = chunk.usage {
82                usage.add(u);
83            }
84            if let Some(r) = chunk.finish_reason {
85                finish = Some(r);
86            }
87        }
88        let tool_calls = parser.finish();
89        Ok(TurnResult {
90            text,
91            usage,
92            finish_reason: finish,
93            tool_calls,
94        })
95    }
96}
atomr_agents_agent/inference.rs

atomr_agents_agent/
inference.rs
