1use futures_core::Stream;
4use serde_json::Value;
5use std::collections::BTreeMap;
6use std::pin::Pin;
7
8use super::error::LlmError;
9use super::types::UsageInfo;
10
11#[derive(Debug, Clone)]
13pub struct ChatResponse {
14 pub content: String,
16
17 pub reasoning_content: Option<String>,
19
20 pub thinking_signature: Option<String>,
22
23 pub tool_calls: Vec<ToolCall>,
25
26 pub usage: UsageInfo,
28
29 pub finish_reason: FinishReason,
31
32 pub raw: Option<Value>,
34}
35
36#[derive(Debug, Clone)]
38pub struct ToolCall {
39 pub id: String,
40 pub name: String,
41 pub arguments: String,
42}
43
44#[derive(Debug, Clone, PartialEq, Eq)]
46pub enum FinishReason {
47 Stop,
49 Length,
51 ToolCalls,
53 ContentFilter,
55 Other(String),
57}
58
59impl FinishReason {
60 #[allow(clippy::should_implement_trait)]
66 pub fn from_str(s: &str) -> Self {
67 match s {
68 "stop" | "end_turn" => Self::Stop,
69 "length" | "max_tokens" => Self::Length,
70 "tool_calls" | "tool_use" => Self::ToolCalls,
71 "content_filter" => Self::ContentFilter,
72 other => Self::Other(other.to_string()),
73 }
74 }
75
76 pub fn as_str(&self) -> &str {
77 match self {
78 Self::Stop => "stop",
79 Self::Length => "length",
80 Self::ToolCalls => "tool_calls",
81 Self::ContentFilter => "content_filter",
82 Self::Other(s) => s,
83 }
84 }
85}
86
87impl std::str::FromStr for FinishReason {
90 type Err = std::convert::Infallible;
91
92 fn from_str(s: &str) -> Result<Self, Self::Err> {
93 Ok(FinishReason::from_str(s))
94 }
95}
96
97#[derive(Clone, Debug)]
99pub enum StreamChunk {
100 Text(String),
102 Thought(String),
104 ThinkingSignature(String),
106 ToolCall(Value),
108 Usage(UsageInfo),
110 Error(String),
112 Stop { finish_reason: Option<String> },
114}
115
116pub struct ChatStream {
120 inner: Pin<Box<dyn Stream<Item = Result<StreamChunk, LlmError>> + Send>>,
121}
122
123impl ChatStream {
124 pub fn new(stream: Pin<Box<dyn Stream<Item = Result<StreamChunk, LlmError>> + Send>>) -> Self {
126 Self { inner: stream }
127 }
128
129 pub fn into_inner(self) -> Pin<Box<dyn Stream<Item = Result<StreamChunk, LlmError>> + Send>> {
133 self.inner
134 }
135
136 pub async fn next(&mut self) -> Option<Result<StreamChunk, LlmError>> {
138 use futures_util::StreamExt;
139 self.inner.next().await
140 }
141
142 pub async fn collect_text(mut self) -> Result<String, LlmError> {
147 let mut text = String::new();
148 while let Some(chunk) = self.next().await {
149 match chunk? {
150 StreamChunk::Text(t) => text.push_str(&t),
151 StreamChunk::Stop { .. } => break,
152 _ => {}
153 }
154 }
155 Ok(text)
156 }
157
158 pub async fn collect_response(mut self) -> Result<ChatResponse, LlmError> {
166 let mut content = String::new();
167 let mut reasoning_content = String::new();
168 let mut thinking_signature = None;
169 let mut tool_call_buf: BTreeMap<usize, ToolCall> = BTreeMap::new();
170 let mut usage = UsageInfo::default();
171 let mut finish_reason = FinishReason::Stop;
172
173 while let Some(chunk) = self.next().await {
174 match chunk? {
175 StreamChunk::Text(t) => content.push_str(&t),
176 StreamChunk::Thought(t) => reasoning_content.push_str(&t),
177 StreamChunk::ThinkingSignature(sig) => {
178 thinking_signature = Some(sig);
179 }
180 StreamChunk::ToolCall(v) => {
181 apply_tool_call_delta(&mut tool_call_buf, &v);
182 }
183 StreamChunk::Usage(u) => {
184 if u.prompt_tokens.is_some() {
186 usage.prompt_tokens = u.prompt_tokens;
187 }
188 if u.completion_tokens.is_some() {
189 usage.completion_tokens = u.completion_tokens;
190 }
191 if u.total_tokens.is_some() {
192 usage.total_tokens = u.total_tokens;
193 }
194 if u.reasoning_tokens.is_some() {
195 usage.reasoning_tokens = u.reasoning_tokens;
196 }
197 }
198 StreamChunk::Error(msg) => {
199 return Err(LlmError::llm(msg));
201 }
202 StreamChunk::Stop {
203 finish_reason: Some(reason),
204 } => {
205 finish_reason = FinishReason::from_str(&reason);
206 break;
207 }
208 StreamChunk::Stop { .. } => break,
209 }
210 }
211
212 let tool_calls: Vec<ToolCall> = tool_call_buf.into_values().collect();
213
214 Ok(ChatResponse {
215 content,
216 reasoning_content: if reasoning_content.is_empty() {
217 None
218 } else {
219 Some(reasoning_content)
220 },
221 thinking_signature,
222 tool_calls,
223 usage,
224 finish_reason,
225 raw: None,
226 })
227 }
228}
229
230fn apply_tool_call_delta(buf: &mut BTreeMap<usize, ToolCall>, value: &Value) {
240 fn apply_one(buf: &mut BTreeMap<usize, ToolCall>, item: &Value) {
242 let index = item.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
243
244 let id = item.get("id").and_then(|i| i.as_str());
245
246 let function = item.get("function");
248 let name = function
249 .and_then(|f| f.get("name"))
250 .and_then(|n| n.as_str())
251 .or_else(|| item.get("name").and_then(|n| n.as_str()));
252 let arguments = function
253 .and_then(|f| f.get("arguments"))
254 .and_then(|a| a.as_str())
255 .or_else(|| item.get("arguments").and_then(|a| a.as_str()));
256
257 if let (Some(id), Some(name)) = (id, name) {
259 let tc = ToolCall {
260 id: id.to_string(),
261 name: name.to_string(),
262 arguments: arguments.unwrap_or("").to_string(),
263 };
264 buf.insert(index, tc);
265 return;
266 }
267
268 if let Some(args_fragment) = arguments
270 && let Some(tc) = buf.get_mut(&index)
271 {
272 tc.arguments.push_str(args_fragment);
273 }
274 }
275
276 let unwrapped = value
279 .get("delta")
280 .and_then(|d| d.get("tool_calls"))
281 .unwrap_or(value);
282
283 if let Some(arr) = unwrapped.as_array() {
284 for item in arr {
285 apply_one(buf, item);
286 }
287 } else {
288 apply_one(buf, unwrapped);
289 }
290}
291
292pub fn extract_tool_calls(value: &Value) -> Option<Vec<ToolCall>> {
298 if let Some(arr) = value.as_array() {
300 let calls: Vec<ToolCall> = arr
301 .iter()
302 .filter_map(|item| {
303 let id = item.get("id")?.as_str()?.to_string();
304 let function = item.get("function")?;
305 let name = function.get("name")?.as_str()?.to_string();
306 let arguments = function
307 .get("arguments")
308 .and_then(|a| a.as_str())
309 .unwrap_or("")
310 .to_string();
311 Some(ToolCall {
312 id,
313 name,
314 arguments,
315 })
316 })
317 .collect();
318 return if calls.is_empty() { None } else { Some(calls) };
319 }
320
321 let id = value.get("id")?.as_str()?.to_string();
323 let name = value.get("name")?.as_str()?.to_string();
324 let arguments = value
325 .get("arguments")
326 .and_then(|a| a.as_str())
327 .unwrap_or("")
328 .to_string();
329 Some(vec![ToolCall {
330 id,
331 name,
332 arguments,
333 }])
334}
335
336pub fn parse_finish_reason(s: &str) -> FinishReason {
340 FinishReason::from_str(s)
341}
342
343#[cfg(test)]
344mod tests {
345 use super::*;
346
347 #[test]
350 fn finish_reason_from_str() {
351 assert_eq!(FinishReason::from_str("stop"), FinishReason::Stop);
352 assert_eq!(FinishReason::from_str("end_turn"), FinishReason::Stop);
353 assert_eq!(FinishReason::from_str("length"), FinishReason::Length);
354 assert_eq!(FinishReason::from_str("max_tokens"), FinishReason::Length);
355 assert_eq!(
356 FinishReason::from_str("tool_calls"),
357 FinishReason::ToolCalls
358 );
359 assert_eq!(FinishReason::from_str("tool_use"), FinishReason::ToolCalls);
360 assert_eq!(
361 FinishReason::from_str("content_filter"),
362 FinishReason::ContentFilter
363 );
364 assert_eq!(
365 FinishReason::from_str("unknown_reason"),
366 FinishReason::Other("unknown_reason".to_string())
367 );
368 }
369
370 #[test]
371 fn finish_reason_as_str() {
372 assert_eq!(FinishReason::Stop.as_str(), "stop");
373 assert_eq!(FinishReason::Length.as_str(), "length");
374 assert_eq!(FinishReason::ToolCalls.as_str(), "tool_calls");
375 assert_eq!(FinishReason::ContentFilter.as_str(), "content_filter");
376 assert_eq!(
378 FinishReason::Other("refusal".to_string()).as_str(),
379 "refusal"
380 );
381 }
382
383 #[test]
384 fn parse_finish_reason_matches_inherent_from_str() {
385 for s in ["stop", "length", "tool_calls", "content_filter", "weird"] {
386 assert_eq!(parse_finish_reason(s), FinishReason::from_str(s));
387 }
388 }
389
390 #[test]
391 fn from_str_trait_parses_infinitely() {
392 let known: FinishReason = "max_tokens".parse().unwrap();
395 assert_eq!(known, FinishReason::Length);
396 let unknown: FinishReason = "totally-new-reason".parse().unwrap();
397 assert_eq!(
398 unknown,
399 FinishReason::Other("totally-new-reason".to_string())
400 );
401 }
402
403 #[tokio::test]
404 async fn into_inner_yields_the_underlying_stream() {
405 use futures_util::StreamExt;
406 let chunks = vec![
407 Ok(StreamChunk::Text("a".into())),
408 Ok(StreamChunk::Stop {
409 finish_reason: Some("stop".into()),
410 }),
411 ];
412 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
413
414 let collected: Vec<_> = stream.into_inner().collect().await;
415 assert_eq!(collected.len(), 2);
416 assert!(matches!(collected[0], Ok(StreamChunk::Text(ref t)) if t == "a"));
417 }
418
419 #[tokio::test]
422 async fn collect_text_basic() {
423 let chunks = vec![
424 Ok(StreamChunk::Text("hello ".into())),
425 Ok(StreamChunk::Text("world".into())),
426 Ok(StreamChunk::Stop {
427 finish_reason: Some("stop".into()),
428 }),
429 ];
430 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
431 let text = stream.collect_text().await.unwrap();
432 assert_eq!(text, "hello world");
433 }
434
435 #[tokio::test]
436 async fn collect_text_ignores_non_text() {
437 let chunks = vec![
438 Ok(StreamChunk::Thought("thinking...".into())),
439 Ok(StreamChunk::Text("visible".into())),
440 Ok(StreamChunk::ToolCall(serde_json::json!({"name": "shell"}))),
441 Ok(StreamChunk::Usage(UsageInfo {
442 prompt_tokens: Some(10),
443 ..Default::default()
444 })),
445 Ok(StreamChunk::Stop {
446 finish_reason: Some("stop".into()),
447 }),
448 ];
449 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
450 let text = stream.collect_text().await.unwrap();
451 assert_eq!(text, "visible");
452 }
453
454 #[tokio::test]
455 async fn collect_text_stops_on_stop_chunk() {
456 let chunks = vec![
457 Ok(StreamChunk::Text("before".into())),
458 Ok(StreamChunk::Stop {
459 finish_reason: Some("stop".into()),
460 }),
461 Ok(StreamChunk::Text("after".into())),
462 ];
463 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
464 let text = stream.collect_text().await.unwrap();
465 assert_eq!(text, "before");
466 }
467
468 #[tokio::test]
469 async fn collect_text_empty_stream() {
470 let chunks: Vec<Result<StreamChunk, LlmError>> = vec![];
471 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
472 let text = stream.collect_text().await.unwrap();
473 assert_eq!(text, "");
474 }
475
476 #[tokio::test]
477 async fn collect_text_error_propagates() {
478 let chunks: Vec<Result<StreamChunk, LlmError>> = vec![
479 Ok(StreamChunk::Text("before".into())),
480 Err(LlmError::llm("stream broke")),
481 ];
482 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
483 let result = stream.collect_text().await;
484 assert!(result.is_err());
485 assert!(result.unwrap_err().to_string().contains("stream broke"));
486 }
487
488 #[tokio::test]
491 async fn collect_response_full_lifecycle() {
492 let chunks = vec![
493 Ok(StreamChunk::Text("Hello ".into())),
494 Ok(StreamChunk::Text("world!".into())),
495 Ok(StreamChunk::ToolCall(serde_json::json!([
496 {
497 "id": "call_1",
498 "function": {
499 "name": "shell",
500 "arguments": "{\"cmd\": \"ls\"}"
501 }
502 }
503 ]))),
504 Ok(StreamChunk::Usage(UsageInfo {
505 prompt_tokens: Some(100),
506 completion_tokens: Some(50),
507 total_tokens: Some(150),
508 reasoning_tokens: None,
509 })),
510 Ok(StreamChunk::Stop {
511 finish_reason: Some("stop".into()),
512 }),
513 ];
514 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
515 let response = stream.collect_response().await.unwrap();
516
517 assert_eq!(response.content, "Hello world!");
518 assert!(response.reasoning_content.is_none());
519 assert_eq!(response.tool_calls.len(), 1);
520 assert_eq!(response.tool_calls[0].id, "call_1");
521 assert_eq!(response.tool_calls[0].name, "shell");
522 assert_eq!(response.tool_calls[0].arguments, "{\"cmd\": \"ls\"}");
523 assert_eq!(response.usage.prompt_tokens, Some(100));
524 assert_eq!(response.usage.completion_tokens, Some(50));
525 assert_eq!(response.usage.total_tokens, Some(150));
526 assert_eq!(response.finish_reason, FinishReason::Stop);
527 }
528
529 #[tokio::test]
530 async fn collect_response_with_thought() {
531 let chunks = vec![
532 Ok(StreamChunk::Thought("thinking...".into())),
533 Ok(StreamChunk::Text("answer".into())),
534 Ok(StreamChunk::Stop {
535 finish_reason: Some("stop".into()),
536 }),
537 ];
538 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
539 let response = stream.collect_response().await.unwrap();
540
541 assert_eq!(response.content, "answer");
542 assert_eq!(response.reasoning_content.as_deref(), Some("thinking..."));
543 }
544
545 #[tokio::test]
546 async fn collect_response_usage_accumulate() {
547 let chunks = vec![
549 Ok(StreamChunk::Usage(UsageInfo {
550 prompt_tokens: Some(20),
551 completion_tokens: Some(0),
552 total_tokens: None,
553 reasoning_tokens: None,
554 })),
555 Ok(StreamChunk::Text("ok".into())),
556 Ok(StreamChunk::Usage(UsageInfo {
557 prompt_tokens: None,
558 completion_tokens: Some(7),
559 total_tokens: None,
560 reasoning_tokens: None,
561 })),
562 Ok(StreamChunk::Stop {
563 finish_reason: Some("end_turn".into()),
564 }),
565 ];
566 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
567 let response = stream.collect_response().await.unwrap();
568
569 assert_eq!(response.usage.prompt_tokens, Some(20));
571 assert_eq!(response.usage.completion_tokens, Some(7));
572 }
573
574 #[tokio::test]
575 async fn collect_response_text_only() {
576 let chunks = vec![
577 Ok(StreamChunk::Text("Just text".into())),
578 Ok(StreamChunk::Stop {
579 finish_reason: Some("stop".into()),
580 }),
581 ];
582 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
583 let response = stream.collect_response().await.unwrap();
584
585 assert_eq!(response.content, "Just text");
586 assert!(response.tool_calls.is_empty());
587 assert_eq!(response.finish_reason, FinishReason::Stop);
588 }
589
590 #[tokio::test]
591 async fn collect_response_no_stop_chunk() {
592 let chunks: Vec<Result<StreamChunk, LlmError>> = vec![Ok(StreamChunk::Text("text".into()))];
593 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
594 let response = stream.collect_response().await.unwrap();
595
596 assert_eq!(response.content, "text");
597 assert_eq!(response.finish_reason, FinishReason::Stop);
598 }
599
600 #[tokio::test]
601 async fn collect_response_tool_calls_format() {
602 let chunks = vec![
603 Ok(StreamChunk::ToolCall(serde_json::json!({
604 "id": "call_abc",
605 "name": "read_file",
606 "arguments": "{\"path\": \"/tmp/test.txt\"}"
607 }))),
608 Ok(StreamChunk::Stop {
609 finish_reason: Some("stop".into()),
610 }),
611 ];
612 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
613 let response = stream.collect_response().await.unwrap();
614
615 assert_eq!(response.tool_calls.len(), 1);
616 assert_eq!(response.tool_calls[0].id, "call_abc");
617 assert_eq!(response.tool_calls[0].name, "read_file");
618 }
619
620 #[tokio::test]
621 async fn collect_response_error_propagates() {
622 let chunks: Vec<Result<StreamChunk, LlmError>> = vec![
623 Ok(StreamChunk::Text("before".into())),
624 Err(LlmError::llm("broken")),
625 ];
626 let stream = ChatStream::new(Box::pin(futures_util::stream::iter(chunks)));
627 let result = stream.collect_response().await;
628 assert!(result.is_err());
629 }
630
631 #[test]
634 fn extract_tool_calls_openai_format() {
635 let value = serde_json::json!([
636 {
637 "id": "call_1",
638 "function": {
639 "name": "shell",
640 "arguments": "{\"cmd\": \"ls\"}"
641 }
642 },
643 {
644 "id": "call_2",
645 "function": {
646 "name": "read_file",
647 "arguments": "{\"path\": \"/tmp\"}"
648 }
649 }
650 ]);
651 let calls = extract_tool_calls(&value).unwrap();
652 assert_eq!(calls.len(), 2);
653 assert_eq!(calls[0].id, "call_1");
654 assert_eq!(calls[0].name, "shell");
655 assert_eq!(calls[1].id, "call_2");
656 assert_eq!(calls[1].name, "read_file");
657 }
658
659 #[test]
660 fn extract_tool_calls_simple_format() {
661 let value = serde_json::json!({
662 "id": "call_1",
663 "name": "shell",
664 "arguments": "{\"cmd\": \"ls\"}"
665 });
666 let calls = extract_tool_calls(&value).unwrap();
667 assert_eq!(calls.len(), 1);
668 assert_eq!(calls[0].name, "shell");
669 }
670
671 #[test]
672 fn extract_tool_calls_invalid_returns_none() {
673 let value = serde_json::json!("just a string");
674 assert!(extract_tool_calls(&value).is_none());
675 }
676
677 #[test]
678 fn extract_tool_calls_empty_array_returns_none() {
679 let value = serde_json::json!([]);
680 assert!(extract_tool_calls(&value).is_none());
681 }
682
683 mod proptest_tests {
686 use super::*;
687 use proptest::prelude::*;
688
689 proptest! {
690 #[test]
691 fn from_str_never_panics(s in ".*") {
692 let _ = FinishReason::from_str(&s);
693 }
694
695 #[test]
696 fn from_str_known_values(
697 variant in proptest::sample::select(vec![
698 ("stop", "Stop"),
699 ("end_turn", "Stop"),
700 ("length", "Length"),
701 ("max_tokens", "Length"),
702 ("tool_calls", "ToolCalls"),
703 ("tool_use", "ToolCalls"),
704 ("content_filter", "ContentFilter"),
705 ])
706 ) {
707 let (s, expected_name) = variant;
708 let fr = FinishReason::from_str(s);
709 match (expected_name, &fr) {
710 ("Stop", FinishReason::Stop) => {}
711 ("Length", FinishReason::Length) => {}
712 ("ToolCalls", FinishReason::ToolCalls) => {}
713 ("ContentFilter", FinishReason::ContentFilter) => {}
714 _ => panic!("from_str({:?}) = {:?}, expected {}", s, fr, expected_name),
715 }
716 }
717
718 #[test]
719 fn from_str_unknown_returns_other(s in "[a-z_]{1,30}") {
720 if ["stop", "end_turn", "length", "max_tokens", "tool_calls", "tool_use", "content_filter"].contains(&s.as_str()) {
722 return Ok(());
723 }
724 let fr = FinishReason::from_str(&s);
725 match fr {
726 FinishReason::Other(inner) => assert_eq!(inner, s),
727 other => panic!("expected Other({:?}), got {:?}", s, other),
728 }
729 }
730 }
731 }
732}