diff --git a/crates/tw-gateway/src/plugin/reply/chat.rs b/crates/tw-gateway/src/plugin/reply/chat.rs index 34dc2acd..e4578ee2 100644 --- a/crates/tw-gateway/src/plugin/reply/chat.rs +++ b/crates/tw-gateway/src/plugin/reply/chat.rs @@ -1,12 +1,19 @@ //! OpenAI Chat Completions 的回答。 //! -//! Chat 的流没有块边界：正文是 `delta.content` 一路；换成推理、开始工具调用、 -//! `finish_reason`、`[DONE]` 都算这一块结束。工具调用按 `index` 分片，OpenAI 一个 -//! 发完再发下一个，所以换了 `index` 或者到了 `finish_reason` 就是上一个收齐了。 +//! Chat 的流没有块边界：每个选项（`n` > 1 时有好几个，帧里按选项的 `index` 分）的正文 +//! 是 `delta.content` 一路；换成推理、开始工具调用、`finish_reason`、`[DONE]` 都算这一块 +//! 结束。工具调用按 `index` 分片，OpenAI 一个发完再发下一个，所以换了 `index` 或者到了 +//! `finish_reason` 就是上一个收齐了。 +//! +//! **每个选项各算各的**：正文一个选项一块（块号就是选项的 `index`），工具调用各攒各的、 +//! 序号各数各的，`finish_reason` 只结束它那一个选项；`[DONE]`、错误帧、用量块把所有选项 +//! 一起收尾。 //! //! 工具调用的分片从原来的帧里摘掉，收齐之后**写成一帧完整的调用**（不变的也是：一帧 -//! 和几帧拼起来是同一个调用），序号按写出去的顺序重新数。全被去掉了的话 -//! `finish_reason` 从 `tool_calls` 改成 `stop`。 +//! 和几帧拼起来是同一个调用），序号按写出去的顺序重新数。一个选项的调用全被去掉了的话， +//! 它的 `finish_reason` 从 `tool_calls` 改成 `stop`。 + +use std::collections::BTreeMap; use serde_json::{Map, Value, json}; @@ -14,21 +21,25 @@ use super::{Call, Chain, Frame, Ids, Out, whole_text}; use crate::error::GatewayError; use crate::plugin::view::{args_text, args_value}; -/// 正文只有一路 -const LANE: u64 = 0; - pub(crate) struct Codec { wants_text: bool, wants_tools: bool, + /// 每个选项一份，按选项的 `index` + choices: BTreeMap, + /// 最近一帧的外壳（id、created、model……），补帧时抄它 + envelope: Map, + ids: Ids, +} + +/// 一个选项的正文和工具调用。正文的块号就是选项的 `index` +#[derive(Default)] +struct Choice { lane_open: bool, calls: Vec, /// 写出去的下一个调用的序号 next_index: u64, - /// 最近一帧的外壳（id、created、model……），补帧时抄它 - envelope: Map, tools_seen: u64, tools_emitted: u64, - ids: Ids, } struct CallBuf { @@ -41,26 +52,34 @@ struct CallBuf { first: Value, } +/// 帧里一个选项的 `index`。没写的按它在数组里的位置 +fn index_of(choice: &Value, pos: usize) -> u64 { + choice + .get("index") + .and_then(Value::as_u64) + .unwrap_or(pos as u64) +} + impl Codec { pub(super) fn new(chain: &Chain) -> Self { Self { wants_text: chain.wants_text(), wants_tools: chain.wants_tools(), - lane_open: false, - calls: Vec::new(), - next_index: 0, + choices: BTreeMap::new(), envelope: Map::new(), - tools_seen: 0, - tools_emitted: 0, ids: Ids::new(), } } - fn chunk(&self, delta: Value) -> Out { + fn at(&mut self, k: u64) -> &mut Choice { + self.choices.entry(k).or_default() + } + + fn chunk(&self, k: u64, delta: Value) -> Out { let mut m = self.envelope.clone(); m.insert( "choices".into(), - json!([{ "index": 0, "delta": delta, "finish_reason": null }]), + json!([{ "index": k, "delta": delta, "finish_reason": null }]), ); Out::new(None, Value::Object(m)) } @@ -68,23 +87,30 @@ impl Codec { async fn close_lane( &mut self, chain: &mut Chain, + k: u64, out: &mut Vec, ) -> Result<(), GatewayError> { - if !self.lane_open { + let c = self.at(k); + if !c.lane_open { return Ok(()); } - self.lane_open = false; - let end = chain.text_end(LANE).await?; + c.lane_open = false; + let end = chain.text_end(k).await?; if !end.is_empty() { - out.push(self.chunk(json!({ "content": end }))); + out.push(self.chunk(k, json!({ "content": end }))); } Ok(()) } - /// 攒着的调用都收齐了：交给插件，写出去 - async fn settle(&mut self, chain: &mut Chain, out: &mut Vec) -> Result<(), GatewayError> { - for c in std::mem::take(&mut self.calls) { - self.tools_seen += 1; + /// 选项 `k` 攒着的调用都收齐了：交给插件，写出去 + async fn settle( + &mut self, + chain: &mut Chain, + k: u64, + out: &mut Vec, + ) -> Result<(), GatewayError> { + for c in std::mem::take(&mut self.at(k).calls) { + self.at(k).tools_seen += 1; let input = if c.custom { Value::String(c.args.clone()) } else { @@ -115,7 +141,10 @@ impl Codec { }; for (id, name, args) in written { let mut entry = c.first.clone(); - entry["index"] = json!(self.next_index); + let ch = self.at(k); + entry["index"] = json!(ch.next_index); + ch.next_index += 1; + ch.tools_emitted += 1; entry["id"] = json!(id); if c.custom { entry["type"] = json!("custom"); @@ -131,14 +160,26 @@ impl Codec { f.insert("arguments".into(), json!(args)); entry["function"] = Value::Object(f); } - self.next_index += 1; - self.tools_emitted += 1; - out.push(self.chunk(json!({ "tool_calls": [entry] }))); + out.push(self.chunk(k, json!({ "tool_calls": [entry] }))); } } Ok(()) } + /// 所有选项一起收尾：扣着的正文、攒着的调用都补出来 + async fn close_all( + &mut self, + chain: &mut Chain, + out: &mut Vec, + ) -> Result<(), GatewayError> { + let ks: Vec = self.choices.keys().copied().collect(); + for k in ks { + self.close_lane(chain, k, out).await?; + self.settle(chain, k, out).await?; + } + Ok(()) + } + pub(super) async fn frame( &mut self, chain: &mut Chain, @@ -147,14 +188,12 @@ impl Codec { ) -> Result<(), GatewayError> { let Some(data) = f.data.as_mut() else { // `[DONE]`：攒着的赶在它前面补出来 - self.close_lane(chain, out).await?; - self.settle(chain, out).await?; + self.close_all(chain, out).await?; out.push(Out::Keep(f)); return Ok(()); }; if data.get("error").is_some_and(|e| !e.is_null()) { - self.close_lane(chain, out).await?; - self.settle(chain, out).await?; + self.close_all(chain, out).await?; out.push(Out::Keep(f)); return Ok(()); } @@ -166,18 +205,52 @@ impl Codec { } } let has_usage = data.get("usage").is_some_and(|u| !u.is_null()); - let Some(choice) = data + let Some(choices) = data .get_mut("choices") .and_then(Value::as_array_mut) - .and_then(|c| c.get_mut(0)) + .filter(|c| !c.is_empty()) else { // 没有 choices 的（流末尾的用量块）：之前攒着的先补出来 - self.close_lane(chain, out).await?; - self.settle(chain, out).await?; + self.close_all(chain, out).await?; out.push(Out::Keep(f)); return Ok(()); }; let mut before: Vec = Vec::new(); + let mut dirty = false; + let mut kept = Vec::with_capacity(choices.len()); + for (pos, mut choice) in std::mem::take(choices).into_iter().enumerate() { + let k = index_of(&choice, pos); + let changed = self.choice(chain, k, &mut choice, &mut before).await?; + dirty |= changed; + // 摘空了的选项不发：没有正文、没有工具调用、没有结束原因，帧上也没有用量 + let empty = choice + .get("delta") + .and_then(Value::as_object) + .is_none_or(|d| d.values().all(|v| v.is_null() || v == "")) + && choice.get("finish_reason").is_none_or(Value::is_null) + && !has_usage; + if !(changed && empty) { + kept.push(choice); + } + } + out.extend(before); + if kept.is_empty() { + return Ok(()); + } + *choices = kept; + f.dirty |= dirty; + out.push(Out::Keep(f)); + Ok(()) + } + + /// 帧里选项 `k` 的那一份。补出来的帧放进 `before`，返回这一份改没改 + async fn choice( + &mut self, + chain: &mut Chain, + k: u64, + choice: &mut Value, + before: &mut Vec, + ) -> Result { let mut dirty = false; let finish = choice .get("finish_reason") @@ -192,7 +265,7 @@ impl Codec { .is_some_and(|t| !t.is_empty()) }); if thinking { - self.close_lane(chain, &mut before).await?; + self.close_lane(chain, k, before).await?; } if self.wants_text && let Some(text) = delta @@ -201,15 +274,15 @@ impl Codec { .filter(|t| !t.is_empty()) .map(str::to_string) { - if !self.calls.is_empty() { - self.settle(chain, &mut before).await?; + if !self.at(k).calls.is_empty() { + self.settle(chain, k, before).await?; } - self.lane_open = true; - let mut got = chain.text(LANE, &text).await?; + self.at(k).lane_open = true; + let mut got = chain.text(k, &text).await?; if finish.is_some() { // 这一帧就是结尾：扣着的接在它自己的正文后面 - self.lane_open = false; - got.push_str(&chain.text_end(LANE).await?); + self.at(k).lane_open = false; + got.push_str(&chain.text_end(k).await?); } if got != text { dirty = true; @@ -224,14 +297,14 @@ impl Codec { && let Some(entries) = delta.get("tool_calls").and_then(Value::as_array).cloned() && !entries.is_empty() { - self.close_lane(chain, &mut before).await?; + self.close_lane(chain, k, before).await?; for (pos, e) in entries.iter().enumerate() { - let k = e.get("index").and_then(Value::as_u64).unwrap_or(pos as u64); - let known = self.calls.iter().any(|c| c.index == k); + let i = e.get("index").and_then(Value::as_u64).unwrap_or(pos as u64); + let known = self.at(k).calls.iter().any(|c| c.index == i); if !known { - // 新的一个调用开始了：前面的都收齐了 - if !self.calls.is_empty() { - self.settle(chain, &mut before).await?; + // 新的一个调用开始了：这个选项前面的都收齐了 + if !self.at(k).calls.is_empty() { + self.settle(chain, k, before).await?; } let custom = e.get("type").and_then(Value::as_str) == Some("custom"); let inner = if custom { @@ -248,8 +321,8 @@ impl Codec { if let Some(o) = first.as_object_mut() { o.remove("index"); } - self.calls.push(CallBuf { - index: k, + self.at(k).calls.push(CallBuf { + index: i, id: e .get("id") .and_then(Value::as_str) @@ -262,9 +335,10 @@ impl Codec { }); } let c = self + .at(k) .calls .iter_mut() - .find(|c| c.index == k) + .find(|c| c.index == i) .expect("just pushed"); let piece = if c.custom { e.pointer("/custom/input") @@ -280,29 +354,15 @@ impl Codec { } } if let Some(reason) = finish { - if self.lane_open { - self.close_lane(chain, &mut before).await?; - } - self.settle(chain, &mut before).await?; - if reason == "tool_calls" && self.tools_seen > 0 && self.tools_emitted == 0 { + self.close_lane(chain, k, before).await?; + self.settle(chain, k, before).await?; + let c = self.at(k); + if reason == "tool_calls" && c.tools_seen > 0 && c.tools_emitted == 0 { choice["finish_reason"] = json!("stop"); dirty = true; } } - out.extend(before); - // 摘空了的帧不发：没有正文、没有工具调用、没有结束原因、没有用量 - let empty = choice - .get("delta") - .and_then(Value::as_object) - .is_none_or(|d| d.values().all(|v| v.is_null() || v == "")) - && choice.get("finish_reason").is_none_or(Value::is_null) - && !has_usage; - if dirty && empty { - return Ok(()); - } - f.dirty |= dirty; - out.push(Out::Keep(f)); - Ok(()) + Ok(dirty) } pub(super) async fn finish( @@ -310,17 +370,26 @@ impl Codec { chain: &mut Chain, out: &mut Vec, ) -> Result<(), GatewayError> { - self.close_lane(chain, out).await?; - self.settle(chain, out).await + self.close_all(chain, out).await } } -/// 整包：`choices[0].message` 的正文和工具调用 +/// 整包：每个选项（`n` > 1 时有好几个）的 `message` 里的正文和工具调用 pub(super) async fn whole(chain: &mut Chain, v: &mut Value) -> Result { - let Some(m) = v - .pointer_mut("/choices/0/message") - .and_then(Value::as_object_mut) - else { + let Some(choices) = v.get_mut("choices").and_then(Value::as_array_mut) else { + return Ok(false); + }; + let mut changed = false; + for (pos, choice) in choices.iter_mut().enumerate() { + let k = index_of(choice, pos); + changed |= whole_choice(chain, k, choice).await?; + } + Ok(changed) +} + +/// 整包里的一个选项。正文的块号是选项的 `index` +async fn whole_choice(chain: &mut Chain, k: u64, choice: &mut Value) -> Result { + let Some(m) = choice.get_mut("message").and_then(Value::as_object_mut) else { return Ok(false); }; let mut changed = false; @@ -328,7 +397,7 @@ pub(super) async fn whole(chain: &mut Chain, v: &mut Value) -> Result Result 1 时有好几个，按候选的 `index` 分）的正文各是一路，块号就是 +//! 候选的 `index`：碰到推理部分、函数调用或者 `finishReason` 就是这个候选的这一块结束 +//! 了，插件扣着的补成一个文字部分。函数调用不用攒，那一个部分就是完整的调用；换出来的 +//! 几个各写成一个 `functionCall` 部分，**推理签名（`thoughtSignature`）留在第一个上**， +//! 去掉的那个的签名挪给同一帧里下一个调用 —— Gemini 下一轮要看到它。 //! -//! 一帧里的部分全被扣下了：没有结束原因就整帧不发，有的话留一个空的文字部分。 +//! 一帧里一个候选的部分全被扣下了：没有结束原因就不发这个候选（一个候选都不剩就整帧 +//! 不发），有的话留一个空的文字部分。 //! 不带 `alt=sse` 的客户端收到的是 JSON 数组，拆帧、拼数组在上一层。 +use std::collections::BTreeSet; + use serde_json::{Map, Value, json}; use super::{Call, Chain, Frame, Ids, Out, whole_text}; use crate::error::GatewayError; -const LANE: u64 = 0; - pub(crate) struct Codec { wants_text: bool, wants_tools: bool, - lane_open: bool, + /// 正文那一块开着的候选（按候选的 `index`） + open: BTreeSet, /// 最近一帧的外壳（`modelVersion`、`responseId`），流结束时补帧抄它 envelope: Map, ids: Ids, @@ -35,25 +38,35 @@ impl Codec { Self { wants_text: chain.wants_text(), wants_tools: chain.wants_tools(), - lane_open: false, + open: BTreeSet::new(), envelope: Map::new(), ids: Ids::new(), } } - async fn end(&mut self, chain: &mut Chain) -> Result { - if !self.lane_open { + async fn end(&mut self, chain: &mut Chain, k: u64) -> Result { + if !self.open.remove(&k) { return Ok(String::new()); } - self.lane_open = false; - chain.text_end(LANE).await + chain.text_end(k).await } - fn chunk(&self, text: String) -> Out { + /// 所有候选一起收尾：扣着的各补成一帧 + async fn end_all(&mut self, chain: &mut Chain, out: &mut Vec) -> Result<(), GatewayError> { + for k in std::mem::take(&mut self.open) { + let end = chain.text_end(k).await?; + if !end.is_empty() { + out.push(self.chunk(k, end)); + } + } + Ok(()) + } + + fn chunk(&self, k: u64, text: String) -> Out { let mut m = self.envelope.clone(); m.insert( "candidates".into(), - json!([{ "content": { "role": "model", "parts": [{ "text": text }] }, "index": 0 }]), + json!([{ "content": { "role": "model", "parts": [{ "text": text }] }, "index": k }]), ); Out::new(None, Value::Object(m)) } @@ -67,10 +80,7 @@ impl Codec { let Some(data) = f.data.as_mut() else { // JSON 数组的 `]`、心跳：之前扣着的先补出来 if f.raw == b"]" { - let end = self.end(chain).await?; - if !end.is_empty() { - out.push(self.chunk(end)); - } + self.end_all(chain, out).await?; } out.push(Out::Keep(f)); return Ok(()); @@ -81,21 +91,55 @@ impl Codec { } } if data.get("error").is_some_and(|e| !e.is_null()) { - let end = self.end(chain).await?; - if !end.is_empty() { - out.push(self.chunk(end)); - } + self.end_all(chain, out).await?; out.push(Out::Keep(f)); return Ok(()); } - let Some(cand) = data + let Some(cands) = data .get_mut("candidates") .and_then(Value::as_array_mut) - .and_then(|c| c.get_mut(0)) + .filter(|c| !c.is_empty()) else { out.push(Out::Keep(f)); return Ok(()); }; + let mut changed = false; + let mut kept = Vec::with_capacity(cands.len()); + for (pos, mut cand) in std::mem::take(cands).into_iter().enumerate() { + let k = cand + .get("index") + .and_then(Value::as_u64) + .unwrap_or(pos as u64); + match self.candidate(chain, k, &mut cand).await? { + // 这个候选的字都扣着：不发 + None => changed = true, + Some(c) => { + changed |= c; + kept.push(cand); + } + } + } + *cands = kept; + if !changed { + out.push(Out::Keep(f)); + return Ok(()); + } + if cands.is_empty() { + return Ok(()); + } + f.dirty = true; + out.push(Out::Keep(f)); + Ok(()) + } + + /// 帧里候选 `k` 的那一份。返回改没改；`None` 是它的部分全被扣下了、又没有结束原因， + /// 这一帧里不发它 + async fn candidate( + &mut self, + chain: &mut Chain, + k: u64, + cand: &mut Value, + ) -> Result, GatewayError> { let finished = field(cand, "finishReason", "finish_reason").is_some_and(|r| !r.is_null()); let parts = cand .pointer("/content/parts") @@ -111,8 +155,8 @@ impl Codec { let thought = p.get("thought").and_then(Value::as_bool) == Some(true); if let Some(t) = text { if thought || !self.wants_text { - if thought && self.lane_open { - let end = self.end(chain).await?; + if thought && self.open.contains(&k) { + let end = self.end(chain, k).await?; if !end.is_empty() { next.push(json!({ "text": end })); changed = true; @@ -121,8 +165,8 @@ impl Codec { next.push(p); continue; } - self.lane_open = true; - let got = chain.text(LANE, &t).await?; + self.open.insert(k); + let got = chain.text(k, &t).await?; if got == t { next.push(p); continue; @@ -141,7 +185,7 @@ impl Codec { next.push(p); continue; }; - let end = self.end(chain).await?; + let end = self.end(chain, k).await?; if !end.is_empty() { next.push(json!({ "text": end })); changed = true; @@ -214,8 +258,8 @@ impl Codec { } } } - if finished && self.lane_open { - let end = self.end(chain).await?; + if finished && self.open.contains(&k) { + let end = self.end(chain, k).await?; if !end.is_empty() { changed = true; match next.last_mut() { @@ -231,22 +275,18 @@ impl Codec { } } if !changed { - out.push(Out::Keep(f)); - return Ok(()); + return Ok(Some(false)); } if next.is_empty() { if !finished { - // 这一帧里的字都扣着：不发 - return Ok(()); + return Ok(None); } next.push(json!({ "text": "" })); } if let Some(c) = cand.get_mut("content") { c["parts"] = Value::Array(next); } - f.dirty = true; - out.push(Out::Keep(f)); - Ok(()) + Ok(Some(true)) } pub(super) async fn finish( @@ -254,18 +294,26 @@ impl Codec { chain: &mut Chain, out: &mut Vec, ) -> Result<(), GatewayError> { - let end = self.end(chain).await?; - if !end.is_empty() { - out.push(self.chunk(end)); - } - Ok(()) + self.end_all(chain, out).await } } -/// 整包：`candidates[0].content.parts` 里的文字和函数调用。每个文字部分算一块 +/// 整包：每个候选（`candidateCount` > 1 时有好几个） pub(super) async fn whole(chain: &mut Chain, v: &mut Value) -> Result { - let Some(parts) = v - .pointer("/candidates/0/content/parts") + let Some(cands) = v.get_mut("candidates").and_then(Value::as_array_mut) else { + return Ok(false); + }; + let mut changed = false; + for cand in cands.iter_mut() { + changed |= whole_candidate(chain, cand).await?; + } + Ok(changed) +} + +/// 整包里一个候选 `content.parts` 里的文字和函数调用。每个文字部分算一块 +async fn whole_candidate(chain: &mut Chain, cand: &mut Value) -> Result { + let Some(parts) = cand + .pointer("/content/parts") .and_then(Value::as_array) .cloned() else { @@ -330,7 +378,7 @@ pub(super) async fn whole(chain: &mut Chain, v: &mut Value) -> Result String { + let chunk = |k: u64, delta: Value, finish: Value| { + data( + json!({"id":"c1","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":k,"delta":delta,"finish_reason":finish}]}), + ) + }; + [ + chunk(0, json!({"role":"assistant","content":""}), Value::Null), + chunk(1, json!({"role":"assistant","content":""}), Value::Null), + chunk(0, json!({"content":text0[0]}), Value::Null), + chunk(1, json!({"content":text1[0]}), Value::Null), + chunk(0, json!({"content":text0[1]}), Value::Null), + chunk(1, json!({"content":text1[1]}), Value::Null), + chunk(0, json!({"tool_calls":[{"index":0,"id":"call_a","type":"function","function":{"name":"shell","arguments":""}}]}), Value::Null), + chunk(1, json!({"tool_calls":[{"index":0,"id":"call_b","type":"function","function":{"name":"shell","arguments":""}}]}), Value::Null), + chunk(0, json!({"tool_calls":[{"index":0,"function":{"arguments":"{\"cmd\":\"ls\"}"}}]}), Value::Null), + chunk(1, json!({"tool_calls":[{"index":0,"function":{"arguments":"{\"cmd\":\"pwd\"}"}}]}), Value::Null), + chunk(0, json!({"tool_calls":[{"index":1,"id":"call_c","type":"function","function":{"name":"read","arguments":"{\"p\":1}"}}]}), Value::Null), + chunk(1, json!({}), json!("tool_calls")), + chunk(0, json!({}), json!("tool_calls")), + data(json!({"id":"c1","object":"chat.completion.chunk","created":1,"model":"m","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2}})), + "data: [DONE]\n\n".to_string(), + ] + .concat() +} + +/// 只读一个选项（帧里 `index` 是 `k` 的那些）的输出 +fn chat_reading_of(out: &str, k: u64) -> (String, Vec, String) { + let only: String = frames(out) + .into_iter() + .flat_map(|(_, v)| { + v["choices"] + .as_array() + .into_iter() + .flatten() + .filter(|c| c["index"] == k) + .map(|c| data(json!({ "choices": [c] }))) + .collect::>() + }) + .collect(); + chat_reading(&only) +} + +/// 工具调用：`shell` 的命令后面加 ` -a`，`read` 去掉 +fn shell_dash_a() -> Double { + tools(|call| { + if call["name"] == "shell" { + let cmd = call["input"]["cmd"].as_str().unwrap_or("?"); + ToolCallOutcome::Replace(vec![ + json!({ "name": "shell", "input": { "cmd": format!("{cmd} -a") } }), + ]) + } else { + ToolCallOutcome::Drop + } + }) +} + +/// `n` > 1：每个选项的正文各攒各的、工具调用各收各的，改过的写回它自己的选项 +#[tokio::test] +async fn chat_choices_are_kept_apart_in_a_stream() { + let input = chat_two_choices(["hel", "lo"], ["wor", "ld"]); + let mut s = Stream::new( + chain_of(Dialect::Chat, vec![upper(), shell_dash_a()], &json!({})).await, + Framing::Sse, + ); + let (out, err) = run(&mut s, &input, 4096).await; + assert!(err.is_none()); + // 换出来的调用没给 id，另发一个：只看序号、名字、参数 + let shape = |calls: Vec| -> Vec<(u64, String, String)> { + calls.into_iter().map(|c| (c.0, c.2, c.3)).collect() + }; + let (text, calls, finish) = chat_reading_of(&out, 0); + assert_eq!(text, "HELLO", "{out}"); + assert_eq!( + shape(calls), + [(0, "shell".into(), "{\"cmd\":\"ls -a\"}".into())], + "{out}" + ); + assert_eq!(finish, "tool_calls"); + let (text, calls, finish) = chat_reading_of(&out, 1); + assert_eq!(text, "WORLD", "{out}"); + assert_eq!( + shape(calls), + [(0, "shell".into(), "{\"cmd\":\"pwd -a\"}".into())], + "{out}" + ); + assert_eq!(finish, "tool_calls"); + assert!(out.ends_with("data: [DONE]\n\n")); +} + +/// 流式插件逐段收、半个密钥先扣着：扣着的那半个只接这个选项后面的字，不会落进另一个选项 +#[tokio::test] +async fn chat_choices_hold_back_apart_in_stream_mode() { + let request = json!({ "messages": [{ "role": "user", "content": format!("key {KEY}") }] }); + let input = chat_two_choices( + [ + &format!("use {}", &KEY[..12]), + &format!("{} now", &KEY[12..]), + ], + ["ab", "c"], + ); + let mut s = Stream::new( + chain_of( + Dialect::Chat, + vec![upper().mode(ReplyMode::Stream), shell_dash_a()], + &request, + ) + .await, + Framing::Sse, + ); + let (out, err) = run(&mut s, &input, 7).await; + assert!(err.is_none()); + assert_eq!( + chat_reading_of(&out, 0).0, + format!("USE {KEY} NOW"), + "{out}" + ); + assert_eq!(chat_reading_of(&out, 1).0, "ABC", "{out}"); +} + // ───────────────────────────────────────────────────────── Responses fn responses_stream() -> String { @@ -1114,6 +1235,77 @@ async fn a_cut_json_array_is_closed_into_a_valid_array() { assert_eq!(got.last().unwrap()["error"]["message"], "cut"); } +/// 只读一个候选（`index` 是 `k` 的那些）的输出：(正文, 调用) +fn gemini_reading_of(chunks: &[Value], k: u64) -> (String, Vec) { + let only: Vec = chunks + .iter() + .flat_map(|c| { + c["candidates"] + .as_array() + .into_iter() + .flatten() + .filter(|x| x["index"] == k) + .map(|x| json!({ "candidates": [x] })) + .collect::>() + }) + .collect(); + let (text, calls, _) = gemini_reading(&only); + (text, calls) +} + +/// 函数调用：`ls` 的目录前面加 `x` +fn ls_x() -> Double { + tools(|call| { + let dir = call["input"]["dir"].as_str().unwrap_or("?"); + ToolCallOutcome::Replace(vec![ + json!({ "name": "ls", "input": { "dir": format!("x{dir}") } }), + ]) + }) +} + +/// `candidateCount` > 1：每个候选的正文各攒各的，最后一帧里两个候选都交给插件 +#[tokio::test] +async fn gemini_candidates_are_kept_apart_in_a_stream() { + let cand = |k: u64, parts: Value| json!({"content":{"role":"model","parts":parts},"index":k}); + let mut last0 = cand(0, json!([{"text":"!"}])); + last0["finishReason"] = json!("STOP"); + let mut last1 = cand( + 1, + json!([{"text":"ld"},{"functionCall":{"name":"ls","args":{"dir":"tmp"}}}]), + ); + last1["finishReason"] = json!("STOP"); + let chunks = [ + json!({"candidates":[cand(0, json!([{"text":"hel"}]))],"modelVersion":"g"}), + json!({"candidates":[cand(1, json!([{"text":"wor"}]))],"modelVersion":"g"}), + json!({"candidates":[cand(0, json!([{"text":"lo"},{"functionCall":{"name":"ls","args":{"dir":"."}}}]))],"modelVersion":"g"}), + json!({"candidates":[last0, last1],"usageMetadata":{"promptTokenCount":1},"modelVersion":"g"}), + ]; + let sse: String = chunks.iter().map(|c| data(c.clone())).collect(); + let mut s = Stream::new( + chain_of(Dialect::Gemini, vec![upper(), ls_x()], &json!({})).await, + Framing::Sse, + ); + let (out, err) = run(&mut s, &sse, 4096).await; + assert!(err.is_none()); + let got: Vec = frames(&out).into_iter().map(|(_, v)| v).collect(); + assert_eq!( + gemini_reading_of(&got, 0), + ( + "HELLO!".into(), + vec![json!({"name":"ls","args":{"dir":"x."}})] + ), + "{out}" + ); + assert_eq!( + gemini_reading_of(&got, 1), + ( + "WORLD".into(), + vec![json!({"name":"ls","args":{"dir":"xtmp"}})] + ), + "{out}" + ); +} + // ───────────────────────────────────────────────────────── 整包 #[tokio::test] @@ -1179,6 +1371,50 @@ async fn whole_bodies_get_block_semantics_in_every_format() { } } +/// 整包里有几个选项（Chat 的 `n`、Gemini 的 `candidateCount`）：每一个都交给插件 +#[tokio::test] +async fn whole_bodies_run_every_choice_through_the_plugins() { + let chat_choice = |k: u64, text: &str| { + json!({"index":k,"message":{"role":"assistant","content":text, + "tool_calls":[{"id":format!("call_{k}"),"type":"function","function":{"name":"Bash","arguments":"{\"command\":\"ls\"}"}}]}, + "finish_reason":"tool_calls"}) + }; + let gemini_cand = |k: u64, text: &str| { + json!({"content":{"role":"model","parts":[ + {"text":text}, + {"functionCall":{"name":"Bash","args":{"command":"ls"}}} + ]},"finishReason":"STOP","index":k}) + }; + let cases = [ + ( + Dialect::Chat, + json!({"id":"c","object":"chat.completion","choices":[chat_choice(0, "hello"), chat_choice(1, "world")]}), + ), + ( + Dialect::Gemini, + json!({"candidates":[gemini_cand(0, "hello"), gemini_cand(1, "world")]}), + ), + ]; + for (d, body) in cases { + let drop_bash = tools(|_| ToolCallOutcome::Drop); + let mut chain = chain_of(d, vec![hold_until_end(), drop_bash], &json!({})).await; + let out = whole(&mut chain, body.to_string().as_bytes()) + .await + .unwrap(); + let v: Value = serde_json::from_slice(&out).unwrap(); + let text = v.to_string(); + assert!( + text.contains("HELLO") && text.contains("WORLD"), + "{d:?}: {text}" + ); + assert!(!text.contains("Bash"), "{d:?}: {text}"); + if d == Dialect::Chat { + assert_eq!(v["choices"][0]["finish_reason"], "stop"); + assert_eq!(v["choices"][1]["finish_reason"], "stop"); + } + } +} + #[tokio::test] async fn reply_runs_are_counted_on_the_plugins() { let state = state();