agent_base/llm/
stream_client.rs

1use std::pin::Pin;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use futures_core::Stream;
6use futures_util::StreamExt;
7use serde_json::Value;
8
9use super::{LlmCapabilities, LlmClient, ReasoningConfig, StreamChunk};
10use crate::types::{AgentResult, ChatMessage, ResponseFormat};
11
12/// Provider-agnostic streaming client trait.
13///
14/// This is the recommended interface for LLM provider integration.
15/// Providers only need to implement [`stream`](StreamClient::stream) —
16/// [`chat`](StreamClient::chat) has a default implementation that collects
17/// text deltas from the stream.
18///
19/// This follows the Rust standard library convention:
20/// [`Iterator`] only requires `next()`, [`std::io::Read`] only requires `read()`.
21///
22/// The older [`LlmClient`] trait is still supported via [`LlmClientAdapter`].
23#[async_trait]
24pub trait StreamClient: Send + Sync {
25    /// Stream LLM response chunks from the provider.
26    ///
27    /// This is the **only required method**. Implementors translate their
28    /// provider's SSE/streaming protocol into [`StreamChunk`] events.
29    async fn stream(
30        &self,
31        messages: &[ChatMessage],
32        tools: &[Value],
33        reasoning: Option<&ReasoningConfig>,
34        response_format: Option<&ResponseFormat>,
35    ) -> AgentResult<Pin<Box<dyn Stream<Item = AgentResult<StreamChunk>> + Send>>>;
36
37    /// Convenience: collect all [`StreamChunk::Text`] deltas into a single string.
38    ///
39    /// The default implementation streams and concatenates text chunks.
40    /// Providers may override this with a dedicated non-streaming API call
41    /// for better latency or cost.
42    async fn chat(
43        &self,
44        messages: &[ChatMessage],
45        tools: &[Value],
46        reasoning: Option<&ReasoningConfig>,
47        response_format: Option<&ResponseFormat>,
48    ) -> AgentResult<String> {
49        let mut stream = self
50            .stream(messages, tools, reasoning, response_format)
51            .await?;
52        let mut text = String::new();
53        while let Some(chunk) = stream.next().await {
54            match chunk? {
55                StreamChunk::Text(t) => text.push_str(&t),
56                StreamChunk::Stop { .. } => break,
57                _ => {}
58            }
59        }
60        Ok(text)
61    }
62
63    /// Return the provider's capabilities.
64    fn capabilities(&self) -> LlmCapabilities;
65
66    /// The model name used by this client (e.g. "claude-sonnet", "gpt-4o").
67    /// Default: "unknown".
68    fn model_name(&self) -> &str {
69        "unknown"
70    }
71}
72
73// ── LlmClient → StreamClient adapter ──
74
75/// Bridges an [`LlmClient`] to the [`StreamClient`] trait.
76///
77/// Wraps an `Arc<dyn LlmClient>` so existing provider implementations
78/// can be used with the new [`StreamClient`]-based engine.
79pub struct LlmClientAdapter {
80    inner: Arc<dyn LlmClient>,
81}
82
83impl LlmClientAdapter {
84    /// Wrap an existing [`LlmClient`] for use as a [`StreamClient`].
85    pub fn new(client: Arc<dyn LlmClient>) -> Self {
86        Self { inner: client }
87    }
88
89    /// Get a reference to the inner [`LlmClient`].
90    pub fn inner(&self) -> &Arc<dyn LlmClient> {
91        &self.inner
92    }
93}
94
95#[async_trait]
96impl StreamClient for LlmClientAdapter {
97    async fn stream(
98        &self,
99        messages: &[ChatMessage],
100        tools: &[Value],
101        reasoning: Option<&ReasoningConfig>,
102        response_format: Option<&ResponseFormat>,
103    ) -> AgentResult<Pin<Box<dyn Stream<Item = AgentResult<StreamChunk>> + Send>>> {
104        self.inner
105            .chat_stream(messages, tools, reasoning, response_format)
106            .await
107    }
108
109    /// Delegate to the inner client's `chat()` for efficiency — avoids
110    /// streaming when only the final text is needed.
111    async fn chat(
112        &self,
113        messages: &[ChatMessage],
114        tools: &[Value],
115        reasoning: Option<&ReasoningConfig>,
116        response_format: Option<&ResponseFormat>,
117    ) -> AgentResult<String> {
118        let result = self
119            .inner
120            .chat(messages, tools, reasoning, response_format)
121            .await?;
122        // Extract text from the chat completion response
123        let text = result
124            .get("choices")
125            .and_then(|c| c.as_array())
126            .and_then(|choices| choices.first())
127            .and_then(|choice| choice.get("message"))
128            .and_then(|msg| msg.get("content"))
129            .and_then(|c| c.as_str())
130            .unwrap_or("")
131            .to_string();
132        Ok(text)
133    }
134
135    fn capabilities(&self) -> LlmCapabilities {
136        self.inner.capabilities()
137    }
138
139    fn model_name(&self) -> &str {
140        self.inner.model_name()
141    }
142}
143
144/// Convenience: wrap an `Arc<dyn LlmClient>` as an `Arc<dyn StreamClient>`.
145///
146/// This is the primary migration path for existing [`LlmClient`] implementations.
147/// Use this when you have a legacy client and need to pass it to APIs that
148/// expect [`StreamClient`].
149pub fn adapt(client: Arc<dyn LlmClient>) -> Arc<dyn StreamClient> {
150    Arc::new(LlmClientAdapter::new(client))
151}
152
153// ── Tests ──
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158    use crate::llm::StreamChunk;
159    use crate::types::AgentError;
160    use async_trait::async_trait;
161    use futures_util::StreamExt;
162    use std::pin::Pin;
163    use std::sync::Mutex;
164
165    // ── Mock clients ──
166
167    /// Mock LlmClient that returns canned JSON responses for `chat()`
168    /// and a canned stream for `chat_stream()`.
169    struct MockLlmClient {
170        chat_response: Mutex<Option<Value>>,
171        stream_chunks: Mutex<Option<Vec<AgentResult<StreamChunk>>>>,
172        caps: LlmCapabilities,
173        model: String,
174    }
175
176    impl MockLlmClient {
177        fn new() -> Self {
178            Self {
179                chat_response: Mutex::new(None),
180                stream_chunks: Mutex::new(None),
181                caps: LlmCapabilities {
182                    supports_streaming: true,
183                    supports_tools: true,
184                    supports_vision: false,
185                    supports_thinking: true,
186                    max_context_tokens: Some(128_000),
187                    max_output_tokens: Some(16_384),
188                },
189                model: "mock-model".into(),
190            }
191        }
192
193        fn with_chat_response(response: Value) -> Self {
194            Self {
195                chat_response: Mutex::new(Some(response)),
196                ..Self::new()
197            }
198        }
199
200        fn with_stream(chunks: Vec<AgentResult<StreamChunk>>) -> Self {
201            Self {
202                stream_chunks: Mutex::new(Some(chunks)),
203                ..Self::new()
204            }
205        }
206    }
207
208    #[async_trait]
209    impl LlmClient for MockLlmClient {
210        async fn chat(
211            &self,
212            _messages: &[ChatMessage],
213            _tools: &[Value],
214            _reasoning: Option<&ReasoningConfig>,
215            _response_format: Option<&ResponseFormat>,
216        ) -> AgentResult<Value> {
217            self.chat_response
218                .lock()
219                .unwrap()
220                .take()
221                .ok_or_else(|| AgentError::internal("no chat response set"))
222        }
223
224        async fn chat_stream(
225            &self,
226            _messages: &[ChatMessage],
227            _tools: &[Value],
228            _reasoning: Option<&ReasoningConfig>,
229            _response_format: Option<&ResponseFormat>,
230        ) -> AgentResult<Pin<Box<dyn Stream<Item = AgentResult<StreamChunk>> + Send>>> {
231            let chunks: Vec<AgentResult<StreamChunk>> = self
232                .stream_chunks
233                .lock()
234                .unwrap()
235                .take()
236                .unwrap_or_default();
237            Ok(Box::pin(futures_util::stream::iter(chunks)))
238        }
239
240        fn capabilities(&self) -> LlmCapabilities {
241            self.caps.clone()
242        }
243
244        fn model_name(&self) -> &str {
245            &self.model
246        }
247    }
248
249    /// Minimal StreamClient that yields a canned stream for testing the
250    /// default `chat()` implementation.
251    struct StubStreamClient {
252        chunks: Mutex<Option<Vec<AgentResult<StreamChunk>>>>,
253        caps: LlmCapabilities,
254    }
255
256    impl StubStreamClient {
257        fn with_chunks(chunks: Vec<StreamChunk>) -> Self {
258            Self {
259                chunks: Mutex::new(Some(chunks.into_iter().map(Ok).collect())),
260                caps: LlmCapabilities::default(),
261            }
262        }
263    }
264
265    #[async_trait]
266    impl StreamClient for StubStreamClient {
267        async fn stream(
268            &self,
269            _messages: &[ChatMessage],
270            _tools: &[Value],
271            _reasoning: Option<&ReasoningConfig>,
272            _response_format: Option<&ResponseFormat>,
273        ) -> AgentResult<Pin<Box<dyn Stream<Item = AgentResult<StreamChunk>> + Send>>> {
274            let chunks: Vec<AgentResult<StreamChunk>> =
275                self.chunks.lock().unwrap().take().unwrap_or_default();
276            Ok(Box::pin(futures_util::stream::iter(chunks)))
277        }
278
279        fn capabilities(&self) -> LlmCapabilities {
280            self.caps.clone()
281        }
282    }
283
284    // ── adapt() ──
285
286    #[test]
287    fn adapt_wraps_llm_client() {
288        let client = Arc::new(MockLlmClient::new());
289        let stream_client = adapt(client);
290        // adapt() returns an Arc<dyn StreamClient> — just verify it builds.
291        let _ = stream_client;
292    }
293
294    // ── LlmClientAdapter::chat() ──
295
296    #[tokio::test]
297    async fn adapter_chat_extracts_content() {
298        let response = serde_json::json!({
299            "choices": [{
300                "message": {
301                    "content": "hello from mock"
302                }
303            }]
304        });
305        let client = Arc::new(MockLlmClient::with_chat_response(response));
306        let adapter = LlmClientAdapter::new(client);
307        let text = adapter
308            .chat(&[], &[], None, None)
309            .await
310            .expect("chat should succeed");
311        assert_eq!(text, "hello from mock");
312    }
313
314    #[tokio::test]
315    async fn adapter_chat_handles_missing_content() {
316        // Response with no choices array
317        let response = serde_json::json!({"error": "something went wrong"});
318        let client = Arc::new(MockLlmClient::with_chat_response(response));
319        let adapter = LlmClientAdapter::new(client);
320        let text = adapter
321            .chat(&[], &[], None, None)
322            .await
323            .expect("chat should succeed with empty string");
324        assert_eq!(text, "");
325    }
326
327    #[tokio::test]
328    async fn adapter_chat_handles_empty_choices() {
329        let response = serde_json::json!({"choices": []});
330        let client = Arc::new(MockLlmClient::with_chat_response(response));
331        let adapter = LlmClientAdapter::new(client);
332        let text = adapter
333            .chat(&[], &[], None, None)
334            .await
335            .expect("chat should succeed with empty string");
336        assert_eq!(text, "");
337    }
338
339    #[tokio::test]
340    async fn adapter_chat_handles_null_content() {
341        // Tool-call response: content is null
342        let response = serde_json::json!({
343            "choices": [{
344                "message": {
345                    "content": null,
346                    "tool_calls": [{"id": "call_1", "function": {"name": "shell"}}]
347                }
348            }]
349        });
350        let client = Arc::new(MockLlmClient::with_chat_response(response));
351        let adapter = LlmClientAdapter::new(client);
352        let text = adapter
353            .chat(&[], &[], None, None)
354            .await
355            .expect("chat should succeed with empty string");
356        assert_eq!(text, "");
357    }
358
359    // ── LlmClientAdapter::stream() ──
360
361    #[tokio::test]
362    async fn adapter_stream_delegates_to_inner() {
363        let chunks = vec![
364            Ok(StreamChunk::Text("hello ".into())),
365            Ok(StreamChunk::Text("world".into())),
366            Ok(StreamChunk::Stop {
367                finish_reason: Some("stop".into()),
368            }),
369        ];
370        let client = Arc::new(MockLlmClient::with_stream(chunks));
371        let adapter = LlmClientAdapter::new(client);
372        let mut stream = adapter
373            .stream(&[], &[], None, None)
374            .await
375            .expect("stream should succeed");
376        let mut texts = Vec::new();
377        while let Some(chunk) = stream.next().await {
378            if let Ok(StreamChunk::Text(t)) = chunk {
379                texts.push(t);
380            }
381        }
382        assert_eq!(texts, vec!["hello ", "world"]);
383    }
384
385    // ── LlmClientAdapter::capabilities() ──
386
387    #[test]
388    fn adapter_capabilities_delegates_to_inner() {
389        let client = Arc::new(MockLlmClient::new());
390        let caps = client.capabilities();
391        let adapter = LlmClientAdapter::new(client);
392        assert_eq!(
393            adapter.capabilities().max_context_tokens,
394            caps.max_context_tokens
395        );
396        assert_eq!(
397            adapter.capabilities().supports_streaming,
398            caps.supports_streaming
399        );
400    }
401
402    // ── LlmClientAdapter::model_name() ──
403
404    #[test]
405    fn adapter_model_name_delegates_to_inner() {
406        let client = Arc::new(MockLlmClient::new());
407        let adapter = LlmClientAdapter::new(client);
408        assert_eq!(adapter.model_name(), "mock-model");
409    }
410
411    #[test]
412    fn default_model_name_is_unknown() {
413        let client = StubStreamClient::with_chunks(vec![]);
414        assert_eq!(client.model_name(), "unknown");
415    }
416
417    // ── LlmClientAdapter::inner() ──
418
419    #[test]
420    fn adapter_inner_returns_reference() {
421        let client = Arc::new(MockLlmClient::new());
422        let adapter = LlmClientAdapter::new(client);
423        let inner: &Arc<dyn LlmClient> = adapter.inner();
424        assert_eq!(inner.model_name(), "mock-model");
425    }
426
427    // ── Default StreamClient::chat() ──
428
429    #[tokio::test]
430    async fn default_chat_collects_text_deltas() {
431        let chunks = vec![
432            StreamChunk::Text("part1".into()),
433            StreamChunk::Text("part2".into()),
434            StreamChunk::Text("part3".into()),
435            StreamChunk::Stop {
436                finish_reason: Some("stop".into()),
437            },
438        ];
439        let client = StubStreamClient::with_chunks(chunks);
440        let text = client
441            .chat(&[], &[], None, None)
442            .await
443            .expect("chat should succeed");
444        assert_eq!(text, "part1part2part3");
445    }
446
447    #[tokio::test]
448    async fn default_chat_stops_on_stop_chunk() {
449        // Text after Stop should be ignored
450        let chunks = vec![
451            StreamChunk::Text("before".into()),
452            StreamChunk::Stop {
453                finish_reason: Some("stop".into()),
454            },
455            StreamChunk::Text("after".into()),
456        ];
457        let client = StubStreamClient::with_chunks(chunks);
458        let text = client
459            .chat(&[], &[], None, None)
460            .await
461            .expect("chat should succeed");
462        assert_eq!(text, "before");
463    }
464
465    #[tokio::test]
466    async fn default_chat_ignores_non_text_chunks() {
467        let chunks = vec![
468            StreamChunk::Thought("thinking...".into()),
469            StreamChunk::Text("visible".into()),
470            StreamChunk::ToolCall(serde_json::json!({"name": "shell"})),
471            StreamChunk::Stop {
472                finish_reason: Some("stop".into()),
473            },
474        ];
475        let client = StubStreamClient::with_chunks(chunks);
476        let text = client
477            .chat(&[], &[], None, None)
478            .await
479            .expect("chat should succeed");
480        assert_eq!(text, "visible");
481    }
482
483    #[tokio::test]
484    async fn default_chat_handles_stream_error() {
485        let chunks = vec![
486            Ok(StreamChunk::Text("before".into())),
487            Err(AgentError::internal("stream broke")),
488        ];
489        let client = StubStreamClient {
490            chunks: Mutex::new(Some(chunks)),
491            caps: LlmCapabilities::default(),
492        };
493        let result = client.chat(&[], &[], None, None).await;
494        assert!(result.is_err());
495        assert!(result.unwrap_err().to_string().contains("stream broke"));
496    }
497
498    #[tokio::test]
499    async fn default_chat_empty_stream_returns_empty() {
500        let client = StubStreamClient::with_chunks(vec![]);
501        let text = client
502            .chat(&[], &[], None, None)
503            .await
504            .expect("chat should succeed");
505        assert_eq!(text, "");
506    }
507
508    // ── Send + Sync ──
509
510    #[test]
511    fn llm_client_adapter_is_send_and_sync() {
512        fn assert_send_sync<T: Send + Sync>() {}
513        assert_send_sync::<LlmClientAdapter>();
514    }
515}
agent_base/llm/stream_client.rs

agent_base/llm/
stream_client.rs
