b3nw/zerogravity
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: completions API improvements, gemini endpoint, response types

Browse Source
This commit is contained in:
Nikketryhard
2026-02-15 17:08:53 -06:00
parent afa96b88a5
commit ca9f808ee3
8 changed files with 1031 additions and 742 deletions
Download Patch File Download Diff File Expand all files Collapse all files

614
src/api/completions.rs
View File

@@ -15,45 +15,77 @@ use super::types::*;
use super::util::{err_response, now_unix};
use super::AppState;
/// Extract a conversation/session ID from a flexible JSON value.
/// Accepts a plain string or an object with an "id" field.
fn extract_conversation_id(conv: &Option<serde_json::Value>) -> Option<String> {
match conv {
Some(serde_json::Value::String(s)) => Some(s.clone()),
Some(obj) => obj["id"].as_str().map(|s| s.to_string()),
None => None,
}
}
/// System fingerprint for completions responses (derived from crate version at compile time).
fn system_fingerprint() -> String {
format!("fp_{}", env!("CARGO_PKG_VERSION").replace('.', ""))
}
/// Build a streaming chunk JSON with all required OpenAI fields.
/// Includes system_fingerprint, service_tier, and logprobs:null in choices.
fn chunk_json(
id: &str, model: &str,
choices: serde_json::Value,
usage: Option<serde_json::Value>,
) -> String {
let mut obj = serde_json::json!({
"id": id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model,
"system_fingerprint": system_fingerprint(),
"service_tier": "default",
"choices": choices,
});
if let Some(u) = usage {
obj["usage"] = u;
}
serde_json::to_string(&obj).unwrap_or_default()
}
/// Build a single choice for a streaming chunk (delta + finish_reason + logprobs).
fn chunk_choice(index: u32, delta: serde_json::Value, finish_reason: Option<&str>) -> serde_json::Value {
serde_json::json!({
"index": index,
"delta": delta,
"logprobs": serde_json::Value::Null,
"finish_reason": finish_reason,
})
}
// ─── Finish reason mapping ───────────────────────────────────────────────────
/// Map Google's finishReason → OpenAI's finish_reason.
/// Google: STOP, MAX_TOKENS, SAFETY, RECITATION, OTHER, BLOCKLIST, PROHIBITED_CONTENT
/// OpenAI: stop, length, content_filter, tool_calls (handled separately)
fn google_to_openai_finish_reason(stop_reason: Option<&str>) -> &'static str {
match stop_reason {
Some("MAX_TOKENS") => "length",
Some("SAFETY") | Some("RECITATION") | Some("BLOCKLIST") | Some("PROHIBITED_CONTENT") => "content_filter",
_ => "stop",
}
}
// ─── Input extraction ────────────────────────────────────────────────────────
/// Extract user text from Chat Completions messages array.
///
/// When tool results are present, builds the full conversation including
/// tool call results so the model can continue after tool use.
/// Builds the full conversation context including all messages (system, user,
/// assistant, tool) so the model has complete history — matching how OpenAI
/// sends the entire messages array to the model.
fn extract_chat_input(messages: &[CompletionMessage]) -> String {
let has_tool_results = messages.iter().any(|m| m.role == "tool");
if has_tool_results {
// Build full conversation context including tool results
return build_conversation_with_tools(messages);
}
// Simple path: no tools, just extract system + last user message
let mut system_parts = Vec::new();
let mut user_parts = Vec::new();
for msg in messages {
let text = extract_message_text(&msg.content);
if text.is_empty() {
continue;
}
match msg.role.as_str() {
"system" | "developer" => system_parts.push(text),
"user" => user_parts.push(text),
_ => {}
}
}
let mut result = String::new();
if !system_parts.is_empty() {
result.push_str(&system_parts.join("\n"));
result.push_str("\n\n");
}
if let Some(last) = user_parts.last() {
result.push_str(last);
}
result.trim().to_string()
// Always build the full conversation — we used to only take the last user
// message which broke multi-turn conversations via the messages array.
build_conversation_with_tools(messages)
}
/// Extract text content from a message's content field (string or array).
@@ -179,18 +211,36 @@ pub(crate) async fn handle_completions(
// Store generation parameters for MITM injection
{
use crate::mitm::store::GenerationParams;
let (response_mime_type, response_schema) = match body.response_format.as_ref() {
Some(rf) => match rf.format_type.as_str() {
"json_object" | "json" => (Some("application/json".to_string()), None),
"json_schema" => {
let schema = rf.json_schema.as_ref().and_then(|js| js.schema.clone());
(Some("application/json".to_string()), schema)
}
_ => (None, None),
},
None => (None, None),
};
let gp = GenerationParams {
temperature: body.temperature,
top_p: body.top_p,
top_k: None, // OpenAI doesn't have top_k
max_output_tokens: body.max_tokens.or(body.max_completion_tokens),
stop_sequences: None, // TODO: body.stop
stop_sequences: body.stop.clone().map(|s| s.into_vec()),
frequency_penalty: body.frequency_penalty,
presence_penalty: body.presence_penalty,
reasoning_effort: body.reasoning_effort.clone(),
response_mime_type,
response_schema,
google_search: body.web_search,
};
// Only store if at least one param is set
if gp.temperature.is_some() || gp.top_p.is_some() || gp.max_output_tokens.is_some()
|| gp.frequency_penalty.is_some() || gp.presence_penalty.is_some()
|| gp.reasoning_effort.is_some() || gp.stop_sequences.is_some()
|| gp.response_mime_type.is_some() || gp.response_schema.is_some()
|| gp.google_search
{
state.mitm_store.set_generation_params(gp).await;
} else {
@@ -216,8 +266,28 @@ pub(crate) async fn handle_completions(
);
}
// Fresh cascade per request
let cascade_id = match state.backend.create_cascade().await {
let n = (body.n.max(1)).min(5); // Cap at 5 to prevent abuse
if n > 1 && body.stream {
warn!("n={n} requested with streaming — streaming only supports n=1, ignoring n");
}
// Session/conversation: reuse cascade if conversation ID provided
let session_id_str = extract_conversation_id(&body.conversation);
// Helper to create a cascade (reuses session or creates fresh)
let create_cascade = |state: Arc<AppState>, session_id: Option<String>| async move {
if let Some(ref sid) = session_id {
state
.sessions
.get_or_create(Some(sid), || state.backend.create_cascade())
.await
.map(|sr| sr.cascade_id)
} else {
state.backend.create_cascade().await
}
};
let cascade_id = match create_cascade(Arc::clone(&state), session_id_str.clone()).await {
Ok(cid) => cid,
Err(e) => {
return err_response(
@@ -228,7 +298,7 @@ pub(crate) async fn handle_completions(
}
};
// Send message
// Send message on primary cascade
state.mitm_store.set_active_cascade(&cascade_id).await;
match state
.backend
@@ -275,7 +345,7 @@ pub(crate) async fn handle_completions(
include_usage,
)
.await
} else {
} else if n <= 1 {
chat_completions_sync(
state,
completion_id,
@@ -284,6 +354,108 @@ pub(crate) async fn handle_completions(
body.timeout,
)
.await
} else {
// n > 1: fire additional (n-1) parallel cascades
let mut extra_cascade_ids = Vec::with_capacity((n - 1) as usize);
for _ in 1..n {
match state.backend.create_cascade().await {
Ok(cid) => {
// Send the same message on each extra cascade
match state.backend.send_message(&cid, &user_text, model.model_enum).await {
Ok((200, _)) => {
let bg = Arc::clone(&state.backend);
let cid2 = cid.clone();
tokio::spawn(async move { let _ = bg.update_annotations(&cid2).await; });
extra_cascade_ids.push(cid);
}
_ => {} // Skip failed cascades
}
}
Err(_) => {} // Skip failed cascade creation
}
}
// Poll all cascades in parallel
let mut handles = Vec::with_capacity(n as usize);
let all_cascade_ids: Vec<String> = std::iter::once(cascade_id.clone())
.chain(extra_cascade_ids)
.collect();
for cid in &all_cascade_ids {
let st = Arc::clone(&state);
let cid = cid.clone();
let timeout = body.timeout;
handles.push(tokio::spawn(async move {
let result = poll_for_response(&st, &cid, timeout).await;
let mitm = match st.mitm_store.take_usage(&cid).await {
Some(u) => Some(u),
None => st.mitm_store.take_usage("_latest").await,
};
(result, mitm)
}));
}
let mut choices = Vec::with_capacity(n as usize);
let mut total_prompt = 0u64;
let mut total_completion = 0u64;
let mut total_cached = 0u64;
let mut total_thinking = 0u64;
for (i, handle) in handles.into_iter().enumerate() {
if let Ok((result, mitm)) = handle.await {
let finish_reason = google_to_openai_finish_reason(
mitm.as_ref().and_then(|u| u.stop_reason.as_deref()),
);
let (pt, ct, cached, thinking) = if let Some(ref mu) = mitm {
(mu.input_tokens, mu.output_tokens, mu.cache_read_input_tokens, mu.thinking_output_tokens)
} else if let Some(u) = &result.usage {
(u.input_tokens, u.output_tokens, 0, 0)
} else {
(0, 0, 0, 0)
};
total_prompt += pt;
total_completion += ct;
total_cached += cached;
total_thinking += thinking;
let mut message = serde_json::json!({
"role": "assistant",
"content": result.text,
});
if let Some(ref thinking_text) = result.thinking {
message["reasoning_content"] = serde_json::json!(thinking_text);
}
choices.push(serde_json::json!({
"index": i,
"message": message,
"logprobs": serde_json::Value::Null,
"finish_reason": finish_reason,
}));
}
}
Json(serde_json::json!({
"id": completion_id,
"object": "chat.completion",
"created": now_unix(),
"model": model_name,
"system_fingerprint": system_fingerprint(),
"service_tier": "default",
"choices": choices,
"usage": {
"prompt_tokens": total_prompt,
"completion_tokens": total_completion,
"total_tokens": total_prompt + total_completion,
"prompt_tokens_details": {
"cached_tokens": total_cached,
},
"completion_tokens_details": {
"reasoning_tokens": total_thinking,
},
},
}))
.into_response()
}
}
@@ -307,21 +479,26 @@ async fn chat_completions_stream(
state.mitm_store.clear_response_async().await;
// Initial role chunk
yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {"role": "assistant", "content": ""},
"finish_reason": serde_json::Value::Null,
}],
})).unwrap_or_default()));
yield Ok::<_, std::convert::Infallible>(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({"role": "assistant", "content": ""}), None)]),
None,
)));
let mut keepalive_counter: u64 = 0;
let mut last_thinking_len: usize = 0;
// Helper: build usage JSON from MITM tokens
let build_usage = |pt: u64, ct: u64, crt: u64, tt: u64| -> serde_json::Value {
serde_json::json!({
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
})
};
while start.elapsed().as_secs() < timeout {
// ── Check for MITM-captured function calls FIRST ──
// This runs independently of LS steps — the MITM captures tool calls
@@ -346,49 +523,28 @@ async fn chat_completions_stream(
},
}));
}
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {"tool_calls": tool_calls},
"finish_reason": serde_json::Value::Null,
}],
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({"tool_calls": tool_calls}), None)]),
None,
)));
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {},
"finish_reason": "tool_calls",
}],
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({}), Some("tool_calls"))]),
None,
)));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([]),
Some(build_usage(pt, ct, crt, tt)),
)));
}
yield Ok(Event::default().data("[DONE]"));
return;
@@ -413,17 +569,11 @@ async fn chat_completions_stream(
let delta = &tc[last_thinking_len..];
last_thinking_len = tc.len();
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {"reasoning_content": delta},
"finish_reason": serde_json::Value::Null,
}],
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({"reasoning_content": delta}), None)]),
None,
)));
}
}
@@ -436,17 +586,11 @@ async fn chat_completions_stream(
};
if !delta.is_empty() {
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {"content": delta},
"finish_reason": serde_json::Value::Null,
}],
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({"content": delta}), None)]),
None,
)));
last_text = text;
}
}
@@ -454,37 +598,24 @@ async fn chat_completions_stream(
// Check if MITM response is complete
if state.mitm_store.is_response_complete() && !last_text.is_empty() {
debug!("Completions: MITM response complete (bypass), text length={}", last_text.len());
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {},
"finish_reason": "stop",
}],
})).unwrap_or_default()));
// Take usage FIRST so we can read stop_reason for finish_reason
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let fr = google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({}), Some(fr))]),
None,
)));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([]),
Some(build_usage(pt, ct, crt, tt)),
)));
}
yield Ok(Event::default().data("[DONE]"));
return;
@@ -514,48 +645,27 @@ async fn chat_completions_stream(
},
}));
}
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {"tool_calls": tool_calls},
"finish_reason": serde_json::Value::Null,
}],
})).unwrap_or_default()));
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {},
"finish_reason": "tool_calls",
}],
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({"tool_calls": tool_calls}), None)]),
None,
)));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({}), Some("tool_calls"))]),
None,
)));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([]),
Some(build_usage(pt, ct, crt, tt)),
)));
}
yield Ok(Event::default().data("[DONE]"));
return;
@@ -587,17 +697,11 @@ async fn chat_completions_stream(
let delta = &tc[last_thinking_len..];
last_thinking_len = tc.len();
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {"reasoning_content": delta},
"finish_reason": serde_json::Value::Null,
}],
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({"reasoning_content": delta}), None)]),
None,
)));
}
}
@@ -611,17 +715,11 @@ async fn chat_completions_stream(
};
if !delta.is_empty() {
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {"content": delta},
"finish_reason": serde_json::Value::Null,
}],
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({"content": delta}), None)]),
None,
)));
last_text = text.to_string();
}
}
@@ -629,37 +727,23 @@ async fn chat_completions_stream(
// Done check: need DONE status AND non-empty text
if is_response_done(steps) && !last_text.is_empty() {
debug!("Completions stream done, text length={}", last_text.len());
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {},
"finish_reason": "stop",
}],
})).unwrap_or_default()));
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let fr = google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({}), Some(fr))]),
None,
)));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([]),
Some(build_usage(pt, ct, crt, tt)),
)));
}
yield Ok(Event::default().data("[DONE]"));
return;
@@ -673,37 +757,23 @@ async fn chat_completions_stream(
let run_status = td["status"].as_str().unwrap_or("");
if run_status.contains("IDLE") && !last_text.is_empty() {
debug!("Completions IDLE, text length={}", last_text.len());
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {},
"finish_reason": "stop",
}],
})).unwrap_or_default()));
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let fr = google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({}), Some(fr))]),
None,
)));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([]),
Some(build_usage(pt, ct, crt, tt)),
)));
}
yield Ok(Event::default().data("[DONE]"));
return;
@@ -728,37 +798,23 @@ async fn chat_completions_stream(
// Timeout
warn!("Completions stream timeout after {}s", timeout);
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {"content": if last_text.is_empty() { "[Timeout waiting for response]" } else { "" }},
"finish_reason": "stop",
}],
})).unwrap_or_default()));
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let fr = google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([chunk_choice(0, serde_json::json!({"content": if last_text.is_empty() { "[Timeout waiting for response]" } else { "" }}), Some(fr))]),
None,
)));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
yield Ok(Event::default().data(chunk_json(
&completion_id, &model_name,
serde_json::json!([]),
Some(build_usage(pt, ct, crt, tt)),
)));
}
yield Ok(Event::default().data("[DONE]"));
};
@@ -789,7 +845,10 @@ async fn chat_completions_sync(
Some(u) => Some(u),
None => state.mitm_store.take_usage("_latest").await,
};
let (prompt_tokens, completion_tokens, cached_tokens, thinking_tokens) = if let Some(mitm_usage) = mitm {
let finish_reason = google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
let (prompt_tokens, completion_tokens, cached_tokens, thinking_tokens) = if let Some(ref mitm_usage) = mitm {
(mitm_usage.input_tokens, mitm_usage.output_tokens, mitm_usage.cache_read_input_tokens, mitm_usage.thinking_output_tokens)
} else if let Some(u) = &result.usage {
(u.input_tokens, u.output_tokens, 0, 0)
@@ -811,10 +870,13 @@ async fn chat_completions_sync(
"object": "chat.completion",
"created": now_unix(),
"model": model_name,
"system_fingerprint": system_fingerprint(),
"service_tier": "default",
"choices": [{
"index": 0,
"message": message,
"finish_reason": "stop",
"logprobs": serde_json::Value::Null,
"finish_reason": finish_reason,
}],
"usage": {
"prompt_tokens": prompt_tokens,