1use std::time::Duration;
2
3use async_trait::async_trait;
4
5use crate::error::{KernelError, Result};
6use crate::llm::tool::{ToolCall, ToolDefinition};
7use crate::llm::types::{
8 LLMRequest, LLMResponse, LLMStream, ModelConfig, ResponseFormat, StreamEvent, TokenUsage,
9};
10
11fn redact_http_body(body: &str) -> String {
18 #[cfg(feature = "safety")]
19 {
20 crate::safety::sanitize::mask_secrets(body)
21 }
22 #[cfg(not(feature = "safety"))]
23 {
24 body.to_string()
25 }
26}
27
28fn openai_tools(tools: &[ToolDefinition]) -> Vec<serde_json::Value> {
30 tools
31 .iter()
32 .map(|t| {
33 serde_json::json!({
34 "type": "function",
35 "function": {
36 "name": t.name,
37 "description": t.description,
38 "parameters": t.input_schema,
39 }
40 })
41 })
42 .collect()
43}
44
45fn openai_response_format(rf: &ResponseFormat) -> Option<serde_json::Value> {
48 match rf {
49 ResponseFormat::Text => None,
50 ResponseFormat::Json => Some(serde_json::json!({ "type": "json_object" })),
51 ResponseFormat::JsonSchema { schema } => Some(serde_json::json!({
52 "type": "json_schema",
53 "json_schema": { "name": "response", "schema": schema, "strict": true }
54 })),
55 }
56}
57
58fn anthropic_tools(tools: &[ToolDefinition]) -> Vec<serde_json::Value> {
60 tools
61 .iter()
62 .map(|t| {
63 serde_json::json!({
64 "name": t.name,
65 "description": t.description,
66 "input_schema": t.input_schema,
67 })
68 })
69 .collect()
70}
71
72fn anthropic_output_config(rf: &ResponseFormat) -> Option<serde_json::Value> {
76 match rf {
77 ResponseFormat::JsonSchema { schema } => Some(serde_json::json!({
78 "format": { "type": "json_schema", "schema": schema }
79 })),
80 ResponseFormat::Json | ResponseFormat::Text => None,
81 }
82}
83
84fn http_client() -> Result<reqwest::Client> {
86 reqwest::Client::builder()
87 .connect_timeout(Duration::from_secs(10))
88 .timeout(Duration::from_secs(120))
89 .build()
90 .map_err(|e| KernelError::Config(format!("Failed to build HTTP client: {}", e)))
91}
92
93fn check_rate_limit(resp: &reqwest::Response) -> Result<()> {
95 if resp.status().as_u16() == 429 {
96 let retry = resp
97 .headers()
98 .get("retry-after")
99 .and_then(|v| v.to_str().ok())
100 .and_then(|v| v.parse().ok())
101 .unwrap_or(60);
102 return Err(KernelError::RateLimited(retry));
103 }
104 Ok(())
105}
106
107#[async_trait]
109pub trait LLMClient: Send + Sync {
110 async fn complete(&self, request: LLMRequest) -> Result<LLMResponse>;
122 fn model_name(&self) -> &str;
124
125 async fn stream_complete(&self, request: LLMRequest) -> Result<LLMStream>;
132}
133
134pub struct OpenAIClient {
136 api_key: String,
137 model: String,
138 base_url: String,
139 client: reqwest::Client,
140}
141
142impl OpenAIClient {
143 pub fn new(config: &ModelConfig) -> Result<Self> {
145 let api_key = std::env::var(&config.api_key_env).map_err(|_| {
146 KernelError::Config(format!(
147 "Environment variable {} not set",
148 config.api_key_env
149 ))
150 })?;
151 Ok(Self {
152 api_key,
153 model: config.model.clone(),
154 base_url: config
155 .base_url
156 .clone()
157 .unwrap_or_else(|| "https://api.openai.com/v1".into()),
158 client: http_client()?,
159 })
160 }
161
162 pub fn from_key(model: impl Into<String>, api_key: impl Into<String>) -> Result<Self> {
176 Ok(Self {
177 api_key: api_key.into(),
178 model: model.into(),
179 base_url: "https://api.openai.com/v1".into(),
180 client: http_client()?,
181 })
182 }
183
184 pub fn from_key_with_client(
190 model: impl Into<String>,
191 api_key: impl Into<String>,
192 client: reqwest::Client,
193 ) -> Self {
194 Self {
195 api_key: api_key.into(),
196 model: model.into(),
197 base_url: "https://api.openai.com/v1".into(),
198 client,
199 }
200 }
201
202 pub fn from_key_with_base_url(
208 model: impl Into<String>,
209 api_key: impl Into<String>,
210 base_url: impl Into<String>,
211 client: reqwest::Client,
212 ) -> Self {
213 Self {
214 api_key: api_key.into(),
215 model: model.into(),
216 base_url: base_url.into(),
217 client,
218 }
219 }
220}
221
222#[derive(serde::Serialize)]
223struct OpenAIChatRequest {
224 model: String,
225 messages: Vec<OpenAIChatMessage>,
226 temperature: f32,
227 #[serde(skip_serializing_if = "Option::is_none")]
228 max_tokens: Option<u32>,
229 #[serde(skip_serializing_if = "std::ops::Not::not")]
230 stream: bool,
231 #[serde(skip_serializing_if = "Option::is_none")]
232 tools: Option<Vec<serde_json::Value>>,
233 #[serde(skip_serializing_if = "Option::is_none")]
234 response_format: Option<serde_json::Value>,
235}
236
237#[derive(serde::Serialize)]
238struct OpenAIChatMessage {
239 role: String,
240 content: String,
241}
242
243#[derive(serde::Deserialize)]
244struct OpenAIChatResponse {
245 #[serde(default)]
246 id: Option<String>,
247 #[serde(default)]
248 created: Option<u64>,
249 choices: Vec<OpenAIChoice>,
250 model: String,
251 usage: Option<OpenAIUsage>,
252}
253
254#[derive(serde::Deserialize)]
255struct OpenAIChoice {
256 message: OpenAIRespMessage,
257 #[serde(default)]
258 finish_reason: Option<String>,
259}
260
261#[derive(serde::Deserialize)]
268struct OpenAIRespMessage {
269 #[serde(default)]
270 content: Option<String>,
271 #[serde(default, alias = "reasoning")]
273 reasoning_content: Option<String>,
274 #[serde(default)]
275 tool_calls: Vec<OpenAIToolCall>,
276}
277
278#[derive(serde::Deserialize)]
279struct OpenAIToolCall {
280 id: String,
281 function: OpenAIFunctionCall,
282}
283
284#[derive(serde::Deserialize)]
285struct OpenAIFunctionCall {
286 name: String,
287 #[serde(default)]
288 arguments: String,
289}
290
291#[derive(serde::Deserialize)]
297struct OpenAIUsage {
298 #[serde(default)]
299 prompt_tokens: u32,
300 #[serde(default)]
301 completion_tokens: u32,
302 #[serde(default)]
303 total_tokens: u32,
304 #[serde(default)]
306 completion_tokens_details: Option<OpenAICompletionTokensDetails>,
307}
308
309#[derive(serde::Deserialize)]
311struct OpenAICompletionTokensDetails {
312 #[serde(default)]
313 reasoning_tokens: Option<u32>,
314}
315
316#[async_trait]
317impl LLMClient for OpenAIClient {
318 async fn complete(&self, request: LLMRequest) -> Result<LLMResponse> {
319 let model = request.model.clone().unwrap_or_else(|| self.model.clone());
320 let temperature = request.temperature;
321 let max_tokens = request.max_tokens;
322 let tools = request
323 .tools
324 .as_deref()
325 .map(openai_tools)
326 .filter(|t| !t.is_empty());
327 let response_format = request
328 .response_format
329 .as_ref()
330 .and_then(openai_response_format);
331 let messages: Vec<_> = request
332 .into_openai_messages()
333 .into_iter()
334 .map(|(role, content)| OpenAIChatMessage { role, content })
335 .collect();
336
337 let body = OpenAIChatRequest {
338 model,
339 messages,
340 temperature,
341 max_tokens,
342 stream: false,
343 tools,
344 response_format,
345 };
346
347 let resp = self
348 .client
349 .post(format!("{}/chat/completions", self.base_url))
350 .header("Authorization", format!("Bearer {}", self.api_key))
351 .json(&body)
352 .send()
353 .await
354 .map_err(|e| KernelError::LlmApi(e.to_string()))?;
355
356 check_rate_limit(&resp)?;
357
358 let status = resp.status();
359
360 if !status.is_success() {
361 let text = resp.text().await.unwrap_or_default();
362 return Err(KernelError::Http {
363 status: status.as_u16(),
364 message: redact_http_body(&text),
365 });
366 }
367
368 let chat_resp: OpenAIChatResponse = resp
369 .json()
370 .await
371 .map_err(|e| KernelError::LlmApi(e.to_string()))?;
372
373 let id = chat_resp.id;
374 let created = chat_resp.created;
375 let first = chat_resp.choices.into_iter().next();
376 let finish_reason = first.as_ref().and_then(|c| c.finish_reason.clone());
377 let (content, reasoning, tool_calls) = match first {
378 Some(c) => {
379 let raw_content = c.message.content.unwrap_or_default();
380 let reasoning = c.message.reasoning_content;
381 let content = promote_reasoning_into_content(raw_content, reasoning.as_deref());
382 let calls = c
383 .message
384 .tool_calls
385 .into_iter()
386 .map(|tc| ToolCall {
387 id: tc.id,
388 name: tc.function.name,
389 arguments: tc.function.arguments,
390 })
391 .collect();
392 (content, reasoning, calls)
393 }
394 None => (String::new(), None, Vec::new()),
395 };
396
397 let usage = chat_resp.usage.map(|u| TokenUsage {
398 prompt_tokens: u.prompt_tokens,
399 completion_tokens: u.completion_tokens,
400 total_tokens: u.total_tokens,
401 reasoning_tokens: u.completion_tokens_details.and_then(|d| d.reasoning_tokens),
402 });
403
404 Ok(LLMResponse {
405 content,
406 reasoning,
407 model: chat_resp.model,
408 usage: usage.unwrap_or_default(),
409 tool_calls,
410 finish_reason,
411 id,
412 created,
413 })
414 }
415
416 fn model_name(&self) -> &str {
417 &self.model
418 }
419
420 async fn stream_complete(&self, request: LLMRequest) -> Result<LLMStream> {
421 let model = request.model.clone().unwrap_or_else(|| self.model.clone());
422 let temperature = request.temperature;
423 let max_tokens = request.max_tokens;
424 let messages: Vec<_> = request
425 .into_openai_messages()
426 .into_iter()
427 .map(|(role, content)| OpenAIChatMessage { role, content })
428 .collect();
429
430 let body = OpenAIChatRequest {
431 model,
432 messages,
433 temperature,
434 max_tokens,
435 stream: true,
436 tools: None,
439 response_format: None,
440 };
441
442 let resp = self
443 .client
444 .post(format!("{}/chat/completions", self.base_url))
445 .header("Authorization", format!("Bearer {}", self.api_key))
446 .json(&body)
447 .send()
448 .await
449 .map_err(|e| KernelError::LlmApi(e.to_string()))?;
450
451 check_rate_limit(&resp)?;
452
453 let status = resp.status();
454 if !status.is_success() {
455 let text = resp.text().await.unwrap_or_default();
456 return Err(KernelError::Http {
457 status: status.as_u16(),
458 message: redact_http_body(&text),
459 });
460 }
461
462 let (tx, rx) = tokio::sync::mpsc::channel::<Result<StreamEvent>>(16);
463
464 tokio::spawn(async move {
465 let mut stream = std::pin::pin!(resp.bytes_stream());
466 let mut buffer: Vec<u8> = Vec::new();
467
468 use tokio_stream::StreamExt;
469
470 while let Some(chunk) = stream.next().await {
471 let chunk = match chunk {
472 Ok(c) => c,
473 Err(e) => {
474 let _ = tx.send(Err(KernelError::LlmApi(e.to_string()))).await;
475 return;
476 }
477 };
478
479 for line in drain_sse_lines(&mut buffer, &chunk) {
480 if let Some(data) = parse_sse_line(&line)
481 && let Some(event) = parse_openai_sse(data)
482 {
483 let is_done = matches!(event, StreamEvent::Done);
484 if tx.send(Ok(event)).await.is_err() || is_done {
485 return;
486 }
487 }
488 }
489 }
490 let _ = tx.send(Ok(StreamEvent::Done)).await;
491 });
492
493 Ok(Box::pin(tokio_stream::wrappers::ReceiverStream::new(rx)))
494 }
495}
496
497fn parse_sse_line(line: &str) -> Option<&str> {
500 line.strip_prefix("data: ").filter(|d| *d != "[DONE]")
501}
502
503fn drain_sse_lines(buffer: &mut Vec<u8>, chunk: &[u8]) -> Vec<String> {
513 buffer.extend_from_slice(chunk);
514 let mut lines = Vec::new();
515 while let Some(pos) = buffer.iter().position(|&b| b == b'\n') {
516 let line: Vec<u8> = buffer.drain(..=pos).collect();
517 lines.push(String::from_utf8_lossy(&line).trim_end().to_string());
518 }
519 lines
520}
521
522fn promote_reasoning_into_content(raw_content: String, reasoning: Option<&str>) -> String {
543 if raw_content.is_empty() {
544 reasoning.unwrap_or_default().to_string()
545 } else {
546 raw_content
547 }
548}
549
550fn parse_openai_sse(data: &str) -> Option<StreamEvent> {
552 let v: serde_json::Value = serde_json::from_str(data).ok()?;
553
554 if let Some(rc) = v
557 .get("choices")?
558 .get(0)?
559 .get("delta")?
560 .get("reasoning_content")
561 .and_then(|c| c.as_str())
562 && !rc.is_empty()
563 {
564 return Some(StreamEvent::ReasoningDelta {
565 content: rc.to_string(),
566 });
567 }
568
569 if let Some(content) = v
571 .get("choices")?
572 .get(0)?
573 .get("delta")?
574 .get("content")
575 .and_then(|c| c.as_str())
576 && !content.is_empty()
577 {
578 return Some(StreamEvent::Delta {
579 content: content.to_string(),
580 });
581 }
582
583 if let Some(usage) = v.get("usage").and_then(|u| {
585 Some(TokenUsage {
586 prompt_tokens: u.get("prompt_tokens")?.as_u64()? as u32,
587 completion_tokens: u.get("completion_tokens")?.as_u64()? as u32,
588 total_tokens: u.get("total_tokens")?.as_u64()? as u32,
589 reasoning_tokens: u
590 .get("completion_tokens_details")
591 .and_then(|d| d.get("reasoning_tokens"))
592 .and_then(|r| r.as_u64())
593 .map(|n| n as u32),
594 })
595 }) {
596 return Some(StreamEvent::Usage(usage));
597 }
598
599 if v.get("choices")?
601 .get(0)?
602 .get("finish_reason")
603 .and_then(|r| r.as_str())
604 .is_some()
605 {
606 return Some(StreamEvent::Done);
607 }
608
609 None
610}
611
612fn parse_anthropic_sse(event_type: &str, data: &str) -> Option<StreamEvent> {
614 let v: serde_json::Value = serde_json::from_str(data).ok()?;
615
616 match event_type {
617 "content_block_delta" => {
618 let delta = v.get("delta")?;
619 match delta.get("type").and_then(|t| t.as_str()) {
621 Some("thinking_delta") => {
622 let text = delta.get("thinking")?.as_str()?;
623 if !text.is_empty() {
624 return Some(StreamEvent::ReasoningDelta {
625 content: text.to_string(),
626 });
627 }
628 None
629 }
630 _ => {
631 let text = delta.get("text")?.as_str()?;
633 if !text.is_empty() {
634 return Some(StreamEvent::Delta {
635 content: text.to_string(),
636 });
637 }
638 None
639 }
640 }
641 }
642 "message_delta" => {
643 let usage = v.get("usage").and_then(|u| {
644 Some(TokenUsage {
645 prompt_tokens: 0,
646 completion_tokens: u.get("output_tokens")?.as_u64()? as u32,
647 total_tokens: 0,
648 reasoning_tokens: None,
649 })
650 });
651 if let Some(usage) = usage {
652 return Some(StreamEvent::Usage(usage));
653 }
654 Some(StreamEvent::Done)
655 }
656 "message_stop" => Some(StreamEvent::Done),
657 _ => None,
658 }
659}
660
661pub struct AnthropicClient {
663 api_key: String,
664 model: String,
665 base_url: String,
666 client: reqwest::Client,
667}
668
669impl AnthropicClient {
670 pub fn new(config: &ModelConfig) -> Result<Self> {
672 let api_key = std::env::var(&config.api_key_env).map_err(|_| {
673 KernelError::Config(format!(
674 "Environment variable {} not set",
675 config.api_key_env
676 ))
677 })?;
678 Ok(Self {
679 api_key,
680 model: config.model.clone(),
681 base_url: config
682 .base_url
683 .clone()
684 .unwrap_or_else(|| "https://api.anthropic.com/v1".into()),
685 client: http_client()?,
686 })
687 }
688
689 pub fn from_key(model: impl Into<String>, api_key: impl Into<String>) -> Result<Self> {
695 Ok(Self {
696 api_key: api_key.into(),
697 model: model.into(),
698 base_url: "https://api.anthropic.com/v1".into(),
699 client: http_client()?,
700 })
701 }
702
703 pub fn from_key_with_client(
705 model: impl Into<String>,
706 api_key: impl Into<String>,
707 client: reqwest::Client,
708 ) -> Self {
709 Self {
710 api_key: api_key.into(),
711 model: model.into(),
712 base_url: "https://api.anthropic.com/v1".into(),
713 client,
714 }
715 }
716
717 pub fn from_key_with_base_url(
722 model: impl Into<String>,
723 api_key: impl Into<String>,
724 base_url: impl Into<String>,
725 client: reqwest::Client,
726 ) -> Self {
727 Self {
728 api_key: api_key.into(),
729 model: model.into(),
730 base_url: base_url.into(),
731 client,
732 }
733 }
734}
735
736#[derive(serde::Serialize)]
737struct AnthropicRequest {
738 model: String,
739 max_tokens: u32,
740 temperature: f32,
741 #[serde(skip_serializing_if = "Option::is_none")]
742 system: Option<String>,
743 messages: Vec<AnthropicMessage>,
744 #[serde(skip_serializing_if = "std::ops::Not::not")]
745 stream: bool,
746 #[serde(skip_serializing_if = "Option::is_none")]
747 tools: Option<Vec<serde_json::Value>>,
748 #[serde(skip_serializing_if = "Option::is_none")]
749 output_config: Option<serde_json::Value>,
750}
751
752#[derive(serde::Serialize)]
753struct AnthropicMessage {
754 role: String,
755 content: String,
756}
757
758#[derive(serde::Deserialize)]
759struct AnthropicResponse {
760 #[serde(default)]
761 id: Option<String>,
762 content: Vec<AnthropicContentBlock>,
763 model: String,
764 #[serde(default)]
765 stop_reason: Option<String>,
766 usage: AnthropicUsage,
767}
768
769#[derive(serde::Deserialize)]
772struct AnthropicContentBlock {
773 #[serde(rename = "type")]
774 block_type: String,
775 #[serde(default)]
776 text: Option<String>,
777 #[serde(default)]
779 thinking: Option<String>,
780 #[serde(default)]
781 id: Option<String>,
782 #[serde(default)]
783 name: Option<String>,
784 #[serde(default)]
785 input: Option<serde_json::Value>,
786}
787
788#[derive(serde::Deserialize)]
789struct AnthropicUsage {
790 input_tokens: u32,
791 output_tokens: u32,
792}
793
794#[async_trait]
795impl LLMClient for AnthropicClient {
796 async fn complete(&self, request: LLMRequest) -> Result<LLMResponse> {
797 let model = request.model.clone().unwrap_or_else(|| self.model.clone());
798 let max_tokens = request.max_tokens.unwrap_or(4096);
799 let temperature = request.temperature;
800 let system = request.system.clone();
801 let tools = request
802 .tools
803 .as_deref()
804 .map(anthropic_tools)
805 .filter(|t| !t.is_empty());
806 let output_config = request
807 .response_format
808 .as_ref()
809 .and_then(anthropic_output_config);
810 let messages: Vec<AnthropicMessage> = request
811 .into_anthropic_messages()
812 .into_iter()
813 .map(|(role, content)| AnthropicMessage { role, content })
814 .collect();
815
816 let body = AnthropicRequest {
817 model,
818 max_tokens,
819 temperature,
820 system,
821 messages,
822 stream: false,
823 tools,
824 output_config,
825 };
826
827 let resp = self
828 .client
829 .post(format!("{}/messages", self.base_url))
830 .header("x-api-key", &self.api_key)
831 .header("anthropic-version", "2023-06-01")
832 .header("content-type", "application/json")
833 .json(&body)
834 .send()
835 .await
836 .map_err(|e| KernelError::LlmApi(e.to_string()))?;
837
838 check_rate_limit(&resp)?;
839
840 let status = resp.status();
841
842 if !status.is_success() {
843 let text = resp.text().await.unwrap_or_default();
844 return Err(KernelError::Http {
845 status: status.as_u16(),
846 message: redact_http_body(&text),
847 });
848 }
849
850 let chat_resp: AnthropicResponse = resp
851 .json()
852 .await
853 .map_err(|e| KernelError::LlmApi(e.to_string()))?;
854
855 let mut content = String::new();
856 let mut reasoning = String::new();
857 let mut tool_calls = Vec::new();
858 for block in chat_resp.content {
859 match block.block_type.as_str() {
860 "text" => {
861 if let Some(t) = block.text {
862 content.push_str(&t);
863 }
864 }
865 "thinking" => {
866 if let Some(t) = block.thinking {
867 reasoning.push_str(&t);
868 }
869 }
870 "tool_use" => {
871 if let (Some(id), Some(name)) = (block.id, block.name) {
872 let arguments = block
873 .input
874 .map(|v| v.to_string())
875 .unwrap_or_else(|| "{}".to_string());
876 tool_calls.push(ToolCall {
877 id,
878 name,
879 arguments,
880 });
881 }
882 }
883 _ => {}
884 }
885 }
886
887 Ok(LLMResponse {
888 content,
889 reasoning: if reasoning.is_empty() {
890 None
891 } else {
892 Some(reasoning)
893 },
894 model: chat_resp.model,
895 usage: TokenUsage {
896 prompt_tokens: chat_resp.usage.input_tokens,
897 completion_tokens: chat_resp.usage.output_tokens,
898 total_tokens: chat_resp.usage.input_tokens + chat_resp.usage.output_tokens,
899 reasoning_tokens: None,
900 },
901 tool_calls,
902 finish_reason: chat_resp.stop_reason,
903 id: chat_resp.id,
904 created: None,
905 })
906 }
907
908 fn model_name(&self) -> &str {
909 &self.model
910 }
911
912 async fn stream_complete(&self, request: LLMRequest) -> Result<LLMStream> {
913 let model = request.model.clone().unwrap_or_else(|| self.model.clone());
914 let max_tokens = request.max_tokens.unwrap_or(4096);
915 let temperature = request.temperature;
916 let system = request.system.clone();
917 let messages: Vec<AnthropicMessage> = request
918 .into_anthropic_messages()
919 .into_iter()
920 .map(|(role, content)| AnthropicMessage { role, content })
921 .collect();
922
923 let body = AnthropicRequest {
924 model,
925 max_tokens,
926 temperature,
927 system,
928 messages,
929 stream: true,
930 tools: None,
933 output_config: None,
934 };
935
936 let resp = self
937 .client
938 .post(format!("{}/messages", self.base_url))
939 .header("x-api-key", &self.api_key)
940 .header("anthropic-version", "2023-06-01")
941 .header("content-type", "application/json")
942 .json(&body)
943 .send()
944 .await
945 .map_err(|e| KernelError::LlmApi(e.to_string()))?;
946
947 check_rate_limit(&resp)?;
948
949 let status = resp.status();
950 if !status.is_success() {
951 let text = resp.text().await.unwrap_or_default();
952 return Err(KernelError::Http {
953 status: status.as_u16(),
954 message: redact_http_body(&text),
955 });
956 }
957
958 let (tx, rx) = tokio::sync::mpsc::channel::<Result<StreamEvent>>(16);
959
960 tokio::spawn(async move {
961 let mut stream = std::pin::pin!(resp.bytes_stream());
962 let mut buffer: Vec<u8> = Vec::new();
963 let mut current_event = String::new();
964
965 use tokio_stream::StreamExt;
966
967 while let Some(chunk) = stream.next().await {
968 let chunk = match chunk {
969 Ok(c) => c,
970 Err(e) => {
971 let _ = tx.send(Err(KernelError::LlmApi(e.to_string()))).await;
972 return;
973 }
974 };
975
976 for line in drain_sse_lines(&mut buffer, &chunk) {
977 if let Some(evt) = line.strip_prefix("event: ") {
978 current_event = evt.to_string();
979 } else if let Some(data) = line.strip_prefix("data: ") {
980 if data == "[DONE]" {
981 let _ = tx.send(Ok(StreamEvent::Done)).await;
982 return;
983 }
984 if let Some(event) = parse_anthropic_sse(¤t_event, data) {
985 let is_done = matches!(event, StreamEvent::Done);
986 if tx.send(Ok(event)).await.is_err() || is_done {
987 return;
988 }
989 }
990 current_event.clear();
991 }
992 }
993 }
994 let _ = tx.send(Ok(StreamEvent::Done)).await;
995 });
996
997 Ok(Box::pin(tokio_stream::wrappers::ReceiverStream::new(rx)))
998 }
999}
1000
1001#[cfg(test)]
1002mod tests {
1003 use super::*;
1004
1005 #[test]
1006 fn parse_sse_line_extracts_data() {
1007 assert_eq!(
1008 parse_sse_line("data: {\"id\":\"1\"}"),
1009 Some("{\"id\":\"1\"}")
1010 );
1011 }
1012
1013 #[test]
1014 fn parse_sse_line_skips_done() {
1015 assert_eq!(parse_sse_line("data: [DONE]"), None);
1016 }
1017
1018 #[test]
1019 fn parse_sse_line_skips_non_data() {
1020 assert_eq!(parse_sse_line("event: ping"), None);
1021 assert_eq!(parse_sse_line(""), None);
1022 }
1023
1024 #[test]
1025 fn drain_sse_lines_reassembles_multibyte_split_across_chunks() {
1026 let full = "data: 안녕\n".as_bytes().to_vec();
1028 let (first, rest) = full.split_at(7);
1030
1031 let mut buffer = Vec::new();
1032 assert!(drain_sse_lines(&mut buffer, first).is_empty());
1035
1036 let lines = drain_sse_lines(&mut buffer, rest);
1037 assert_eq!(lines, vec!["data: 안녕".to_string()]);
1038 assert!(!lines[0].contains('\u{FFFD}'));
1040 }
1041
1042 #[test]
1043 fn drain_sse_lines_handles_multiple_lines_and_keeps_partial_tail() {
1044 let mut buffer = Vec::new();
1045 let lines = drain_sse_lines(&mut buffer, b"event: ping\r\ndata: {}\npartial");
1046 assert_eq!(
1047 lines,
1048 vec!["event: ping".to_string(), "data: {}".to_string()]
1049 );
1050 let lines = drain_sse_lines(&mut buffer, b" tail\n");
1052 assert_eq!(lines, vec!["partial tail".to_string()]);
1053 }
1054
1055 #[test]
1056 fn openai_delta_extraction() {
1057 let data = r#"{"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}"#;
1058 let event = parse_openai_sse(data).unwrap();
1059 match event {
1060 StreamEvent::Delta { content } => assert_eq!(content, "Hello"),
1061 _ => panic!("expected Delta, got {:?}", event),
1062 }
1063 }
1064
1065 #[test]
1066 fn openai_usage_extraction() {
1067 let data = r#"{"id":"chatcmpl-1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}"#;
1068 let event = parse_openai_sse(data).unwrap();
1069 match event {
1070 StreamEvent::Usage(usage) => {
1071 assert_eq!(usage.prompt_tokens, 10);
1072 assert_eq!(usage.completion_tokens, 5);
1073 assert_eq!(usage.total_tokens, 15);
1074 }
1075 _ => panic!("expected Usage, got {:?}", event),
1076 }
1077 }
1078
1079 #[test]
1080 fn openai_finish_reason_is_done() {
1081 let data =
1082 r#"{"id":"chatcmpl-1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}"#;
1083 let event = parse_openai_sse(data).unwrap();
1084 assert!(matches!(event, StreamEvent::Done));
1085 }
1086
1087 #[test]
1088 fn openai_empty_delta_skipped() {
1089 let data = r#"{"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":""},"finish_reason":null}]}"#;
1090 assert!(parse_openai_sse(data).is_none());
1091 }
1092
1093 #[test]
1094 fn anthropic_content_block_delta() {
1095 let data = r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}"#;
1096 let event = parse_anthropic_sse("content_block_delta", data).unwrap();
1097 match event {
1098 StreamEvent::Delta { content } => assert_eq!(content, "Hello"),
1099 _ => panic!("expected Delta, got {:?}", event),
1100 }
1101 }
1102
1103 #[test]
1104 fn anthropic_message_delta_usage() {
1105 let data = r#"{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":5}}"#;
1106 let event = parse_anthropic_sse("message_delta", data).unwrap();
1107 match event {
1108 StreamEvent::Usage(usage) => assert_eq!(usage.completion_tokens, 5),
1109 _ => panic!("expected Usage, got {:?}", event),
1110 }
1111 }
1112
1113 #[test]
1114 fn anthropic_message_stop() {
1115 let event = parse_anthropic_sse("message_stop", r#"{"type":"message_stop"}"#).unwrap();
1116 assert!(matches!(event, StreamEvent::Done));
1117 }
1118
1119 #[test]
1120 fn anthropic_unknown_event_ignored() {
1121 assert!(parse_anthropic_sse("ping", "{}").is_none());
1122 }
1123
1124 fn sample_tool() -> ToolDefinition {
1125 ToolDefinition {
1126 name: "get_weather".into(),
1127 description: "Get weather".into(),
1128 input_schema: serde_json::json!({
1129 "type": "object",
1130 "properties": { "location": { "type": "string" } },
1131 "required": ["location"]
1132 }),
1133 }
1134 }
1135
1136 #[test]
1137 fn openai_tools_use_function_wrapper() {
1138 let out = openai_tools(&[sample_tool()]);
1139 assert_eq!(out.len(), 1);
1140 assert_eq!(out[0]["type"], "function");
1141 assert_eq!(out[0]["function"]["name"], "get_weather");
1142 assert_eq!(out[0]["function"]["parameters"]["required"][0], "location");
1144 }
1145
1146 #[test]
1147 fn openai_response_format_maps_each_variant() {
1148 assert!(openai_response_format(&ResponseFormat::Text).is_none());
1149 assert_eq!(
1150 openai_response_format(&ResponseFormat::Json).unwrap()["type"],
1151 "json_object"
1152 );
1153 let schema = serde_json::json!({"type": "object"});
1154 let js = openai_response_format(&ResponseFormat::JsonSchema { schema }).unwrap();
1155 assert_eq!(js["type"], "json_schema");
1156 assert_eq!(js["json_schema"]["strict"], true);
1157 }
1158
1159 #[test]
1160 fn anthropic_tools_use_input_schema_key() {
1161 let out = anthropic_tools(&[sample_tool()]);
1162 assert_eq!(out[0]["name"], "get_weather");
1163 assert_eq!(out[0]["input_schema"]["type"], "object");
1164 assert!(out[0].get("function").is_none());
1165 }
1166
1167 #[test]
1168 fn anthropic_output_config_only_for_json_schema() {
1169 assert!(anthropic_output_config(&ResponseFormat::Text).is_none());
1170 assert!(anthropic_output_config(&ResponseFormat::Json).is_none());
1171 let schema = serde_json::json!({"type": "object"});
1172 let cfg = anthropic_output_config(&ResponseFormat::JsonSchema { schema }).unwrap();
1173 assert_eq!(cfg["format"]["type"], "json_schema");
1174 }
1175
1176 #[test]
1177 fn openai_request_serializes_tools_and_format() {
1178 let body = OpenAIChatRequest {
1179 model: "gpt-4o".into(),
1180 messages: vec![OpenAIChatMessage {
1181 role: "user".into(),
1182 content: "hi".into(),
1183 }],
1184 temperature: 0.7,
1185 max_tokens: None,
1186 stream: false,
1187 tools: Some(openai_tools(&[sample_tool()])),
1188 response_format: Some(serde_json::json!({ "type": "json_object" })),
1189 };
1190 let json = serde_json::to_value(&body).unwrap();
1191 assert_eq!(json["tools"][0]["function"]["name"], "get_weather");
1192 assert_eq!(json["response_format"]["type"], "json_object");
1193 assert!(json.get("max_tokens").is_none());
1195 }
1196
1197 #[test]
1198 fn openai_response_parses_tool_calls() {
1199 let raw = r#"{
1200 "id": "chatcmpl-1",
1201 "created": 1700000000,
1202 "model": "gpt-4o",
1203 "choices": [{
1204 "index": 0,
1205 "message": {
1206 "role": "assistant",
1207 "content": null,
1208 "tool_calls": [{
1209 "id": "call_abc",
1210 "type": "function",
1211 "function": { "name": "get_weather", "arguments": "{\"location\":\"Paris\"}" }
1212 }]
1213 },
1214 "finish_reason": "tool_calls"
1215 }],
1216 "usage": { "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15 }
1217 }"#;
1218 let resp: OpenAIChatResponse = serde_json::from_str(raw).unwrap();
1219 assert_eq!(resp.id.as_deref(), Some("chatcmpl-1"));
1220 let choice = resp.choices.into_iter().next().unwrap();
1221 assert_eq!(choice.finish_reason.as_deref(), Some("tool_calls"));
1222 assert!(choice.message.content.is_none());
1223 assert_eq!(choice.message.tool_calls.len(), 1);
1224 assert_eq!(choice.message.tool_calls[0].function.name, "get_weather");
1225 }
1226
1227 #[test]
1228 fn openai_response_parses_glm47_reasoning_content() {
1229 let raw = r#"{
1231 "id":"chatcmpl-1","created":1700000000,"model":"glm-4.7",
1232 "choices":[{"index":0,"message":{"role":"assistant","content":null,
1233 "reasoning_content":"thinking... {\"rating\":\"Buy\"}"},"finish_reason":"stop"}],
1234 "usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15,
1235 "completion_tokens_details":{"reasoning_tokens":3}}
1236 }"#;
1237 let resp: OpenAIChatResponse = serde_json::from_str(raw).unwrap();
1238 let choice = resp.choices.into_iter().next().unwrap();
1239 assert!(choice.message.content.is_none());
1240 assert!(choice.message.reasoning_content.is_some());
1241 assert_eq!(
1242 resp.usage
1243 .unwrap()
1244 .completion_tokens_details
1245 .unwrap()
1246 .reasoning_tokens,
1247 Some(3)
1248 );
1249 }
1250
1251 #[test]
1252 fn openai_complete_promotes_reasoning_when_content_empty() {
1253 let raw = r#"{
1257 "id":"chatcmpl-1","created":1700000000,"model":"glm-4.7",
1258 "choices":[{"index":0,"message":{"role":"assistant","content":"",
1259 "reasoning_content":"사고 {\"rating\":\"Sell\",\"key_thesis\":\"약세\"}"},"finish_reason":"stop"}]
1260 }"#;
1261 let resp: OpenAIChatResponse = serde_json::from_str(raw).unwrap();
1262 let first = resp.choices.into_iter().next().unwrap();
1263 let raw_content = first.message.content.unwrap_or_default();
1264 let reasoning = first.message.reasoning_content.clone();
1265 let content = promote_reasoning_into_content(raw_content, reasoning.as_deref());
1266 assert!(
1267 content.contains("\"rating\":\"Sell\""),
1268 "promoted content must contain JSON: {content}"
1269 );
1270 assert_eq!(
1271 reasoning.as_deref(),
1272 Some("사고 {\"rating\":\"Sell\",\"key_thesis\":\"약세\"}")
1273 );
1274 }
1275
1276 #[test]
1277 fn openai_complete_keeps_content_when_present() {
1278 let raw = r#"{
1280 "id":"chatcmpl-1","created":1700000000,"model":"gpt-4o",
1281 "choices":[{"index":0,"message":{"role":"assistant","content":"answer",
1282 "reasoning_content":"thought"},"finish_reason":"stop"}]
1283 }"#;
1284 let resp: OpenAIChatResponse = serde_json::from_str(raw).unwrap();
1285 let first = resp.choices.into_iter().next().unwrap();
1286 let raw_content = first.message.content.clone().unwrap_or_default();
1287 let reasoning = first.message.reasoning_content.as_deref();
1288 let content = promote_reasoning_into_content(raw_content.clone(), reasoning);
1289 assert_eq!(content, "answer");
1290 assert_eq!(reasoning, Some("thought"));
1291 }
1292
1293 #[test]
1294 fn promote_reasoning_into_content_edge_cases() {
1295 assert_eq!(
1298 promote_reasoning_into_content(String::new(), Some("chain-of-thought")),
1299 "chain-of-thought"
1300 );
1301 assert_eq!(
1303 promote_reasoning_into_content("ans".into(), Some("cot")),
1304 "ans"
1305 );
1306 assert_eq!(promote_reasoning_into_content(String::new(), None), "");
1308 assert_eq!(promote_reasoning_into_content("x".into(), None), "x");
1309 }
1310
1311 #[test]
1312 fn openai_response_parses_deepseek_reasoning_alias() {
1313 let raw = r#"{"model":"deepseek-r1","choices":[{"message":{"content":"ans","reasoning":"thought"}}]}"#;
1315 let resp: OpenAIChatResponse = serde_json::from_str(raw).unwrap();
1316 assert_eq!(
1317 resp.choices[0].message.reasoning_content.as_deref(),
1318 Some("thought")
1319 );
1320 }
1321
1322 #[test]
1323 fn openai_sse_reasoning_delta_extracted() {
1324 let data = r#"{"choices":[{"index":0,"delta":{"reasoning_content":"thinking"},"finish_reason":null}]}"#;
1325 let event = parse_openai_sse(data).unwrap();
1326 match event {
1327 StreamEvent::ReasoningDelta { content } => assert_eq!(content, "thinking"),
1328 other => panic!("expected ReasoningDelta, got {other:?}"),
1329 }
1330 }
1331
1332 #[test]
1333 fn openai_sse_content_delta_still_works_alongside_reasoning() {
1334 let data = r#"{"choices":[{"index":0,"delta":{"content":"answer"},"finish_reason":null}]}"#;
1336 let event = parse_openai_sse(data).unwrap();
1337 assert!(matches!(event, StreamEvent::Delta { .. }));
1338 }
1339
1340 #[test]
1341 fn openai_sse_usage_carries_reasoning_tokens() {
1342 let data = r#"{"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],
1344 "usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3,
1345 "completion_tokens_details":{"reasoning_tokens":7}}}"#;
1346 let event = parse_openai_sse(data).unwrap();
1347 match event {
1348 StreamEvent::Usage(u) => assert_eq!(u.reasoning_tokens, Some(7)),
1349 other => panic!("expected Usage, got {other:?}"),
1350 }
1351 }
1352
1353 #[test]
1354 fn anthropic_response_parses_tool_use_block() {
1355 let raw = r#"{
1356 "id": "msg_1",
1357 "model": "claude-sonnet-4-6",
1358 "stop_reason": "tool_use",
1359 "content": [
1360 { "type": "text", "text": "Let me check." },
1361 { "type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": { "location": "Paris" } }
1362 ],
1363 "usage": { "input_tokens": 12, "output_tokens": 8 }
1364 }"#;
1365 let resp: AnthropicResponse = serde_json::from_str(raw).unwrap();
1366 assert_eq!(resp.stop_reason.as_deref(), Some("tool_use"));
1367 assert_eq!(resp.content.len(), 2);
1368 assert_eq!(resp.content[0].block_type, "text");
1369 assert_eq!(resp.content[1].block_type, "tool_use");
1370 assert_eq!(resp.content[1].name.as_deref(), Some("get_weather"));
1371 assert_eq!(resp.content[1].input.as_ref().unwrap()["location"], "Paris");
1372 }
1373
1374 #[test]
1375 fn anthropic_response_parses_thinking_block() {
1376 let raw = r#"{
1378 "id": "msg_2",
1379 "model": "claude-sonnet-4-6",
1380 "stop_reason": "end_turn",
1381 "content": [
1382 { "type": "thinking", "thinking": "step by step..." },
1383 { "type": "text", "text": "Final answer." }
1384 ],
1385 "usage": { "input_tokens": 5, "output_tokens": 9 }
1386 }"#;
1387 let resp: AnthropicResponse = serde_json::from_str(raw).unwrap();
1388 let mut reasoning = String::new();
1389 let mut content = String::new();
1390 for block in resp.content {
1391 match block.block_type.as_str() {
1392 "text" => {
1393 if let Some(t) = block.text {
1394 content.push_str(&t);
1395 }
1396 }
1397 "thinking" => {
1398 if let Some(t) = block.thinking {
1399 reasoning.push_str(&t);
1400 }
1401 }
1402 _ => {}
1403 }
1404 }
1405 assert_eq!(content, "Final answer.");
1406 assert_eq!(reasoning, "step by step...");
1407 }
1408
1409 #[test]
1410 fn anthropic_sse_thinking_delta_extracted() {
1411 let data = r#"{"delta":{"type":"thinking_delta","thinking":"a thought"}}"#;
1412 let event = parse_anthropic_sse("content_block_delta", data).unwrap();
1413 match event {
1414 StreamEvent::ReasoningDelta { content } => assert_eq!(content, "a thought"),
1415 other => panic!("expected ReasoningDelta, got {other:?}"),
1416 }
1417 }
1418}