mirror of
https://github.com/block/buzz.git
synced 2026-10-27 12:40:30 +00:00
feat(agent): make Gemini and MLflow-route models usable through databricks_v2 (#3569)
## Summary
Makes Gemini — and every other non-Claude, non-GPT-5 model on the
Databricks MLflow route (`databricks_v2`) — usable in an agent loop.
These are the non-`benchmarks/` changes from
`benchmark/harness-accounting-and-solo`, lifted onto a clean base off
`main` so they can land independently while the harness work continues.
Two defects made these models unusable, one fatal and one silent. Both
live only in `openai_body` / `parse_openai`, which is the
least-exercised of the three `databricks_v2` sub-routes — the
`luna`/`sol` conditions run the Responses route and the `opus`
conditions run the Anthropic route, so **this change is inert for every
model already in use** and only lights up the MLflow path.
## Why a third route at all
`databricks_v2_route_for_model` buckets by model family: `claude*` →
Anthropic Messages, `gpt-5`/code-names → OpenAI Responses, **everything
else → MLflow chat-completions**. Gemini, Qwen, gpt-oss, and friends all
fall through to that third pair — and both bugs below live only there.
## D1 — dropped thought signatures (fatal)
Gemini returns a `thoughtSignature` on every tool call and **requires it
echoed back**. `openai_body` reserialized each call as `{id, type,
function}` only, dropping the field, so the next request 400'd:
```
HTTP 400 Function call is missing a thought_signature in functionCall parts.
```
For a coding agent this fires on the **first** tool call, so the model
never completes a single turn.
**Position is load-bearing.** A four-shape replay probe against the live
gateway established that the signature must sit as a *sibling* of
`function` — nesting it inside `function{}` fails with the *same* 400 as
omitting it. A fix that "preserves the field" without preserving its
position passes a unit test and still 400s.
The fix: `ToolCall` gains `provider_extra: Map<String, Value>`.
`parse_openai` captures every top-level wire key except the three we
model (`id`, `type`, `function`); `openai_body` re-emits them beside
`function`. Keeping *whatever we did not model*, rather than naming
`thoughtSignature`, means the next provider with an opaque per-call
token needs no change here. The Responses and Anthropic replay shapes
are fully modelled, so they pass `Default::default()` and stay
**byte-identical** to before.
### D1b — duplicate tool-call ids (same root cause)
Gemini returns the **function name** as the id, so two parallel calls to
one function arrive sharing an id — and that id is what pairs a
`role:"tool"` result back to its call, making two results
indistinguishable. `dedupe_provider_ids` suffixes collisions
(`get_weather`, `get_weather-2`). Safe because both halves of the
pairing (the assistant `tool_calls[].id` and the result's
`tool_call_id`) are re-emitted from this same value; the provider never
sees its original id again.
## D2 — block-array content discarded (silent, worse than a crash)
`parse_openai` read `content` with `as_str()`, which returns `""` for
anything that isn't a JSON string. Gemini (and Qwen35, gpt-oss) send an
array of typed blocks:
```json
"content": [
{"type": "reasoning", "summary": [{"type": "summary_text", "text": "…"}]},
{"type": "text", "text": "391"}
]
```
So the model answered and the answer was thrown away — no error, no
warning, just a turn that looked like the model had said nothing. On a
benchmark this reads as "Gemini is bad at the task" rather than "buzz
dropped the reply."
`openai_content_parts` now accepts either shape — string as before, or a
block array where `text` blocks concatenate into text and `reasoning`
blocks into reasoning (Gemini nests the prose one level down under
`summary`). Message-level `reasoning_content` / `reasoning` still win
when present, so DeepSeek and vLLM-style hosts are unchanged; block
reasoning is the last fallback.
## Also: a turn-start log line (`buzz-acp` `pool.rs`)
Small, independent observability change that also rides in the
non-benchmark delta: `run_prompt_task` now emits a `pool::prompt` "turn
starting" line, labelled by the same `prompt_label` helper as
`log_stop_reason`, so a log reads as start/stop pairs. An unpaired start
is the only durable evidence that a turn was entered and never returned
— without it, a stalled agent and an agent nobody woke leave identical
(zero-completion) logs.
## Interaction with #3538
#3538 (already merged) rewrote `databricks_v2_route_for_model` to route
by boundary-aware model-family segments. That change and this one touch
**different functions** in `llm.rs` — routing vs. body/parse — and
compose cleanly; the family routing decides *which* pair runs, and this
fixes the MLflow pair it can now select.
## Testing
- `cargo fmt --all -- --check`, `cargo clippy -p buzz-agent -p buzz-acp
--all-targets -- -D warnings` — clean.
- `cargo test -p buzz-agent -p buzz-acp` — all green (304 + 632 lib
tests plus integration suites, 0 failures). Five new tests cover:
block-array text extraction, plain-string regression, passthrough
capture (and non-duplication of the modelled keys), replay position
(`thoughtSignature` beside `function`, not inside it), and id
de-duplication.
- Wire evidence: the four-shape replay table and the reasoning-effort
probe were run against `block-lakehouse-staging` (recorded in the design
doc).
## Relationship to the benchmark branch
The full design write-up (four-shape replay table, position-matters
analysis, effort verification, and open pricing item) lives in
`docs/08-gemini-provider-fixes.md` on
`benchmark/harness-accounting-and-solo`. The benchmark manifests and
endpoint-config entries that exercise these models are separable and
stay on that branch.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Signed-off-by: Atish Patel <atish@squareup.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1897,6 +1897,18 @@ pub async fn run_prompt_task(
|
||||
None => prompt_sections.iter().map(String::as_str).collect(),
|
||||
};
|
||||
|
||||
// Turn start, labelled exactly as `log_stop_reason` labels the end, so a
|
||||
// log reads as start/stop pairs. Purely observational: an unpaired start is
|
||||
// the only durable evidence that a turn was entered and never returned, and
|
||||
// without it a stalled agent and an agent nobody woke leave identical logs —
|
||||
// zero completions either way, so anything reading them afterwards has to
|
||||
// guess which happened.
|
||||
tracing::info!(
|
||||
target: "pool::prompt",
|
||||
"turn starting for {}",
|
||||
prompt_label(&source)
|
||||
);
|
||||
|
||||
// When control_rx is Some (channel tasks), wrap the prompt in select! so
|
||||
// the main loop can cancel, interrupt, or rotate it. Heartbeats
|
||||
// (control_rx=None) take the simple await path — they are not controllable.
|
||||
@@ -3131,12 +3143,19 @@ fn classify_control_cancel_failure(
|
||||
}
|
||||
}
|
||||
|
||||
/// Log a stop reason at the appropriate tracing level.
|
||||
fn log_stop_reason(source: &PromptSource, stop_reason: &StopReason) {
|
||||
let label = match source {
|
||||
/// How a turn's source is named in the `pool::prompt` log lines.
|
||||
///
|
||||
/// Shared by the turn-start and turn-stop lines so a log can be read as pairs.
|
||||
fn prompt_label(source: &PromptSource) -> String {
|
||||
match source {
|
||||
PromptSource::Channel(cid) => format!("channel {cid}"),
|
||||
PromptSource::Heartbeat => "heartbeat".to_string(),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/// Log a stop reason at the appropriate tracing level.
|
||||
fn log_stop_reason(source: &PromptSource, stop_reason: &StopReason) {
|
||||
let label = prompt_label(source);
|
||||
match stop_reason {
|
||||
StopReason::EndTurn => {
|
||||
tracing::info!(target: "pool::prompt", "turn complete for {label}: end_turn");
|
||||
|
||||
@@ -696,6 +696,9 @@ pub(crate) fn push_hook_outputs_as_tool_results(
|
||||
provider_id: provider_id.clone(),
|
||||
name: tool_name,
|
||||
arguments: serde_json::json!({}),
|
||||
// Synthesised locally, so there is no provider wire form to
|
||||
// preserve.
|
||||
provider_extra: Default::default(),
|
||||
}],
|
||||
});
|
||||
history.push(HistoryItem::ToolResult(ToolResult {
|
||||
|
||||
@@ -3,7 +3,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
use reqwest::Client;
|
||||
use serde_json::{json, Value};
|
||||
use serde_json::{json, Map, Value};
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::time::Instant;
|
||||
|
||||
@@ -848,11 +848,24 @@ fn openai_body(
|
||||
let calls: Vec<Value> = tool_calls
|
||||
.iter()
|
||||
.map(|c| {
|
||||
json!({
|
||||
"id": c.provider_id, "type": "function",
|
||||
"function": { "name": c.name,
|
||||
"arguments": serde_json::to_string(&c.arguments)
|
||||
.unwrap_or_else(|_| "{}".into()) } })
|
||||
let mut call = serde_json::Map::new();
|
||||
call.insert("id".into(), json!(c.provider_id));
|
||||
call.insert("type".into(), json!("function"));
|
||||
// Provider-owned fields go back beside `function`,
|
||||
// which is where the provider put them. Position is
|
||||
// load-bearing, not cosmetic: Gemini rejects a
|
||||
// `thoughtSignature` nested inside `function{}` with
|
||||
// the same 400 it gives for one that is missing.
|
||||
for (k, v) in &c.provider_extra {
|
||||
call.insert(k.clone(), v.clone());
|
||||
}
|
||||
call.insert(
|
||||
"function".into(),
|
||||
json!({ "name": c.name,
|
||||
"arguments": serde_json::to_string(&c.arguments)
|
||||
.unwrap_or_else(|_| "{}".into()) }),
|
||||
);
|
||||
Value::Object(call)
|
||||
})
|
||||
.collect();
|
||||
msg.insert("tool_calls".into(), Value::Array(calls));
|
||||
@@ -1120,10 +1133,15 @@ fn parse_responses(v: Value) -> Result<LlmResponse, AgentError> {
|
||||
let args: Value = serde_json::from_str(raw).map_err(|e| {
|
||||
AgentError::Llm(format!("function_call.arguments not valid JSON: {e}"))
|
||||
})?;
|
||||
// No passthrough on this route: `responses_body` replays a
|
||||
// function call as `{call_id, name, arguments}` and the Responses
|
||||
// API asks for nothing else, so an empty map keeps the request
|
||||
// byte-identical to before.
|
||||
tool_calls.push(make_tool_call(
|
||||
str_field(item, "call_id"),
|
||||
str_field(item, "name"),
|
||||
args,
|
||||
Default::default(),
|
||||
)?);
|
||||
}
|
||||
Some("reasoning") => {
|
||||
@@ -1301,6 +1319,76 @@ fn str_field(v: &Value, key: &str) -> String {
|
||||
v.get(key).and_then(Value::as_str).unwrap_or("").to_owned()
|
||||
}
|
||||
|
||||
/// Append `part` to `buf` on its own line, ignoring empties.
|
||||
fn push_part(buf: &mut String, part: &str) {
|
||||
if part.is_empty() {
|
||||
return;
|
||||
}
|
||||
if !buf.is_empty() {
|
||||
buf.push('\n');
|
||||
}
|
||||
buf.push_str(part);
|
||||
}
|
||||
|
||||
/// Split an OpenAI-shaped `message.content` into `(text, reasoning)`.
|
||||
///
|
||||
/// Standard OpenAI sends a string. Several models on the Databricks MLflow route
|
||||
/// — Gemini, Qwen35, gpt-oss — send an array of typed blocks instead, and
|
||||
/// `as_str()` yields nothing for an array, so their entire answer was being
|
||||
/// discarded: no error, no warning, just a turn that looked like the model had
|
||||
/// said nothing. `parse_anthropic` already walks a block array; this gives
|
||||
/// `parse_openai` the same tolerance.
|
||||
fn openai_content_parts(content: Option<&Value>) -> (String, String) {
|
||||
let mut text = String::new();
|
||||
let mut reasoning = String::new();
|
||||
match content {
|
||||
Some(Value::String(s)) => text.push_str(s),
|
||||
Some(Value::Array(blocks)) => {
|
||||
for b in blocks {
|
||||
match b.get("type").and_then(Value::as_str) {
|
||||
Some("text") => push_part(&mut text, &str_field(b, "text")),
|
||||
Some("reasoning") => match b.get("summary").and_then(Value::as_array) {
|
||||
// Gemini nests the prose one level down under `summary`.
|
||||
Some(summary) => {
|
||||
for s in summary {
|
||||
push_part(&mut reasoning, &str_field(s, "text"));
|
||||
}
|
||||
}
|
||||
None => push_part(&mut reasoning, &str_field(b, "text")),
|
||||
},
|
||||
// An untyped block carrying text is still the model talking;
|
||||
// treating it as text loses nothing and keeps one more
|
||||
// provider out of the silent-empty-answer failure mode.
|
||||
_ => push_part(&mut text, &str_field(b, "text")),
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
(text, reasoning)
|
||||
}
|
||||
|
||||
/// Make `provider_id` unique across one assistant turn's tool calls.
|
||||
///
|
||||
/// Gemini returns the function name as the id, so two parallel calls to the same
|
||||
/// function arrive sharing one id — and that id is what pairs a `role:"tool"`
|
||||
/// result back to its call, leaving two results indistinguishable. Rewriting is
|
||||
/// safe because both halves of that pairing are re-emitted from this same value;
|
||||
/// the provider never sees its original id again.
|
||||
fn dedupe_provider_ids(calls: &mut [ToolCall]) {
|
||||
let mut seen: BTreeSet<String> = BTreeSet::new();
|
||||
for c in calls.iter_mut() {
|
||||
if seen.contains(&c.provider_id) {
|
||||
let mut n = 2;
|
||||
while seen.contains(&format!("{}-{n}", c.provider_id)) {
|
||||
n += 1;
|
||||
}
|
||||
c.provider_id = format!("{}-{n}", c.provider_id);
|
||||
}
|
||||
seen.insert(c.provider_id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_anthropic(v: Value) -> Result<LlmResponse, AgentError> {
|
||||
let stop = map_stop(v.get("stop_reason").and_then(Value::as_str));
|
||||
let mut tool_calls = Vec::new();
|
||||
@@ -1323,10 +1411,12 @@ fn parse_anthropic(v: Value) -> Result<LlmResponse, AgentError> {
|
||||
reasoning.push_str(t);
|
||||
}
|
||||
}
|
||||
// Anthropic's replay shape is fully modelled, so nothing to keep.
|
||||
Some("tool_use") => tool_calls.push(make_tool_call(
|
||||
str_field(b, "id"),
|
||||
str_field(b, "name"),
|
||||
b.get("input").cloned().unwrap_or(Value::Null),
|
||||
Default::default(),
|
||||
)?),
|
||||
_ => {}
|
||||
}
|
||||
@@ -1358,17 +1448,23 @@ fn parse_openai(v: Value) -> Result<LlmResponse, AgentError> {
|
||||
let msg = choice
|
||||
.get("message")
|
||||
.ok_or_else(|| AgentError::Llm("missing message".into()))?;
|
||||
let text = str_field(msg, "content");
|
||||
let (text, block_reasoning) = openai_content_parts(msg.get("content"));
|
||||
// DeepSeek and vLLM-style OpenAI-compat hosts expose reasoning tokens on the
|
||||
// message object. Prefer `reasoning_content` (DeepSeek's field name); fall
|
||||
// back to `reasoning` (some other providers). Both are absent for standard
|
||||
// OpenAI responses, which leaves this empty without any special-casing.
|
||||
// back to `reasoning` (some other providers), and last to reasoning blocks
|
||||
// found inside `content`. All three are absent for standard OpenAI
|
||||
// responses, which leaves this empty without any special-casing.
|
||||
let reasoning = {
|
||||
let rc = str_field(msg, "reasoning_content");
|
||||
if rc.is_empty() {
|
||||
let rc = if rc.is_empty() {
|
||||
str_field(msg, "reasoning")
|
||||
} else {
|
||||
rc
|
||||
};
|
||||
if rc.is_empty() {
|
||||
block_reasoning
|
||||
} else {
|
||||
rc
|
||||
}
|
||||
};
|
||||
let mut tool_calls = Vec::new();
|
||||
@@ -1380,13 +1476,25 @@ fn parse_openai(v: Value) -> Result<LlmResponse, AgentError> {
|
||||
let raw = f.get("arguments").and_then(Value::as_str).unwrap_or("{}");
|
||||
let args: Value = serde_json::from_str(raw)
|
||||
.map_err(|e| AgentError::Llm(format!("tool_call.arguments not valid JSON: {e}")))?;
|
||||
// Everything on the wire object we do not model, kept for replay.
|
||||
let extra = tc
|
||||
.as_object()
|
||||
.map(|o| {
|
||||
o.iter()
|
||||
.filter(|(k, _)| !matches!(k.as_str(), "id" | "type" | "function"))
|
||||
.map(|(k, v)| (k.clone(), v.clone()))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
tool_calls.push(make_tool_call(
|
||||
str_field(tc, "id"),
|
||||
str_field(f, "name"),
|
||||
args,
|
||||
extra,
|
||||
)?);
|
||||
}
|
||||
}
|
||||
dedupe_provider_ids(&mut tool_calls);
|
||||
let input_tokens = openai_chat_input_tokens(&v);
|
||||
let output_tokens = sum_usage(&v, &["completion_tokens"]);
|
||||
let cached_input_tokens = openai_chat_cached_tokens(&v);
|
||||
@@ -1401,7 +1509,12 @@ fn parse_openai(v: Value) -> Result<LlmResponse, AgentError> {
|
||||
})
|
||||
}
|
||||
|
||||
fn make_tool_call(id: String, name: String, args: Value) -> Result<ToolCall, AgentError> {
|
||||
fn make_tool_call(
|
||||
id: String,
|
||||
name: String,
|
||||
args: Value,
|
||||
provider_extra: Map<String, Value>,
|
||||
) -> Result<ToolCall, AgentError> {
|
||||
if id.is_empty() || name.is_empty() {
|
||||
return Err(AgentError::Llm("tool_call missing id or name".into()));
|
||||
}
|
||||
@@ -1418,6 +1531,7 @@ fn make_tool_call(id: String, name: String, args: Value) -> Result<ToolCall, Age
|
||||
provider_id: id,
|
||||
name,
|
||||
arguments,
|
||||
provider_extra,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2365,6 +2479,7 @@ mod tests {
|
||||
provider_id: "toolu_1".into(),
|
||||
name: "dev__view_image".into(),
|
||||
arguments: serde_json::json!({"source":"x.png"}),
|
||||
provider_extra: Default::default(),
|
||||
}],
|
||||
},
|
||||
HistoryItem::ToolResult(ToolResult {
|
||||
@@ -2414,6 +2529,7 @@ mod tests {
|
||||
provider_id: "call_abc".into(),
|
||||
name: "dev__shell".into(),
|
||||
arguments: serde_json::json!({"command": "ls"}),
|
||||
provider_extra: Default::default(),
|
||||
}],
|
||||
},
|
||||
HistoryItem::ToolResult(ToolResult {
|
||||
@@ -2512,6 +2628,7 @@ mod tests {
|
||||
provider_id: "call_x".into(),
|
||||
name: "t".into(),
|
||||
arguments: serde_json::json!({}),
|
||||
provider_extra: Default::default(),
|
||||
}],
|
||||
},
|
||||
];
|
||||
@@ -2790,11 +2907,13 @@ mod tests {
|
||||
provider_id: "toolu_a".into(),
|
||||
name: "dev__view_image".into(),
|
||||
arguments: serde_json::json!({"source": "a.png"}),
|
||||
provider_extra: Default::default(),
|
||||
},
|
||||
ToolCall {
|
||||
provider_id: "toolu_b".into(),
|
||||
name: "dev__view_image".into(),
|
||||
arguments: serde_json::json!({"source": "b.png"}),
|
||||
provider_extra: Default::default(),
|
||||
},
|
||||
],
|
||||
},
|
||||
@@ -3819,6 +3938,92 @@ mod tests {
|
||||
assert_eq!(parse_anthropic(v).unwrap().input_tokens, None);
|
||||
}
|
||||
|
||||
/// A Gemini reply as the Databricks MLflow route actually returns it:
|
||||
/// block-array `content`, and a `thoughtSignature` beside `function`.
|
||||
fn gemini_choice() -> Value {
|
||||
json!({"choices": [{"finish_reason": "tool_calls", "message": {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "reasoning", "summary": [{"type": "summary_text", "text": "weighing it"}]},
|
||||
{"type": "text", "text": "391"}
|
||||
],
|
||||
"tool_calls": [{
|
||||
"id": "get_weather", "type": "function", "thoughtSignature": "SIG-A",
|
||||
"function": {"name": "get_weather", "arguments": "{\"city\":\"Paris\"}"}
|
||||
}]
|
||||
}}]})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_openai_reads_text_out_of_a_block_array() {
|
||||
// Before this, `as_str()` on the array yielded "" and the model's answer
|
||||
// was discarded with no error at all.
|
||||
let r = parse_openai(gemini_choice()).unwrap();
|
||||
assert_eq!(r.text, "391");
|
||||
assert_eq!(r.reasoning, "weighing it");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_openai_still_reads_a_plain_string_content() {
|
||||
let v = json!({"choices": [{"finish_reason": "stop", "message": {
|
||||
"role": "assistant", "content": "plain"}}]});
|
||||
let r = parse_openai(v).unwrap();
|
||||
assert_eq!(r.text, "plain");
|
||||
assert_eq!(r.reasoning, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_openai_keeps_unmodelled_tool_call_fields() {
|
||||
let r = parse_openai(gemini_choice()).unwrap();
|
||||
let extra = &r.tool_calls[0].provider_extra;
|
||||
assert_eq!(extra.get("thoughtSignature"), Some(&json!("SIG-A")));
|
||||
// `id`/`type`/`function` are modelled, so they must not be duplicated
|
||||
// into the passthrough — they would be re-emitted twice.
|
||||
assert!(!extra.contains_key("id"));
|
||||
assert!(!extra.contains_key("type"));
|
||||
assert!(!extra.contains_key("function"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_body_replays_the_signature_beside_function_not_inside_it() {
|
||||
// Position is what the gateway checks: nested inside `function{}` it is
|
||||
// rejected with the same 400 as a missing signature.
|
||||
let r = parse_openai(gemini_choice()).unwrap();
|
||||
let history = vec![HistoryItem::Assistant {
|
||||
text: r.text.clone(),
|
||||
tool_calls: r.tool_calls.clone(),
|
||||
}];
|
||||
let body = openai_body(
|
||||
&cfg(Provider::DatabricksV2),
|
||||
"sys",
|
||||
&history,
|
||||
&[],
|
||||
"databricks-gemini-3-6-flash",
|
||||
None,
|
||||
);
|
||||
let call = &body["messages"][1]["tool_calls"][0];
|
||||
assert_eq!(call["thoughtSignature"], json!("SIG-A"));
|
||||
assert!(call["function"].get("thoughtSignature").is_none());
|
||||
assert_eq!(call["function"]["name"], json!("get_weather"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_openai_makes_duplicate_tool_call_ids_unique() {
|
||||
// Gemini returns the function name as the id, so parallel calls to one
|
||||
// function collide and their results become indistinguishable.
|
||||
let v = json!({"choices": [{"finish_reason": "tool_calls", "message": {
|
||||
"role": "assistant", "content": "",
|
||||
"tool_calls": [
|
||||
{"id": "get_weather", "type": "function",
|
||||
"function": {"name": "get_weather", "arguments": "{\"city\":\"Paris\"}"}},
|
||||
{"id": "get_weather", "type": "function",
|
||||
"function": {"name": "get_weather", "arguments": "{\"city\":\"Rome\"}"}}
|
||||
]}}]});
|
||||
let r = parse_openai(v).unwrap();
|
||||
assert_eq!(r.tool_calls[0].provider_id, "get_weather");
|
||||
assert_eq!(r.tool_calls[1].provider_id, "get_weather-2");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_openai_uses_prompt_tokens() {
|
||||
let v = serde_json::json!({
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
/// Byte-equivalent charged to the handoff/context-pressure gate for a single
|
||||
/// image tool result. The gate maps bytes to tokens at 1 byte/token (see
|
||||
@@ -93,6 +93,13 @@ impl HistoryItem {
|
||||
+ serde_json::to_vec(&c.arguments)
|
||||
.map(|b| b.len())
|
||||
.unwrap_or(0)
|
||||
// `provider_extra` (e.g. a Gemini
|
||||
// `thoughtSignature`) is re-serialized into
|
||||
// every replayed call, so it counts toward the
|
||||
// request body and the context-pressure gate.
|
||||
+ serde_json::to_vec(&c.provider_extra)
|
||||
.map(|b| b.len())
|
||||
.unwrap_or(0)
|
||||
})
|
||||
.sum::<usize>()
|
||||
}
|
||||
@@ -108,6 +115,17 @@ pub struct ToolCall {
|
||||
pub provider_id: String,
|
||||
pub name: String,
|
||||
pub arguments: Value,
|
||||
/// Fields the provider put on the tool call that we do not model, kept so
|
||||
/// the assistant turn can be replayed the way it arrived.
|
||||
///
|
||||
/// Gemini on the Databricks MLflow route returns a `thoughtSignature` per
|
||||
/// call and *requires* it echoed back: replaying without it fails the whole
|
||||
/// request with `Function call is missing a thought_signature in functionCall
|
||||
/// parts`. For an agent loop that lands on the very first tool call, so the
|
||||
/// model is unusable without this. Carrying whatever we did not model,
|
||||
/// rather than naming that one field, means the next provider with an opaque
|
||||
/// per-call token needs no change here.
|
||||
pub provider_extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -353,6 +371,39 @@ mod tests {
|
||||
assert!(item.estimated_bytes() >= 3_118_884);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assistant_size_counts_provider_extra() {
|
||||
// A Gemini `thoughtSignature` rides the wire on every replayed call, so
|
||||
// both size measures must see it — otherwise `truncate_history` and the
|
||||
// handoff gate under-count and let the real request exceed the budget.
|
||||
let mut extra = Map::new();
|
||||
extra.insert("thoughtSignature".into(), Value::String("S".repeat(500)));
|
||||
let with_extra = HistoryItem::Assistant {
|
||||
text: String::new(),
|
||||
tool_calls: vec![ToolCall {
|
||||
provider_id: "id".into(),
|
||||
name: "t".into(),
|
||||
arguments: Value::Null,
|
||||
provider_extra: extra,
|
||||
}],
|
||||
};
|
||||
let without_extra = HistoryItem::Assistant {
|
||||
text: String::new(),
|
||||
tool_calls: vec![ToolCall {
|
||||
provider_id: "id".into(),
|
||||
name: "t".into(),
|
||||
arguments: Value::Null,
|
||||
provider_extra: Map::new(),
|
||||
}],
|
||||
};
|
||||
assert!(with_extra.estimated_bytes() > without_extra.estimated_bytes() + 500);
|
||||
assert_eq!(
|
||||
with_extra.estimated_bytes(),
|
||||
with_extra.context_pressure_bytes(),
|
||||
"provider_extra is text, so both measures must agree"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_content_size_is_identical_for_both_measures() {
|
||||
// Only images diverge; text must size the same under both paths.
|
||||
|
||||
Reference in New Issue
Block a user