mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-24 12:44:18 +08:00
Jason
bf40b0138c
- Add usage_daily_rollups table (schema v6) to aggregate proxy request logs into daily summaries, reducing query overhead for statistics - Add rollup_and_prune DAO that aggregates old detail logs (>N days) into rollup rows and deletes the originals - Update all usage stats queries to UNION detail logs with rollup data - Introduce incremental auto-vacuum for SQLite, with startup and periodic cleanup of old stream_check_logs and request log rollups - Split backup export/import into full vs sync variants: WebDAV sync now skips local-only table data (proxy_request_logs, stream_check_logs, provider_health, proxy_live_backup, usage_daily_rollups) while preserving them on import - Add enable_logging guard to skip request log writes when disabled - Apply cargo fmt formatting fixes across multiple modules
745 lines
40 KiB
Rust
745 lines
40 KiB
Rust
//! 流式响应转换模块
|
||
//!
|
||
//! 实现 OpenAI SSE → Anthropic SSE 格式转换
|
||
|
||
use bytes::Bytes;
|
||
use futures::stream::{Stream, StreamExt};
|
||
use serde::{Deserialize, Serialize};
|
||
use serde_json::json;
|
||
use std::collections::{HashMap, HashSet};
|
||
|
||
/// OpenAI 流式响应数据结构
|
||
#[derive(Debug, Deserialize)]
|
||
struct OpenAIStreamChunk {
|
||
id: String,
|
||
model: String,
|
||
choices: Vec<StreamChoice>,
|
||
#[serde(default)]
|
||
usage: Option<Usage>,
|
||
}
|
||
|
||
#[derive(Debug, Deserialize)]
|
||
struct StreamChoice {
|
||
delta: Delta,
|
||
#[serde(default)]
|
||
finish_reason: Option<String>,
|
||
}
|
||
|
||
#[derive(Debug, Deserialize)]
|
||
struct Delta {
|
||
#[serde(default)]
|
||
content: Option<String>,
|
||
#[serde(default)]
|
||
reasoning: Option<String>, // OpenRouter 的推理内容
|
||
#[serde(default)]
|
||
tool_calls: Option<Vec<DeltaToolCall>>,
|
||
}
|
||
|
||
#[derive(Debug, Deserialize, Serialize)]
|
||
struct DeltaToolCall {
|
||
index: usize,
|
||
#[serde(default)]
|
||
id: Option<String>,
|
||
#[serde(rename = "type", default)]
|
||
call_type: Option<String>,
|
||
#[serde(default)]
|
||
function: Option<DeltaFunction>,
|
||
}
|
||
|
||
#[derive(Debug, Deserialize, Serialize)]
|
||
struct DeltaFunction {
|
||
#[serde(default)]
|
||
name: Option<String>,
|
||
#[serde(default)]
|
||
arguments: Option<String>,
|
||
}
|
||
|
||
/// OpenAI 流式响应的 usage 信息（完整版）
|
||
#[derive(Debug, Deserialize)]
|
||
struct Usage {
|
||
#[serde(default)]
|
||
prompt_tokens: u32,
|
||
#[serde(default)]
|
||
completion_tokens: u32,
|
||
#[serde(default)]
|
||
prompt_tokens_details: Option<PromptTokensDetails>,
|
||
/// Some compatible servers return Anthropic-style cache fields directly
|
||
#[serde(default)]
|
||
cache_read_input_tokens: Option<u32>,
|
||
#[serde(default)]
|
||
cache_creation_input_tokens: Option<u32>,
|
||
}
|
||
|
||
/// Nested token details from OpenAI format
|
||
#[derive(Debug, Deserialize)]
|
||
struct PromptTokensDetails {
|
||
#[serde(default)]
|
||
cached_tokens: u32,
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
struct ToolBlockState {
|
||
anthropic_index: u32,
|
||
id: String,
|
||
name: String,
|
||
started: bool,
|
||
pending_args: String,
|
||
}
|
||
|
||
/// 创建 Anthropic SSE 流
|
||
pub fn create_anthropic_sse_stream(
|
||
stream: impl Stream<Item = Result<Bytes, reqwest::Error>> + Send + 'static,
|
||
) -> impl Stream<Item = Result<Bytes, std::io::Error>> + Send {
|
||
async_stream::stream! {
|
||
let mut buffer = String::new();
|
||
let mut message_id = None;
|
||
let mut current_model = None;
|
||
let mut next_content_index: u32 = 0;
|
||
let mut has_sent_message_start = false;
|
||
let mut current_non_tool_block_type: Option<&'static str> = None;
|
||
let mut current_non_tool_block_index: Option<u32> = None;
|
||
let mut tool_blocks_by_index: HashMap<usize, ToolBlockState> = HashMap::new();
|
||
let mut open_tool_block_indices: HashSet<u32> = HashSet::new();
|
||
|
||
tokio::pin!(stream);
|
||
|
||
while let Some(chunk) = stream.next().await {
|
||
match chunk {
|
||
Ok(bytes) => {
|
||
let text = String::from_utf8_lossy(&bytes);
|
||
buffer.push_str(&text);
|
||
|
||
while let Some(pos) = buffer.find("\n\n") {
|
||
let line = buffer[..pos].to_string();
|
||
buffer = buffer[pos + 2..].to_string();
|
||
|
||
if line.trim().is_empty() {
|
||
continue;
|
||
}
|
||
|
||
for l in line.lines() {
|
||
if let Some(data) = l.strip_prefix("data: ") {
|
||
if data.trim() == "[DONE]" {
|
||
log::debug!("[Claude/OpenRouter] <<< OpenAI SSE: [DONE]");
|
||
let event = json!({"type": "message_stop"});
|
||
let sse_data = format!("event: message_stop\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
log::debug!("[Claude/OpenRouter] >>> Anthropic SSE: message_stop");
|
||
yield Ok(Bytes::from(sse_data));
|
||
continue;
|
||
}
|
||
|
||
if let Ok(chunk) = serde_json::from_str::<OpenAIStreamChunk>(data) {
|
||
log::debug!("[Claude/OpenRouter] <<< SSE chunk received");
|
||
|
||
if message_id.is_none() {
|
||
message_id = Some(chunk.id.clone());
|
||
}
|
||
if current_model.is_none() {
|
||
current_model = Some(chunk.model.clone());
|
||
}
|
||
|
||
if let Some(choice) = chunk.choices.first() {
|
||
if !has_sent_message_start {
|
||
// Build usage with cache tokens if available from first chunk
|
||
let mut start_usage = json!({
|
||
"input_tokens": 0,
|
||
"output_tokens": 0
|
||
});
|
||
if let Some(u) = &chunk.usage {
|
||
start_usage["input_tokens"] = json!(u.prompt_tokens);
|
||
if let Some(cached) = extract_cache_read_tokens(u) {
|
||
start_usage["cache_read_input_tokens"] = json!(cached);
|
||
}
|
||
if let Some(created) = u.cache_creation_input_tokens {
|
||
start_usage["cache_creation_input_tokens"] = json!(created);
|
||
}
|
||
}
|
||
|
||
let event = json!({
|
||
"type": "message_start",
|
||
"message": {
|
||
"id": message_id.clone().unwrap_or_default(),
|
||
"type": "message",
|
||
"role": "assistant",
|
||
"model": current_model.clone().unwrap_or_default(),
|
||
"usage": start_usage
|
||
}
|
||
});
|
||
let sse_data = format!("event: message_start\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
has_sent_message_start = true;
|
||
}
|
||
|
||
// 处理 reasoning（thinking）
|
||
if let Some(reasoning) = &choice.delta.reasoning {
|
||
if current_non_tool_block_type != Some("thinking") {
|
||
if let Some(index) = current_non_tool_block_index.take() {
|
||
let event = json!({
|
||
"type": "content_block_stop",
|
||
"index": index
|
||
});
|
||
let sse_data = format!("event: content_block_stop\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
}
|
||
let index = next_content_index;
|
||
next_content_index += 1;
|
||
let event = json!({
|
||
"type": "content_block_start",
|
||
"index": index,
|
||
"content_block": {
|
||
"type": "thinking",
|
||
"thinking": ""
|
||
}
|
||
});
|
||
let sse_data = format!("event: content_block_start\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
current_non_tool_block_type = Some("thinking");
|
||
current_non_tool_block_index = Some(index);
|
||
}
|
||
|
||
if let Some(index) = current_non_tool_block_index {
|
||
let event = json!({
|
||
"type": "content_block_delta",
|
||
"index": index,
|
||
"delta": {
|
||
"type": "thinking_delta",
|
||
"thinking": reasoning
|
||
}
|
||
});
|
||
let sse_data = format!("event: content_block_delta\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
}
|
||
}
|
||
|
||
// 处理文本内容
|
||
if let Some(content) = &choice.delta.content {
|
||
if !content.is_empty() {
|
||
if current_non_tool_block_type != Some("text") {
|
||
if let Some(index) = current_non_tool_block_index.take() {
|
||
let event = json!({
|
||
"type": "content_block_stop",
|
||
"index": index
|
||
});
|
||
let sse_data = format!("event: content_block_stop\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
}
|
||
|
||
let index = next_content_index;
|
||
next_content_index += 1;
|
||
let event = json!({
|
||
"type": "content_block_start",
|
||
"index": index,
|
||
"content_block": {
|
||
"type": "text",
|
||
"text": ""
|
||
}
|
||
});
|
||
let sse_data = format!("event: content_block_start\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
current_non_tool_block_type = Some("text");
|
||
current_non_tool_block_index = Some(index);
|
||
}
|
||
|
||
if let Some(index) = current_non_tool_block_index {
|
||
let event = json!({
|
||
"type": "content_block_delta",
|
||
"index": index,
|
||
"delta": {
|
||
"type": "text_delta",
|
||
"text": content
|
||
}
|
||
});
|
||
let sse_data = format!("event: content_block_delta\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
}
|
||
}
|
||
}
|
||
|
||
// 处理工具调用
|
||
if let Some(tool_calls) = &choice.delta.tool_calls {
|
||
if let Some(index) = current_non_tool_block_index.take() {
|
||
let event = json!({
|
||
"type": "content_block_stop",
|
||
"index": index
|
||
});
|
||
let sse_data = format!("event: content_block_stop\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
}
|
||
current_non_tool_block_type = None;
|
||
|
||
for tool_call in tool_calls {
|
||
let (
|
||
anthropic_index,
|
||
id,
|
||
name,
|
||
should_start,
|
||
pending_after_start,
|
||
immediate_delta,
|
||
) = {
|
||
let state = tool_blocks_by_index
|
||
.entry(tool_call.index)
|
||
.or_insert_with(|| {
|
||
let index = next_content_index;
|
||
next_content_index += 1;
|
||
ToolBlockState {
|
||
anthropic_index: index,
|
||
id: String::new(),
|
||
name: String::new(),
|
||
started: false,
|
||
pending_args: String::new(),
|
||
}
|
||
});
|
||
|
||
if let Some(id) = &tool_call.id {
|
||
state.id = id.clone();
|
||
}
|
||
if let Some(function) = &tool_call.function {
|
||
if let Some(name) = &function.name {
|
||
state.name = name.clone();
|
||
}
|
||
}
|
||
|
||
let should_start =
|
||
!state.started
|
||
&& !state.id.is_empty()
|
||
&& !state.name.is_empty();
|
||
if should_start {
|
||
state.started = true;
|
||
}
|
||
let pending_after_start = if should_start
|
||
&& !state.pending_args.is_empty()
|
||
{
|
||
Some(std::mem::take(&mut state.pending_args))
|
||
} else {
|
||
None
|
||
};
|
||
let args_delta = tool_call
|
||
.function
|
||
.as_ref()
|
||
.and_then(|f| f.arguments.clone());
|
||
let immediate_delta = if let Some(args) = args_delta {
|
||
if state.started {
|
||
Some(args)
|
||
} else {
|
||
state.pending_args.push_str(&args);
|
||
None
|
||
}
|
||
} else {
|
||
None
|
||
};
|
||
(
|
||
state.anthropic_index,
|
||
state.id.clone(),
|
||
state.name.clone(),
|
||
should_start,
|
||
pending_after_start,
|
||
immediate_delta,
|
||
)
|
||
};
|
||
|
||
if should_start {
|
||
let event = json!({
|
||
"type": "content_block_start",
|
||
"index": anthropic_index,
|
||
"content_block": {
|
||
"type": "tool_use",
|
||
"id": id,
|
||
"name": name
|
||
}
|
||
});
|
||
let sse_data = format!("event: content_block_start\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
open_tool_block_indices.insert(anthropic_index);
|
||
}
|
||
|
||
if let Some(args) = pending_after_start {
|
||
let event = json!({
|
||
"type": "content_block_delta",
|
||
"index": anthropic_index,
|
||
"delta": {
|
||
"type": "input_json_delta",
|
||
"partial_json": args
|
||
}
|
||
});
|
||
let sse_data = format!("event: content_block_delta\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
}
|
||
|
||
if let Some(args) = immediate_delta {
|
||
let event = json!({
|
||
"type": "content_block_delta",
|
||
"index": anthropic_index,
|
||
"delta": {
|
||
"type": "input_json_delta",
|
||
"partial_json": args
|
||
}
|
||
});
|
||
let sse_data = format!("event: content_block_delta\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
}
|
||
}
|
||
}
|
||
|
||
// 处理 finish_reason
|
||
if let Some(finish_reason) = &choice.finish_reason {
|
||
if let Some(index) = current_non_tool_block_index.take() {
|
||
let event = json!({
|
||
"type": "content_block_stop",
|
||
"index": index
|
||
});
|
||
let sse_data = format!("event: content_block_stop\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
}
|
||
current_non_tool_block_type = None;
|
||
|
||
// Late start for blocks that accumulated args before id/name arrived.
|
||
let mut late_tool_starts: Vec<(u32, String, String, String)> =
|
||
Vec::new();
|
||
for (tool_idx, state) in tool_blocks_by_index.iter_mut() {
|
||
if state.started {
|
||
continue;
|
||
}
|
||
let has_payload = !state.pending_args.is_empty()
|
||
|| !state.id.is_empty()
|
||
|| !state.name.is_empty();
|
||
if !has_payload {
|
||
continue;
|
||
}
|
||
let fallback_id = if state.id.is_empty() {
|
||
format!("tool_call_{tool_idx}")
|
||
} else {
|
||
state.id.clone()
|
||
};
|
||
let fallback_name = if state.name.is_empty() {
|
||
"unknown_tool".to_string()
|
||
} else {
|
||
state.name.clone()
|
||
};
|
||
state.started = true;
|
||
let pending = std::mem::take(&mut state.pending_args);
|
||
late_tool_starts.push((
|
||
state.anthropic_index,
|
||
fallback_id,
|
||
fallback_name,
|
||
pending,
|
||
));
|
||
}
|
||
late_tool_starts.sort_unstable_by_key(|(index, _, _, _)| *index);
|
||
for (index, id, name, pending) in late_tool_starts {
|
||
let event = json!({
|
||
"type": "content_block_start",
|
||
"index": index,
|
||
"content_block": {
|
||
"type": "tool_use",
|
||
"id": id,
|
||
"name": name
|
||
}
|
||
});
|
||
let sse_data = format!("event: content_block_start\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
open_tool_block_indices.insert(index);
|
||
if !pending.is_empty() {
|
||
let delta_event = json!({
|
||
"type": "content_block_delta",
|
||
"index": index,
|
||
"delta": {
|
||
"type": "input_json_delta",
|
||
"partial_json": pending
|
||
}
|
||
});
|
||
let delta_sse = format!("event: content_block_delta\ndata: {}\n\n",
|
||
serde_json::to_string(&delta_event).unwrap_or_default());
|
||
yield Ok(Bytes::from(delta_sse));
|
||
}
|
||
}
|
||
|
||
if !open_tool_block_indices.is_empty() {
|
||
let mut tool_indices: Vec<u32> =
|
||
open_tool_block_indices.iter().copied().collect();
|
||
tool_indices.sort_unstable();
|
||
for index in tool_indices {
|
||
let event = json!({
|
||
"type": "content_block_stop",
|
||
"index": index
|
||
});
|
||
let sse_data = format!("event: content_block_stop\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
}
|
||
open_tool_block_indices.clear();
|
||
}
|
||
|
||
let stop_reason = map_stop_reason(Some(finish_reason));
|
||
// Build usage with cache token fields
|
||
let usage_json = chunk.usage.as_ref().map(|u| {
|
||
let mut uj = json!({
|
||
"input_tokens": u.prompt_tokens,
|
||
"output_tokens": u.completion_tokens
|
||
});
|
||
if let Some(cached) = extract_cache_read_tokens(u) {
|
||
uj["cache_read_input_tokens"] = json!(cached);
|
||
}
|
||
if let Some(created) = u.cache_creation_input_tokens {
|
||
uj["cache_creation_input_tokens"] = json!(created);
|
||
}
|
||
uj
|
||
});
|
||
let event = json!({
|
||
"type": "message_delta",
|
||
"delta": {
|
||
"stop_reason": stop_reason,
|
||
"stop_sequence": null
|
||
},
|
||
"usage": usage_json
|
||
});
|
||
let sse_data = format!("event: message_delta\ndata: {}\n\n",
|
||
serde_json::to_string(&event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
Err(e) => {
|
||
log::error!("Stream error: {e}");
|
||
let error_event = json!({
|
||
"type": "error",
|
||
"error": {
|
||
"type": "stream_error",
|
||
"message": format!("Stream error: {e}")
|
||
}
|
||
});
|
||
let sse_data = format!("event: error\ndata: {}\n\n",
|
||
serde_json::to_string(&error_event).unwrap_or_default());
|
||
yield Ok(Bytes::from(sse_data));
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Extract cache_read tokens from Usage, checking both direct field and nested details
|
||
fn extract_cache_read_tokens(usage: &Usage) -> Option<u32> {
|
||
// Direct field takes priority (compatible servers)
|
||
if let Some(v) = usage.cache_read_input_tokens {
|
||
return Some(v);
|
||
}
|
||
// OpenAI standard: prompt_tokens_details.cached_tokens
|
||
usage
|
||
.prompt_tokens_details
|
||
.as_ref()
|
||
.map(|d| d.cached_tokens)
|
||
.filter(|&v| v > 0)
|
||
}
|
||
|
||
/// 映射停止原因
|
||
fn map_stop_reason(finish_reason: Option<&str>) -> Option<String> {
|
||
finish_reason.map(|r| {
|
||
match r {
|
||
"tool_calls" | "function_call" => "tool_use",
|
||
"stop" => "end_turn",
|
||
"length" => "max_tokens",
|
||
"content_filter" => "end_turn",
|
||
other => {
|
||
log::warn!("[Claude/OpenRouter] Unknown finish_reason in streaming: {other}");
|
||
"end_turn"
|
||
}
|
||
}
|
||
.to_string()
|
||
})
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use futures::stream;
|
||
use futures::StreamExt;
|
||
use serde_json::Value;
|
||
use std::collections::HashMap;
|
||
|
||
#[test]
|
||
fn test_map_stop_reason_legacy_and_filtered_values() {
|
||
assert_eq!(
|
||
map_stop_reason(Some("function_call")),
|
||
Some("tool_use".to_string())
|
||
);
|
||
assert_eq!(
|
||
map_stop_reason(Some("content_filter")),
|
||
Some("end_turn".to_string())
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_streaming_tool_calls_routed_by_index() {
|
||
let input = concat!(
|
||
"data: {\"id\":\"chatcmpl_1\",\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_0\",\"type\":\"function\",\"function\":{\"name\":\"first_tool\"}}]}}]}\n\n",
|
||
"data: {\"id\":\"chatcmpl_1\",\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":1,\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"second_tool\"}}]}}]}\n\n",
|
||
"data: {\"id\":\"chatcmpl_1\",\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":1,\"function\":{\"arguments\":\"{\\\"b\\\":2}\"}}]}}]}\n\n",
|
||
"data: {\"id\":\"chatcmpl_1\",\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"a\\\":1}\"}}]}}]}\n\n",
|
||
"data: {\"id\":\"chatcmpl_1\",\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{},\"finish_reason\":\"tool_calls\"}],\"usage\":{\"prompt_tokens\":8,\"completion_tokens\":4}}\n\n",
|
||
"data: [DONE]\n\n"
|
||
);
|
||
|
||
let upstream = stream::iter(vec![Ok(Bytes::from(input.as_bytes().to_vec()))]);
|
||
let converted = create_anthropic_sse_stream(upstream);
|
||
let chunks: Vec<_> = converted.collect().await;
|
||
|
||
let merged = chunks
|
||
.into_iter()
|
||
.map(|chunk| String::from_utf8_lossy(chunk.unwrap().as_ref()).to_string())
|
||
.collect::<String>();
|
||
|
||
let events: Vec<Value> = merged
|
||
.split("\n\n")
|
||
.filter_map(|block| {
|
||
let data = block.lines().find_map(|line| line.strip_prefix("data: "))?;
|
||
serde_json::from_str::<Value>(data).ok()
|
||
})
|
||
.collect();
|
||
|
||
let mut tool_index_by_call: HashMap<String, u64> = HashMap::new();
|
||
for event in &events {
|
||
if event.get("type").and_then(|v| v.as_str()) == Some("content_block_start")
|
||
&& event
|
||
.pointer("/content_block/type")
|
||
.and_then(|v| v.as_str())
|
||
== Some("tool_use")
|
||
{
|
||
if let (Some(call_id), Some(index)) = (
|
||
event.pointer("/content_block/id").and_then(|v| v.as_str()),
|
||
event.get("index").and_then(|v| v.as_u64()),
|
||
) {
|
||
tool_index_by_call.insert(call_id.to_string(), index);
|
||
}
|
||
}
|
||
}
|
||
|
||
assert_eq!(tool_index_by_call.len(), 2);
|
||
assert_ne!(
|
||
tool_index_by_call.get("call_0"),
|
||
tool_index_by_call.get("call_1")
|
||
);
|
||
|
||
let deltas: Vec<(u64, String)> = events
|
||
.iter()
|
||
.filter(|event| {
|
||
event.get("type").and_then(|v| v.as_str()) == Some("content_block_delta")
|
||
&& event.pointer("/delta/type").and_then(|v| v.as_str())
|
||
== Some("input_json_delta")
|
||
})
|
||
.filter_map(|event| {
|
||
let index = event.get("index").and_then(|v| v.as_u64())?;
|
||
let partial_json = event
|
||
.pointer("/delta/partial_json")
|
||
.and_then(|v| v.as_str())?
|
||
.to_string();
|
||
Some((index, partial_json))
|
||
})
|
||
.collect();
|
||
|
||
assert_eq!(deltas.len(), 2);
|
||
let second_idx = deltas
|
||
.iter()
|
||
.find_map(|(index, payload)| (payload == "{\"b\":2}").then_some(*index))
|
||
.unwrap();
|
||
let first_idx = deltas
|
||
.iter()
|
||
.find_map(|(index, payload)| (payload == "{\"a\":1}").then_some(*index))
|
||
.unwrap();
|
||
|
||
assert_eq!(second_idx, *tool_index_by_call.get("call_1").unwrap());
|
||
assert_eq!(first_idx, *tool_index_by_call.get("call_0").unwrap());
|
||
|
||
assert!(events.iter().any(|event| {
|
||
event.get("type").and_then(|v| v.as_str()) == Some("message_delta")
|
||
&& event.pointer("/delta/stop_reason").and_then(|v| v.as_str()) == Some("tool_use")
|
||
}));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_streaming_delays_tool_start_until_id_and_name_ready() {
|
||
let input = concat!(
|
||
"data: {\"id\":\"chatcmpl_2\",\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"a\\\":\"}}]}}]}\n\n",
|
||
"data: {\"id\":\"chatcmpl_2\",\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_0\",\"type\":\"function\",\"function\":{\"name\":\"first_tool\"}}]}}]}\n\n",
|
||
"data: {\"id\":\"chatcmpl_2\",\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"1}\"}}]}}]}\n\n",
|
||
"data: {\"id\":\"chatcmpl_2\",\"model\":\"gpt-4o\",\"choices\":[{\"delta\":{},\"finish_reason\":\"tool_calls\"}],\"usage\":{\"prompt_tokens\":6,\"completion_tokens\":2}}\n\n",
|
||
"data: [DONE]\n\n"
|
||
);
|
||
|
||
let upstream = stream::iter(vec![Ok(Bytes::from(input.as_bytes().to_vec()))]);
|
||
let converted = create_anthropic_sse_stream(upstream);
|
||
let chunks: Vec<_> = converted.collect().await;
|
||
let merged = chunks
|
||
.into_iter()
|
||
.map(|chunk| String::from_utf8_lossy(chunk.unwrap().as_ref()).to_string())
|
||
.collect::<String>();
|
||
|
||
let events: Vec<Value> = merged
|
||
.split("\n\n")
|
||
.filter_map(|block| {
|
||
let data = block.lines().find_map(|line| line.strip_prefix("data: "))?;
|
||
serde_json::from_str::<Value>(data).ok()
|
||
})
|
||
.collect();
|
||
|
||
let starts: Vec<&Value> = events
|
||
.iter()
|
||
.filter(|event| {
|
||
event.get("type").and_then(|v| v.as_str()) == Some("content_block_start")
|
||
&& event
|
||
.pointer("/content_block/type")
|
||
.and_then(|v| v.as_str())
|
||
== Some("tool_use")
|
||
})
|
||
.collect();
|
||
assert_eq!(starts.len(), 1);
|
||
assert_eq!(
|
||
starts[0]
|
||
.pointer("/content_block/id")
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or(""),
|
||
"call_0"
|
||
);
|
||
assert_eq!(
|
||
starts[0]
|
||
.pointer("/content_block/name")
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or(""),
|
||
"first_tool"
|
||
);
|
||
|
||
let deltas: Vec<&str> = events
|
||
.iter()
|
||
.filter(|event| {
|
||
event.get("type").and_then(|v| v.as_str()) == Some("content_block_delta")
|
||
&& event.pointer("/delta/type").and_then(|v| v.as_str())
|
||
== Some("input_json_delta")
|
||
})
|
||
.filter_map(|event| {
|
||
event
|
||
.pointer("/delta/partial_json")
|
||
.and_then(|v| v.as_str())
|
||
})
|
||
.collect();
|
||
assert!(deltas.contains(&"{\"a\":"));
|
||
assert!(deltas.contains(&"1}"));
|
||
}
|
||
}
|