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    images: Option<Arc<dyn crate::images::ImageResolver>>,
67}
68
69impl LlmClient {
70    /// Build a client with a fresh, default circuit breaker.
71    pub fn new(config: LlmConfig) -> Result<Self> {
72        Self::with_breaker(config, Arc::new(CircuitBreaker::default()))
73    }
74
75    /// Build a client sharing an existing circuit breaker (e.g. one breaker
76    /// across several clients hitting the same provider).
77    pub fn with_breaker(config: LlmConfig, breaker: Arc<CircuitBreaker>) -> Result<Self> {
78        let http = reqwest::Client::builder().build()?;
79        Ok(Self {
80            http,
81            base_url: config.base_url.trim_end_matches('/').into(),
82            api_key: config.api_key.into(),
83            breaker,
84            timeout: config.timeout,
85            images: None,
86        })
87    }
88
89    /// Bind a caller-aware image resolver. Text-only clients need no resolver.
90    pub fn with_image_resolver(mut self, resolver: Arc<dyn crate::images::ImageResolver>) -> Self {
91        self.images = Some(resolver);
92        self
93    }
94
95    /// Access the shared circuit breaker (health checks, metrics).
96    pub fn breaker(&self) -> &Arc<CircuitBreaker> {
97        &self.breaker
98    }
99
100    /// Execute one auditable streaming provider attempt. Retry and durable
101    /// attempt accounting belong exclusively to `af-agent-runtime`.
102    pub async fn complete_stream_single_attempt<F>(
103        &self,
104        request: &CompletionRequest,
105        on_delta: F,
106    ) -> Result<CompletionResponse>
107    where
108        F: FnMut(&str, bool) + Send,
109    {
110        let cancellation = tokio_util::sync::CancellationToken::new();
111        self.complete_stream_single_attempt_controlled(
112            request,
113            &cancellation,
114            Instant::now() + self.timeout,
115            on_delta,
116        )
117        .await
118    }
119
120    /// Execute one provider attempt bounded by caller cancellation and deadline.
121    pub async fn complete_stream_single_attempt_controlled<F>(
122        &self,
123        request: &CompletionRequest,
124        cancellation: &tokio_util::sync::CancellationToken,
125        deadline: Instant,
126        mut on_delta: F,
127    ) -> Result<CompletionResponse>
128    where
129        F: FnMut(&str, bool) + Send,
130    {
131        if self.breaker.is_open() {
132            return Err(LlmError::CircuitOpen);
133        }
134        let url = format!("{}/chat/completions", self.base_url);
135        let result = tokio::select! {
136            biased;
137            _ = cancellation.cancelled() => Err(LlmError::Cancelled),
138            _ = tokio::time::sleep_until(deadline.into()) => Err(LlmError::DeadlineExceeded),
139            result = self.stream_collect_once(&url, request, &mut on_delta) => result,
140        };
141        match &result {
142            Ok(_) => self.breaker.record_success(),
143            Err(LlmError::Cancelled | LlmError::DeadlineExceeded | LlmError::InvalidInput(_)) => {}
144            Err(_) => self.breaker.record_failure(),
145        }
146        result
147    }
148
149    async fn stream_collect_once<F>(
150        &self,
151        url: &str,
152        request: &CompletionRequest,
153        on_delta: &mut F,
154    ) -> Result<CompletionResponse>
155    where
156        F: FnMut(&str, bool) + Send,
157    {
158        let payload = crate::images::provider_request(request, self.images.as_deref()).await?;
159        let mut req = self.http.post(url).json(&payload);
160        if let Some(attempt_id) = &request.provider_attempt_id {
161            req = req
162                .header("Idempotency-Key", attempt_id)
163                .header("X-Agent-Factory-Attempt-Id", attempt_id);
164        }
165        if !self.api_key.is_empty() {
166            req = req.bearer_auth(self.api_key.as_ref());
167        }
168        let response = req.send().await?;
169        let status = response.status();
170        if !status.is_success() {
171            let body = response.text().await.unwrap_or_default();
172            return Err(LlmError::Api {
173                status: status.as_u16(),
174                body: if request
175                    .messages
176                    .iter()
177                    .any(|message| !message.images.is_empty())
178                {
179                    "[multimodal provider error redacted]".into()
180                } else {
181                    redact_and_truncate(&body, self.api_key.as_ref())
182                },
183            });
184        }
185
186        let mut assembler = StreamAssembler::default();
187        let mut decoder = SseDecoder::default();
188        let mut bytes = response.bytes_stream();
189        while let Some(chunk) = bytes.next().await {
190            for data in decoder.push(&chunk?)? {
191                if data == "[DONE]" {
192                    decoder.finish()?;
193                    return assembler.finish();
194                }
195                if let Some(delta) = assembler.apply_json(&data)? {
196                    on_delta(&delta.content, delta.has_tool_calls);
197                }
198            }
199        }
200        decoder.finish()?;
201        Err(LlmError::StreamProtocol(
202            "stream ended before [DONE]".into(),
203        ))
204    }
205}
206
207const MAX_ERROR_BODY_BYTES: usize = 4 * 1024;
208
209fn redact_and_truncate(body: &str, secret: &str) -> String {
210    let redacted = if secret.is_empty() {
211        body.to_owned()
212    } else {
213        body.replace(secret, "[REDACTED]")
214    };
215    if redacted.len() <= MAX_ERROR_BODY_BYTES {
216        return redacted;
217    }
218    let mut end = MAX_ERROR_BODY_BYTES;
219    while !redacted.is_char_boundary(end) {
220        end -= 1;
221    }
222    format!("{}...[truncated]", &redacted[..end])
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228
229    #[test]
230    fn prepared_streaming_request_is_wire_stable() {
231        let mut request = CompletionRequest::new("model", vec![]).stream(true);
232        request.stream_options = Some(crate::StreamOptions {
233            include_usage: true,
234        });
235        let wire = serde_json::to_value(&request).unwrap();
236        assert_eq!(wire["stream"], true);
237        assert_eq!(wire["stream_options"]["include_usage"], true);
238    }
239
240    #[test]
241    fn config_debug_and_provider_errors_hide_credentials() {
242        let config = LlmConfig::new("https://provider.invalid/v1", "top-secret");
243        assert!(!format!("{config:?}").contains("top-secret"));
244
245        let body = format!("token=top-secret {}", "界".repeat(MAX_ERROR_BODY_BYTES));
246        let safe = redact_and_truncate(&body, "top-secret");
247        assert!(!safe.contains("top-secret"));
248        assert!(safe.len() <= MAX_ERROR_BODY_BYTES + "...[truncated]".len());
249    }
250
251    #[tokio::test]
252    async fn caller_cancellation_wins_before_network_io() {
253        let client =
254            LlmClient::new(LlmConfig::new("https://provider.invalid/v1", "secret")).unwrap();
255        let cancellation = tokio_util::sync::CancellationToken::new();
256        cancellation.cancel();
257        let result = client
258            .complete_stream_single_attempt_controlled(
259                &CompletionRequest::new("model", vec![]),
260                &cancellation,
261                Instant::now() + Duration::from_secs(1),
262                |_, _| {},
263            )
264            .await;
265        assert!(matches!(result, Err(LlmError::Cancelled)));
266        assert_eq!(client.breaker().status().failure_count, 0);
267    }
268}
af_llm/client.rs

af_llm/
client.rs
