1use bytes::Bytes;
4use futures::{Stream, StreamExt};
5use reqwest::Client;
6use serde::Deserialize;
7use serde_json::Value as JsonValue;
8use std::future::Future;
9use std::pin::Pin;
10use std::sync::Arc;
11
12use super::openai_responses_shared::parse_streaming_json;
13use super::shared_client;
14use super::sse::split_complete_lines;
15use crate::{
16 Api, AssistantMessage, ContentBlock, Context, Message, MessageContent, Model, Provider,
17 ProviderEvent, StopReason, StreamOptions, StreamResult, TextContent, ThinkingContent, ToolCall,
18 ToolResultMessage, Usage, error::ProviderError,
19};
20
21fn is_zai(model: &Model) -> bool {
23 model.provider.eq_ignore_ascii_case("zai") || model.base_url.contains("api.z.ai")
24}
25
26#[derive(Clone)]
28pub struct OpenAiProvider {
29 client: &'static Client,
30 api_key: Option<String>,
31 base_url: Option<String>,
32 extra_headers: Vec<(String, String)>,
34}
35
36impl OpenAiProvider {
37 pub fn new() -> Self {
42 Self {
43 client: shared_client(),
44 api_key: None,
45 base_url: None,
46 extra_headers: Vec::new(),
47 }
48 }
49
50 pub fn with_api_key(api_key: impl Into<String>) -> Self {
52 Self {
53 client: shared_client(),
54 api_key: Some(api_key.into()),
55 base_url: None,
56 extra_headers: Vec::new(),
57 }
58 }
59
60 pub fn with_base_url(base_url: &str) -> Self {
64 Self {
65 client: shared_client(),
66 api_key: None,
67 base_url: Some(base_url.to_string()),
68 extra_headers: Vec::new(),
69 }
70 }
71
72 pub fn with_base_url_and_key(base_url: &str, api_key: Option<String>) -> Self {
76 Self {
77 client: shared_client(),
78 api_key,
79 base_url: Some(base_url.to_string()),
80 extra_headers: Vec::new(),
81 }
82 }
83
84 pub fn with_config(
88 base_url: &str,
89 api_key: Option<String>,
90 extra_headers: Vec<(String, String)>,
91 ) -> Self {
92 Self {
93 client: shared_client(),
94 api_key,
95 base_url: Some(base_url.to_string()),
96 extra_headers,
97 }
98 }
99}
100
101impl Default for OpenAiProvider {
102 fn default() -> Self {
103 Self::new()
104 }
105}
106
107impl Provider for OpenAiProvider {
108 fn stream<'a>(
109 &'a self,
110 model: &'a Model,
111 context: &'a Context,
112 options: Option<StreamOptions>,
113 ) -> Pin<Box<dyn Future<Output = StreamResult> + Send + 'a>> {
114 Box::pin(async move {
115 let options = options.unwrap_or_default();
116
117 let effective_base_url = self.base_url.as_deref().unwrap_or(&model.base_url);
119 let url = format!("{}/chat/completions", effective_base_url);
120
121 let api_key = options
123 .api_key
124 .as_ref()
125 .or(self.api_key.as_ref())
126 .ok_or_else(|| ProviderError::MissingApiKey)?;
127
128 let normalized = normalize_messages(&context.messages, &model.provider, &model.id);
130 let messages = build_messages_from_normalized(&context.system_prompt, &normalized)?;
131
132 let mut body = serde_json::json!({
134 "model": model.id,
135 "messages": messages,
136 "stream": true,
137 "stream_options": { "include_usage": true },
138 });
139
140 if let Some(temp) = options.temperature {
142 body["temperature"] = serde_json::json!(temp);
143 }
144
145 if let Some(max) = options.max_tokens {
146 body["max_completion_tokens"] = serde_json::json!(max);
149 body["max_tokens"] = serde_json::json!(max);
150 }
151
152 if !context.tools.is_empty() {
154 body["tools"] = build_tools(&context.tools)?;
155 }
156
157 if model.reasoning {
162 let openai_opts = options
163 .provider_options
164 .as_ref()
165 .and_then(|po| po.openai.as_ref());
166
167 let effort = openai_opts
168 .and_then(|o| o.reasoning_effort.clone())
169 .or_else(|| {
170 options
171 .thinking_level
172 .as_ref()
173 .and_then(|l| l.as_str().map(String::from))
174 });
175
176 if let Some(effort_str) = effort {
177 body["reasoning_effort"] = serde_json::json!(effort_str);
178 }
179 }
180
181 if is_zai(model) {
185 if model.reasoning {
186 body["enable_thinking"] = serde_json::json!(true);
187 }
188 if !context.tools.is_empty() {
189 body["tool_stream"] = serde_json::json!(true);
190 }
191 }
192
193 tracing::info!(
194 "Sending request to {} model={} body_len={} enable_thinking={} tool_stream={}",
195 url,
196 model.id,
197 body.to_string().len(),
198 body.get("enable_thinking").is_some(),
199 body.get("tool_stream").is_some()
200 );
201 tracing::debug!("Request body: {}", body.to_string());
202
203 let mut headers = reqwest::header::HeaderMap::new();
205 headers.insert(
206 reqwest::header::AUTHORIZATION,
207 format!("Bearer {}", api_key).parse().map_err(|e| {
208 ProviderError::InvalidResponse(format!("invalid bearer header: {e}"))
209 })?,
210 );
211 headers.insert(
212 reqwest::header::CONTENT_TYPE,
213 "application/json".parse().map_err(|e| {
214 ProviderError::InvalidResponse(format!("invalid header value: {e}"))
215 })?,
216 );
217
218 for (k, v) in &self.extra_headers {
220 if let (Ok(name), Ok(value)) = (
221 k.parse::<reqwest::header::HeaderName>(),
222 v.parse::<reqwest::header::HeaderValue>(),
223 ) {
224 headers.insert(name, value);
225 }
226 }
227
228 for (k, v) in &options.headers {
230 if let (Ok(name), Ok(value)) = (
231 k.parse::<reqwest::header::HeaderName>(),
232 v.parse::<reqwest::header::HeaderValue>(),
233 ) {
234 headers.insert(name, value);
235 }
236 }
237
238 let response = self
240 .client
241 .post(&url)
242 .headers(headers)
243 .json(&body)
244 .send()
245 .await
246 .map_err(ProviderError::RequestFailed)?;
247
248 if !response.status().is_success() {
249 let status = response.status();
250 let body: String = response.text().await.unwrap_or_default();
251 return Err(ProviderError::HttpError(
252 crate::error::HttpErrorDetail::new(status.as_u16(), body),
253 ));
254 }
255
256 let provider_name = model.provider.clone();
258 let model_id = model.id.clone();
259
260 let start_event = ProviderEvent::Start {
262 partial: Arc::new(AssistantMessage::new(
263 Api::OpenAiCompletions,
264 &provider_name,
265 &model_id,
266 )),
267 };
268
269 let stream = response
280 .bytes_stream()
281 .scan(
282 (
283 Vec::new(),
284 std::collections::HashMap::<usize, (String, String, String)>::new(),
285 std::collections::HashMap::<String, usize>::new(), false,
287 AssistantMessage::new(Api::OpenAiCompletions, &provider_name, &model_id),
288 ),
289 move |(
290 pending_bytes,
291 pending_tc,
292 tc_id_to_index,
293 thinking_started,
294 accumulated_output,
295 ),
296 chunk: Result<Bytes, reqwest::Error>| {
297 let events = match chunk {
298 Ok(bytes) => {
299 let mut combined =
301 Vec::with_capacity(pending_bytes.len() + bytes.len());
302 combined.extend_from_slice(pending_bytes);
303 combined.extend_from_slice(&bytes);
304
305 let (text, trailing) = split_complete_lines(&combined);
309 *pending_bytes = trailing;
310
311 tracing::debug!(
312 "parse_sse_events input: {} bytes, {} lines",
313 text.len(),
314 text.lines().count()
315 );
316 let raw_events = parse_sse_events(
317 &text,
318 &provider_name,
319 &model_id,
320 accumulated_output,
321 );
322 tracing::debug!("parse_sse_events output: {} events", raw_events.len());
323
324 let mut processed = Vec::new();
326 for event in raw_events {
327 match &event {
328 ProviderEvent::ThinkingDelta { content_index, .. } => {
329 if !*thinking_started {
331 *thinking_started = true;
332 processed.push(ProviderEvent::ThinkingStart {
333 content_index: *content_index,
334 partial: Arc::new(AssistantMessage::new(
335 Api::OpenAiCompletions,
336 &provider_name,
337 &model_id,
338 )),
339 });
340 }
341 processed.push(event);
342 }
343 ProviderEvent::ToolCallStart {
344 content_index,
345 tool_call_id,
346 tool_name,
347 ..
348 } => {
349 let entry =
350 pending_tc.entry(*content_index).or_insert_with(|| {
351 (String::new(), String::new(), String::new())
352 });
353 if let Some(id) = tool_call_id
354 && !id.is_empty()
355 {
356 entry.0 = id.clone();
357 tc_id_to_index.insert(id.clone(), *content_index);
358 }
359 if let Some(name) = tool_name
360 && !name.is_empty()
361 {
362 entry.1 = name.clone();
363 }
364 processed.push(event);
365 }
366 ProviderEvent::ToolCallDelta {
367 content_index,
368 delta,
369 ..
370 } => {
371 let idx = if pending_tc.contains_key(content_index) {
373 *content_index
374 } else {
375 tc_id_to_index
377 .values()
378 .copied()
379 .find(|i| *i == *content_index)
380 .unwrap_or(*content_index)
381 };
382 let entry = pending_tc.entry(idx).or_insert_with(|| {
383 (String::new(), String::new(), String::new())
384 });
385 tracing::debug!(
386 "[TC-DELTA] idx={}, delta_len={}, accumulated_len={}",
387 idx,
388 delta.len(),
389 entry.2.len() + delta.len()
390 );
391 entry.2.push_str(delta);
392 processed.push(event);
393 }
394 ProviderEvent::ToolCallEnd { .. } => {
395 processed.push(event);
397 }
398 ProviderEvent::Done { reason, .. } => {
399 if matches!(reason, StopReason::ToolUse) {
401 let mut indices: Vec<usize> =
402 pending_tc.keys().copied().collect();
403 indices.sort();
404 for idx in indices {
405 let (id, name, arguments) = &pending_tc[&idx];
406 tracing::debug!(
407 "[TC-END] idx={}, id={}, name={}, args_len={}",
408 idx,
409 id.len(),
410 name.len(),
411 arguments.len()
412 );
413 let args_value = parse_streaming_json(arguments);
414 processed.push(ProviderEvent::ToolCallEnd {
415 content_index: idx,
416 tool_call: crate::ToolCall {
417 content_type:
418 crate::messages::ToolCallType::ToolCall,
419 id: id.clone(),
420 name: name.clone(),
421 arguments: args_value,
422 thought_signature: None,
423 },
424 partial: Arc::new(AssistantMessage::new(
425 Api::OpenAiCompletions,
426 &provider_name,
427 &model_id,
428 )),
429 });
430 }
431 }
432 pending_tc.clear();
436 tc_id_to_index.clear();
437 processed.push(event);
438 }
439 _ => {
440 processed.push(event);
441 }
442 }
443 }
444 processed
445 }
446 Err(e) => {
447 vec![ProviderEvent::Error {
448 reason: StopReason::Error,
449 error: create_error_message(
450 &e.to_string(),
451 &provider_name,
452 &model_id,
453 ),
454 }]
455 }
456 };
457 async move { Some(futures::stream::iter(events)) }
459 },
460 )
461 .flatten();
462
463 let stream_with_start = futures::stream::once(async move { start_event }).chain(stream);
465 Ok(Box::pin(stream_with_start) as Pin<Box<dyn Stream<Item = ProviderEvent> + Send>>)
466 })
467 }
468}
469
470fn build_messages_from_normalized(
476 system_prompt: &Option<String>,
477 messages: &[Message],
478) -> Result<Vec<JsonValue>, ProviderError> {
479 let mut result = Vec::new();
480
481 if let Some(prompt) = system_prompt {
483 result.push(serde_json::json!({
484 "role": "system",
485 "content": prompt,
486 }));
487 }
488
489 for msg in messages {
491 match msg {
492 Message::User(u) => {
493 let content: String = match &u.content {
494 MessageContent::Text(s) => s.clone(),
495 MessageContent::Blocks(blocks) => blocks_to_content(blocks)?.to_string(),
496 };
497 result.push(serde_json::json!({
498 "role": "user",
499 "content": content,
500 }));
501 }
502 Message::Assistant(a) => {
503 let mut text_parts = Vec::new();
505 let mut tool_calls = Vec::new();
506 for block in &a.content {
507 match block {
508 ContentBlock::Text(t) => {
509 text_parts.push(t.text.clone());
510 }
511 ContentBlock::Thinking(_) => {
512 }
514 ContentBlock::ToolCall(tc) => {
515 tool_calls.push(serde_json::json!({
516 "id": tc.id,
517 "type": "function",
518 "function": {
519 "name": tc.name,
520 "arguments": tc.arguments.to_string(),
521 },
522 }));
523 }
524 ContentBlock::Image(_) | ContentBlock::Unknown(_) => {}
525 }
526 }
527
528 let mut msg_obj = serde_json::json!({
529 "role": "assistant",
530 "content": text_parts.join(""),
531 });
532 if !tool_calls.is_empty() {
533 msg_obj["tool_calls"] = serde_json::json!(tool_calls);
534 }
535 result.push(msg_obj);
536 }
537 Message::ToolResult(t) => {
538 let result_text: String = t
539 .content
540 .iter()
541 .filter_map(|b| b.as_text())
542 .collect::<Vec<_>>()
543 .join("");
544 result.push(serde_json::json!({
545 "role": "tool",
546 "tool_call_id": t.tool_call_id,
547 "content": result_text,
548 }));
549 }
550 }
551 }
552
553 Ok(result)
554}
555
556fn blocks_to_content(blocks: &[ContentBlock]) -> Result<JsonValue, ProviderError> {
558 if blocks.len() == 1
559 && let Some(text) = blocks[0].as_text()
560 {
561 return Ok(JsonValue::String(text.to_string()));
562 }
563
564 let items: Result<Vec<_>, _> = blocks
565 .iter()
566 .map(|block| match block {
567 ContentBlock::Text(t) => Ok(serde_json::json!({
568 "type": "text",
569 "text": t.text,
570 })),
571 ContentBlock::ToolCall(tc) => Ok(serde_json::json!({
572 "type": "function",
573 "id": tc.id,
574 "function": {
575 "name": tc.name,
576 "arguments": tc.arguments.to_string(),
577 },
578 })),
579 ContentBlock::Thinking(th) => Ok(serde_json::json!({
580 "type": "thinking",
581 "thinking": th.thinking,
582 })),
583 ContentBlock::Image(img) => Ok(serde_json::json!({
584 "type": "image_url",
585 "image_url": {
586 "url": format!("data:{};base64,{}", img.mime_type, img.data),
587 },
588 })),
589 ContentBlock::Unknown(_) => Err(ProviderError::InvalidResponse(
590 "Unknown content block type".into(),
591 )),
592 })
593 .collect();
594
595 Ok(serde_json::json!(items?))
596}
597
598fn build_tools(tools: &[crate::Tool]) -> Result<JsonValue, ProviderError> {
600 let items: Vec<_> = tools
601 .iter()
602 .map(|tool| {
603 serde_json::json!({
604 "type": "function",
605 "function": {
606 "name": tool.name,
607 "description": tool.description,
608 "parameters": tool.parameters,
609 },
610 })
611 })
612 .collect();
613
614 Ok(serde_json::json!(items))
615}
616
617fn parse_sse_events(
628 text: &str,
629 _provider: &str,
630 _model_id: &str,
631 output: &mut AssistantMessage,
632) -> Vec<ProviderEvent> {
633 let mut events = Vec::with_capacity(text.len() / 80);
641
642 let mut accumulated_usage = Usage::default();
643
644 for line in text.split('\n') {
645 let line = line.trim_end_matches('\r');
646 if line.is_empty() {
647 continue;
648 }
649
650 if !line.starts_with("data: ") {
652 continue;
653 }
654
655 let data = &line[6..]; if data == "[DONE]" {
659 break;
660 }
661
662 if data.is_empty() {
663 continue;
664 }
665
666 let chunk = match serde_json::from_str::<SSEChunk>(data) {
667 Ok(c) => c,
668 Err(_) => continue,
669 };
670
671 if let Some(chunk_usage) = &chunk.usage {
677 accumulated_usage.input = chunk_usage.prompt_tokens.max(accumulated_usage.input);
678 accumulated_usage.output = chunk_usage.completion_tokens.max(accumulated_usage.output);
679 accumulated_usage.cache_read = chunk_usage
680 .prompt_tokens_details
681 .as_ref()
682 .map(|d| d.cached_tokens)
683 .unwrap_or(0)
684 .max(accumulated_usage.cache_read);
685 accumulated_usage.total_tokens =
686 chunk_usage.total_tokens.max(accumulated_usage.total_tokens);
687 }
688
689 for choice in &chunk.choices {
690 if let Some(delta) = &choice.delta {
691 if let Some(content) = &delta.content {
692 let last_text_idx = output
694 .content
695 .iter()
696 .rposition(|b| matches!(b, ContentBlock::Text(_)));
697 if let Some(idx) = last_text_idx
698 && let ContentBlock::Text(t) = &mut output.content[idx]
699 {
700 t.text.push_str(content);
701 } else {
702 output
703 .content
704 .push(ContentBlock::Text(TextContent::new(content.clone())));
705 }
706 events.push(ProviderEvent::TextDelta {
707 content_index: choice.index,
708 delta: content.clone(),
709 partial: Arc::new(output.clone()),
710 });
711 }
712
713 if let Some(ref reasoning) = delta.reasoning_content
715 && !reasoning.is_empty()
716 {
717 let last_think_idx = output
719 .content
720 .iter()
721 .rposition(|b| matches!(b, ContentBlock::Thinking(_)));
722 if let Some(idx) = last_think_idx
723 && let ContentBlock::Thinking(t) = &mut output.content[idx]
724 {
725 t.thinking.push_str(reasoning);
726 } else {
727 output
728 .content
729 .push(ContentBlock::Thinking(ThinkingContent::new(
730 reasoning.clone(),
731 )));
732 }
733 events.push(ProviderEvent::ThinkingDelta {
734 content_index: choice.index,
735 delta: reasoning.clone(),
736 partial: Arc::new(output.clone()),
737 });
738 }
739
740 if let Some(tool_calls) = &delta.tool_calls {
741 for tc in tool_calls {
742 let tc_index = tc.index.unwrap_or(choice.index);
743
744 if tc.id.is_some()
746 || tc.function.as_ref().and_then(|f| f.name.as_ref()).is_some()
747 {
748 events.push(ProviderEvent::ToolCallStart {
749 content_index: tc_index,
750 tool_call_id: tc.id.clone(),
751 tool_name: tc.function.as_ref().and_then(|f| f.name.clone()),
752 partial: Arc::new(output.clone()),
753 });
754 }
755
756 if let Some(func) = &tc.function {
758 events.push(ProviderEvent::ToolCallDelta {
759 content_index: tc_index,
760 delta: func.arguments.clone().unwrap_or_default(),
761 partial: Arc::new(output.clone()),
762 });
763 }
764 }
765 }
766 }
767
768 if choice.finish_reason.is_some() {
769 let reason = match choice.finish_reason.as_deref() {
770 Some("stop") | Some("end") => StopReason::Stop,
771 Some("length") => StopReason::Length,
772 Some("tool_calls") | Some("function_call") => StopReason::ToolUse,
773 Some("content_filter") => StopReason::Error,
774 Some(unknown) => {
775 tracing::warn!("Unknown finish_reason: '{}', treating as Error", unknown);
776 StopReason::Error
777 }
778 None => StopReason::Stop,
779 };
780 tracing::info!("finish_reason={:?} → {:?}", choice.finish_reason, reason);
781
782 let mut done_msg = output.clone();
783 done_msg.stop_reason = reason;
784 done_msg.usage = accumulated_usage.clone();
785 events.push(ProviderEvent::Done {
786 reason,
787 message: done_msg,
788 });
789 }
790 }
791 }
792
793 events
794}
795
796fn create_error_message(msg: &str, provider: &str, model_id: &str) -> AssistantMessage {
798 let mut message = AssistantMessage::new(Api::OpenAiCompletions, provider, model_id);
799 message.stop_reason = StopReason::Error;
800 message.error_message = Some(msg.to_string());
801 message
802}
803
804#[derive(Debug, Deserialize)]
806struct SSEChunk {
808 _id: Option<String>,
809 #[serde(rename = "model")]
810 _model: Option<String>,
811 choices: Vec<Choice>,
812 usage: Option<UsageInfo>,
813}
814
815#[derive(Debug, Deserialize)]
816struct Choice {
818 index: usize,
819 delta: Option<Delta>,
820 finish_reason: Option<String>,
821}
822
823#[derive(Debug, Deserialize)]
824struct Delta {
825 content: Option<String>,
826 reasoning_content: Option<String>,
827 tool_calls: Option<Vec<ToolCallDelta>>,
828}
829
830#[derive(Debug, Deserialize)]
831struct ToolCallDelta {
833 index: Option<usize>,
834 id: Option<String>,
835 #[serde(rename = "type")]
836 _type_: Option<String>,
837 function: Option<FunctionDelta>,
838}
839
840#[derive(Debug, Deserialize)]
841struct FunctionDelta {
843 name: Option<String>,
844 arguments: Option<String>,
845}
846
847#[derive(Debug, Deserialize, Clone)]
848struct UsageInfo {
849 prompt_tokens: usize,
850 completion_tokens: usize,
851 total_tokens: usize,
852 #[serde(rename = "prompt_tokens_details")]
853 prompt_tokens_details: Option<PromptTokensDetails>,
854}
855
856#[derive(Debug, Deserialize, Clone)]
857struct PromptTokensDetails {
858 #[serde(rename = "cached_tokens")]
859 cached_tokens: usize,
860}
861
862pub fn normalize_messages(messages: &[Message], provider: &str, model_id: &str) -> Vec<Message> {
882 let provider_lower = provider.to_lowercase();
883 let model_lower = model_id.to_lowercase();
884
885 let is_anthropic = provider_lower == "anthropic"
886 || provider_lower.contains("vertex-anthropic")
887 || provider_lower.contains("google-vertex") && model_lower.contains("claude");
888
889 let is_mistral = provider_lower == "mistral"
890 || model_lower.contains("mistral")
891 || model_lower.contains("devstral");
892
893 let is_deepseek = model_lower.contains("deepseek");
894
895 let is_claude = model_lower.contains("claude");
896
897 let needs_tool_id_scrub = is_mistral || is_claude;
898
899 let messages: Vec<Message> = messages.to_vec();
900
901 let messages = if is_deepseek {
903 messages
904 .iter()
905 .map(|msg| match msg {
906 Message::Assistant(a) => {
907 let has_reasoning = a
908 .content
909 .iter()
910 .any(|b| matches!(b, ContentBlock::Thinking(_)));
911 if has_reasoning {
912 Message::Assistant(a.clone())
913 } else {
914 let mut new_content = a.content.clone();
915 new_content
916 .push(ContentBlock::Thinking(ThinkingContent::new(String::new())));
917 let mut new_msg = AssistantMessage::new(a.api, &a.provider, &a.model);
918 new_msg.content = new_content;
919 new_msg.usage = a.usage.clone();
920 new_msg.stop_reason = a.stop_reason;
921 new_msg.response_id = a.response_id.clone();
922 new_msg.timestamp = a.timestamp;
923 Message::Assistant(new_msg)
924 }
925 }
926 _ => msg.clone(),
927 })
928 .collect()
929 } else {
930 messages
931 };
932
933 let messages = if is_anthropic {
935 messages
936 .iter()
937 .filter_map(|msg| match msg {
938 Message::User(u) => {
939 let filtered = filter_empty_content(&u.content);
940 filtered.as_ref()?;
941 #[allow(clippy::expect_used)]
945 Some(Message::User(crate::UserMessage {
946 role: u.role,
947 content: filtered.expect("checked above"),
948 timestamp: u.timestamp,
949 }))
950 }
951 Message::Assistant(a) => {
952 let filtered: Vec<ContentBlock> = a
953 .content
954 .iter()
955 .filter(|b| !is_empty_content_block(b))
956 .cloned()
957 .collect();
958 if filtered.is_empty() {
959 return None; }
961 let mut new_msg = AssistantMessage::new(a.api, &a.provider, &a.model);
962 new_msg.content = filtered;
963 new_msg.usage = a.usage.clone();
964 new_msg.stop_reason = a.stop_reason;
965 new_msg.response_id = a.response_id.clone();
966 new_msg.timestamp = a.timestamp;
967 Some(Message::Assistant(new_msg))
968 }
969 Message::ToolResult(t) => Some(Message::ToolResult(t.clone())),
970 })
971 .collect()
972 } else {
973 messages
974 };
975
976 let messages = if is_anthropic {
978 messages
979 .iter()
980 .flat_map(|msg| match msg {
981 Message::Assistant(a) => {
982 let parts = &a.content;
983 let has_tool = parts.iter().any(|b| matches!(b, ContentBlock::ToolCall(_)));
984 let has_non_tool = parts
985 .iter()
986 .any(|b| !matches!(b, ContentBlock::ToolCall(_)));
987
988 if has_tool && has_non_tool {
989 let non_tools: Vec<ContentBlock> = parts
991 .iter()
992 .filter(|b| !matches!(b, ContentBlock::ToolCall(_)))
993 .cloned()
994 .collect();
995 let tools: Vec<ContentBlock> = parts
996 .iter()
997 .filter(|b| matches!(b, ContentBlock::ToolCall(_)))
998 .cloned()
999 .collect();
1000
1001 let mut msg1 = AssistantMessage::new(a.api, &a.provider, &a.model);
1002 msg1.content = non_tools;
1003 msg1.usage = a.usage.clone();
1004 msg1.stop_reason = a.stop_reason;
1005 msg1.response_id = a.response_id.clone();
1006 msg1.timestamp = a.timestamp;
1007
1008 let mut msg2 = AssistantMessage::new(a.api, &a.provider, &a.model);
1009 msg2.content = tools;
1010 msg2.usage = a.usage.clone();
1011 msg2.stop_reason = a.stop_reason;
1012 msg2.response_id = a.response_id.clone();
1013 msg2.timestamp = a.timestamp;
1014
1015 vec![Message::Assistant(msg1), Message::Assistant(msg2)]
1016 } else {
1017 vec![Message::Assistant(a.clone())]
1018 }
1019 }
1020 _ => vec![msg.clone()],
1021 })
1022 .collect()
1023 } else {
1024 messages
1025 };
1026
1027 if needs_tool_id_scrub {
1030 messages
1031 .iter()
1032 .map(|msg| match msg {
1033 Message::Assistant(a) => {
1034 let new_content: Vec<ContentBlock> = a
1035 .content
1036 .iter()
1037 .map(|block| match block {
1038 ContentBlock::ToolCall(tc) => ContentBlock::ToolCall(ToolCall::new(
1039 scrub_tool_id(&tc.id, is_mistral, is_anthropic),
1040 tc.name.clone(),
1041 tc.arguments.clone(),
1042 )),
1043 _ => block.clone(),
1044 })
1045 .collect();
1046 let mut new_msg = AssistantMessage::new(a.api, &a.provider, &a.model);
1047 new_msg.content = new_content;
1048 new_msg.usage = a.usage.clone();
1049 new_msg.stop_reason = a.stop_reason;
1050 new_msg.response_id = a.response_id.clone();
1051 new_msg.timestamp = a.timestamp;
1052 Message::Assistant(new_msg)
1053 }
1054 Message::ToolResult(t) => Message::ToolResult(ToolResultMessage::new(
1055 scrub_tool_id(&t.tool_call_id, is_mistral, is_anthropic),
1056 &t.tool_name,
1057 t.content.clone(),
1058 )),
1059 _ => msg.clone(),
1060 })
1061 .collect()
1062 } else {
1063 messages
1064 }
1065}
1066
1067fn is_empty_content_block(block: &ContentBlock) -> bool {
1069 match block {
1070 ContentBlock::Text(t) => t.text.trim().is_empty(),
1071 ContentBlock::Thinking(t) => t.thinking.trim().is_empty(),
1072 _ => false,
1073 }
1074}
1075
1076fn filter_empty_content(content: &MessageContent) -> Option<MessageContent> {
1079 match content {
1080 MessageContent::Text(s) => {
1081 if s.trim().is_empty() {
1082 None
1083 } else {
1084 Some(MessageContent::Text(s.clone()))
1085 }
1086 }
1087 MessageContent::Blocks(blocks) => {
1088 let filtered: Vec<ContentBlock> = blocks
1089 .iter()
1090 .filter(|b| !is_empty_content_block(b))
1091 .cloned()
1092 .collect();
1093 if filtered.is_empty() {
1094 None
1095 } else {
1096 Some(MessageContent::Blocks(filtered))
1097 }
1098 }
1099 }
1100}
1101
1102fn scrub_tool_id(id: &str, is_mistral: bool, is_anthropic: bool) -> String {
1106 if is_mistral {
1107 let alphanumeric: String = id.chars().filter(|c| c.is_alphanumeric()).take(9).collect();
1108 if alphanumeric.len() < 9 {
1109 format!("{}{}", alphanumeric, "0".repeat(9 - alphanumeric.len()))
1110 } else {
1111 alphanumeric
1112 }
1113 } else if is_anthropic {
1114 id.chars()
1117 .map(|c| {
1118 if c.is_alphanumeric() || c == '_' || c == '-' {
1119 c
1120 } else {
1121 '_'
1122 }
1123 })
1124 .take(64)
1125 .collect()
1126 } else {
1127 id.chars()
1128 .filter(|c| c.is_alphanumeric() || *c == '_')
1129 .collect()
1130 }
1131}
1132
1133#[cfg(test)]
1134mod tests {
1135 use super::*;
1136
1137 const PROVIDER: &str = "openai";
1138 const MODEL: &str = "gpt-4o";
1139
1140 fn parse_sse(sse: &str) -> Vec<ProviderEvent> {
1141 let mut output = AssistantMessage::new(Api::OpenAiCompletions, PROVIDER, MODEL);
1142 parse_sse_events(sse, PROVIDER, MODEL, &mut output)
1143 }
1144
1145 #[test]
1148 fn parse_single_text_event() {
1149 let sse = "data: {\"id\":\"chatcmpl-1\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hello\"}}]}\n\n";
1150 let events = parse_sse(sse);
1151 assert_eq!(events.len(), 1);
1152 match &events[0] {
1153 ProviderEvent::TextDelta {
1154 delta,
1155 content_index,
1156 ..
1157 } => {
1158 assert_eq!(delta, "Hello");
1159 assert_eq!(*content_index, 0);
1160 }
1161 other => panic!("expected TextDelta, got {other:?}"),
1162 }
1163 }
1164
1165 #[test]
1166 fn parse_multiple_text_events() {
1167 let sse = concat!(
1168 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hel\"}}]}\n",
1169 "\n",
1170 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"lo!\"}}]}\n",
1171 "\n"
1172 );
1173 let events = parse_sse(sse);
1174 assert_eq!(events.len(), 2);
1175 let texts: Vec<&str> = events
1176 .iter()
1177 .filter_map(|e| match e {
1178 ProviderEvent::TextDelta { delta, .. } => Some(delta.as_str()),
1179 _ => None,
1180 })
1181 .collect();
1182 assert_eq!(texts, vec!["Hel", "lo!"]);
1183 }
1184
1185 #[test]
1186 fn parse_done_terminator() {
1187 let sse = concat!(
1188 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"X\"}}]}\n",
1189 "\n",
1190 "data: [DONE]\n",
1191 "\n",
1192 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"NEVER\"}}]}\n"
1193 );
1194 let events = parse_sse(sse);
1195 assert_eq!(events.len(), 1);
1197 match &events[0] {
1198 ProviderEvent::TextDelta { delta, .. } => assert_eq!(delta, "X"),
1199 other => panic!("expected TextDelta, got {other:?}"),
1200 }
1201 }
1202
1203 #[test]
1206 fn parse_finish_reason_stop() {
1207 let sse = "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":null,\"finish_reason\":\"stop\"}]}\n\n";
1208 let events = parse_sse(sse);
1209 assert_eq!(events.len(), 1);
1210 match &events[0] {
1211 ProviderEvent::Done { reason, .. } => assert!(matches!(reason, StopReason::Stop)),
1212 other => panic!("expected Done, got {other:?}"),
1213 }
1214 }
1215
1216 #[test]
1217 fn parse_finish_reason_length() {
1218 let sse = "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":null,\"finish_reason\":\"length\"}]}\n\n";
1219 let events = parse_sse(sse);
1220 match &events[0] {
1221 ProviderEvent::Done { reason, .. } => assert!(matches!(reason, StopReason::Length)),
1222 other => panic!("expected Done with Length, got {other:?}"),
1223 }
1224 }
1225
1226 #[test]
1227 fn parse_finish_reason_tool_calls() {
1228 let sse = "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":null,\"finish_reason\":\"tool_calls\"}]}\n\n";
1229 let events = parse_sse(sse);
1230 match &events[0] {
1231 ProviderEvent::Done { reason, .. } => assert!(matches!(reason, StopReason::ToolUse)),
1232 other => panic!("expected Done with ToolUse, got {other:?}"),
1233 }
1234 }
1235
1236 #[test]
1239 fn parse_tool_call_deltas() {
1240 let sse = concat!(
1241 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"get_weather\",\"arguments\":\"\"}}]}}]}\n",
1242 "\n",
1243 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"city\\\":\\\"SF\\\"}\"}}]}}]}\n",
1244 "\n"
1245 );
1246 let events = parse_sse(sse);
1247 assert_eq!(events.len(), 3);
1250 let starts: Vec<&str> = events
1251 .iter()
1252 .filter_map(|e| match e {
1253 ProviderEvent::ToolCallStart { tool_name, .. } => tool_name.as_deref(),
1254 _ => None,
1255 })
1256 .collect();
1257 assert_eq!(starts, vec!["get_weather"]);
1258 let deltas: Vec<&str> = events
1259 .iter()
1260 .filter_map(|e| match e {
1261 ProviderEvent::ToolCallDelta { delta, .. } => Some(delta.as_str()),
1262 _ => None,
1263 })
1264 .collect();
1265 assert_eq!(deltas, vec!["", "{\"city\":\"SF\"}"]);
1266 }
1267
1268 #[test]
1269 fn parse_tool_call_with_no_arguments_field() {
1270 let sse = "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"name\":\"run\"}}]}}]}\n\n";
1272 let events = parse_sse(sse);
1273 assert_eq!(events.len(), 2);
1274 match &events[0] {
1275 ProviderEvent::ToolCallStart { tool_name, .. } => {
1276 assert_eq!(tool_name.as_deref(), Some("run"));
1277 }
1278 other => panic!("expected ToolCallStart, got {other:?}"),
1279 }
1280 match &events[1] {
1281 ProviderEvent::ToolCallDelta { delta, .. } => assert_eq!(delta, ""),
1282 other => panic!("expected ToolCallDelta, got {other:?}"),
1283 }
1284 }
1285
1286 #[test]
1289 fn parse_usage_in_chunk() {
1290 let sse = concat!(
1293 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":8,\"total_tokens\":18,\"prompt_tokens_details\":{\"cached_tokens\":3}}}\n",
1294 "\n",
1295 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":null,\"finish_reason\":\"stop\"}]}\n"
1296 );
1297 let events = parse_sse(sse);
1298 assert_eq!(events.len(), 2);
1300 match &events[1] {
1301 ProviderEvent::Done { message, .. } => {
1302 assert_eq!(message.usage.input, 10);
1303 assert_eq!(message.usage.output, 8);
1304 assert_eq!(message.usage.total_tokens, 18);
1305 assert_eq!(message.usage.cache_read, 3);
1306 }
1307 other => panic!("expected Done, got {other:?}"),
1308 }
1309 }
1310
1311 #[test]
1312 fn parse_usage_without_cache_details() {
1313 let sse = concat!(
1315 "data: {\"id\":\"c\",\"choices\":[],\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":2,\"total_tokens\":7}}\n",
1316 "\n",
1317 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":null,\"finish_reason\":\"stop\"}]}\n"
1318 );
1319 let events = parse_sse(sse);
1320 match &events[0] {
1321 ProviderEvent::Done { message, .. } => {
1322 assert_eq!(message.usage.input, 5);
1323 assert_eq!(message.usage.output, 2);
1324 assert_eq!(message.usage.cache_read, 0);
1325 }
1326 other => panic!("expected Done, got {other:?}"),
1327 }
1328 }
1329
1330 #[test]
1333 fn parse_empty_input() {
1334 let events = parse_sse("");
1335 assert!(events.is_empty());
1336 }
1337
1338 #[test]
1339 fn parse_only_empty_lines() {
1340 let events = parse_sse("\n\n\n");
1341 assert!(events.is_empty());
1342 }
1343
1344 #[test]
1345 fn parse_malformed_json_after_data() {
1346 let sse = "data: {not json at all}\ndata: also bad\ndata: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"}}]}\n";
1347 let events = parse_sse(sse);
1348 assert_eq!(events.len(), 1);
1350 match &events[0] {
1351 ProviderEvent::TextDelta { delta, .. } => assert_eq!(delta, "ok"),
1352 other => panic!("expected TextDelta, got {other:?}"),
1353 }
1354 }
1355
1356 #[test]
1357 fn parse_empty_data_line() {
1358 let sse = "data: \ndata: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"X\"}}]}\n";
1359 let events = parse_sse(sse);
1360 assert_eq!(events.len(), 1);
1361 }
1362
1363 #[test]
1364 fn parse_non_data_lines_ignored() {
1365 let sse = "event: ping\nid: 42\nretry: 5000\ndata: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Y\"}}]}\n";
1366 let events = parse_sse(sse);
1367 assert_eq!(events.len(), 1);
1368 }
1369
1370 #[test]
1371 fn parse_carriage_return_line_endings() {
1372 let sse = "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"CR\"}}]}\r\n\r\n";
1373 let events = parse_sse(sse);
1374 assert_eq!(events.len(), 1);
1375 match &events[0] {
1376 ProviderEvent::TextDelta { delta, .. } => assert_eq!(delta, "CR"),
1377 other => panic!("expected TextDelta, got {other:?}"),
1378 }
1379 }
1380
1381 #[test]
1384 fn parse_full_stream_with_text_tool_and_done() {
1385 let sse = concat!(
1386 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Let me\"}}]}\n",
1387 "\n",
1388 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" check\"}}]}\n",
1389 "\n",
1390 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"search\",\"arguments\":\"{\\\"q\\\":\\\"rust\\\"}\"}}]}}]}\n",
1391 "\n",
1392 "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":null,\"finish_reason\":\"tool_calls\"}]}\n",
1393 "\n",
1394 "data: [DONE]\n"
1395 );
1396 let events = parse_sse(sse);
1397 assert_eq!(events.len(), 5); let mut text_count = 0;
1400 let mut tc_start_count = 0;
1401 let mut tc_delta_count = 0;
1402 let mut done_count = 0;
1403 for e in &events {
1404 match e {
1405 ProviderEvent::TextDelta { .. } => text_count += 1,
1406 ProviderEvent::ToolCallStart { .. } => tc_start_count += 1,
1407 ProviderEvent::ToolCallDelta { .. } => tc_delta_count += 1,
1408 ProviderEvent::Done { reason, .. } => {
1409 done_count += 1;
1410 assert!(matches!(reason, StopReason::ToolUse));
1411 }
1412 other => panic!("unexpected event: {other:?}"),
1413 }
1414 }
1415 assert_eq!(text_count, 2);
1416 assert_eq!(tc_start_count, 1);
1417 assert_eq!(tc_delta_count, 1);
1418 assert_eq!(done_count, 1);
1419 }
1420}