mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-08-02 18:41:35 +08:00
refactor(proxy): extract take_sse_block helper with CRLF delimiter support
Replace inline `buffer.find("\n\n")` SSE splitting logic across streaming,
streaming_responses, response_handler, and response_processor with a shared
`take_sse_block` function that handles both `\n\n` and `\r\n\r\n` delimiters.
This commit is contained in:
@@ -2,7 +2,7 @@
|
|||||||
//!
|
//!
|
||||||
//! 实现 OpenAI SSE → Anthropic SSE 格式转换
|
//! 实现 OpenAI SSE → Anthropic SSE 格式转换
|
||||||
|
|
||||||
use crate::proxy::sse::strip_sse_field;
|
use crate::proxy::sse::{strip_sse_field, take_sse_block};
|
||||||
use bytes::Bytes;
|
use bytes::Bytes;
|
||||||
use futures::stream::{Stream, StreamExt};
|
use futures::stream::{Stream, StreamExt};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
@@ -110,10 +110,7 @@ pub fn create_anthropic_sse_stream<E: std::error::Error + Send + 'static>(
|
|||||||
let text = String::from_utf8_lossy(&bytes);
|
let text = String::from_utf8_lossy(&bytes);
|
||||||
buffer.push_str(&text);
|
buffer.push_str(&text);
|
||||||
|
|
||||||
while let Some(pos) = buffer.find("\n\n") {
|
while let Some(line) = take_sse_block(&mut buffer) {
|
||||||
let line = buffer[..pos].to_string();
|
|
||||||
buffer = buffer[pos + 2..].to_string();
|
|
||||||
|
|
||||||
if line.trim().is_empty() {
|
if line.trim().is_empty() {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,7 +9,7 @@
|
|||||||
//! 与 Chat Completions 的 delta chunk 模型完全不同,需要独立的状态机处理。
|
//! 与 Chat Completions 的 delta chunk 模型完全不同,需要独立的状态机处理。
|
||||||
|
|
||||||
use super::transform_responses::{build_anthropic_usage_from_responses, map_responses_stop_reason};
|
use super::transform_responses::{build_anthropic_usage_from_responses, map_responses_stop_reason};
|
||||||
use crate::proxy::sse::strip_sse_field;
|
use crate::proxy::sse::{strip_sse_field, take_sse_block};
|
||||||
use bytes::Bytes;
|
use bytes::Bytes;
|
||||||
use futures::stream::{Stream, StreamExt};
|
use futures::stream::{Stream, StreamExt};
|
||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
@@ -122,10 +122,7 @@ pub fn create_anthropic_sse_stream_from_responses<E: std::error::Error + Send +
|
|||||||
buffer.push_str(&text);
|
buffer.push_str(&text);
|
||||||
|
|
||||||
// SSE 事件由 \n\n 分隔
|
// SSE 事件由 \n\n 分隔
|
||||||
while let Some(pos) = buffer.find("\n\n") {
|
while let Some(block) = take_sse_block(&mut buffer) {
|
||||||
let block = buffer[..pos].to_string();
|
|
||||||
buffer = buffer[pos + 2..].to_string();
|
|
||||||
|
|
||||||
if block.trim().is_empty() {
|
if block.trim().is_empty() {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -5,7 +5,7 @@
|
|||||||
use super::session::ProxySession;
|
use super::session::ProxySession;
|
||||||
use super::usage::parser::TokenUsage;
|
use super::usage::parser::TokenUsage;
|
||||||
use super::ProxyError;
|
use super::ProxyError;
|
||||||
use crate::proxy::sse::strip_sse_field;
|
use crate::proxy::sse::{strip_sse_field, take_sse_block};
|
||||||
use bytes::Bytes;
|
use bytes::Bytes;
|
||||||
use futures::stream::{Stream, StreamExt};
|
use futures::stream::{Stream, StreamExt};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
@@ -86,10 +86,7 @@ impl StreamHandler {
|
|||||||
buffer.push_str(&text);
|
buffer.push_str(&text);
|
||||||
|
|
||||||
// 提取完整事件
|
// 提取完整事件
|
||||||
while let Some(pos) = buffer.find("\n\n") {
|
while let Some(event_text) = take_sse_block(&mut buffer) {
|
||||||
let event_text = buffer[..pos].to_string();
|
|
||||||
buffer = buffer[pos + 2..].to_string();
|
|
||||||
|
|
||||||
for line in event_text.lines() {
|
for line in event_text.lines() {
|
||||||
if let Some(data) = strip_sse_field(line, "data") {
|
if let Some(data) = strip_sse_field(line, "data") {
|
||||||
if data.trim() != "[DONE]" {
|
if data.trim() != "[DONE]" {
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ use super::{
|
|||||||
handler_context::{RequestContext, StreamingTimeoutConfig},
|
handler_context::{RequestContext, StreamingTimeoutConfig},
|
||||||
hyper_client::ProxyResponse,
|
hyper_client::ProxyResponse,
|
||||||
server::ProxyState,
|
server::ProxyState,
|
||||||
sse::strip_sse_field,
|
sse::{strip_sse_field, take_sse_block},
|
||||||
usage::parser::TokenUsage,
|
usage::parser::TokenUsage,
|
||||||
ProxyError,
|
ProxyError,
|
||||||
};
|
};
|
||||||
@@ -623,10 +623,7 @@ pub fn create_logged_passthrough_stream(
|
|||||||
buffer.push_str(&text);
|
buffer.push_str(&text);
|
||||||
|
|
||||||
// 尝试解析并记录完整的 SSE 事件
|
// 尝试解析并记录完整的 SSE 事件
|
||||||
while let Some(pos) = buffer.find("\n\n") {
|
while let Some(event_text) = take_sse_block(&mut buffer) {
|
||||||
let event_text = buffer[..pos].to_string();
|
|
||||||
buffer = buffer[pos + 2..].to_string();
|
|
||||||
|
|
||||||
if !event_text.trim().is_empty() {
|
if !event_text.trim().is_empty() {
|
||||||
// 提取 data 部分并尝试解析为 JSON
|
// 提取 data 部分并尝试解析为 JSON
|
||||||
for line in event_text.lines() {
|
for line in event_text.lines() {
|
||||||
|
|||||||
@@ -4,9 +4,27 @@ pub(crate) fn strip_sse_field<'a>(line: &'a str, field: &str) -> Option<&'a str>
|
|||||||
.or_else(|| line.strip_prefix(&format!("{field}:")))
|
.or_else(|| line.strip_prefix(&format!("{field}:")))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
pub(crate) fn take_sse_block(buffer: &mut String) -> Option<String> {
|
||||||
|
let mut best: Option<(usize, usize)> = None;
|
||||||
|
|
||||||
|
for (delimiter, len) in [("\r\n\r\n", 4usize), ("\n\n", 2usize)] {
|
||||||
|
if let Some(pos) = buffer.find(delimiter) {
|
||||||
|
if best.is_none_or(|(best_pos, _)| pos < best_pos) {
|
||||||
|
best = Some((pos, len));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let (pos, len) = best?;
|
||||||
|
let block = buffer[..pos].to_string();
|
||||||
|
buffer.drain(..pos + len);
|
||||||
|
Some(block)
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::strip_sse_field;
|
use super::{strip_sse_field, take_sse_block};
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn strip_sse_field_accepts_optional_space() {
|
fn strip_sse_field_accepts_optional_space() {
|
||||||
@@ -28,4 +46,26 @@ mod tests {
|
|||||||
);
|
);
|
||||||
assert_eq!(strip_sse_field("id:1", "data"), None);
|
assert_eq!(strip_sse_field("id:1", "data"), None);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn take_sse_block_supports_lf_delimiters() {
|
||||||
|
let mut buffer = "data: {\"ok\":true}\n\nrest".to_string();
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
take_sse_block(&mut buffer),
|
||||||
|
Some("data: {\"ok\":true}".to_string())
|
||||||
|
);
|
||||||
|
assert_eq!(buffer, "rest");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn take_sse_block_supports_crlf_delimiters() {
|
||||||
|
let mut buffer = "data: {\"ok\":true}\r\n\r\nrest".to_string();
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
take_sse_block(&mut buffer),
|
||||||
|
Some("data: {\"ok\":true}".to_string())
|
||||||
|
);
|
||||||
|
assert_eq!(buffer, "rest");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user