af_llm/
client.rs

1//! Async LLM client — port of `agent_core/llm/client.py::llm_completion`.
2//!
3//! Talks to any OpenAI-compatible `/chat/completions` endpoint. In production
4//! this points at the LiteLLM proxy, which owns model routing and fallback; the
5//! client itself owns one bounded streaming provider attempt and a shared
6//! circuit breaker. The event-sourced Agent runtime is the sole retry owner.
7
8use std::sync::Arc;
9use std::time::Duration;
10
11use futures::stream::StreamExt;
12
13use crate::circuit_breaker::CircuitBreaker;
14use crate::error::{LlmError, Result};
15use crate::stream::{SseDecoder, StreamAssembler};
16use crate::types::{CompletionRequest, CompletionResponse};
17
18/// Default per-request timeout. DeepSeek's server-side timeout is 900s, which
19/// caused 15-minute hangs that exhausted DB connections; 30s is generous for
20/// summarization and agent flows can override per call.
21pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30);
22
23/// Configuration for [`LlmClient`].
24#[derive(Debug, Clone)]
25pub struct LlmConfig {
26    /// Base URL of the OpenAI-compatible API, e.g. `http://litellm:4000/v1`.
27    pub base_url: String,
28    /// Bearer token. May be empty when the proxy handles auth upstream.
29    pub api_key: String,
30    /// Per-request timeout.
31    pub timeout: Duration,
32}
33
34impl LlmConfig {
35    pub fn new(base_url: impl Into<String>, api_key: impl Into<String>) -> Self {
36        Self {
37            base_url: base_url.into(),
38            api_key: api_key.into(),
39            timeout: DEFAULT_TIMEOUT,
40        }
41    }
42}
43
44/// Cheap to clone (everything shareable lives behind `Arc`). Share one instance
45/// across tasks so they all observe the same circuit-breaker state.
46#[derive(Clone)]
47pub struct LlmClient {
48    http: reqwest::Client,
49    base_url: Arc<str>,
50    api_key: Arc<str>,
51    breaker: Arc<CircuitBreaker>,
52}
53
54impl LlmClient {
55    /// Build a client with a fresh, default circuit breaker.
56    pub fn new(config: LlmConfig) -> Result<Self> {
57        Self::with_breaker(config, Arc::new(CircuitBreaker::default()))
58    }
59
60    /// Build a client sharing an existing circuit breaker (e.g. one breaker
61    /// across several clients hitting the same provider).
62    pub fn with_breaker(config: LlmConfig, breaker: Arc<CircuitBreaker>) -> Result<Self> {
63        let http = reqwest::Client::builder().timeout(config.timeout).build()?;
64        Ok(Self {
65            http,
66            base_url: config.base_url.trim_end_matches('/').into(),
67            api_key: config.api_key.into(),
68            breaker,
69        })
70    }
71
72    /// Access the shared circuit breaker (health checks, metrics).
73    pub fn breaker(&self) -> &Arc<CircuitBreaker> {
74        &self.breaker
75    }
76
77    /// Execute one auditable streaming provider attempt. Retry and durable
78    /// attempt accounting belong exclusively to `af-agent-runtime`.
79    pub async fn complete_stream_single_attempt<F>(
80        &self,
81        request: &CompletionRequest,
82        mut on_delta: F,
83    ) -> Result<CompletionResponse>
84    where
85        F: FnMut(&str, bool) + Send,
86    {
87        if self.breaker.is_open() {
88            return Err(LlmError::CircuitOpen);
89        }
90        let url = format!("{}/chat/completions", self.base_url);
91        let result = self.stream_collect_once(&url, request, &mut on_delta).await;
92        match &result {
93            Ok(_) => self.breaker.record_success(),
94            Err(_) => self.breaker.record_failure(),
95        }
96        result
97    }
98
99    async fn stream_collect_once<F>(
100        &self,
101        url: &str,
102        request: &CompletionRequest,
103        on_delta: &mut F,
104    ) -> Result<CompletionResponse>
105    where
106        F: FnMut(&str, bool) + Send,
107    {
108        let mut req = self.http.post(url).json(request);
109        if let Some(attempt_id) = &request.provider_attempt_id {
110            req = req
111                .header("Idempotency-Key", attempt_id)
112                .header("X-Agent-Factory-Attempt-Id", attempt_id);
113        }
114        if !self.api_key.is_empty() {
115            req = req.bearer_auth(self.api_key.as_ref());
116        }
117        let response = req.send().await?;
118        let status = response.status();
119        if !status.is_success() {
120            return Err(LlmError::Api {
121                status: status.as_u16(),
122                body: response.text().await.unwrap_or_default(),
123            });
124        }
125
126        let mut assembler = StreamAssembler::default();
127        let mut decoder = SseDecoder::default();
128        let mut bytes = response.bytes_stream();
129        while let Some(chunk) = bytes.next().await {
130            for data in decoder.push(&chunk?)? {
131                if data == "[DONE]" {
132                    decoder.finish()?;
133                    return assembler.finish();
134                }
135                if let Some(delta) = assembler.apply_json(&data)? {
136                    on_delta(&delta.content, delta.has_tool_calls);
137                }
138            }
139        }
140        decoder.finish()?;
141        Err(LlmError::StreamProtocol(
142            "stream ended before [DONE]".into(),
143        ))
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    #[test]
152    fn prepared_streaming_request_is_wire_stable() {
153        let mut request = CompletionRequest::new("model", vec![]).stream(true);
154        request.stream_options = Some(crate::StreamOptions {
155            include_usage: true,
156        });
157        let wire = serde_json::to_value(&request).unwrap();
158        assert_eq!(wire["stream"], true);
159        assert_eq!(wire["stream_options"]["include_usage"], true);
160    }
161}
af_llm/client.rs

af_llm/
client.rs
