Skip to main content

af_agent/
model.rs

1//! The chat-model abstraction the agent loop runs against.
2//!
3//! Decoupling the loop from a concrete client means: (a) the loop is unit-
4//! testable with a scripted mock — no network; (b) any backend (the real
5//! [`LlmClient`], a local model, a replay harness) plugs in by implementing one
6//! method.
7
8use std::sync::Arc;
9
10use af_llm::{CompletionRequest, CompletionResponse, LlmClient, LlmError};
11use async_trait::async_trait;
12use tokio::sync::mpsc::UnboundedSender;
13
14/// Anything that can turn a chat-completion request into a response.
15#[async_trait]
16pub trait ChatModel: Send + Sync {
17    /// Stream cumulative `(content, has_tool_calls)` updates and return the
18    /// canonical terminal response. This is the only model invocation path.
19    async fn complete_streaming(
20        &self,
21        request: &CompletionRequest,
22        delta_tx: UnboundedSender<(String, bool)>,
23    ) -> Result<CompletionResponse, LlmError>;
24}
25
26/// The production model is the real LLM client.
27#[async_trait]
28impl ChatModel for LlmClient {
29    async fn complete_streaming(
30        &self,
31        request: &CompletionRequest,
32        delta_tx: UnboundedSender<(String, bool)>,
33    ) -> Result<CompletionResponse, LlmError> {
34        self.complete_stream_single_attempt(request, |content, has_tools| {
35            let _ = delta_tx.send((content.to_string(), has_tools));
36        })
37        .await
38    }
39}
40
41/// A type-erased model, so callers (e.g. an HTTP service holding one `Agent` in
42/// shared state) can pick the backend at runtime — real client vs stub vs
43/// replay — without the loop being generic over every concrete type.
44#[async_trait]
45impl ChatModel for Arc<dyn ChatModel> {
46    async fn complete_streaming(
47        &self,
48        request: &CompletionRequest,
49        delta_tx: UnboundedSender<(String, bool)>,
50    ) -> Result<CompletionResponse, LlmError> {
51        (**self).complete_streaming(request, delta_tx).await
52    }
53}