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::fmt;
9use std::sync::Arc;
10use std::time::{Duration, Instant};
11
12use futures::stream::StreamExt;
13
14use crate::circuit_breaker::CircuitBreaker;
15use crate::error::{LlmError, Result};
16use crate::stream::{SseDecoder, StreamAssembler};
17use crate::types::{CompletionRequest, CompletionResponse};
18
19/// Default per-request timeout. DeepSeek's server-side timeout is 900s, which
20/// caused 15-minute hangs that exhausted DB connections; 30s is generous for
21/// summarization and agent flows can override per call.
22pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30);
23
24/// Configuration for [`LlmClient`].
25#[derive(Clone)]
26pub struct LlmConfig {
27    /// Base URL of the OpenAI-compatible API, e.g. `http://litellm:4000/v1`.
28    pub base_url: String,
29    /// Bearer token. May be empty when the proxy handles auth upstream.
30    pub api_key: String,
31    /// Per-request timeout.
32    pub timeout: Duration,
33}
34
35impl fmt::Debug for LlmConfig {
36    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
37        formatter
38            .debug_struct("LlmConfig")
39            .field("base_url", &self.base_url)
40            .field("api_key", &"[REDACTED]")
41            .field("timeout", &self.timeout)
42            .finish()
43    }
44}
45
46impl LlmConfig {
47    /// Client for an OpenAI-compatible `base_url` using bearer `api_key`, with default retry, timeout and breaker settings.
48    pub fn new(base_url: impl Into<String>, api_key: impl Into<String>) -> Self {
49        Self {
50            base_url: base_url.into(),
51            api_key: api_key.into(),
52            timeout: DEFAULT_TIMEOUT,
53        }
54    }
55}
56
57/// Cheap to clone (everything shareable lives behind `Arc`). Share one instance
58/// across tasks so they all observe the same circuit-breaker state.
59#[derive(Clone)]
60pub struct LlmClient {
61    http: reqwest::Client,
62    base_url: Arc<str>,
63    api_key: Arc<str>,
64    breaker: Arc<CircuitBreaker>,
65    timeout: Duration,
66}
67
68impl LlmClient {
69    /// Build a client with a fresh, default circuit breaker.
70    pub fn new(config: LlmConfig) -> Result<Self> {
71        Self::with_breaker(config, Arc::new(CircuitBreaker::default()))
72    }
73
74    /// Build a client sharing an existing circuit breaker (e.g. one breaker
75    /// across several clients hitting the same provider).
76    pub fn with_breaker(config: LlmConfig, breaker: Arc<CircuitBreaker>) -> Result<Self> {
77        let http = reqwest::Client::builder().build()?;
78        Ok(Self {
79            http,
80            base_url: config.base_url.trim_end_matches('/').into(),
81            api_key: config.api_key.into(),
82            breaker,
83            timeout: config.timeout,
84        })
85    }
86
87    /// Access the shared circuit breaker (health checks, metrics).
88    pub fn breaker(&self) -> &Arc<CircuitBreaker> {
89        &self.breaker
90    }
91
92    /// Execute one auditable streaming provider attempt. Retry and durable
93    /// attempt accounting belong exclusively to `af-agent-runtime`.
94    pub async fn complete_stream_single_attempt<F>(
95        &self,
96        request: &CompletionRequest,
97        on_delta: F,
98    ) -> Result<CompletionResponse>
99    where
100        F: FnMut(&str, bool) + Send,
101    {
102        let cancellation = tokio_util::sync::CancellationToken::new();
103        self.complete_stream_single_attempt_controlled(
104            request,
105            &cancellation,
106            Instant::now() + self.timeout,
107            on_delta,
108        )
109        .await
110    }
111
112    /// Execute one provider attempt bounded by caller cancellation and deadline.
113    pub async fn complete_stream_single_attempt_controlled<F>(
114        &self,
115        request: &CompletionRequest,
116        cancellation: &tokio_util::sync::CancellationToken,
117        deadline: Instant,
118        mut on_delta: F,
119    ) -> Result<CompletionResponse>
120    where
121        F: FnMut(&str, bool) + Send,
122    {
123        if self.breaker.is_open() {
124            return Err(LlmError::CircuitOpen);
125        }
126        let url = format!("{}/chat/completions", self.base_url);
127        let result = tokio::select! {
128            biased;
129            _ = cancellation.cancelled() => Err(LlmError::Cancelled),
130            _ = tokio::time::sleep_until(deadline.into()) => Err(LlmError::DeadlineExceeded),
131            result = self.stream_collect_once(&url, request, &mut on_delta) => result,
132        };
133        match &result {
134            Ok(_) => self.breaker.record_success(),
135            Err(LlmError::Cancelled | LlmError::DeadlineExceeded) => {}
136            Err(_) => self.breaker.record_failure(),
137        }
138        result
139    }
140
141    async fn stream_collect_once<F>(
142        &self,
143        url: &str,
144        request: &CompletionRequest,
145        on_delta: &mut F,
146    ) -> Result<CompletionResponse>
147    where
148        F: FnMut(&str, bool) + Send,
149    {
150        let mut req = self.http.post(url).json(request);
151        if let Some(attempt_id) = &request.provider_attempt_id {
152            req = req
153                .header("Idempotency-Key", attempt_id)
154                .header("X-Agent-Factory-Attempt-Id", attempt_id);
155        }
156        if !self.api_key.is_empty() {
157            req = req.bearer_auth(self.api_key.as_ref());
158        }
159        let response = req.send().await?;
160        let status = response.status();
161        if !status.is_success() {
162            let body = response.text().await.unwrap_or_default();
163            return Err(LlmError::Api {
164                status: status.as_u16(),
165                body: redact_and_truncate(&body, self.api_key.as_ref()),
166            });
167        }
168
169        let mut assembler = StreamAssembler::default();
170        let mut decoder = SseDecoder::default();
171        let mut bytes = response.bytes_stream();
172        while let Some(chunk) = bytes.next().await {
173            for data in decoder.push(&chunk?)? {
174                if data == "[DONE]" {
175                    decoder.finish()?;
176                    return assembler.finish();
177                }
178                if let Some(delta) = assembler.apply_json(&data)? {
179                    on_delta(&delta.content, delta.has_tool_calls);
180                }
181            }
182        }
183        decoder.finish()?;
184        Err(LlmError::StreamProtocol(
185            "stream ended before [DONE]".into(),
186        ))
187    }
188}
189
190const MAX_ERROR_BODY_BYTES: usize = 4 * 1024;
191
192fn redact_and_truncate(body: &str, secret: &str) -> String {
193    let redacted = if secret.is_empty() {
194        body.to_owned()
195    } else {
196        body.replace(secret, "[REDACTED]")
197    };
198    if redacted.len() <= MAX_ERROR_BODY_BYTES {
199        return redacted;
200    }
201    let mut end = MAX_ERROR_BODY_BYTES;
202    while !redacted.is_char_boundary(end) {
203        end -= 1;
204    }
205    format!("{}...[truncated]", &redacted[..end])
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[test]
213    fn prepared_streaming_request_is_wire_stable() {
214        let mut request = CompletionRequest::new("model", vec![]).stream(true);
215        request.stream_options = Some(crate::StreamOptions {
216            include_usage: true,
217        });
218        let wire = serde_json::to_value(&request).unwrap();
219        assert_eq!(wire["stream"], true);
220        assert_eq!(wire["stream_options"]["include_usage"], true);
221    }
222
223    #[test]
224    fn config_debug_and_provider_errors_hide_credentials() {
225        let config = LlmConfig::new("https://provider.invalid/v1", "top-secret");
226        assert!(!format!("{config:?}").contains("top-secret"));
227
228        let body = format!("token=top-secret {}", "界".repeat(MAX_ERROR_BODY_BYTES));
229        let safe = redact_and_truncate(&body, "top-secret");
230        assert!(!safe.contains("top-secret"));
231        assert!(safe.len() <= MAX_ERROR_BODY_BYTES + "...[truncated]".len());
232    }
233
234    #[tokio::test]
235    async fn caller_cancellation_wins_before_network_io() {
236        let client =
237            LlmClient::new(LlmConfig::new("https://provider.invalid/v1", "secret")).unwrap();
238        let cancellation = tokio_util::sync::CancellationToken::new();
239        cancellation.cancel();
240        let result = client
241            .complete_stream_single_attempt_controlled(
242                &CompletionRequest::new("model", vec![]),
243                &cancellation,
244                Instant::now() + Duration::from_secs(1),
245                |_, _| {},
246            )
247            .await;
248        assert!(matches!(result, Err(LlmError::Cancelled)));
249        assert_eq!(client.breaker().status().failure_count, 0);
250    }
251}
af_llm/client.rs

af_llm/
client.rs
