agent-loop: new crate — LLM loop kernel + Model clients (phase 0)
Extract the LLM agent loop into a standalone workspace crate with zero deps on skald-core/core-api (blueprint project-loop.md, D13-D15): - kernel: round loop, model fallback with rebuild, parallel tool fan-out (ordered id alloc / bounded concurrent exec / ordered record), streaming deltas drained before outcomes, sticky cancellation - models: OpenAiModel/AnthropicModel/OllamaModel/LmStudioModel ported from llm-client onto the Model trait; ModelError carries the HTTP status; is_retriable default = the 401/403/404/422 rule - DTL as crate protocol (ToolRendering Inline/DeferredToolReference/ SystemToolBlock; Anthropic conversions + Kimi system+tools passthrough), host catalog behind ActivationSource/ToolActivator - HistoryStore durability contract + InMemoryStore; LinearAssembler with well-formed projection (incl. DTL injection, summary, crash survivors) - LoopManager singleton (broadcast bus + live registry), one live loop per conversation, orphan-marking on start_turn - 32 tests green (kernel §13 suite, assembler DTL, SSE/Anthropic ports), clippy clean
This commit is contained in:
@@ -0,0 +1,775 @@
|
||||
//! Anthropic client (`/v1/messages`). Ported from `llm-client/src/anthropic.rs`
|
||||
//! onto the `Model` trait — including the DTL conversions (blueprint §4.10):
|
||||
//! `defer_loading`, `_tool_references` → `tool_reference` blocks, and the
|
||||
//! `cache_control` breakpoint moved onto the last non-deferred tool.
|
||||
|
||||
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 super::{SseDecoder, error_response_body, headers_to_json, redact_key};
|
||||
use crate::APP_NAME;
|
||||
use crate::model::{
|
||||
Model, ModelError, ModelRequest, ModelResponse, NamedModel, RawMeta, StreamDelta, ToolCall,
|
||||
Usage,
|
||||
};
|
||||
|
||||
const DEFAULT_BASE_URL: &str = "https://api.anthropic.com";
|
||||
const ANTHROPIC_VERSION: &str = "2023-06-01";
|
||||
|
||||
pub struct AnthropicModel {
|
||||
base_url: String,
|
||||
api_key: String,
|
||||
default_model: String,
|
||||
/// Extra top-level request-body keys merged into every request (e.g. the
|
||||
/// `thinking` config for extended reasoning).
|
||||
extra_body: Option<Value>,
|
||||
app_name: String,
|
||||
http: reqwest::Client,
|
||||
}
|
||||
|
||||
impl AnthropicModel {
|
||||
pub fn new(api_key: impl Into<String>, default_model: impl Into<String>) -> Self {
|
||||
Self::with_extra_body(api_key, default_model, None)
|
||||
}
|
||||
|
||||
pub fn with_base_url(
|
||||
base_url: impl Into<String>,
|
||||
api_key: impl Into<String>,
|
||||
default_model: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
base_url: base_url.into(),
|
||||
api_key: api_key.into(),
|
||||
default_model: default_model.into(),
|
||||
extra_body: None,
|
||||
app_name: APP_NAME.to_string(),
|
||||
http: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Like `new` but with extra request-body keys (e.g. `{"thinking": {...}}`).
|
||||
pub fn with_extra_body(
|
||||
api_key: impl Into<String>,
|
||||
default_model: impl Into<String>,
|
||||
extra_body: Option<Value>,
|
||||
) -> Self {
|
||||
Self {
|
||||
base_url: DEFAULT_BASE_URL.to_string(),
|
||||
api_key: api_key.into(),
|
||||
default_model: default_model.into(),
|
||||
extra_body,
|
||||
app_name: APP_NAME.to_string(),
|
||||
http: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_app_name(mut self, app_name: impl Into<String>) -> Self {
|
||||
self.app_name = app_name.into();
|
||||
self
|
||||
}
|
||||
|
||||
/// Merges `extra_body` (then the request's own `extras`) into `body` and
|
||||
/// enforces Anthropic's extended-thinking constraints: when `thinking` is
|
||||
/// enabled, `temperature` is not allowed and `max_tokens` must be strictly
|
||||
/// greater than `budget_tokens`.
|
||||
fn apply_extra(&self, body: &mut Value, req_extras: &Value) {
|
||||
for extra in [self.extra_body.as_ref(), Some(req_extras).filter(|v| v.is_object())]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
{
|
||||
let Some(extra) = extra.as_object() else { continue };
|
||||
let Some(obj) = body.as_object_mut() else { return };
|
||||
for (k, v) in extra {
|
||||
obj.insert(k.clone(), v.clone());
|
||||
}
|
||||
}
|
||||
let Some(obj) = body.as_object_mut() else { return };
|
||||
if obj.get("thinking").map(|t| t["type"] == json!("enabled")).unwrap_or(false) {
|
||||
obj.remove("temperature");
|
||||
let budget = obj["thinking"]["budget_tokens"].as_i64().unwrap_or(0);
|
||||
let cur_max = obj.get("max_tokens").and_then(|v| v.as_i64()).unwrap_or(4096);
|
||||
if budget > 0 && cur_max <= budget {
|
||||
obj.insert("max_tokens".to_string(), json!(budget + 4096));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Converts OpenAI-format tool definitions to Anthropic format.
|
||||
/// OpenAI: { "type": "function", "function": { "name", "description", "parameters" } }
|
||||
/// Anthropic: { "name", "description", "input_schema" }
|
||||
///
|
||||
/// DTL (`DeferredToolReference`): a top-level `defer_loading: true` on the
|
||||
/// OpenAI tool object is carried through. When any tool is deferred, the
|
||||
/// cache breakpoint is placed on the last **non-deferred** tool — a
|
||||
/// deferred tool cannot carry `cache_control` (the API 400s).
|
||||
fn convert_tools(tools: &[Value]) -> Vec<Value> {
|
||||
let has_deferred = tools.iter().any(|t| t["defer_loading"].as_bool() == Some(true));
|
||||
let mut out: Vec<Value> = tools
|
||||
.iter()
|
||||
.filter_map(|t| {
|
||||
let func = &t["function"];
|
||||
let name = func["name"].as_str()?;
|
||||
let mut tool = json!({
|
||||
"name": name,
|
||||
"description": func["description"].as_str().unwrap_or(""),
|
||||
"input_schema": func["parameters"],
|
||||
});
|
||||
if t["defer_loading"].as_bool() == Some(true) {
|
||||
tool["defer_loading"] = json!(true);
|
||||
}
|
||||
Some(tool)
|
||||
})
|
||||
.collect();
|
||||
if has_deferred
|
||||
&& let Some(t) = out.iter_mut().rev().find(|t| t["defer_loading"].as_bool() != Some(true))
|
||||
{
|
||||
t["cache_control"] = json!({ "type": "ephemeral" });
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Converts OpenAI-format messages to Anthropic format: system extracted
|
||||
/// separately; assistant tool_calls → tool_use blocks; consecutive `tool`
|
||||
/// messages grouped into one user message of tool_result blocks.
|
||||
fn convert_messages(messages: &[Value]) -> Vec<Value> {
|
||||
let mut out: Vec<Value> = Vec::new();
|
||||
let mut i = 0;
|
||||
|
||||
while i < messages.len() {
|
||||
let msg = &messages[i];
|
||||
let role = msg["role"].as_str().unwrap_or("");
|
||||
|
||||
match role {
|
||||
"system" => { i += 1; }
|
||||
|
||||
"user" => {
|
||||
out.push(json!({
|
||||
"role": "user",
|
||||
"content": convert_user_content(&msg["content"]),
|
||||
}));
|
||||
i += 1;
|
||||
}
|
||||
|
||||
"assistant" => {
|
||||
if let Some(tool_calls) = msg["tool_calls"].as_array() {
|
||||
let mut content: Vec<Value> = Vec::new();
|
||||
|
||||
let text = msg["content"].as_str().unwrap_or("");
|
||||
if !text.is_empty() {
|
||||
content.push(json!({ "type": "text", "text": text }));
|
||||
}
|
||||
|
||||
for tc in tool_calls {
|
||||
let id = tc["id"].as_str().unwrap_or("");
|
||||
let name = tc["function"]["name"].as_str().unwrap_or("");
|
||||
let args_str = tc["function"]["arguments"].as_str().unwrap_or("{}");
|
||||
let input: Value = serde_json::from_str(args_str)
|
||||
.unwrap_or(Value::Object(Default::default()));
|
||||
|
||||
content.push(json!({
|
||||
"type": "tool_use",
|
||||
"id": id,
|
||||
"name": name,
|
||||
"input": input,
|
||||
}));
|
||||
}
|
||||
|
||||
out.push(json!({ "role": "assistant", "content": content }));
|
||||
} else {
|
||||
out.push(json!({
|
||||
"role": "assistant",
|
||||
"content": msg["content"].as_str().unwrap_or(""),
|
||||
}));
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
|
||||
"tool" => {
|
||||
// Group consecutive tool results into a single user message.
|
||||
let mut results: Vec<Value> = Vec::new();
|
||||
while i < messages.len() && messages[i]["role"].as_str() == Some("tool") {
|
||||
let tm = &messages[i];
|
||||
// DTL (`DeferredToolReference`): a tool result carrying
|
||||
// `_tool_references` becomes a content array of
|
||||
// `tool_reference` blocks, which the API expands into
|
||||
// the deferred tools' full definitions.
|
||||
let content: Value = match tm["_tool_references"].as_array() {
|
||||
Some(refs) if !refs.is_empty() => Value::Array(
|
||||
refs.iter()
|
||||
.filter_map(|r| r.as_str())
|
||||
.map(|name| json!({ "type": "tool_reference", "tool_name": name }))
|
||||
.collect(),
|
||||
),
|
||||
_ => Value::String(tm["content"].as_str().unwrap_or("").to_string()),
|
||||
};
|
||||
results.push(json!({
|
||||
"type": "tool_result",
|
||||
"tool_use_id": tm["tool_call_id"].as_str().unwrap_or(""),
|
||||
"content": content,
|
||||
}));
|
||||
i += 1;
|
||||
}
|
||||
out.push(json!({ "role": "user", "content": results }));
|
||||
}
|
||||
|
||||
_ => { i += 1; }
|
||||
}
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
/// Shared `/v1/messages` body (the caller adds `stream` on top).
|
||||
fn tools_body(&self, system: Option<Value>, messages: Vec<Value>, tools: Vec<Value>, req: &ModelRequest) -> Value {
|
||||
let max_tokens = req.max_tokens.unwrap_or(4096);
|
||||
let mut body = json!({
|
||||
"model": req.model,
|
||||
"max_tokens": max_tokens,
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
});
|
||||
|
||||
if let Some(sys) = system { body["system"] = sys; }
|
||||
if let Some(t) = req.temperature { body["temperature"] = t.into(); }
|
||||
self.apply_extra(&mut body, &req.extras);
|
||||
body
|
||||
}
|
||||
|
||||
/// Collects ALL system-role messages into the single `system` parameter.
|
||||
/// Structured content (a text-block array with `cache_control`) is kept
|
||||
/// in array form so the cache breakpoint survives.
|
||||
fn merged_system(messages: &[Value]) -> Option<Value> {
|
||||
let sys: Vec<&Value> = messages
|
||||
.iter()
|
||||
.filter(|m| m["role"].as_str() == Some("system"))
|
||||
.collect();
|
||||
if sys.is_empty() { return None; }
|
||||
|
||||
if !sys.iter().any(|m| m["content"].is_array()) {
|
||||
let parts: Vec<&str> = sys.iter().filter_map(|m| m["content"].as_str()).collect();
|
||||
return if parts.is_empty() { None } else { Some(Value::String(parts.join("\n\n---\n\n"))) };
|
||||
}
|
||||
|
||||
let mut blocks: Vec<Value> = Vec::new();
|
||||
for m in &sys {
|
||||
match &m["content"] {
|
||||
Value::String(s) if !s.is_empty() => blocks.push(json!({ "type": "text", "text": s })),
|
||||
Value::Array(arr) => {
|
||||
for b in arr {
|
||||
if b["type"].as_str() == Some("text") {
|
||||
blocks.push(b.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
if blocks.is_empty() { None } else { Some(Value::Array(blocks)) }
|
||||
}
|
||||
|
||||
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",
|
||||
})
|
||||
}
|
||||
|
||||
/// Sends the request WITHOUT `error_for_status`, so the caller can read
|
||||
/// the error body and attach the payload to the `ModelError`.
|
||||
async fn send_request(&self, body: &Value) -> Result<reqwest::Response, ModelError> {
|
||||
self.http
|
||||
.post(self.url())
|
||||
.header("x-api-key", &self.api_key)
|
||||
.header("anthropic-version", ANTHROPIC_VERSION)
|
||||
.header("X-Title", &self.app_name)
|
||||
.json(body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(ModelError::from_reqwest)
|
||||
}
|
||||
|
||||
/// Joined `thinking` blocks of a content array (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")) }
|
||||
}
|
||||
|
||||
/// The buffered path.
|
||||
async fn buffered(&self, req: &ModelRequest) -> Result<ModelResponse, ModelError> {
|
||||
let system = Self::merged_system(&req.messages);
|
||||
let anthropic_messages = Self::convert_messages(&req.messages);
|
||||
let anthropic_tools = Self::convert_tools(&req.tools);
|
||||
let body = self.tools_body(system, anthropic_messages, anthropic_tools, req);
|
||||
|
||||
debug!(model = %req.model, tools = req.tools.len(), "anthropic: sending request");
|
||||
trace!(body = %body, "anthropic: 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();
|
||||
let resp_text = http_resp.text().await.map_err(ModelError::from_reqwest)?;
|
||||
if !status.is_success() {
|
||||
return Err(ModelError {
|
||||
status: Some(status.as_u16()),
|
||||
message: format!("anthropic: HTTP {status} from {url}\nbody: {resp_text}", url = self.url()),
|
||||
raw: Some(RawMeta {
|
||||
request_headers: Some(request_headers),
|
||||
request_body: Some(request_body),
|
||||
response_headers: Some(response_headers),
|
||||
response_body: Some(error_response_body(resp_text)),
|
||||
}),
|
||||
});
|
||||
}
|
||||
let resp: Value = serde_json::from_str(&resp_text).map_err(|e| {
|
||||
ModelError::new(None, format!("anthropic: failed to parse response JSON: {e}\nbody: {resp_text}"))
|
||||
})?;
|
||||
|
||||
let raw = RawMeta {
|
||||
request_headers: Some(request_headers),
|
||||
request_body: Some(request_body),
|
||||
response_headers: Some(response_headers),
|
||||
response_body: Some(resp.clone()),
|
||||
};
|
||||
|
||||
let stop_reason = resp["stop_reason"].as_str().unwrap_or("");
|
||||
let mut usage = Usage {
|
||||
input_tokens: resp["usage"]["input_tokens"].as_u64().map(|n| n as u32),
|
||||
output_tokens: resp["usage"]["output_tokens"].as_u64().map(|n| n as u32),
|
||||
cache_read: resp["usage"]["cache_read_input_tokens"].as_u64().map(|n| n as u32),
|
||||
cache_write: resp["usage"]["cache_creation_input_tokens"].as_u64().map(|n| n as u32),
|
||||
cost_usd: None,
|
||||
truncated: stop_reason == "max_tokens",
|
||||
};
|
||||
let content_blocks = resp["content"].as_array().cloned().unwrap_or_default();
|
||||
info!(model = %req.model, ?usage.input_tokens, ?usage.output_tokens, stop_reason, "anthropic: response received");
|
||||
if usage.truncated {
|
||||
warn!(model = %req.model, ?usage.output_tokens, "anthropic: response truncated (max_tokens reached)");
|
||||
}
|
||||
|
||||
let has_tool_use = content_blocks.iter().any(|b| b["type"].as_str() == Some("tool_use"));
|
||||
let reasoning = Self::reasoning_of(&content_blocks);
|
||||
|
||||
// Anthropic sometimes returns stop_reason "end_turn" even when
|
||||
// tool_use blocks are present — check the blocks directly.
|
||||
let mut resp_out = if stop_reason == "tool_use" || has_tool_use {
|
||||
let text: String = content_blocks
|
||||
.iter()
|
||||
.filter(|b| b["type"].as_str() == Some("text"))
|
||||
.filter_map(|b| b["text"].as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
usage.truncated = false;
|
||||
let calls: Vec<ToolCall> = content_blocks
|
||||
.iter()
|
||||
.filter(|b| b["type"].as_str() == Some("tool_use"))
|
||||
.map(|b| ToolCall {
|
||||
id: b["id"].as_str().unwrap_or("").to_string(),
|
||||
name: b["name"].as_str().unwrap_or("").to_string(),
|
||||
arguments: b["input"].clone(),
|
||||
})
|
||||
.collect();
|
||||
ModelResponse::ToolCalls { content: text, calls, reasoning, usage, raw: None }
|
||||
} else {
|
||||
let content = content_blocks
|
||||
.iter()
|
||||
.find(|b| b["type"].as_str() == Some("text"))
|
||||
.and_then(|b| b["text"].as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
ModelResponse::Message { content, reasoning, usage, raw: None }
|
||||
};
|
||||
match &mut resp_out {
|
||||
ModelResponse::Message { raw: r, .. } | ModelResponse::ToolCalls { raw: r, .. } => {
|
||||
*r = Some(raw)
|
||||
}
|
||||
}
|
||||
Ok(resp_out)
|
||||
}
|
||||
|
||||
/// SSE streaming path: Anthropic streams typed events (`message_start` /
|
||||
/// `content_block_*` / `message_delta`); text and thinking deltas are
|
||||
/// forwarded best-effort while blocks accumulate into the same
|
||||
/// `ModelResponse` the buffered path returns.
|
||||
#[allow(clippy::result_large_err)]
|
||||
async fn stream_chat(
|
||||
&self,
|
||||
req: &ModelRequest,
|
||||
delta_tx: &mpsc::Sender<StreamDelta>,
|
||||
emitted: &mut bool,
|
||||
) -> Result<ModelResponse, ModelError> {
|
||||
let system = Self::merged_system(&req.messages);
|
||||
let anthropic_messages = Self::convert_messages(&req.messages);
|
||||
let anthropic_tools = Self::convert_tools(&req.tools);
|
||||
let mut body = self.tools_body(system, anthropic_messages, anthropic_tools, req);
|
||||
body["stream"] = json!(true);
|
||||
|
||||
debug!(model = %req.model, tools = req.tools.len(), "anthropic: sending streaming request");
|
||||
trace!(body = %body, "anthropic: streaming 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.map_err(ModelError::from_reqwest)?;
|
||||
return Err(ModelError {
|
||||
status: Some(status.as_u16()),
|
||||
message: format!("anthropic: HTTP {status} from {url}\nbody: {resp_text}", url = self.url()),
|
||||
raw: Some(RawMeta {
|
||||
request_headers: Some(request_headers),
|
||||
request_body: Some(request_body),
|
||||
response_headers: Some(response_headers),
|
||||
response_body: Some(error_response_body(resp_text)),
|
||||
}),
|
||||
});
|
||||
}
|
||||
|
||||
/// 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| -> Result<(), ModelError> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
"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(ModelError::new(None, format!("anthropic: stream error event: {payload}")));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(())
|
||||
};
|
||||
|
||||
while let Some(chunk) = byte_stream.next().await {
|
||||
let chunk = chunk.map_err(ModelError::from_reqwest)?;
|
||||
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 usage_struct = Usage {
|
||||
input_tokens: usage["input_tokens"].as_u64().map(|n| n as u32),
|
||||
output_tokens: usage["output_tokens"].as_u64().map(|n| n as u32),
|
||||
cache_read: usage["cache_read_input_tokens"].as_u64().map(|n| n as u32),
|
||||
cache_write: usage["cache_creation_input_tokens"].as_u64().map(|n| n as u32),
|
||||
cost_usd: None,
|
||||
truncated: stop == "max_tokens",
|
||||
};
|
||||
info!(model = %req.model, ?usage_struct.input_tokens, ?usage_struct.output_tokens, stop_reason = stop, "anthropic: streaming response completed");
|
||||
if usage_struct.truncated {
|
||||
warn!(model = %req.model, "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 = text_of("thinking");
|
||||
let reasoning = if reasoning_text.is_empty() { None } else { Some(reasoning_text) };
|
||||
let tool_blocks: Vec<&Block> = blocks.values().filter(|b| b.kind == "tool_use").collect();
|
||||
|
||||
// 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 = RawMeta {
|
||||
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,
|
||||
})),
|
||||
};
|
||||
|
||||
let mut resp_out = 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();
|
||||
ModelResponse::ToolCalls { content: text_of("text"), calls, reasoning, usage: usage_struct, raw: None }
|
||||
} else {
|
||||
ModelResponse::Message { content: text_of("text"), reasoning, usage: usage_struct, raw: None }
|
||||
};
|
||||
match &mut resp_out {
|
||||
ModelResponse::Message { raw: r, .. } | ModelResponse::ToolCalls { raw: r, .. } => {
|
||||
*r = Some(raw)
|
||||
}
|
||||
}
|
||||
Ok(resp_out)
|
||||
}
|
||||
}
|
||||
|
||||
impl NamedModel for AnthropicModel {
|
||||
fn default_model(&self) -> &str { &self.default_model }
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Model for AnthropicModel {
|
||||
async fn complete(
|
||||
&self,
|
||||
req: &ModelRequest,
|
||||
deltas: Option<mpsc::Sender<StreamDelta>>,
|
||||
) -> Result<ModelResponse, ModelError> {
|
||||
match deltas {
|
||||
None => self.buffered(req).await,
|
||||
Some(delta_tx) => {
|
||||
let mut emitted = false;
|
||||
match self.stream_chat(req, &delta_tx, &mut emitted).await {
|
||||
Ok(ok) => Ok(ok),
|
||||
// Pre-stream failure (nothing shown yet): retry buffered.
|
||||
// A mid-stream failure propagates to the fallback logic.
|
||||
Err(e) if !emitted => {
|
||||
debug!(model = %req.model, error = %e, "anthropic: streaming failed before any delta; retrying buffered");
|
||||
self.buffered(req).await
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// User content arrives either as a plain string or as an OpenAI-style parts
|
||||
/// array (text + `image_url` data URLs + `file` PDF parts). Strings pass
|
||||
/// through; parts become Anthropic blocks. Unknown parts are dropped with a
|
||||
/// warning.
|
||||
fn convert_user_content(content: &Value) -> Value {
|
||||
let Some(parts) = content.as_array() else {
|
||||
return Value::String(content.as_str().unwrap_or("").to_string());
|
||||
};
|
||||
let mut blocks = Vec::new();
|
||||
for p in parts {
|
||||
match p["type"].as_str().unwrap_or("") {
|
||||
"text" => blocks.push(json!({
|
||||
"type": "text",
|
||||
"text": p["text"].as_str().unwrap_or(""),
|
||||
})),
|
||||
"image_url" => {
|
||||
if let Some(block) = parse_data_image(&p["image_url"]) {
|
||||
blocks.push(block);
|
||||
}
|
||||
}
|
||||
"file" => {
|
||||
if let Some(block) = parse_data_document(&p["file"]) {
|
||||
blocks.push(block);
|
||||
}
|
||||
}
|
||||
other => tracing::warn!(part_type = other, "dropping content part unsupported by Anthropic"),
|
||||
}
|
||||
}
|
||||
Value::Array(blocks)
|
||||
}
|
||||
|
||||
/// `{"url": "data:<mime>;base64,<data>"}` → an Anthropic base64 image block.
|
||||
fn parse_data_image(image_url: &Value) -> Option<Value> {
|
||||
let url = image_url["url"].as_str().or_else(|| image_url.as_str())?;
|
||||
let (mime, data) = url.strip_prefix("data:")?.split_once(";base64,")?;
|
||||
Some(json!({
|
||||
"type": "image",
|
||||
"source": { "type": "base64", "media_type": mime, "data": data },
|
||||
}))
|
||||
}
|
||||
|
||||
/// `{"file_data": "data:application/pdf;base64,<data>"}` → an Anthropic
|
||||
/// base64 `document` block (the native PDF input).
|
||||
fn parse_data_document(file: &Value) -> Option<Value> {
|
||||
let url = file["file_data"].as_str()?;
|
||||
let (mime, data) = url.strip_prefix("data:")?.split_once(";base64,")?;
|
||||
Some(json!({
|
||||
"type": "document",
|
||||
"source": { "type": "base64", "media_type": mime, "data": data },
|
||||
}))
|
||||
}
|
||||
|
||||
#[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!(
|
||||
AnthropicModel::reasoning_of(&blocks),
|
||||
Some("first\nsecond".to_string())
|
||||
);
|
||||
assert_eq!(AnthropicModel::reasoning_of(&[]), None);
|
||||
assert_eq!(
|
||||
AnthropicModel::reasoning_of(&[json!({"type": "text", "text": "a"})]),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_tools_carries_defer_loading_and_moves_cache_control() {
|
||||
let tools = vec![
|
||||
json!({"type":"function","function":{"name":"a","description":"","parameters":{}}}),
|
||||
json!({"type":"function","function":{"name":"b","description":"","parameters":{}},"defer_loading":true}),
|
||||
json!({"type":"function","function":{"name":"c","description":"","parameters":{}},"defer_loading":true}),
|
||||
];
|
||||
let out = AnthropicModel::convert_tools(&tools);
|
||||
assert_eq!(out[0]["cache_control"], json!({"type": "ephemeral"}));
|
||||
assert!(out[0].get("defer_loading").is_none());
|
||||
assert_eq!(out[1]["defer_loading"], json!(true));
|
||||
assert!(out[1].get("cache_control").is_none());
|
||||
assert_eq!(out[2]["defer_loading"], json!(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_messages_tool_references_become_blocks() {
|
||||
let messages = vec![
|
||||
json!({"role":"assistant","content":"","tool_calls":[
|
||||
{"id":"t1","type":"function","function":{"name":"activate_tools","arguments":"{\"groups\":[\"gmail\"]}"}}
|
||||
]}),
|
||||
json!({"role":"tool","tool_call_id":"t1","content":"ok","_tool_references":["mcp__gmail__send"]}),
|
||||
];
|
||||
let out = AnthropicModel::convert_messages(&messages);
|
||||
assert_eq!(out.len(), 2);
|
||||
let results = out[1]["content"].as_array().unwrap();
|
||||
assert_eq!(
|
||||
results[0]["content"],
|
||||
json!([{ "type": "tool_reference", "tool_name": "mcp__gmail__send" }])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_content_string_passthrough() {
|
||||
let v = convert_user_content(&json!("hello"));
|
||||
assert_eq!(v, json!("hello"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_content_parts_become_anthropic_blocks() {
|
||||
let v = convert_user_content(&json!([
|
||||
{ "type": "text", "text": "what is this?" },
|
||||
{ "type": "image_url", "image_url": { "url": "data:image/png;base64,QUJD" } },
|
||||
]));
|
||||
assert_eq!(v, json!([
|
||||
{ "type": "text", "text": "what is this?" },
|
||||
{ "type": "image", "source": { "type": "base64", "media_type": "image/png", "data": "QUJD" } },
|
||||
]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_content_drops_video_and_non_data_urls() {
|
||||
let v = convert_user_content(&json!([
|
||||
{ "type": "text", "text": "t" },
|
||||
{ "type": "video_url", "video_url": { "url": "data:video/mp4;base64,QUJD" } },
|
||||
{ "type": "image_url", "image_url": { "url": "https://example.com/x.png" } },
|
||||
]));
|
||||
assert_eq!(v, json!([{ "type": "text", "text": "t" }]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_content_file_part_becomes_document_block() {
|
||||
let v = convert_user_content(&json!([
|
||||
{ "type": "text", "text": "read this" },
|
||||
{ "type": "file", "file": { "filename": "a.pdf", "file_data": "data:application/pdf;base64,QUJD" } },
|
||||
]));
|
||||
assert_eq!(v, json!([
|
||||
{ "type": "text", "text": "read this" },
|
||||
{ "type": "document", "source": { "type": "base64", "media_type": "application/pdf", "data": "QUJD" } },
|
||||
]));
|
||||
|
||||
let v = convert_user_content(&json!([
|
||||
{ "type": "file", "file": { "filename": "a.pdf", "file_data": "https://example.com/a.pdf" } },
|
||||
]));
|
||||
assert_eq!(v, json!([]));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
//! LM Studio client — a thin wrapper over [`OpenAiModel`] defaulting to
|
||||
//! `http://localhost:1234/v1` with no API key. (LM Studio can also be served
|
||||
//! by a YAML-declared provider; this client is kept for explicit use.)
|
||||
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use super::openai::OpenAiModel;
|
||||
use crate::model::{Model, ModelError, ModelRequest, ModelResponse, NamedModel, StreamDelta};
|
||||
|
||||
pub struct LmStudioModel {
|
||||
inner: OpenAiModel,
|
||||
}
|
||||
|
||||
impl LmStudioModel {
|
||||
/// `base_url` defaults to `http://localhost:1234/v1` if `None`.
|
||||
pub fn new(base_url: Option<impl Into<String>>, default_model: impl Into<String>) -> Self {
|
||||
let url = base_url
|
||||
.map(|u| u.into())
|
||||
.unwrap_or_else(|| "http://localhost:1234/v1".to_string());
|
||||
Self { inner: OpenAiModel::new(url, "", default_model) }
|
||||
}
|
||||
}
|
||||
|
||||
impl NamedModel for LmStudioModel {
|
||||
fn default_model(&self) -> &str { self.inner.default_model() }
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Model for LmStudioModel {
|
||||
/// LM Studio is OpenAI-compatible: everything forwards to the inner
|
||||
/// client (its pre-delta buffered retry covers local builds rejecting
|
||||
/// `stream_options`).
|
||||
async fn complete(
|
||||
&self,
|
||||
req: &ModelRequest,
|
||||
deltas: Option<mpsc::Sender<StreamDelta>>,
|
||||
) -> Result<ModelResponse, ModelError> {
|
||||
self.inner.complete(req, deltas).await
|
||||
}
|
||||
|
||||
fn is_retriable(&self, err: &ModelError) -> bool { self.inner.is_retriable(err) }
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
//! Shipped `Model` clients (blueprint D13): OpenAI-compatible, Anthropic,
|
||||
//! Ollama, LM Studio — plus the shared SSE decoder and HTTP helpers.
|
||||
//!
|
||||
//! All clients are stateless (connection config only) and share the same
|
||||
//! failure policy: if a stream dies BEFORE any delta, the client retries
|
||||
//! buffered on the same model (providers rejecting `stream` keep working); a
|
||||
//! mid-stream failure propagates to the caller's fallback logic.
|
||||
|
||||
pub mod anthropic;
|
||||
pub mod lm_studio;
|
||||
pub mod ollama;
|
||||
pub mod openai;
|
||||
mod sse;
|
||||
|
||||
pub use anthropic::AnthropicModel;
|
||||
pub use lm_studio::LmStudioModel;
|
||||
pub use ollama::OllamaModel;
|
||||
pub use openai::OpenAiModel;
|
||||
pub(crate) use sse::SseDecoder;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
/// Converts a reqwest `HeaderMap` into a JSON object (for payload logging).
|
||||
pub(crate) fn headers_to_json(headers: &reqwest::header::HeaderMap) -> Value {
|
||||
let map: serde_json::Map<String, Value> = headers
|
||||
.iter()
|
||||
.map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("<binary>").into()))
|
||||
.collect();
|
||||
Value::Object(map)
|
||||
}
|
||||
|
||||
/// Raw error body → JSON for the payload log: parsed JSON when the provider
|
||||
/// returned JSON, else the raw text wrapped as a JSON string so a non-JSON
|
||||
/// body (HTML gateway page) is still preserved verbatim.
|
||||
pub(crate) fn error_response_body(text: String) -> Value {
|
||||
serde_json::from_str::<Value>(&text).unwrap_or(Value::String(text))
|
||||
}
|
||||
|
||||
/// Redacted preview of an API key: first 7 chars + "***".
|
||||
pub(crate) fn redact_key(key: &str) -> String {
|
||||
if key.len() > 7 { format!("{}***", &key[..7]) } else { "***".to_string() }
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
//! Ollama client (native `/api/chat` endpoint). Ported from
|
||||
//! `llm-client/src/ollama.rs`. No streaming, no tool support — tool-call
|
||||
//! messages are flattened to text, mirroring the previous default behavior.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::model::{Model, ModelError, ModelRequest, ModelResponse, NamedModel, StreamDelta, Usage};
|
||||
|
||||
/// Ollama client. Defaults to `http://localhost:11434`. No API key required.
|
||||
pub struct OllamaModel {
|
||||
base_url: String,
|
||||
default_model: String,
|
||||
http: reqwest::Client,
|
||||
}
|
||||
|
||||
impl OllamaModel {
|
||||
/// `base_url` defaults to `http://localhost:11434` if `None`.
|
||||
pub fn new(base_url: Option<impl Into<String>>, default_model: impl Into<String>) -> Self {
|
||||
let url = base_url
|
||||
.map(|u| u.into())
|
||||
.unwrap_or_else(|| "http://localhost:11434".to_string());
|
||||
Self { base_url: url, default_model: default_model.into(), http: reqwest::Client::new() }
|
||||
}
|
||||
}
|
||||
|
||||
impl NamedModel for OllamaModel {
|
||||
fn default_model(&self) -> &str { &self.default_model }
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Model for OllamaModel {
|
||||
async fn complete(
|
||||
&self,
|
||||
req: &ModelRequest,
|
||||
_deltas: Option<mpsc::Sender<StreamDelta>>,
|
||||
) -> Result<ModelResponse, ModelError> {
|
||||
// Flatten to plain text messages: tool results and assistant
|
||||
// tool_calls are dropped (no native tool support on this path).
|
||||
let msgs: Vec<Value> = req
|
||||
.messages
|
||||
.iter()
|
||||
.filter_map(|m| {
|
||||
let role = m["role"].as_str()?;
|
||||
if !matches!(role, "system" | "user" | "assistant") {
|
||||
return None;
|
||||
}
|
||||
let content = m["content"].as_str().unwrap_or("").to_string();
|
||||
Some(json!({ "role": role, "content": content }))
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut options_obj = json!({});
|
||||
if let Some(t) = req.temperature { options_obj["temperature"] = t.into(); }
|
||||
if let Some(n) = req.max_tokens { options_obj["num_predict"] = n.into(); }
|
||||
|
||||
let body = json!({
|
||||
"model": req.model,
|
||||
"messages": msgs,
|
||||
"stream": false,
|
||||
"options": options_obj,
|
||||
});
|
||||
|
||||
let url = format!("{}/api/chat", self.base_url.trim_end_matches('/'));
|
||||
|
||||
let http_resp = self
|
||||
.http
|
||||
.post(&url)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(ModelError::from_reqwest)?;
|
||||
|
||||
let status = http_resp.status();
|
||||
if !status.is_success() {
|
||||
let resp_text = http_resp.text().await.map_err(ModelError::from_reqwest)?;
|
||||
return Err(ModelError::new(
|
||||
Some(status.as_u16()),
|
||||
format!("ollama: HTTP {status} from {url}\nbody: {resp_text}"),
|
||||
));
|
||||
}
|
||||
|
||||
let resp: Value = http_resp.json().await.map_err(ModelError::from_reqwest)?;
|
||||
|
||||
let content = resp["message"]["content"]
|
||||
.as_str()
|
||||
.ok_or_else(|| ModelError::new(None, "ollama: missing content in response"))?
|
||||
.to_string();
|
||||
|
||||
Ok(ModelResponse::Message {
|
||||
content,
|
||||
reasoning: None,
|
||||
usage: Usage {
|
||||
input_tokens: resp["prompt_eval_count"].as_u64().map(|n| n as u32),
|
||||
output_tokens: resp["eval_count"].as_u64().map(|n| n as u32),
|
||||
..Usage::default()
|
||||
},
|
||||
raw: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,459 @@
|
||||
//! OpenAI-compatible client (OpenAI, OpenRouter, Moonshot/Kimi, and every
|
||||
//! provider declared via YAML). Ported from `llm-client/src/openai.rs` onto
|
||||
//! the `Model` trait.
|
||||
//!
|
||||
//! Kimi's `SystemToolBlock` DTL needs NO client code: messages are passed
|
||||
//! through verbatim and the endpoint speaks the `{role:"system", tools:[…]}`
|
||||
//! convention natively.
|
||||
|
||||
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 super::{SseDecoder, error_response_body, headers_to_json, redact_key};
|
||||
use crate::APP_NAME;
|
||||
use crate::model::{
|
||||
Model, ModelError, ModelRequest, ModelResponse, NamedModel, RawMeta, StreamDelta, ToolCall,
|
||||
Usage,
|
||||
};
|
||||
|
||||
/// OpenAI ChatGPT client (also compatible with any OpenAI-spec endpoint).
|
||||
pub struct OpenAiModel {
|
||||
base_url: String,
|
||||
api_key: String,
|
||||
default_model: String,
|
||||
extra_params: Option<Value>,
|
||||
/// When true, Anthropic-compatible prompt-caching hints are injected
|
||||
/// (OpenRouter routing to Anthropic models).
|
||||
enable_prompt_cache: bool,
|
||||
app_name: String,
|
||||
http: reqwest::Client,
|
||||
}
|
||||
|
||||
impl OpenAiModel {
|
||||
/// Minimal constructor: base URL + key + default model name (used as the
|
||||
/// selector id by `SingleModel`).
|
||||
pub fn new(
|
||||
base_url: impl Into<String>,
|
||||
api_key: impl Into<String>,
|
||||
default_model: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::with_options(base_url, api_key, default_model, None, false)
|
||||
}
|
||||
|
||||
pub fn with_options(
|
||||
base_url: impl Into<String>,
|
||||
api_key: impl Into<String>,
|
||||
default_model: impl Into<String>,
|
||||
extra_params: Option<Value>,
|
||||
enable_prompt_cache: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
base_url: base_url.into(),
|
||||
api_key: api_key.into(),
|
||||
default_model: default_model.into(),
|
||||
extra_params,
|
||||
enable_prompt_cache,
|
||||
app_name: APP_NAME.to_string(),
|
||||
http: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Override the `X-Title` header (OpenRouter rankings).
|
||||
pub fn with_app_name(mut self, app_name: impl Into<String>) -> Self {
|
||||
self.app_name = app_name.into();
|
||||
self
|
||||
}
|
||||
|
||||
/// Merges extra top-level object keys into `body` (later maps win).
|
||||
fn merge_extra(body: &mut Value, extra: Option<&Value>) {
|
||||
if let Some(Value::Object(extra)) = extra
|
||||
&& let Some(b) = body.as_object_mut()
|
||||
{
|
||||
for (k, v) in extra {
|
||||
b.insert(k.clone(), v.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn url(&self) -> String {
|
||||
format!("{}/chat/completions", self.base_url.trim_end_matches('/'))
|
||||
}
|
||||
|
||||
/// Shared request body for the buffered and the streaming path.
|
||||
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 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, req: &ModelRequest) -> Value {
|
||||
if let Some(t) = req.max_tokens { body["max_tokens"] = t.into(); }
|
||||
if let Some(t) = req.temperature { body["temperature"] = t.into(); }
|
||||
Self::merge_extra(&mut body, self.extra_params.as_ref());
|
||||
Self::merge_extra(&mut body, Some(&req.extras));
|
||||
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) -> Result<reqwest::Response, ModelError> {
|
||||
let mut req = self
|
||||
.http
|
||||
.post(self.url())
|
||||
.bearer_auth(&self.api_key)
|
||||
.header("X-Title", &self.app_name);
|
||||
if self.enable_prompt_cache {
|
||||
req = req.header("anthropic-beta", "prompt-caching-2024-07-31");
|
||||
}
|
||||
req.json(body).send().await.map_err(ModelError::from_reqwest)
|
||||
}
|
||||
|
||||
/// The buffered path.
|
||||
async fn buffered(&self, req: &ModelRequest) -> Result<ModelResponse, ModelError> {
|
||||
let body = self.finalize_body(self.base_body(&req.model, &req.messages, &req.tools), req);
|
||||
|
||||
debug!(model = %req.model, tools = req.tools.len(), prompt_cache = self.enable_prompt_cache, "openai: sending request");
|
||||
trace!(body = %body, "openai: 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();
|
||||
let resp_text = http_resp.text().await.map_err(ModelError::from_reqwest)?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(ModelError {
|
||||
status: Some(status.as_u16()),
|
||||
message: format!("openai: HTTP {status} from {url}\nbody: {resp_text}", url = self.url()),
|
||||
raw: Some(RawMeta {
|
||||
request_headers: Some(request_headers),
|
||||
request_body: Some(request_body),
|
||||
response_headers: Some(response_headers),
|
||||
response_body: Some(error_response_body(resp_text)),
|
||||
}),
|
||||
});
|
||||
}
|
||||
|
||||
let resp: Value = serde_json::from_str(&resp_text).map_err(|e| {
|
||||
ModelError::new(None, format!("openai: failed to parse response JSON: {e}\nbody: {resp_text}"))
|
||||
})?;
|
||||
let response_body: Value = serde_json::from_str(&resp_text).unwrap_or(Value::Null);
|
||||
|
||||
let raw = RawMeta {
|
||||
request_headers: Some(request_headers),
|
||||
request_body: Some(request_body),
|
||||
response_headers: Some(response_headers),
|
||||
response_body: Some(response_body),
|
||||
};
|
||||
|
||||
Ok(parse_turn(&resp, &req.model).with_raw(raw))
|
||||
}
|
||||
|
||||
/// SSE streaming path. Accumulates fragments into the same `ModelResponse`
|
||||
/// the buffered path returns, forwarding deltas best-effort. `emitted`
|
||||
/// tracks whether any delta was pushed, distinguishing a pre-stream
|
||||
/// failure (safe to retry buffered) from a mid-stream one.
|
||||
async fn stream_chat(
|
||||
&self,
|
||||
req: &ModelRequest,
|
||||
delta_tx: &mpsc::Sender<StreamDelta>,
|
||||
emitted: &mut bool,
|
||||
) -> Result<ModelResponse, ModelError> {
|
||||
let mut body = self.base_body(&req.model, &req.messages, &req.tools);
|
||||
body["stream"] = json!(true);
|
||||
body["stream_options"] = json!({ "include_usage": true });
|
||||
let body = self.finalize_body(body, req);
|
||||
|
||||
debug!(model = %req.model, tools = req.tools.len(), prompt_cache = self.enable_prompt_cache, "openai: sending streaming request");
|
||||
trace!(body = %body, "openai: streaming 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.map_err(ModelError::from_reqwest)?;
|
||||
return Err(ModelError {
|
||||
status: Some(status.as_u16()),
|
||||
message: format!("openai: HTTP {status} from {url}\nbody: {resp_text}", url = self.url()),
|
||||
raw: Some(RawMeta {
|
||||
request_headers: Some(request_headers),
|
||||
request_body: Some(request_body),
|
||||
response_headers: Some(response_headers),
|
||||
response_body: Some(error_response_body(resp_text)),
|
||||
}),
|
||||
});
|
||||
}
|
||||
|
||||
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();
|
||||
|
||||
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()));
|
||||
}
|
||||
// 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.map_err(ModelError::from_reqwest)?;
|
||||
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 = usage.as_ref()
|
||||
.and_then(|u| u["prompt_tokens_details"]["cached_tokens"].as_u64())
|
||||
.map(|n| n as u32);
|
||||
let cost_usd = usage.as_ref().and_then(|u| u["cost"].as_f64());
|
||||
let reasoning_content = if reasoning.is_empty() { None } else { Some(reasoning) };
|
||||
info!(model = %req.model, ?input_tokens, ?output_tokens, finish_reason = finish, "openai: streaming response completed");
|
||||
if finish == "length" {
|
||||
warn!(model = %req.model, ?output_tokens, "openai: response truncated (max_tokens reached)");
|
||||
}
|
||||
|
||||
let usage_struct = Usage {
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cache_read,
|
||||
cache_write: None,
|
||||
cost_usd,
|
||||
truncated: finish == "length",
|
||||
};
|
||||
|
||||
// Reassemble the streamed message for the payload log (buffered shape).
|
||||
let logged_tool_calls: Vec<Value> = tool_calls.iter()
|
||||
.map(|(_idx, (id, name, args))| json!({
|
||||
"id": id,
|
||||
"type": "function",
|
||||
"function": { "name": name, "arguments": args },
|
||||
}))
|
||||
.collect();
|
||||
let mut logged_message = json!({ "role": "assistant", "content": content.clone() });
|
||||
if let Some(rc) = &reasoning_content {
|
||||
logged_message["reasoning_content"] = rc.clone().into();
|
||||
}
|
||||
if !logged_tool_calls.is_empty() {
|
||||
logged_message["tool_calls"] = Value::Array(logged_tool_calls);
|
||||
}
|
||||
let raw = RawMeta {
|
||||
request_headers: Some(request_headers),
|
||||
request_body: Some(request_body),
|
||||
response_headers: Some(response_headers),
|
||||
response_body: Some(json!({
|
||||
"streamed": true,
|
||||
"choices": [{ "finish_reason": finish, "message": logged_message }],
|
||||
"usage": usage,
|
||||
})),
|
||||
};
|
||||
|
||||
let mut resp = 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();
|
||||
ModelResponse::ToolCalls { content, calls, reasoning: reasoning_content, usage: usage_struct, raw: None }
|
||||
} else {
|
||||
ModelResponse::Message { content, reasoning: reasoning_content, usage: usage_struct, raw: None }
|
||||
};
|
||||
set_raw(&mut resp, raw);
|
||||
Ok(resp)
|
||||
}
|
||||
}
|
||||
|
||||
impl NamedModel for OpenAiModel {
|
||||
fn default_model(&self) -> &str { &self.default_model }
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Model for OpenAiModel {
|
||||
async fn complete(
|
||||
&self,
|
||||
req: &ModelRequest,
|
||||
deltas: Option<mpsc::Sender<StreamDelta>>,
|
||||
) -> Result<ModelResponse, ModelError> {
|
||||
match deltas {
|
||||
None => self.buffered(req).await,
|
||||
Some(delta_tx) => {
|
||||
let mut emitted = false;
|
||||
match self.stream_chat(req, &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. A mid-stream
|
||||
// failure instead propagates to the fallback logic.
|
||||
Err(e) if !emitted => {
|
||||
debug!(model = %req.model, error = %e, "openai: streaming failed before any delta; retrying buffered");
|
||||
self.buffered(req).await
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── response parsing (shared by buffered and tests) ──
|
||||
|
||||
trait WithRaw {
|
||||
fn with_raw(self, raw: RawMeta) -> ModelResponse;
|
||||
}
|
||||
|
||||
impl WithRaw for ModelResponse {
|
||||
fn with_raw(mut self, raw: RawMeta) -> ModelResponse {
|
||||
set_raw(&mut self, raw);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
fn set_raw(resp: &mut ModelResponse, raw: RawMeta) {
|
||||
match resp {
|
||||
ModelResponse::Message { raw: r, .. } | ModelResponse::ToolCalls { raw: r, .. } => {
|
||||
*r = Some(raw)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a buffered OpenAI response body into a `ModelResponse`.
|
||||
fn parse_turn(resp: &Value, model: &str) -> ModelResponse {
|
||||
let usage = Usage {
|
||||
input_tokens: resp["usage"]["prompt_tokens"].as_u64().map(|n| n as u32),
|
||||
output_tokens: resp["usage"]["completion_tokens"].as_u64().map(|n| n as u32),
|
||||
cache_read: resp["usage"]["prompt_tokens_details"]["cached_tokens"].as_u64().map(|n| n as u32),
|
||||
cache_write: None,
|
||||
cost_usd: resp["usage"]["cost"].as_f64(),
|
||||
truncated: false,
|
||||
};
|
||||
|
||||
let choice = &resp["choices"][0];
|
||||
let message = &choice["message"];
|
||||
let finish = choice["finish_reason"].as_str().unwrap_or("stop");
|
||||
if finish == "length" {
|
||||
warn!(model = %model, "openai: response truncated (max_tokens reached)");
|
||||
}
|
||||
|
||||
let reasoning_content = message["reasoning_content"].as_str()
|
||||
.or_else(|| message["reasoning"].as_str())
|
||||
.map(str::to_string);
|
||||
|
||||
let tool_calls_array = message["tool_calls"].as_array().filter(|a| !a.is_empty());
|
||||
|
||||
// Some models (e.g. Qwen via OpenRouter) return finish_reason "stop" even
|
||||
// when tool_calls are present, so check the array directly.
|
||||
if finish == "tool_calls" || tool_calls_array.is_some() {
|
||||
let content = message["content"].as_str().unwrap_or("").to_string();
|
||||
let calls = tool_calls_array
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.map(|tc| ToolCall {
|
||||
id: tc["id"].as_str().unwrap_or("").to_string(),
|
||||
name: tc["function"]["name"].as_str().unwrap_or("").to_string(),
|
||||
arguments: tc["function"]["arguments"]
|
||||
.as_str()
|
||||
.and_then(|s| serde_json::from_str(s).ok())
|
||||
.unwrap_or(Value::Object(Default::default())),
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
ModelResponse::ToolCalls { content, calls, reasoning: reasoning_content, usage, raw: None }
|
||||
} else {
|
||||
// content can be null for thinking models or finish_reason="length".
|
||||
let content = match message["content"].as_str() {
|
||||
Some(s) => s.to_string(),
|
||||
None => {
|
||||
warn!(finish_reason = finish, raw_message = %message, "openai: response has null content");
|
||||
String::new()
|
||||
}
|
||||
};
|
||||
let mut usage = usage;
|
||||
usage.truncated = finish == "length";
|
||||
ModelResponse::Message { content, reasoning: reasoning_content, usage, raw: None }
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
//! 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).
|
||||
//!
|
||||
//! Ported verbatim from `llm-client`.
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct SseDecoder {
|
||||
buf: Vec<u8>,
|
||||
}
|
||||
|
||||
impl SseDecoder {
|
||||
pub(crate) fn new() -> Self { Self::default() }
|
||||
|
||||
pub(crate) 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(crate) 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()) }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SseDecoder;
|
||||
|
||||
#[test]
|
||||
fn sse_decoder_buffers_partial_lines_across_chunks() {
|
||||
let mut dec = SseDecoder::new();
|
||||
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() {
|
||||
// "€" 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);
|
||||
let mut dec = SseDecoder::new();
|
||||
let (first, second) = (dec.feed(a), dec.feed(b));
|
||||
assert!(first.is_empty());
|
||||
assert_eq!(second.len(), 1);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user