@@ -676,14 +676,9 @@ buf = buf[nl + 1..].to_string(); tracing::debug!("stream line read: {}", line); - // Standard SSE uses "data: " but some providers (e.g. Kimi) - // omit the space and send "data:". Accept both. - let Some(data) = line - .strip_prefix("data: ") - .or_else(|| line.strip_prefix("data:")) - else { + let Some(data) = sse_data_payload(&line) else { continue; }; if data == "[DONE]" { tracing::debug!("stream line received [DONE]"); @@ -902,57 +897,9 @@ tracing::info!("non-stream response status: {}", v.status()); let text = Self::error_for_status_with_body(v).await?.text().await?; if text.trim().starts_with("data:") { - let mut content = String::new(); - let mut completion_tokens = 0; - let mut prompt_tokens = 0; - let mut total_tokens = 0; - for line in text.lines() { - let line = line.trim(); - if line.starts_with("data:") { - let data_part = line[5..].trim(); - if data_part == "[DONE]" { - continue; - } - if let Ok(val) = serde_json::from_str::(data_part) { - if let Some(choices) = val["choices"].as_array() { - if let Some(choice) = choices.get(0) { - if let Some(delta) = choice.get("delta") { - if let Some(chunk) = delta["content"].as_str() { - content.push_str(chunk); - } - } - if let Some(message) = choice.get("message") { - if let Some(text_content) = message["content"].as_str() { - content.push_str(text_content); - } - } - } - } - if let Some(usage) = val.get("usage") { - if let Some(ct) = usage["completion_tokens"].as_u64() { - completion_tokens = ct as usize; - } - if let Some(pt) = usage["prompt_tokens"].as_u64() { - prompt_tokens = pt as usize; - } - if let Some(tt) = usage["total_tokens"].as_u64() { - total_tokens = tt as usize; - } - } - } - } - } - Ok(( - content, - Usage { - prompt_tokens, - completion_tokens, - total_tokens, - ..Default::default() - }, - )) + Ok(complete_from_sse_text(&text)) } else { let v: Value = serde_json::from_str(&text)?; let content = v["choices"][0]["message"]["content"] .as_str() @@ -1018,38 +965,9 @@ Ok(None) } fn strip_media_urls(text: &str) -> String { - let mut result = String::new(); - let mut current_idx = 0; - while let Some(start_idx) = text[current_idx..].find("data:") { - let abs_start = current_idx + start_idx; - if let Some(comma_offset) = text[abs_start..].find(";base64,") { - if comma_offset < 100 { - let in_between = &text[abs_start + 5..abs_start + comma_offset]; - if !in_between.contains("data:") { - let base64_start = abs_start + comma_offset + 8; - let mut base64_end = base64_start; - for c in text[base64_start..].chars() { - if c.is_ascii_alphanumeric() || c == '+' || c == '/' || c == '=' { - base64_end += c.len_utf8(); - } else { - break; - } - } - if base64_end > base64_start { - result.push_str(&text[current_idx..abs_start]); - current_idx = base64_end; - continue; - } - } - } - } - result.push_str(&text[current_idx..abs_start + 5]); - current_idx = abs_start + 5; - } - result.push_str(&text[current_idx..]); - result + rewrite_media_data_urls(text, |_| Some(String::new())) } async fn embed_with_config( &self, @@ -1502,49 +1420,139 @@ parse_rerank_results(&v) } } -fn replace_media_urls_with_placeholders(text: &str, limit_bytes: usize) -> String { - let mut result = String::new(); - let mut current_idx = 0; - while let Some(start_idx) = text[current_idx..].find("data:") { - let abs_start = current_idx + start_idx; - if let Some(comma_offset) = text[abs_start..].find(";base64,") { - if comma_offset < 100 { - let mime = &text[abs_start + 5..abs_start + comma_offset]; - if !mime.contains("data:") { - let base64_start = abs_start + comma_offset + 8; - let mut base64_end = base64_start; - for c in text[base64_start..].chars() { - if c.is_ascii_alphanumeric() || c == '+' || c == '/' || c == '=' { - base64_end += c.len_utf8(); - } else { - break; - } +fn sse_data_payload(line: &str) -> Option<&str> { + let line = line.trim(); + // Standard SSE uses "data: " but some providers omit the space. + line.strip_prefix("data: ") + .or_else(|| line.strip_prefix("data:")) + .map(str::trim) +} + +fn complete_from_sse_text(text: &str) -> (String, Usage) { + let mut content = String::new(); + let mut usage = Usage::default(); + + for data in text.lines().filter_map(sse_data_payload) { + if data == "[DONE]" { + continue; + } + let Ok(val) = serde_json::from_str::(data) else { + continue; + }; + if let Some(choices) = val["choices"].as_array() { + if let Some(choice) = choices.first() { + if let Some(delta) = choice.get("delta") { + if let Some(chunk) = delta["content"].as_str() { + content.push_str(chunk); } - if base64_end > base64_start { - let b64_len = base64_end - base64_start; - let approx_size = (b64_len * 3) / 4; - if approx_size > limit_bytes { - result.push_str(&text[current_idx..abs_start]); - result.push_str(&format!( - "[Media attachment: {}, approximate size: {} bytes]", - mime, approx_size - )); - current_idx = base64_end; - continue; - } + } + if let Some(message) = choice.get("message") { + if let Some(text_content) = message["content"].as_str() { + content.push_str(text_content); } } } } - result.push_str(&text[current_idx..abs_start + 5]); - current_idx = abs_start + 5; + if let Some(value) = val.get("usage") { + usage.prompt_tokens = value["prompt_tokens"].as_u64().unwrap_or(0) as usize; + usage.completion_tokens = value["completion_tokens"].as_u64().unwrap_or(0) as usize; + usage.total_tokens = value["total_tokens"].as_u64().unwrap_or(0) as usize; + } } + + (content, usage) +} + +struct MediaDataUrl<'a> { + start: usize, + end: usize, + mime: &'a str, + approx_size: usize, +} + +struct MediaScan<'a> { + data_start: usize, + media: Option>, +} + +fn next_media_scan(text: &str, current_idx: usize) -> Option> { + let data_start = current_idx + text[current_idx..].find("data:")?; + Some(MediaScan { + data_start, + media: parse_media_data_url_at(text, data_start), + }) +} + +fn parse_media_data_url_at(text: &str, start: usize) -> Option> { + let comma_offset = text[start..].find(";base64,")?; + if comma_offset >= 100 { + return None; + } + let mime = &text[start + "data:".len()..start + comma_offset]; + if mime.contains("data:") { + return None; + } + + let base64_start = start + comma_offset + ";base64,".len(); + let mut base64_end = base64_start; + for ch in text[base64_start..].chars() { + if ch.is_ascii_alphanumeric() || ch == '+' || ch == '/' || ch == '=' { + base64_end += ch.len_utf8(); + } else { + break; + } + } + if base64_end == base64_start { + return None; + } + + let base64_len = base64_end - base64_start; + Some(MediaDataUrl { + start, + end: base64_end, + mime, + approx_size: (base64_len * 3) / 4, + }) +} + +fn rewrite_media_data_urls( + text: &str, + mut replacement: impl FnMut(&MediaDataUrl<'_>) -> Option, +) -> String { + let mut result = String::new(); + let mut current_idx = 0; + while let Some(scan) = next_media_scan(text, current_idx) { + if let Some(media) = scan.media { + if let Some(value) = replacement(&media) { + result.push_str(&text[current_idx..media.start]); + result.push_str(&value); + current_idx = media.end; + continue; + } + } + + result.push_str(&text[current_idx..scan.data_start + "data:".len()]); + current_idx = scan.data_start + "data:".len(); + } result.push_str(&text[current_idx..]); result } +fn replace_media_urls_with_placeholders(text: &str, limit_bytes: usize) -> String { + rewrite_media_data_urls(text, |media| { + if media.approx_size > limit_bytes { + Some(format!( + "[Media attachment: {}, approximate size: {} bytes]", + media.mime, media.approx_size + )) + } else { + None + } + }) +} + fn prepare_messages_for_api(messages: &[Message]) -> Vec { let mut prepared = Vec::new(); for mut message in messages.iter().cloned() { if message.normalize_for_api() { @@ -1624,8 +1632,41 @@ "check this: and this right here" ); } + #[test] + fn media_placeholders_preserve_small_attachments() { + let input = concat!( + "small data:image/png;base64,YWJj ", + "large data:audio/mp3;base64,YWJjZGVm" + ); + + assert_eq!( + replace_media_urls_with_placeholders(input, 4), + concat!( + "small data:image/png;base64,YWJj ", + "large [Media attachment: audio/mp3, approximate size: 6 bytes]" + ) + ); + } + + #[test] + fn complete_from_sse_text_accepts_sse_spacing_variants() { + let text = concat!( + "data:{\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n", + "data: {\"choices\":[{\"message\":{\"content\":\"lo\"}}],", + "\"usage\":{\"prompt_tokens\":2,\"completion_tokens\":3,\"total_tokens\":5}}\n", + "data: [DONE]\n" + ); + + let (content, usage) = complete_from_sse_text(text); + + assert_eq!(content, "hello"); + assert_eq!(usage.prompt_tokens, 2); + assert_eq!(usage.completion_tokens, 3); + assert_eq!(usage.total_tokens, 5); + } + async fn spawn_http_response(status: &str, content_type: &str, body: &str) -> String { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let status = status.to_string();