token streaming & reasoning display: live SSE tokens frontend to back
Nightly Build / build (push) Successful in 6m59s
Nightly Build / build (push) Successful in 6m59s
- Add StreamDelta(SseDecoder) framing shared by OpenAI/Anthropic - OpenAiClient: stream=true + reasoning_content deltas, index-based tool_calls accumulation, usage from final chunk - AnthropicClient: message_start/content_block_*/message_delta events, thinking_delta->reasoning, input_json_delta->tool input - TokenDelta ServerEvent variant wired through ChatHub + WS broadcast - Frontend throttled flush (~15 Hz), pending bubble mutate-in-place, reasoning as collapsed-by-default <details> - Drop streaming bubble on error/llm_failed/model_fallback - i18n: chat.reasoning key added to en/fr/it
This commit is contained in:
@@ -5,9 +5,11 @@ edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
core-api = { path = "../core-api" }
|
||||
reqwest = { version = "0.13", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json"] }
|
||||
reqwest = { version = "0.13", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json", "stream"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
async-trait = "0.1"
|
||||
anyhow = "1"
|
||||
tracing = "0.1"
|
||||
tokio = { version = "1", features = ["sync"] }
|
||||
futures-util = "0.3"
|
||||
|
||||
@@ -1,8 +1,12 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures_util::StreamExt;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::{debug, info, trace, warn};
|
||||
|
||||
use crate::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, Role, ToolCall, headers_to_json, redact_key};
|
||||
use crate::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, Role, SseDecoder, StreamDelta, ToolCall, headers_to_json, redact_key};
|
||||
|
||||
const DEFAULT_BASE_URL: &str = "https://api.anthropic.com";
|
||||
const ANTHROPIC_VERSION: &str = "2023-06-01";
|
||||
@@ -157,6 +161,242 @@ impl AnthropicClient {
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
/// Assembles the `/v1/messages` request body shared by the buffered and the
|
||||
/// streaming path (the caller adds `stream` on top).
|
||||
fn tools_body(&self, system: Option<String>, messages: Vec<Value>, tools: Vec<Value>, options: &ChatOptions) -> Value {
|
||||
let max_tokens = options.max_tokens.unwrap_or(4096);
|
||||
let mut body = json!({
|
||||
"model": options.model,
|
||||
"max_tokens": max_tokens,
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
});
|
||||
|
||||
if let Some(sys) = system { body["system"] = sys.into(); }
|
||||
if let Some(t) = options.temperature { body["temperature"] = t.into(); }
|
||||
self.apply_extra(&mut body);
|
||||
body
|
||||
}
|
||||
|
||||
/// Collects ALL system-role messages (main prompt, mid-conversation
|
||||
/// summary, tail_reminder) into a single `system:` string. The Anthropic
|
||||
/// API only accepts a single system parameter.
|
||||
fn merged_system(messages: &[Value]) -> Option<String> {
|
||||
let parts: Vec<&str> = messages
|
||||
.iter()
|
||||
.filter(|m| m["role"].as_str() == Some("system"))
|
||||
.filter_map(|m| m["content"].as_str())
|
||||
.collect();
|
||||
if parts.is_empty() { None } else { Some(parts.join("\n\n---\n\n")) }
|
||||
}
|
||||
|
||||
fn url(&self) -> String {
|
||||
format!("{}/v1/messages", self.base_url.trim_end_matches('/'))
|
||||
}
|
||||
|
||||
fn logged_headers(&self) -> Value {
|
||||
json!({
|
||||
"x-api-key": redact_key(&self.api_key),
|
||||
"anthropic-version": ANTHROPIC_VERSION,
|
||||
"content-type": "application/json",
|
||||
})
|
||||
}
|
||||
|
||||
async fn send_request(&self, body: &Value) -> reqwest::Result<reqwest::Response> {
|
||||
self.http
|
||||
.post(self.url())
|
||||
.header("x-api-key", &self.api_key)
|
||||
.header("anthropic-version", ANTHROPIC_VERSION)
|
||||
.header("X-Title", core_api::APP_NAME)
|
||||
.json(body)
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()
|
||||
}
|
||||
|
||||
/// Joined `thinking` blocks of a content array, if any (extended thinking).
|
||||
fn reasoning_of(content_blocks: &[Value]) -> Option<String> {
|
||||
let parts: Vec<&str> = content_blocks
|
||||
.iter()
|
||||
.filter(|b| b["type"].as_str() == Some("thinking"))
|
||||
.filter_map(|b| b["thinking"].as_str())
|
||||
.collect();
|
||||
if parts.is_empty() { None } else { Some(parts.join("\n")) }
|
||||
}
|
||||
|
||||
/// SSE streaming path behind `chat_with_tools_raw_streaming`. Anthropic
|
||||
/// streams typed events (`message_start` / `content_block_*` /
|
||||
/// `message_delta` / `message_stop`); text and thinking deltas are
|
||||
/// forwarded to `delta_tx` best-effort while the blocks are accumulated
|
||||
/// into the same `LlmTurn` the buffered path returns.
|
||||
async fn stream_chat(
|
||||
&self,
|
||||
messages: &[Value],
|
||||
tools: &[Value],
|
||||
options: &ChatOptions,
|
||||
delta_tx: &mpsc::Sender<StreamDelta>,
|
||||
emitted: &mut bool,
|
||||
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
||||
let system = Self::merged_system(messages);
|
||||
let anthropic_messages = Self::convert_messages(messages);
|
||||
let anthropic_tools = Self::convert_tools(tools);
|
||||
let mut body = self.tools_body(system, anthropic_messages, anthropic_tools, options);
|
||||
body["stream"] = json!(true);
|
||||
|
||||
debug!(model = %options.model, tools = tools.len(), "anthropic: sending streaming chat_with_tools request");
|
||||
trace!(body = %body, "anthropic: streaming chat_with_tools request body");
|
||||
|
||||
let request_body = body.clone();
|
||||
let request_headers = self.logged_headers();
|
||||
|
||||
let http_resp = self.send_request(&body).await?;
|
||||
let response_headers = headers_to_json(http_resp.headers());
|
||||
|
||||
/// One content block being accumulated by index.
|
||||
#[derive(Default)]
|
||||
struct Block {
|
||||
kind: String, // "text" | "thinking" | "tool_use"
|
||||
buf: String, // text/thinking content or input_json fragments
|
||||
id: String,
|
||||
name: String,
|
||||
}
|
||||
|
||||
let mut blocks: BTreeMap<u64, Block> = BTreeMap::new();
|
||||
let mut stop_reason: Option<String> = None;
|
||||
let mut usage = json!({});
|
||||
let mut sse = SseDecoder::new();
|
||||
let mut byte_stream = http_resp.bytes_stream();
|
||||
|
||||
let mut handle_payload = |payload: &str, emitted: &mut bool| -> anyhow::Result<()> {
|
||||
let Ok(v) = serde_json::from_str::<Value>(payload) else { return Ok(()) };
|
||||
match v["type"].as_str().unwrap_or("") {
|
||||
"message_start" => {
|
||||
if let Some(u) = v["message"]["usage"].as_object() {
|
||||
for (k, val) in u { usage[k.clone()] = val.clone(); }
|
||||
}
|
||||
}
|
||||
"content_block_start" => {
|
||||
let idx = v["index"].as_u64().unwrap_or(0);
|
||||
let cb = &v["content_block"];
|
||||
let block = blocks.entry(idx).or_default();
|
||||
block.kind = cb["type"].as_str().unwrap_or("").to_string();
|
||||
block.id = cb["id"].as_str().unwrap_or("").to_string();
|
||||
block.name = cb["name"].as_str().unwrap_or("").to_string();
|
||||
}
|
||||
"content_block_delta" => {
|
||||
let idx = v["index"].as_u64().unwrap_or(0);
|
||||
let delta = &v["delta"];
|
||||
match delta["type"].as_str().unwrap_or("") {
|
||||
"text_delta" => {
|
||||
if let Some(t) = delta["text"].as_str().filter(|t| !t.is_empty()) {
|
||||
blocks.entry(idx).or_default().buf.push_str(t);
|
||||
*emitted = true;
|
||||
let _ = delta_tx.try_send(StreamDelta::Text(t.to_string()));
|
||||
}
|
||||
}
|
||||
"thinking_delta" => {
|
||||
if let Some(t) = delta["thinking"].as_str().filter(|t| !t.is_empty()) {
|
||||
blocks.entry(idx).or_default().buf.push_str(t);
|
||||
*emitted = true;
|
||||
let _ = delta_tx.try_send(StreamDelta::Reasoning(t.to_string()));
|
||||
}
|
||||
}
|
||||
"input_json_delta" => {
|
||||
if let Some(j) = delta["partial_json"].as_str() {
|
||||
blocks.entry(idx).or_default().buf.push_str(j);
|
||||
}
|
||||
}
|
||||
// signature_delta and unknown deltas carry no displayable text.
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
"message_delta" => {
|
||||
if let Some(sr) = v["delta"]["stop_reason"].as_str() {
|
||||
stop_reason = Some(sr.to_string());
|
||||
}
|
||||
if let Some(u) = v["usage"].as_object() {
|
||||
for (k, val) in u { usage[k.clone()] = val.clone(); }
|
||||
}
|
||||
}
|
||||
"error" => {
|
||||
return Err(anyhow::anyhow!("anthropic: stream error event: {payload}"));
|
||||
}
|
||||
// content_block_stop / message_stop / ping: nothing to accumulate.
|
||||
_ => {}
|
||||
}
|
||||
Ok(())
|
||||
};
|
||||
|
||||
while let Some(chunk) = byte_stream.next().await {
|
||||
let chunk = chunk?;
|
||||
for payload in sse.feed(&chunk) {
|
||||
handle_payload(&payload, emitted)?;
|
||||
}
|
||||
}
|
||||
for payload in sse.finish() {
|
||||
handle_payload(&payload, emitted)?;
|
||||
}
|
||||
|
||||
let stop = stop_reason.as_deref().unwrap_or("");
|
||||
let input_tokens = usage["input_tokens"].as_u64().map(|n| n as u32);
|
||||
let output_tokens = usage["output_tokens"].as_u64().map(|n| n as u32);
|
||||
let cache_read_tokens = usage["cache_read_input_tokens"].as_u64().map(|n| n as u32);
|
||||
let cache_creation_tokens = usage["cache_creation_input_tokens"].as_u64().map(|n| n as u32);
|
||||
info!(model = %options.model, ?input_tokens, ?output_tokens, stop_reason = stop, "anthropic: streaming response completed");
|
||||
if stop == "max_tokens" {
|
||||
warn!(model = %options.model, ?output_tokens, "anthropic: response truncated (max_tokens reached)");
|
||||
}
|
||||
|
||||
let text_of = |kind: &str| -> String {
|
||||
blocks.values()
|
||||
.filter(|b| b.kind == kind)
|
||||
.map(|b| b.buf.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
};
|
||||
let reasoning = text_of("thinking");
|
||||
let reasoning_content = if reasoning.is_empty() { None } else { Some(reasoning) };
|
||||
let tool_blocks: Vec<&Block> = blocks.values().filter(|b| b.kind == "tool_use").collect();
|
||||
|
||||
let turn = if !tool_blocks.is_empty() {
|
||||
let calls = tool_blocks
|
||||
.iter()
|
||||
.map(|b| ToolCall {
|
||||
id: b.id.clone(),
|
||||
name: b.name.clone(),
|
||||
arguments: serde_json::from_str(&b.buf).unwrap_or(Value::Object(Default::default())),
|
||||
})
|
||||
.collect();
|
||||
LlmTurn::ToolCalls { content: text_of("text"), calls, input_tokens, output_tokens, reasoning_content, cache_read_tokens, cache_creation_tokens, cost: None }
|
||||
} else {
|
||||
let truncated = stop == "max_tokens";
|
||||
LlmTurn::Message(ChatResponse {
|
||||
content: text_of("text"), input_tokens, output_tokens, truncated,
|
||||
reasoning_content, cache_read_tokens, cache_creation_tokens, cost: None,
|
||||
})
|
||||
};
|
||||
|
||||
// Buffered-shaped response body for the payload log.
|
||||
let content_log: Vec<Value> = blocks.values().map(|b| match b.kind.as_str() {
|
||||
"tool_use" => json!({"type": "tool_use", "id": b.id, "name": b.name, "input": serde_json::from_str::<Value>(&b.buf).unwrap_or(json!({}))}),
|
||||
"thinking" => json!({"type": "thinking", "thinking": b.buf}),
|
||||
_ => json!({"type": "text", "text": b.buf}),
|
||||
}).collect();
|
||||
let raw_meta = LlmRawMeta {
|
||||
request_headers: Some(request_headers),
|
||||
request_body: Some(request_body),
|
||||
response_headers: Some(response_headers),
|
||||
response_body: Some(json!({
|
||||
"streamed": true,
|
||||
"content": content_log,
|
||||
"stop_reason": stop,
|
||||
"usage": usage,
|
||||
})),
|
||||
};
|
||||
|
||||
Ok((turn, Some(raw_meta)))
|
||||
}
|
||||
}
|
||||
|
||||
/// User content arrives either as a plain string or as an OpenAI-style parts
|
||||
@@ -274,7 +514,7 @@ impl ChatbotClient for AnthropicClient {
|
||||
|
||||
let content = resp["content"]
|
||||
.as_array()
|
||||
.and_then(|arr| arr.first())
|
||||
.and_then(|arr| arr.iter().find(|b| b["type"].as_str() == Some("text")))
|
||||
.and_then(|block| block["text"].as_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("Missing content in Anthropic response"))?
|
||||
.to_string();
|
||||
@@ -304,58 +544,22 @@ impl ChatbotClient for AnthropicClient {
|
||||
tools: &[Value],
|
||||
options: &ChatOptions,
|
||||
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
||||
// Collect ALL system-role messages (main prompt, mid-conversation
|
||||
// summary, tail_reminder) and merge them into a single `system:`
|
||||
// string. The Anthropic API only accepts a single system parameter;
|
||||
// mid-conversation system messages generated by build_openai_messages
|
||||
// are intentionally used for injecting compaction summaries and tail
|
||||
// reminders — they must not be silently dropped.
|
||||
let system: Option<String> = {
|
||||
let parts: Vec<&str> = messages
|
||||
.iter()
|
||||
.filter(|m| m["role"].as_str() == Some("system"))
|
||||
.filter_map(|m| m["content"].as_str())
|
||||
.collect();
|
||||
if parts.is_empty() { None } else { Some(parts.join("\n\n---\n\n")) }
|
||||
};
|
||||
|
||||
// Mid-conversation system messages (compaction summaries, tail
|
||||
// reminders) are merged into the single `system:` parameter — they
|
||||
// must not be silently dropped.
|
||||
let system = Self::merged_system(messages);
|
||||
let anthropic_messages = Self::convert_messages(messages);
|
||||
let anthropic_tools = Self::convert_tools(tools);
|
||||
let body = self.tools_body(system, anthropic_messages, anthropic_tools, options);
|
||||
|
||||
let max_tokens = options.max_tokens.unwrap_or(4096);
|
||||
let mut body = json!({
|
||||
"model": options.model,
|
||||
"max_tokens": max_tokens,
|
||||
"messages": anthropic_messages,
|
||||
"tools": anthropic_tools,
|
||||
});
|
||||
|
||||
if let Some(sys) = system { body["system"] = sys.into(); }
|
||||
if let Some(t) = options.temperature { body["temperature"] = t.into(); }
|
||||
self.apply_extra(&mut body);
|
||||
|
||||
let url = format!("{}/v1/messages", self.base_url.trim_end_matches('/'));
|
||||
debug!(model = %options.model, tools = tools.len(), "anthropic: sending chat_with_tools request");
|
||||
trace!(body = %body, "anthropic: chat_with_tools request body");
|
||||
|
||||
// Capture request metadata for logging.
|
||||
let request_body = body.clone();
|
||||
let request_headers = json!({
|
||||
"x-api-key": redact_key(&self.api_key),
|
||||
"anthropic-version": ANTHROPIC_VERSION,
|
||||
"content-type": "application/json",
|
||||
});
|
||||
let request_headers = self.logged_headers();
|
||||
|
||||
let http_resp = self
|
||||
.http
|
||||
.post(&url)
|
||||
.header("x-api-key", &self.api_key)
|
||||
.header("anthropic-version", ANTHROPIC_VERSION)
|
||||
.header("X-Title", core_api::APP_NAME)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
let http_resp = self.send_request(&body).await?;
|
||||
|
||||
let response_headers = headers_to_json(http_resp.headers());
|
||||
let resp_text = http_resp.text().await?;
|
||||
@@ -383,6 +587,7 @@ impl ChatbotClient for AnthropicClient {
|
||||
}
|
||||
|
||||
let has_tool_use = content_blocks.iter().any(|b| b["type"].as_str() == Some("tool_use"));
|
||||
let reasoning_content = Self::reasoning_of(&content_blocks);
|
||||
|
||||
// Check content blocks directly: Anthropic sometimes returns stop_reason "end_turn"
|
||||
// even when tool_use blocks are present, so stop_reason alone is not reliable.
|
||||
@@ -404,7 +609,7 @@ impl ChatbotClient for AnthropicClient {
|
||||
})
|
||||
.collect();
|
||||
|
||||
LlmTurn::ToolCalls { content: text, calls, input_tokens, output_tokens, reasoning_content: None, cache_read_tokens, cache_creation_tokens, cost }
|
||||
LlmTurn::ToolCalls { content: text, calls, input_tokens, output_tokens, reasoning_content, cache_read_tokens, cache_creation_tokens, cost }
|
||||
} else {
|
||||
let content = content_blocks
|
||||
.iter()
|
||||
@@ -414,17 +619,55 @@ impl ChatbotClient for AnthropicClient {
|
||||
.to_string();
|
||||
|
||||
let truncated = stop_reason == "max_tokens";
|
||||
LlmTurn::Message(ChatResponse { content, input_tokens, output_tokens, truncated, reasoning_content: None, cache_read_tokens, cache_creation_tokens, cost })
|
||||
LlmTurn::Message(ChatResponse { content, input_tokens, output_tokens, truncated, reasoning_content, cache_read_tokens, cache_creation_tokens, cost })
|
||||
};
|
||||
|
||||
Ok((turn, Some(raw_meta)))
|
||||
}
|
||||
|
||||
async fn chat_with_tools_raw_streaming(
|
||||
&self,
|
||||
messages: &[Value],
|
||||
tools: &[Value],
|
||||
options: &ChatOptions,
|
||||
delta_tx: mpsc::Sender<StreamDelta>,
|
||||
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
||||
let mut emitted = false;
|
||||
match self.stream_chat(messages, tools, options, &delta_tx, &mut emitted).await {
|
||||
Ok(ok) => Ok(ok),
|
||||
// Pre-stream failure (nothing shown yet): retry buffered. A
|
||||
// mid-stream failure propagates to the model-fallback logic.
|
||||
Err(e) if !emitted => {
|
||||
debug!(model = %options.model, error = %e, "anthropic: streaming failed before any delta; retrying buffered");
|
||||
self.chat_with_tools_raw(messages, tools, options).await
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn reasoning_of_joins_thinking_blocks() {
|
||||
let blocks = vec![
|
||||
json!({"type": "thinking", "thinking": "first"}),
|
||||
json!({"type": "text", "text": "answer"}),
|
||||
json!({"type": "thinking", "thinking": "second"}),
|
||||
];
|
||||
assert_eq!(
|
||||
AnthropicClient::reasoning_of(&blocks),
|
||||
Some("first\nsecond".to_string())
|
||||
);
|
||||
assert_eq!(AnthropicClient::reasoning_of(&[]), None);
|
||||
assert_eq!(
|
||||
AnthropicClient::reasoning_of(&[json!({"type": "text", "text": "a"})]),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_content_string_passthrough() {
|
||||
let v = convert_user_content(&json!("hello"));
|
||||
|
||||
@@ -6,11 +6,54 @@ pub mod openai;
|
||||
// Re-export the trait and all associated types from core-api so existing
|
||||
// callers that import from `llm_client` continue to work unchanged.
|
||||
pub use core_api::chatbot::{
|
||||
ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, Role, ToolCall,
|
||||
ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, Role, StreamDelta,
|
||||
ToolCall,
|
||||
};
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
/// Incremental SSE decoder: feed raw response bytes, get back the payload of
|
||||
/// every complete `data:` line seen (`[DONE]` included — callers decide).
|
||||
/// Buffers partial lines across chunks; `event:` lines and comments are
|
||||
/// skipped (both OpenAI and Anthropic put the event type inside the JSON).
|
||||
#[derive(Default)]
|
||||
pub struct SseDecoder {
|
||||
buf: Vec<u8>,
|
||||
}
|
||||
|
||||
impl SseDecoder {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn feed(&mut self, bytes: &[u8]) -> Vec<String> {
|
||||
self.buf.extend_from_slice(bytes);
|
||||
let mut out = Vec::new();
|
||||
while let Some(pos) = self.buf.iter().position(|&b| b == b'\n') {
|
||||
let line: Vec<u8> = self.buf.drain(..=pos).collect();
|
||||
if let Some(payload) = parse_sse_line(&line) {
|
||||
out.push(payload);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Flush a trailing line not terminated by `\n` at end-of-stream.
|
||||
pub fn finish(&mut self) -> Vec<String> {
|
||||
let rest = std::mem::take(&mut self.buf);
|
||||
parse_sse_line(&rest).into_iter().collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// A complete SSE line is valid UTF-8 (a multibyte sequence never contains a
|
||||
/// `\n` byte), but decode lossily anyway — a corrupt line is skipped, not fatal.
|
||||
fn parse_sse_line(line: &[u8]) -> Option<String> {
|
||||
let line = String::from_utf8_lossy(line);
|
||||
let line = line.trim_end_matches('\r').trim();
|
||||
let data = line.strip_prefix("data:")?.trim_start();
|
||||
if data.is_empty() { None } else { Some(data.to_string()) }
|
||||
}
|
||||
|
||||
/// Converts a reqwest `HeaderMap` into a `serde_json::Value` object.
|
||||
pub fn headers_to_json(headers: &reqwest::header::HeaderMap) -> Value {
|
||||
let map: serde_json::Map<String, Value> = headers
|
||||
@@ -72,3 +115,41 @@ pub fn http_status(err: &anyhow::Error) -> Option<u16> {
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SseDecoder;
|
||||
|
||||
#[test]
|
||||
fn sse_decoder_buffers_partial_lines_across_chunks() {
|
||||
let mut dec = SseDecoder::new();
|
||||
// A payload split mid-JSON across two chunks yields one complete line.
|
||||
assert!(dec.feed(br#"data: {"a": 1"#).is_empty());
|
||||
assert_eq!(dec.feed(b"}\r\n").len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sse_decoder_skips_events_comments_and_keeps_done() {
|
||||
let mut dec = SseDecoder::new();
|
||||
let out = dec.feed(b"event: message_start\n: ping\n\ndata: {\"type\":\"ping\"}\ndata: [DONE]\n");
|
||||
assert_eq!(out, vec!["{\"type\":\"ping\"}".to_string(), "[DONE]".to_string()]);
|
||||
assert!(dec.finish().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sse_decoder_finish_flushes_unterminated_tail() {
|
||||
let mut dec = SseDecoder::new();
|
||||
assert!(dec.feed(b"data: tail-without-newline").is_empty());
|
||||
assert_eq!(dec.finish(), vec!["tail-without-newline".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sse_decoder_handles_multibyte_split() {
|
||||
let mut dec = SseDecoder::new();
|
||||
// "€" is 3 bytes in UTF-8; split across the chunk boundary.
|
||||
let payload = "data: {\"t\":\"€\"}\n".as_bytes();
|
||||
let (a, b) = payload.split_at(12);
|
||||
assert!(dec.feed(a).is_empty());
|
||||
assert_eq!(dec.feed(b), vec!["{\"t\":\"€\"}".to_string()]);
|
||||
}
|
||||
}
|
||||
|
||||
+229
-41
@@ -1,8 +1,12 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures_util::StreamExt;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::{debug, info, trace, warn};
|
||||
|
||||
use crate::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, Role, ToolCall, headers_to_json, redact_key};
|
||||
use crate::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, Role, SseDecoder, StreamDelta, ToolCall, headers_to_json, redact_key};
|
||||
use core_api::APP_NAME;
|
||||
|
||||
/// OpenAI ChatGPT client (also compatible with any OpenAI-spec endpoint).
|
||||
@@ -45,6 +49,205 @@ impl OpenAiClient {
|
||||
fn url(&self) -> String {
|
||||
format!("{}/chat/completions", self.base_url.trim_end_matches('/'))
|
||||
}
|
||||
|
||||
/// Shared request body for the buffered and the streaming path. Caller adds
|
||||
/// `max_tokens`/`temperature`/`extra_params` afterwards via `finalize_body`.
|
||||
fn base_body(&self, model: &str, messages: &[Value], tools: &[Value]) -> Value {
|
||||
let mut body = json!({
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
});
|
||||
|
||||
if !tools.is_empty() {
|
||||
// When prompt caching is enabled, tag the last tool with cache_control
|
||||
// so the entire tools array is included in the Anthropic KV cache prefix.
|
||||
let tools_value: Value = if self.enable_prompt_cache {
|
||||
let mut tagged = tools.to_vec();
|
||||
if let Some(last) = tagged.last_mut() {
|
||||
last["cache_control"] = json!({"type": "ephemeral"});
|
||||
}
|
||||
tagged.into()
|
||||
} else {
|
||||
tools.into()
|
||||
};
|
||||
body["tools"] = tools_value;
|
||||
body["tool_choice"] = "auto".into();
|
||||
}
|
||||
body
|
||||
}
|
||||
|
||||
fn finalize_body(&self, mut body: Value, options: &ChatOptions) -> Value {
|
||||
if let Some(t) = options.max_tokens { body["max_tokens"] = t.into(); }
|
||||
if let Some(t) = options.temperature { body["temperature"] = t.into(); }
|
||||
self.apply_extra(&mut body);
|
||||
body
|
||||
}
|
||||
|
||||
/// Request metadata for logging (shared by buffered and streaming paths).
|
||||
fn logged_headers(&self) -> Value {
|
||||
let mut logged_headers = json!({
|
||||
"authorization": format!("Bearer {}", redact_key(&self.api_key)),
|
||||
"content-type": "application/json",
|
||||
});
|
||||
if self.enable_prompt_cache {
|
||||
logged_headers["anthropic-beta"] = "prompt-caching-2024-07-31".into();
|
||||
}
|
||||
logged_headers
|
||||
}
|
||||
|
||||
async fn send_request(&self, body: &Value) -> reqwest::Result<reqwest::Response> {
|
||||
let mut req = self.http.post(self.url()).bearer_auth(&self.api_key).header("X-Title", APP_NAME);
|
||||
if self.enable_prompt_cache {
|
||||
req = req.header("anthropic-beta", "prompt-caching-2024-07-31");
|
||||
}
|
||||
req.json(body).send().await
|
||||
}
|
||||
|
||||
/// SSE streaming path behind `chat_with_tools_raw_streaming`. Accumulates
|
||||
/// content/reasoning/tool-call fragments into the same `LlmTurn` the
|
||||
/// buffered path would return, while forwarding text/reasoning deltas to
|
||||
/// `delta_tx` (try_send, best-effort). `emitted` tracks whether any delta
|
||||
/// was pushed, so the caller can distinguish a pre-stream failure (safe to
|
||||
/// retry buffered) from a mid-stream one (partial output already shown).
|
||||
async fn stream_chat(
|
||||
&self,
|
||||
messages: &[Value],
|
||||
tools: &[Value],
|
||||
options: &ChatOptions,
|
||||
delta_tx: &mpsc::Sender<StreamDelta>,
|
||||
emitted: &mut bool,
|
||||
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
||||
let mut body = self.base_body(&options.model, messages, tools);
|
||||
body["stream"] = json!(true);
|
||||
body["stream_options"] = json!({ "include_usage": true });
|
||||
let body = self.finalize_body(body, options);
|
||||
|
||||
debug!(model = %options.model, tools = tools.len(), prompt_cache = self.enable_prompt_cache, "openai: sending streaming chat_with_tools request");
|
||||
trace!(body = %body, "openai: streaming chat_with_tools request body");
|
||||
|
||||
let request_body = body.clone();
|
||||
let request_headers = self.logged_headers();
|
||||
|
||||
let http_resp = self.send_request(&body).await?;
|
||||
|
||||
let response_headers = headers_to_json(http_resp.headers());
|
||||
let status = http_resp.status();
|
||||
if !status.is_success() {
|
||||
let resp_text = http_resp.text().await?;
|
||||
return Err(crate::LlmError {
|
||||
status: Some(status.as_u16()),
|
||||
message: format!(
|
||||
"openai: HTTP {status} from {url}\nbody: {resp_text}",
|
||||
url = self.url(),
|
||||
),
|
||||
}.into());
|
||||
}
|
||||
|
||||
let mut content = String::new();
|
||||
let mut reasoning = String::new();
|
||||
// index → (id, name, arguments fragment buffer)
|
||||
let mut tool_calls: BTreeMap<u64, (String, String, String)> = BTreeMap::new();
|
||||
let mut finish_reason: Option<String> = None;
|
||||
let mut usage: Option<Value> = None;
|
||||
let mut sse = SseDecoder::new();
|
||||
let mut byte_stream = http_resp.bytes_stream();
|
||||
|
||||
// One SSE `data:` payload. Fragments update the accumulators; text and
|
||||
// reasoning also go out as deltas. Unparseable chunks are skipped —
|
||||
// the assembled turn stays consistent.
|
||||
let mut handle_payload = |payload: &str, emitted: &mut bool| {
|
||||
if payload == "[DONE]" {
|
||||
return;
|
||||
}
|
||||
let Ok(v) = serde_json::from_str::<Value>(payload) else { return };
|
||||
if let Some(u) = v.get("usage").filter(|u| !u.is_null()) {
|
||||
usage = Some(u.clone());
|
||||
}
|
||||
let Some(choice) = v["choices"].as_array().and_then(|a| a.first()) else { return };
|
||||
if let Some(fr) = choice["finish_reason"].as_str() {
|
||||
finish_reason = Some(fr.to_string());
|
||||
}
|
||||
let delta = &choice["delta"];
|
||||
if let Some(t) = delta["content"].as_str().filter(|t| !t.is_empty()) {
|
||||
content.push_str(t);
|
||||
*emitted = true;
|
||||
let _ = delta_tx.try_send(StreamDelta::Text(t.to_string()));
|
||||
}
|
||||
// Same normalization as the buffered path: DeepSeek uses
|
||||
// `reasoning_content`, MiniMax M3 and others `reasoning`.
|
||||
if let Some(t) = delta["reasoning_content"].as_str()
|
||||
.or_else(|| delta["reasoning"].as_str())
|
||||
.filter(|t| !t.is_empty())
|
||||
{
|
||||
reasoning.push_str(t);
|
||||
*emitted = true;
|
||||
let _ = delta_tx.try_send(StreamDelta::Reasoning(t.to_string()));
|
||||
}
|
||||
if let Some(tc_arr) = delta["tool_calls"].as_array() {
|
||||
for tc in tc_arr {
|
||||
let idx = tc["index"].as_u64().unwrap_or(0);
|
||||
let entry = tool_calls.entry(idx).or_default();
|
||||
if let Some(id) = tc["id"].as_str() { entry.0 = id.to_string(); }
|
||||
if let Some(n) = tc["function"]["name"].as_str() { entry.1 = n.to_string(); }
|
||||
if let Some(a) = tc["function"]["arguments"].as_str() { entry.2.push_str(a); }
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
while let Some(chunk) = byte_stream.next().await {
|
||||
let chunk = chunk?;
|
||||
for payload in sse.feed(&chunk) {
|
||||
handle_payload(&payload, emitted);
|
||||
}
|
||||
}
|
||||
for payload in sse.finish() {
|
||||
handle_payload(&payload, emitted);
|
||||
}
|
||||
|
||||
let finish = finish_reason.as_deref().unwrap_or("stop");
|
||||
let input_tokens = usage.as_ref().and_then(|u| u["prompt_tokens"].as_u64()).map(|n| n as u32);
|
||||
let output_tokens = usage.as_ref().and_then(|u| u["completion_tokens"].as_u64()).map(|n| n as u32);
|
||||
let cache_read_tokens = usage.as_ref()
|
||||
.and_then(|u| u["prompt_tokens_details"]["cached_tokens"].as_u64())
|
||||
.map(|n| n as u32);
|
||||
let cost = usage.as_ref().and_then(|u| u["cost"].as_f64());
|
||||
let reasoning_content = if reasoning.is_empty() { None } else { Some(reasoning) };
|
||||
info!(model = %options.model, ?input_tokens, ?output_tokens, finish_reason = finish, "openai: streaming response completed");
|
||||
if finish == "length" {
|
||||
warn!(model = %options.model, ?output_tokens, "openai: response truncated (max_tokens reached)");
|
||||
}
|
||||
|
||||
let turn = if !tool_calls.is_empty() {
|
||||
let calls = tool_calls
|
||||
.into_values()
|
||||
.map(|(id, name, args)| ToolCall {
|
||||
id,
|
||||
name,
|
||||
arguments: serde_json::from_str(&args).unwrap_or(Value::Object(Default::default())),
|
||||
})
|
||||
.collect();
|
||||
LlmTurn::ToolCalls { content, calls, input_tokens, output_tokens, reasoning_content, cache_read_tokens, cache_creation_tokens: None, cost }
|
||||
} else {
|
||||
let truncated = finish == "length";
|
||||
LlmTurn::Message(ChatResponse { content, input_tokens, output_tokens, truncated, reasoning_content, cache_read_tokens, cache_creation_tokens: None, cost })
|
||||
};
|
||||
|
||||
// Synthesize a buffered-shaped response body for the payload log, so a
|
||||
// streamed call leaves the same debugging trail as a buffered one.
|
||||
let response_body = json!({
|
||||
"streamed": true,
|
||||
"choices": [{ "finish_reason": finish }],
|
||||
"usage": usage,
|
||||
});
|
||||
let raw_meta = LlmRawMeta {
|
||||
request_headers: Some(request_headers),
|
||||
request_body: Some(request_body),
|
||||
response_headers: Some(response_headers),
|
||||
response_body: Some(response_body),
|
||||
};
|
||||
|
||||
Ok((turn, Some(raw_meta)))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -123,53 +326,16 @@ impl ChatbotClient for OpenAiClient {
|
||||
tools: &[Value],
|
||||
options: &ChatOptions,
|
||||
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
||||
let mut body = json!({
|
||||
"model": options.model,
|
||||
"messages": messages,
|
||||
});
|
||||
|
||||
if !tools.is_empty() {
|
||||
// When prompt caching is enabled, tag the last tool with cache_control
|
||||
// so the entire tools array is included in the Anthropic KV cache prefix.
|
||||
let tools_value: Value = if self.enable_prompt_cache {
|
||||
let mut tagged = tools.to_vec();
|
||||
if let Some(last) = tagged.last_mut() {
|
||||
last["cache_control"] = json!({"type": "ephemeral"});
|
||||
}
|
||||
tagged.into()
|
||||
} else {
|
||||
tools.into()
|
||||
};
|
||||
body["tools"] = tools_value;
|
||||
body["tool_choice"] = "auto".into();
|
||||
}
|
||||
|
||||
if let Some(t) = options.max_tokens { body["max_tokens"] = t.into(); }
|
||||
if let Some(t) = options.temperature { body["temperature"] = t.into(); }
|
||||
self.apply_extra(&mut body);
|
||||
let body = self.finalize_body(self.base_body(&options.model, messages, tools), options);
|
||||
|
||||
debug!(model = %options.model, tools = tools.len(), prompt_cache = self.enable_prompt_cache, "openai: sending chat_with_tools request");
|
||||
trace!(body = %body, "openai: chat_with_tools request body");
|
||||
|
||||
// Capture request metadata for logging.
|
||||
let mut logged_headers = json!({
|
||||
"authorization": format!("Bearer {}", redact_key(&self.api_key)),
|
||||
"content-type": "application/json",
|
||||
});
|
||||
if self.enable_prompt_cache {
|
||||
logged_headers["anthropic-beta"] = "prompt-caching-2024-07-31".into();
|
||||
}
|
||||
let request_body = body.clone();
|
||||
let request_headers = logged_headers;
|
||||
let request_headers = self.logged_headers();
|
||||
|
||||
let mut req = self.http.post(self.url()).bearer_auth(&self.api_key).header("X-Title", APP_NAME);
|
||||
if self.enable_prompt_cache {
|
||||
req = req.header("anthropic-beta", "prompt-caching-2024-07-31");
|
||||
}
|
||||
let http_resp = req
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
let http_resp = self.send_request(&body).await?;
|
||||
|
||||
let response_headers = headers_to_json(http_resp.headers());
|
||||
let status = http_resp.status();
|
||||
@@ -262,4 +428,26 @@ impl ChatbotClient for OpenAiClient {
|
||||
|
||||
Ok((turn, Some(raw_meta)))
|
||||
}
|
||||
|
||||
async fn chat_with_tools_raw_streaming(
|
||||
&self,
|
||||
messages: &[Value],
|
||||
tools: &[Value],
|
||||
options: &ChatOptions,
|
||||
delta_tx: mpsc::Sender<StreamDelta>,
|
||||
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
|
||||
let mut emitted = false;
|
||||
match self.stream_chat(messages, tools, options, &delta_tx, &mut emitted).await {
|
||||
Ok(ok) => Ok(ok),
|
||||
// Nothing was ever streamed: some OpenAI-compatible providers reject
|
||||
// `stream`/`stream_options` outright — retry buffered so they keep
|
||||
// working exactly as before. A mid-stream failure (deltas already
|
||||
// shown) instead propagates to the model-fallback logic.
|
||||
Err(e) if !emitted => {
|
||||
debug!(model = %options.model, error = %e, "openai: streaming failed before any delta; retrying buffered");
|
||||
self.chat_with_tools_raw(messages, tools, options).await
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user