Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 93 additions & 5 deletions crates/utopia-llm/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,11 @@ pub struct ToolCall {
pub name: String,
/// JSON 字符串参数(协议原样透传)
pub arguments: String,
/// 端点挂在调用上的私货,原样带回。Gemini 3 在每个 function call 上给一个
/// `extra_content.google.thought_signature`,同一回合把结果送回去时**必须**
/// 原样附上,否则 400「Function call is missing a thought_signature」;上一回合
/// 的可以不带。内容不解读,不是对象就当没有。
pub extra_content: Option<serde_json::Value>,
}

/// 工具对话的一个 assistant 回合:文本与工具调用至少其一。
Expand All @@ -38,11 +43,17 @@ impl AssistantTurn {
msg["tool_calls"] = json!(self
.tool_calls
.iter()
.map(|c| json!({
"id": c.id,
"type": "function",
"function": { "name": c.name, "arguments": c.arguments },
}))
.map(|c| {
let mut call = json!({
"id": c.id,
"type": "function",
"function": { "name": c.name, "arguments": c.arguments },
});
if let Some(extra) = &c.extra_content {
call["extra_content"] = extra.clone();
}
call
})
.collect::<Vec<_>>());
}
msg
Expand All @@ -54,6 +65,11 @@ pub fn tool_result_message(tool_call_id: &str, content: &str) -> serde_json::Val
json!({ "role": "tool", "tool_call_id": tool_call_id, "content": content })
}

/// 一条工具调用上端点挂的 `extra_content`:是对象才算,别的形状不猜。
fn extra_content_of(call: &serde_json::Value) -> Option<serde_json::Value> {
call.get("extra_content").filter(|v| v.is_object()).cloned()
}

/// 带工具的流式回合事件。
#[derive(Debug)]
pub enum ToolStreamItem {
Expand Down Expand Up @@ -842,6 +858,7 @@ impl LlmClient {
.as_str()
.unwrap_or("{}")
.to_string(),
extra_content: extra_content_of(c),
})
})
.collect()
Expand Down Expand Up @@ -928,9 +945,13 @@ impl LlmClient {
id: String::new(),
name: String::new(),
arguments: String::new(),
extra_content: None,
});
}
let slot = &mut calls[idx];
if let Some(extra) = extra_content_of(tc) {
slot.extra_content = Some(extra);
}
if let Some(id) = tc["id"].as_str() {
slot.id.push_str(id);
}
Expand Down Expand Up @@ -2928,4 +2949,71 @@ data: [DONE]
.unwrap();
assert_eq!(out, vec![vec![1.0, 10.0], vec![2.0, 20.0]]);
}

/// Gemini 3 的形状:tool_call 不带 `index`,带 `extra_content.google.thought_signature`,
/// finish_reason 说 `stop`。两种传输都要把 `extra_content` 留住,回灌时原样带回
#[tokio::test]
async fn a_tool_call_keeps_its_extra_content_in_both_transports() {
use futures_util::TryStreamExt;
let extra = json!({"google": {"thought_signature": "Eq0CCqoC"}});
let call = json!({
"id": "call_170866", "type": "function",
"function": {"name": "search_kb", "arguments": "{\"query\":\"lore\"}"},
"extra_content": extra
});
let body = json!({"choices":[{"message":{"role":"assistant","tool_calls":[call]},"finish_reason":"stop"}]})
.to_string();
let (addr, server) = an_http_response("200 OK", "application/json", &body).await;
let turn = client_at(addr)
.chat_tools_with(&[], None, None)
.await
.unwrap();
server.await.unwrap();
assert_eq!(turn.tool_calls[0].extra_content, Some(extra.clone()));
assert_eq!(turn.to_message()["tool_calls"][0]["extra_content"], extra);

let frame =
json!({"choices":[{"delta":{"role":"assistant","tool_calls":[call]},"index":0}]});
let done =
json!({"choices":[{"delta":{"role":"assistant"},"finish_reason":"stop","index":0}]});
let sse = format!("data: {frame}\n\ndata: {done}\n\ndata: [DONE]\n\n");
let (addr, server) = an_http_response("200 OK", "text/event-stream", &sse).await;
let stream = client_at(addr)
.chat_tools_stream_with(&[], None, None)
.await
.unwrap();
let items: Vec<ToolStreamItem> = stream.try_collect().await.unwrap();
server.await.unwrap();
let [ToolStreamItem::Turn(turn)] = items.as_slice() else {
panic!("expected one completed turn: {items:?}");
};
assert_eq!(turn.tool_calls[0].extra_content, Some(extra.clone()));
assert_eq!(turn.to_message()["tool_calls"][0]["extra_content"], extra);
}

/// 没有 `extra_content`、或它不是对象的调用,回灌时也不长出这个键
#[tokio::test]
async fn a_tool_call_without_extra_content_echoes_none() {
for call in [
json!({"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}),
json!({"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}, "extra_content": null}),
json!({"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}, "extra_content": "sig"}),
] {
let body = json!({"choices":[{"message":{"role":"assistant","tool_calls":[call]},"finish_reason":"tool_calls"}]})
.to_string();
let (addr, server) = an_http_response("200 OK", "application/json", &body).await;
let turn = client_at(addr)
.chat_tools_with(&[], None, None)
.await
.unwrap();
server.await.unwrap();
assert_eq!(turn.tool_calls[0].extra_content, None, "{call}");
assert!(
turn.to_message()["tool_calls"][0]
.get("extra_content")
.is_none(),
"{call}"
);
}
}
}
102 changes: 85 additions & 17 deletions crates/utopia-server/src/api/rig_model.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,11 +104,14 @@ impl CompletionModel for RigModel {
.tool_calls
.iter()
.map(|c| {
Ok(RawStreamingChoice::ToolCall(RawStreamingToolCall::new(
call_id(&c.id),
c.name.clone(),
args_value(&c.arguments),
)))
Ok(RawStreamingChoice::ToolCall(
RawStreamingToolCall::new(
call_id(&c.id),
c.name.clone(),
args_value(&c.arguments),
)
.with_additional_params(c.extra_content.clone()),
))
})
.collect();
v.push(Ok(RawStreamingChoice::FinalResponse(
Expand Down Expand Up @@ -300,14 +303,22 @@ pub(super) fn push_message(out: &mut Vec<Value>, m: &Message) {
for c in content {
match c {
AssistantContent::Text(t) => text.push_str(&t.text),
AssistantContent::ToolCall(tc) => calls.push(json!({
"id": tc.id.as_str(),
"type": "function",
"function": {
"name": tc.function.name,
"arguments": args_string(&tc.function.arguments),
// 端点挂在调用上的 `extra_content`(Gemini 3 的 thought_signature)
// 跟着回去:同一回合里缺了它,端点拒收整个请求
AssistantContent::ToolCall(tc) => {
let mut call = json!({
"id": tc.id.as_str(),
"type": "function",
"function": {
"name": tc.function.name,
"arguments": args_string(&tc.function.arguments),
}
});
if let Some(extra) = &tc.additional_params {
call["extra_content"] = extra.clone();
}
})),
calls.push(call);
}
// 推理块与图片不回灌
_ => {}
}
Expand Down Expand Up @@ -370,11 +381,12 @@ fn choice_of(turn: &AssistantTurn) -> Vec<AssistantContent> {
out.push(AssistantContent::text(text));
}
for c in &turn.tool_calls {
out.push(AssistantContent::tool_call(
call_id(&c.id),
c.name.clone(),
args_value(&c.arguments),
));
let mut call =
AssistantContent::tool_call(call_id(&c.id), c.name.clone(), args_value(&c.arguments));
if let AssistantContent::ToolCall(tc) = &mut call {
tc.additional_params = c.extra_content.clone();
}
out.push(call);
}
out
}
Expand Down Expand Up @@ -510,6 +522,7 @@ mod tests {
id: "c9".into(),
name: "find_entities".into(),
arguments: "{\"name\":\"Acme\"}".into(),
extra_content: None,
}],
};
let choice = choice_of(&turn);
Expand All @@ -524,4 +537,59 @@ mod tests {
other => panic!("expected a tool call, got {other:?}"),
}
}

/// Gemini 3 把 thought_signature 挂在调用的 `extra_content` 上,同一回合回灌时
/// 必须原样带回:rig 这边装在 `additional_params`,上线时还原成 `extra_content`
#[test]
fn a_tool_call_carries_its_extra_content_from_turn_to_wire() {
let extra = json!({"google": {"thought_signature": "Eq0CCqoC"}});
let turn = AssistantTurn {
finish_reason: Some("stop".into()),
content: None,
tool_calls: vec![utopia_llm::ToolCall {
id: "call_170866".into(),
name: "search_kb".into(),
arguments: "{\"query\":\"lore\"}".into(),
extra_content: Some(extra.clone()),
}],
};
let choice = choice_of(&turn);
let AssistantContent::ToolCall(call) = &choice[0] else {
panic!("expected a tool call, got {choice:?}");
};
assert_eq!(call.additional_params, Some(extra.clone()));

let w = wire(&req(vec![
Message::user("i want to know about lore"),
Message::Assistant {
id: None,
content: choice,
},
Message::tool_result("call_170866", "search_kb", "found"),
]));
assert_eq!(w.messages[2]["tool_calls"][0]["extra_content"], extra);
assert_eq!(
w.messages[2]["tool_calls"][0]["function"]["name"],
"search_kb"
);
}

#[test]
fn a_tool_call_without_extra_content_adds_no_key_on_the_wire() {
let w = wire(&req(vec![
Message::user("q"),
Message::Assistant {
id: None,
content: vec![AssistantContent::tool_call(
"c1",
"search_chunks",
json!({ "query": "x" }),
)],
},
Message::tool_result("c1", "search_chunks", "found"),
]));
assert!(w.messages[2]["tool_calls"][0]
.get("extra_content")
.is_none());
}
}
Loading