From 4145ad72a36100038dd4d783d56f926f712592d7 Mon Sep 17 00:00:00 2001 From: WaylandYang Date: Wed, 9 Sep 2026 19:52:42 +0800 Subject: [PATCH 1/3] The chat loop runs on rig Co-Authored-By: Claude Fable 5.1 Signed-off-by: WaylandYang --- Cargo.lock | 341 +++++++++++- crates/utopia-llm/src/lib.rs | 90 ++- crates/utopia-server/Cargo.toml | 2 + crates/utopia-server/src/api/agent.rs | 445 +++++++++++++++ crates/utopia-server/src/api/chat.rs | 634 +++++++++------------- crates/utopia-server/src/api/mod.rs | 2 + crates/utopia-server/src/api/rig_model.rs | 507 +++++++++++++++++ 7 files changed, 1604 insertions(+), 417 deletions(-) create mode 100644 crates/utopia-server/src/api/agent.rs create mode 100644 crates/utopia-server/src/api/rig_model.rs diff --git a/Cargo.lock b/Cargo.lock index cfb8cb373..f6bdb5b63 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -140,6 +140,12 @@ dependencies = [ "password-hash", ] +[[package]] +name = "as-any" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0f477b951e452a0b6b4a10b53ccd569042d1d01729b519e02074a9c0958a063" + [[package]] name = "assert-json-diff" version = "2.0.2" @@ -696,6 +702,15 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" +[[package]] +name = "convert_case" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "affbf0190ed2caf063e3def54ff444b449371d55c58e513a95ab98eca50adb49" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "cookie" version = "0.18.2" @@ -1006,6 +1021,12 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +[[package]] +name = "data-encoding" +version = "2.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4583a4551df46e2792f82ceeac45e850d2e2d5debba0b91f102385cda5b11f06" + [[package]] name = "datasketches" version = "0.2.0" @@ -1200,6 +1221,12 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "ecb" version = "0.1.2" @@ -1310,7 +1337,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -1343,6 +1370,17 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "eventsource-stream" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74fef4569247a5f429d9156b9d0a2599914385dd189c539334c625d8099d90ab" +dependencies = [ + "futures-core", + "nom 7.1.3", + "pin-project-lite", +] + [[package]] name = "fast-float2" version = "0.2.4" @@ -1568,6 +1606,16 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" +[[package]] +name = "futures-timer" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968" +dependencies = [ + "gloo-timers", + "send_wrapper", +] + [[package]] name = "futures-util" version = "0.3.34" @@ -1651,6 +1699,24 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "43503cc176394dd30a6525f5f36e838339b8b5619be33ed9a7783841580a97b6" +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "gloo-timers" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "482ce8a491a501da4cd806bd190275363d674f2845005c6ddbd5d3e1dd54495d" +dependencies = [ + "futures-channel", + "futures-core", + "js-sys", + "wasm-bindgen", +] + [[package]] name = "group" version = "0.13.0" @@ -2687,7 +2753,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2880,7 +2946,7 @@ dependencies = [ "oxiri", "oxrdf", "ryu-js", - "thiserror 1.0.69", + "thiserror 2.0.20", ] [[package]] @@ -2892,7 +2958,7 @@ dependencies = [ "oxilangtag", "oxiri", "rand 0.9.5", - "thiserror 1.0.69", + "thiserror 2.0.20", ] [[package]] @@ -2905,7 +2971,7 @@ dependencies = [ "oxrdf", "oxrdfxml", "oxttl", - "thiserror 1.0.69", + "thiserror 2.0.20", ] [[package]] @@ -2918,7 +2984,7 @@ dependencies = [ "oxiri", "oxrdf", "quick-xml 0.37.5", - "thiserror 1.0.69", + "thiserror 2.0.20", ] [[package]] @@ -2931,7 +2997,7 @@ dependencies = [ "oxilangtag", "oxiri", "oxrdf", - "thiserror 1.0.69", + "thiserror 2.0.20", ] [[package]] @@ -3125,6 +3191,26 @@ dependencies = [ "siphasher", ] +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -3239,6 +3325,15 @@ dependencies = [ "elliptic-curve", ] +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + [[package]] name = "proc-macro-error-attr3" version = "3.1.1" @@ -3376,7 +3471,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3549,6 +3644,26 @@ dependencies = [ "bitflags", ] +[[package]] +name = "ref-cast" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e440fb4e4b4147295338efb76001ab9e4efc0e5839df2c47fc5ac2381d365c3" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ecd8964f8453721699a1ed72037b0db49ce2f5a5138486ee89bed6f67cdf3a" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.5", +] + [[package]] name = "regex" version = "1.13.1" @@ -3599,6 +3714,7 @@ dependencies = [ "js-sys", "log", "mime", + "mime_guess", "percent-encoding", "pin-project-lite", "quinn", @@ -3632,6 +3748,76 @@ dependencies = [ "subtle", ] +[[package]] +name = "rig-agent" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b66477c0d0a786358d078e7e3541a73089dd9009100d6e2b3fe3b0d730e8550f" +dependencies = [ + "async-stream", + "fastrand", + "futures", + "http", + "indexmap", + "rig-core", + "rig-derive", + "schemars", + "serde", + "serde_json", + "thiserror 2.0.20", + "tokio", + "tracing", + "tracing-futures", +] + +[[package]] +name = "rig-core" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "432d83e0facf16749f91fe729cbffca84437e8062d2f4e92f4f12e903693922d" +dependencies = [ + "as-any", + "async-stream", + "base64 0.22.1", + "bytes", + "eventsource-stream", + "fastrand", + "futures", + "futures-timer", + "glob", + "http", + "indexmap", + "mime", + "mime_guess", + "ordered-float", + "pin-project-lite", + "reqwest", + "rig-derive", + "schemars", + "serde", + "serde_json", + "sha2 0.10.9", + "thiserror 2.0.20", + "tokio", + "tokio-tungstenite", + "tracing", + "tracing-futures", + "url", +] + +[[package]] +name = "rig-derive" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de0a33f1bac45f16e50146c248bcbbfaa44518c7252d274e972c7f4ad71aaba7" +dependencies = [ + "convert_case", + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "ring" version = "0.17.14" @@ -3707,7 +3893,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3765,7 +3951,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3822,6 +4008,31 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "schemars" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a" +dependencies = [ + "dyn-clone", + "ref-cast", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98c67716b46af2f0b8cf752abc930f6f9aecfbf671ecfb531db8a31dbe4e2ba" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn 3.0.5", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -3890,6 +4101,12 @@ version = "1.0.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" +[[package]] +name = "send_wrapper" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd0b0ec5f1c1ca621c432a25813d8d60c88abe6d3e08a3eb9cf37d97a0fe3d73" + [[package]] name = "serde" version = "1.0.229" @@ -3920,6 +4137,17 @@ dependencies = [ "syn 3.0.5", ] +[[package]] +name = "serde_derive_internals" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f852137cce035d6a4df67ccce505ff6b3e9fd3a10e3e52b24dc71e650bb1a9bd" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.5", +] + [[package]] name = "serde_json" version = "1.0.151" @@ -4666,10 +4894,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4852,6 +5080,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-tungstenite" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f72a05e828585856dacd553fba484c242c46e391fb0e58917c942ee9202915c" +dependencies = [ + "futures-util", + "log", + "rustls", + "rustls-pki-types", + "tokio", + "tokio-rustls", + "tungstenite", + "webpki-roots 0.26.11", +] + [[package]] name = "tokio-util" version = "0.7.19" @@ -4866,6 +5110,36 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.13+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.3+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" +dependencies = [ + "winnow", +] + [[package]] name = "tower" version = "0.5.3" @@ -4956,6 +5230,18 @@ dependencies = [ "valuable", ] +[[package]] +name = "tracing-futures" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" +dependencies = [ + "futures", + "futures-task", + "pin-project", + "tracing", +] + [[package]] name = "tracing-log" version = "0.2.0" @@ -4997,6 +5283,24 @@ version = "0.25.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d2df906b07856748fa3f6e0ad0cbaa047052d4a7dd609e231c4f72cee8c36f31" +[[package]] +name = "tungstenite" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c01152af293afb9c7c2a57e4b559c5620b421f6d133261c60dd2d0cdb38e6b8" +dependencies = [ + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.9.5", + "rustls", + "rustls-pki-types", + "sha1", + "thiserror 2.0.20", +] + [[package]] name = "type1-encoding-parser" version = "0.1.1" @@ -5248,6 +5552,8 @@ dependencies = [ "percent-encoding", "quick-xml 0.42.0", "reqwest", + "rig-agent", + "rig-core", "serde", "serde_json", "sha2 0.11.0", @@ -5526,7 +5832,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -5828,6 +6134,15 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + [[package]] name = "wiremock" version = "0.6.5" diff --git a/crates/utopia-llm/src/lib.rs b/crates/utopia-llm/src/lib.rs index ded4092b6..6d3bfd9a0 100644 --- a/crates/utopia-llm/src/lib.rs +++ b/crates/utopia-llm/src/lib.rs @@ -123,6 +123,29 @@ pub fn out_of_credit(err: &anyhow::Error) -> Option<&OutOfCredit> { err.chain().find_map(|e| e.downcast_ref::()) } +/// 端点回答了，而且回的是「不」：欠费与限流之外的所有非 2xx。 +/// +/// **做成类型是为了让状态码留下来。** 对话循环要在「这家不支持 tool calling」 +/// （400，请求形状被拒）与「网断了」「密钥错了」之间做不同的事：前者退成一次性 +/// RAG 还能答，后者退了也答不了。从前这一类只剩一段文本，状态码埋在字符串里， +/// 于是循环把**任何**首轮错误都当成不支持工具——一次网络抖动就退成了 RAG， +/// 然后 RAG 也死在同一个抖动上。 +#[derive(Debug, thiserror::Error)] +#[error("{kind} request failed ({reason}): {detail}")] +pub struct Rejected { + /// 哪一类请求（`LLM` / `Embedding`），只进错误文本 + pub kind: String, + pub status: u16, + /// 状态码带原因短语，如 `400 Bad Request`——错误文本从前就是这么写的 + pub reason: String, + pub detail: String, +} + +/// anyhow 错误链里的 [`Rejected`]，穿透 context 层。 +pub fn rejected(err: &anyhow::Error) -> Option<&Rejected> { + err.chain().find_map(|e| e.downcast_ref::()) +} + /// `Retry-After` 的整数秒形态。 /// /// 规范还允许 HTTP-date，这里**不解析**：为一个很少有人发的头引一个日期库不划算， @@ -169,7 +192,12 @@ fn failure( detail, }); } - anyhow::anyhow!("{kind} request failed ({status}): {detail}") + anyhow::Error::new(Rejected { + kind: kind.to_string(), + status: status.as_u16(), + reason: status.to_string(), + detail, + }) } /// 把一个非成功状态的响应变成错误（#527）。 @@ -285,15 +313,45 @@ impl LlmClient { &self, messages: &[serde_json::Value], tools: &serde_json::Value, + ) -> anyhow::Result { + self.chat_tools_with(messages, Some(tools), None).await + } + + /// 带工具的请求体。`tools` 为 None 就不带工具字段——**不是空数组**： + /// 有的端点见到 `"tools": []` 会 400。`tool_choice` 按 OpenAI 协议原样透传 + /// （`"auto"` / `"none"` / `"required"` / `{"type":"function",...}`）， + /// 也只在给了的时候才写进去。 + fn tools_body( + &self, + messages: &[serde_json::Value], + tools: Option<&serde_json::Value>, + tool_choice: Option<&serde_json::Value>, + stream: bool, + ) -> serde_json::Value { + let mut body = json!({ + "model": self.model, + "messages": messages, + "stream": stream, + }); + if let Some(tools) = tools { + body["tools"] = tools.clone(); + if let Some(choice) = tool_choice { + body["tool_choice"] = choice.clone(); + } + } + body + } + + /// 工具对话（非流式），工具清单与 `tool_choice` 都可选。 + pub async fn chat_tools_with( + &self, + messages: &[serde_json::Value], + tools: Option<&serde_json::Value>, + tool_choice: Option<&serde_json::Value>, ) -> anyhow::Result { let resp = self .request("/chat/completions") - .json(&json!({ - "model": self.model, - "messages": messages, - "tools": tools, - "stream": false, - })) + .json(&self.tools_body(messages, tools, tool_choice, false)) .send() .await .map_err(Unreachable)?; @@ -341,15 +399,21 @@ impl LlmClient { &self, messages: &[serde_json::Value], tools: &serde_json::Value, + ) -> anyhow::Result> + Send + use<>> { + self.chat_tools_stream_with(messages, Some(tools), None) + .await + } + + /// 工具对话（流式），工具清单与 `tool_choice` 都可选；见 [`Self::chat_tools_stream`]。 + pub async fn chat_tools_stream_with( + &self, + messages: &[serde_json::Value], + tools: Option<&serde_json::Value>, + tool_choice: Option<&serde_json::Value>, ) -> anyhow::Result> + Send + use<>> { let resp = self .request("/chat/completions") - .json(&json!({ - "model": self.model, - "messages": messages, - "tools": tools, - "stream": true, - })) + .json(&self.tools_body(messages, tools, tool_choice, true)) .send() .await .map_err(Unreachable)?; diff --git a/crates/utopia-server/Cargo.toml b/crates/utopia-server/Cargo.toml index 1d16872b8..133139dc0 100644 --- a/crates/utopia-server/Cargo.toml +++ b/crates/utopia-server/Cargo.toml @@ -48,6 +48,8 @@ url = "2.5.8" percent-encoding = "2.3.2" base64.workspace = true subtle.workspace = true +rig-agent = "0.42" +rig-core = "0.42" [dev-dependencies] wiremock = "0.6.5" diff --git a/crates/utopia-server/src/api/agent.rs b/crates/utopia-server/src/api/agent.rs new file mode 100644 index 000000000..26a97da57 --- /dev/null +++ b/crates/utopia-server/src/api/agent.rs @@ -0,0 +1,445 @@ +//! 对话循环交给 rig 的 runner（#546）。 +//! +//! 从前 `chat.rs` 里手写着一个 1,100 行的循环，每个策略问题都落成它里面的又一个 +//! 分支；#509（模型说「我去查」然后就结束了一轮）的守卫（#543）是启发式上再叠 +//! 启发式——实测追问四次，模型四次都答 DONE。这里把循环换成 rig 的：工具是 +//! `DynamicTool`，策略是钩子，终止是**结构性的**—— +//! +//! **一轮对话在调过至少一个工具之前不能结束。** 首个请求带 `tool_choice: required`， +//! 直到某个工具跑过为止；`no_evidence_needed` 是给「你好」「说短一点」这类不需要 +//! 库的问题准备的那个工具。于是「请稍等，我将调用工具」不再是一个可能的终态： +//! 模型要么真的调了，要么明说这题不用查。 +//! +//! 端点不认 `required`（400 时 `RigModel` 去掉它重发；或者干脆无视）的那一层， +//! 由 `on_model_turn_finished` 兜一次：首轮只有文字、一个工具没跑，退回去要求它调； +//! 再不调就认了——那是端点的事，日志里能看见。 + +use super::tools::{self, ToolCtx, ToolSink}; +use crate::state::AppState; +use rig_agent::agent::{ + AgentHook, CompletionCallAction, CompletionCallEvent, HookContext, ModelTurnAction, + ModelTurnFinished, RequestPatch, RetryRequest, ToolCall as ToolCallEvent, ToolCallAction, + ToolResultAction, ToolResultEvent, +}; +use rig_agent::tool::{DynamicTool, ToolContext, ToolOutput}; +use rig_core::message::{AssistantContent, Message, ToolChoice}; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use utopia_core::models::DataSourceView; +use uuid::Uuid; + +/// 「这题不用查」的那个工具。它存在是为了让首轮的 `required` 有一个诚实的出口： +/// 不是每个问题都关于库里的数据，而强迫模型对「你好」跑一次检索是另一种错 +pub const NO_EVIDENCE_TOOL: &str = "no_evidence_needed"; + +const NO_EVIDENCE_DESCRIPTION: &str = "Declare that this question does not need the knowledge \ + base: a greeting, a question about this conversation itself, or a request to reword or \ + shorten a previous answer. Never use it for a question about facts, entities, documents, \ + dates or data; for those, call the tool that gathers the evidence. After calling it, answer \ + directly."; + +/// 端点无视了 `required`、首轮只有文字时退回去的那句话。**没有 DONE 出口**： +/// 不查的正当理由只有一个，就是调 `no_evidence_needed` +const MUST_CALL: &str = "(system) Every turn starts with a tool call. Call the tool that gathers \ + the evidence for this question, or call no_evidence_needed if the question is not about the \ + knowledge base at all. Do not describe a plan."; + +/// 弹药耗尽那一轮的系统提示补语；工具同时被撤走，模型只能作答 +const BUDGET_EXHAUSTED: &str = + "\n\n(system) Tool budget exhausted. Answer now from the evidence gathered above."; + +/// 一场对话里工具共用的东西：库、权限、引用清单，以及给界面的轨迹。 +/// +/// rig 并发跑同一轮的多个工具，而引用编号是有状态的（`[3]` 取决于之前引过几个）， +/// 所以 `sink` 是一把异步锁，一个工具跑完另一个再进 +pub struct Shared { + pub state: AppState, + pub kb_id: Uuid, + pub workspace_id: Uuid, + pub mounted_sources: Vec, + pub can_write: bool, + pub actor: Uuid, + /// 工具清单（与 MCP 共用的那一份），`check_call` 的判据从这里取 + pub schema: Value, + /// 日志里写模型名，好按模型统计 + pub model: String, + pub sink: tokio::sync::Mutex, + /// 工具跑完留给界面的一步，按 rig 的 internal_call_id 取； + /// `check_call` 拒掉的调用也在这里留一步 + steps: Mutex>, + /// 任何工具（含 `no_evidence_needed`）跑过一次：`required` 的闸门就过了 + gate_passed: AtomicBool, + /// 端点无视 `required` 时的退回只给一次 + nudged: AtomicBool, +} + +impl Shared { + #[allow(clippy::too_many_arguments)] + pub fn new( + state: AppState, + kb_id: Uuid, + workspace_id: Uuid, + mounted_sources: Vec, + can_write: bool, + actor: Uuid, + schema: Value, + model: String, + ) -> Arc { + Arc::new(Self { + state, + kb_id, + workspace_id, + mounted_sources, + can_write, + actor, + schema, + model, + sink: tokio::sync::Mutex::new(ToolSink::default()), + steps: Mutex::new(HashMap::new()), + gate_passed: AtomicBool::new(false), + nudged: AtomicBool::new(false), + }) + } + + fn keep_step(&self, internal_call_id: &str, step: Value) { + self.steps + .lock() + .expect("steps lock") + .insert(internal_call_id.to_string(), step); + } + + /// 取走这次调用留给界面的那一步（没有 = 未知工具，或不留痕的闸门工具） + pub fn take_step(&self, internal_call_id: &str) -> Option { + self.steps + .lock() + .expect("steps lock") + .remove(internal_call_id) + } + + fn tool_ctx(&self) -> ToolCtx<'_> { + ToolCtx { + state: &self.state, + kb_id: self.kb_id, + workspace_id: self.workspace_id, + mounted_sources: &self.mounted_sources, + can_write: self.can_write, + actor: Some(self.actor), + // 网页端对话不经令牌：说话的就是这个人本人 + via_token: None, + } + } +} + +/// 工具跑完留在 rig 工具上下文里的那一步，`on_tool_result` 从那里取 +#[derive(Clone)] +struct Step(Value); + +/// 工具清单变成 rig 的动态工具：名字、描述、参数 schema 都来自 `tools_schema`， +/// 执行还是 `tools::dispatch`。**清单是唯一的真相**，这里不抄第二份 +pub fn dynamic_tools(shared: &Arc) -> Vec { + let mut out = Vec::new(); + for t in shared.schema.as_array().into_iter().flatten() { + let f = &t["function"]; + let Some(name) = f["name"].as_str() else { + continue; + }; + let owned = shared.clone(); + let tool_name = name.to_string(); + out.push(DynamicTool::new( + name, + f["description"].as_str().unwrap_or_default(), + f["parameters"].clone(), + move |ctx: &mut ToolContext, args: Value| { + let shared = owned.clone(); + let name = tool_name.clone(); + Box::pin(async move { + let tool_ctx = shared.tool_ctx(); + let (result, step) = { + let mut sink = shared.sink.lock().await; + tools::dispatch(&tool_ctx, &mut sink, &name, &args).await + }; + shared.gate_passed.store(true, Ordering::Relaxed); + ctx.insert_result(Step(step)); + Ok(ToolOutput::text(result)) + }) + }, + )); + } + let owned = shared.clone(); + out.push(DynamicTool::new( + NO_EVIDENCE_TOOL, + NO_EVIDENCE_DESCRIPTION, + json!({ + "type": "object", + "properties": { + "reason": { + "type": "string", + "description": "One short phrase: why no evidence is needed." + } + }, + "required": ["reason"] + }), + move |_ctx: &mut ToolContext, args: Value| { + let shared = owned.clone(); + Box::pin(async move { + tracing::info!( + model = shared.model, + reason = args["reason"].as_str().unwrap_or(""), + "模型声明这题不用查" + ); + shared.gate_passed.store(true, Ordering::Relaxed); + Ok(ToolOutput::text( + "Understood. Answer the user directly now.", + )) + }) + }, + )); + out +} + +/// 循环的策略，作为 rig 的钩子。 +#[derive(Clone)] +pub struct Policy { + pub shared: Arc, + /// 系统提示原文：弹药耗尽那一轮要在它后面补一句 + pub preamble: String, + /// 允许的工具轮数；第 `max_rounds + 1` 次请求撤走工具、命令作答 + pub max_rounds: usize, +} + +impl AgentHook for Policy { + fn on_completion_call( + &self, + _ctx: &HookContext, + event: CompletionCallEvent<'_>, + ) -> impl std::future::Future + Send { + let turn = event.turn; + let action = if turn > self.max_rounds { + // 弹药耗尽：撤走工具（`RigModel` 对 None 的处理是根本不带工具字段）， + // 系统提示末尾命令它就现有证据作答 + CompletionCallAction::Patch( + RequestPatch::new() + .tool_choice(ToolChoice::None) + .preamble(format!("{}{BUDGET_EXHAUSTED}", self.preamble)), + ) + } else if !self.shared.gate_passed.load(Ordering::Relaxed) { + // 一个工具都还没跑：这一轮必须调一个 + CompletionCallAction::Patch(RequestPatch::new().tool_choice(ToolChoice::Required)) + } else { + CompletionCallAction::Continue + }; + async move { action } + } + + fn on_model_turn_finished( + &self, + _ctx: &HookContext, + event: ModelTurnFinished<'_>, + ) -> impl std::future::Future + Send { + let has_tool_call = event + .content + .iter() + .any(|c| matches!(c, AssistantContent::ToolCall(_))); + let turn = event.turn; + let shared = self.shared.clone(); + let max_rounds = self.max_rounds; + async move { + if has_tool_call || shared.gate_passed.load(Ordering::Relaxed) || turn > max_rounds { + return ModelTurnAction::Continue; + } + // 只有文字、闸门没过：端点无视了 `required`。退回去一次 + if !shared.nudged.swap(true, Ordering::Relaxed) { + tracing::warn!( + model = shared.model, + turn, + "首轮没有工具调用（端点未执行 tool_choice=required），退回要求调用" + ); + return ModelTurnAction::Retry(RetryRequest::Feedback(MUST_CALL.into())); + } + tracing::warn!( + model = shared.model, + turn, + "退回一次后仍无工具调用，按原文收尾（sources 为空）" + ); + ModelTurnAction::Continue + } + } + + fn on_tool_call( + &self, + _ctx: &HookContext, + event: ToolCallEvent<'_>, + ) -> impl std::future::Future + Send { + // **说不清自己要做什么的调用不执行。** 把话回给模型，让它重来； + // 界面上照样显示成一次没做成的调用 + let action = match super::chat::check_call(&self.shared.schema, event.tool_name, event.args) + { + Ok(_) => ToolCallAction::Run, + Err((message, step)) => { + self.shared.keep_step(event.internal_call_id, step); + ToolCallAction::Skip(message) + } + }; + async move { action } + } + + fn on_tool_result( + &self, + _ctx: &HookContext, + event: ToolResultEvent<'_>, + ) -> impl std::future::Future + Send { + if let Some(Step(step)) = event.tool_context.result::() { + self.shared.keep_step(event.internal_call_id, step.clone()); + } + async { ToolResultAction::Keep } + } +} + +/// 落库的历史（OpenAI 协议的 JSON）变回 rig 的消息。 +/// +/// 最后那一轮的工具往返放在它的结论之前——顺序就是真实顺序，模型读起来 +/// 是「我问了、我查了、我答了」。tool 消息的工具名在协议里没有，从前面那条 +/// 带 tool_calls 的 assistant 消息里按 id 找回来 +pub fn history_messages(turns: &[(String, String)], last_tool_exchange: &[Value]) -> Vec { + let last_assistant = turns.iter().rposition(|(role, _)| role == "assistant"); + let mut out = Vec::new(); + for (i, (role, content)) in turns.iter().enumerate() { + if Some(i) == last_assistant { + out.extend(exchange_messages(last_tool_exchange)); + } + match role.as_str() { + "assistant" => out.push(Message::assistant(content.clone())), + _ => out.push(Message::user(content.clone())), + } + } + out +} + +fn exchange_messages(exchange: &[Value]) -> Vec { + let mut names: HashMap = HashMap::new(); + let mut out = Vec::new(); + for m in exchange { + match m["role"].as_str() { + Some("assistant") => { + let mut content = Vec::new(); + if let Some(text) = m["content"].as_str().filter(|t| !t.is_empty()) { + content.push(AssistantContent::text(text)); + } + for c in m["tool_calls"].as_array().into_iter().flatten() { + let id = c["id"].as_str().unwrap_or_default(); + let name = c["function"]["name"].as_str().unwrap_or_default(); + names.insert(id.to_string(), name.to_string()); + content.push(AssistantContent::tool_call( + id, + name, + super::rig_model::args_value( + c["function"]["arguments"].as_str().unwrap_or("{}"), + ), + )); + } + if !content.is_empty() { + out.push(Message::Assistant { id: None, content }); + } + } + Some("tool") => { + let id = m["tool_call_id"].as_str().unwrap_or_default(); + let name = names.get(id).cloned().unwrap_or_default(); + out.push(Message::tool_result( + id, + name, + m["content"].as_str().unwrap_or_default(), + )); + } + _ => {} + } + } + out +} + +/// 前几轮已经认下的实体，连 id 一起交回去（作为本轮的上下文文档，见 `rig_model::wire` +/// 里它落在哪、以什么角色）。少了它，模型只看得见上一轮的最终答案文字，不知道 +/// 自己拿到过哪些 id，于是从名字重搜一遍；同名歧义时两轮可能落到不同的实体上 +pub fn known_entities_block(entities: &[Value], limit: usize) -> Option { + if entities.is_empty() { + return None; + } + let lines: Vec = entities + .iter() + .take(limit) + .map(|e| { + format!( + "{} | {} | {}", + e["id"].as_str().unwrap_or("?"), + e["name"].as_str().unwrap_or("?"), + e["type"].as_str().unwrap_or("?") + ) + }) + .collect(); + Some(format!( + "Entities already identified earlier in this conversation (id | name | type). \ + Call entity_facts with these ids directly; do not look them up by name again:\n{}", + lines.join("\n") + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 上一轮的工具往返插在它的结论之前，tool 消息找回自己的工具名 + #[test] + fn the_last_exchange_sits_before_its_conclusion() { + let turns = vec![ + ("user".to_string(), "q1".to_string()), + ("assistant".to_string(), "a1".to_string()), + ("user".to_string(), "q2".to_string()), + ]; + let exchange = vec![ + json!({ "role": "assistant", "content": null, "tool_calls": [ + { "id": "c1", "type": "function", + "function": { "name": "search_chunks", "arguments": "{\"query\":\"x\"}" } } + ]}), + json!({ "role": "tool", "tool_call_id": "c1", "content": "found" }), + ]; + let msgs = history_messages(&turns, &exchange); + assert_eq!(msgs.len(), 5); + assert!(matches!(&msgs[0], Message::User { .. })); + match &msgs[1] { + Message::Assistant { content, .. } => match &content[0] { + AssistantContent::ToolCall(tc) => { + assert_eq!(tc.id.as_str(), "c1"); + assert_eq!(tc.function.arguments["query"], "x"); + } + other => panic!("expected the tool call, got {other:?}"), + }, + other => panic!("expected the exchange first, got {other:?}"), + } + match &msgs[2] { + Message::User { content } => match &content[0] { + rig_core::message::UserContent::ToolResult(r) => { + assert_eq!(r.call.as_str(), "c1"); + assert_eq!(r.name, "search_chunks"); + } + other => panic!("expected the tool result, got {other:?}"), + }, + other => panic!("expected the tool result, got {other:?}"), + } + assert!(matches!(&msgs[3], Message::Assistant { .. })); + assert!(matches!(&msgs[4], Message::User { .. })); + } + + #[test] + fn the_entity_block_lists_id_name_type_and_stops_at_the_limit() { + assert!(known_entities_block(&[], 20).is_none()); + let entities = vec![ + json!({ "id": "e1", "name": "Acme", "type": "Organization" }), + json!({ "id": "e2", "name": "Bob", "type": "Person" }), + ]; + let block = known_entities_block(&entities, 1).unwrap(); + assert!(block.contains("e1 | Acme | Organization")); + assert!(!block.contains("e2")); + assert!(block.starts_with("Entities already identified")); + } +} diff --git a/crates/utopia-server/src/api/chat.rs b/crates/utopia-server/src/api/chat.rs index d982d9997..4c03985c2 100644 --- a/crates/utopia-server/src/api/chat.rs +++ b/crates/utopia-server/src/api/chat.rs @@ -2,12 +2,19 @@ //! 事件序列：step*（行动轨迹）| sources（引用清单，随检索增量更新）| delta*（增量文本）→ done | error。 //! 模型不支持 tool-calling 时自动降级为一次性 RAG 注入。 -use super::tools; +use super::agent; +use super::rig_model::{self, RigModel}; use crate::live::Frame; use axum::extract::{Path, Query, State}; use axum::response::sse::{Event, KeepAlive, Sse}; use axum::Json; use futures_util::{Stream, StreamExt}; +use rig_agent::agent::{AgentBuilder, MultiTurnStreamItem, StreamingError}; +use rig_agent::completion::PromptError; +use rig_agent::tool::server::ToolServer; +use rig_core::completion::{CompletionError, Document}; +use rig_core::message::Message; +use rig_core::streaming::{StreamedAssistantContent, StreamedUserContent}; use serde::Deserialize; use serde_json::json; use std::convert::Infallible; @@ -29,66 +36,6 @@ const KNOWN_ENTITY_LIMIT: usize = 20; const MAX_HISTORY: usize = 20; const MAX_ROUNDS: usize = 6; -/// 模型说「我去查」却一个工具都没调就结束了一轮时，追问的那句话（#509）。 -/// -/// 实测 DeepSeek-V3 会答「请稍等，我将进行相关搜索」然后什么都不做；再问一次， -/// 还是「稍等，我正在搜索」。提示词第 1、2 条早就禁了这种叙述，所以这不是缺一句 -/// 指令，是模型没听、而循环把它当成了最终答案。守卫放在循环里，不放在提示词里。 -/// -/// 追问只给一次、不流式、只认三种回复：调工具（照常执行）、一个词 DONE（上一句 -/// 本来就是答案：打招呼、问这场对话、拒答，原样收尾）、其他任何文字（还在说空话， -/// 明说没查到证据）。给它 DONE 这条出口，是为了不让「把那句话说短一点」这种 -/// 本就不需要工具的回答被追问成第二个答案。 -/// -/// 措辞把 DONE 的门开得窄：只有问题**不是关于用户数据**时才许说 DONE。实测还有 -/// 一种更坏的停法——不说「稍等」，直接写「以下是我找到的内容」然后凭记忆作答， -/// 库里 steps、sources 全是 0。它没停住，它在撒谎。对它，「若已答完就说 DONE」 -/// 是一条太宽的出口，所以这里明说：关于数据的事实性回答没有工具就不算答。 -const STALL_NUDGE: &str = "(system) Your last message ended the turn without calling any tool, \ - and it cites nothing. An answer about the user's data that was not gathered with a tool is \ - not an answer, whatever the message says it found: call the tool now. Do not describe a \ - plan. Reply with the single word DONE only if the question was not about the user's data \ - at all: a greeting, a question about this transcript, or a refusal."; - -/// 追问后仍不查时补在答案末尾的话。承诺已经流给用户了，收不回来；能做的是 -/// 让文字和空白的轨迹不再互相矛盾——对一个把「每个回答可追溯」当卖点的产品， -/// 一句叙述了从未发生的查证的回答比「不知道」更糟。 -const NO_EVIDENCE_NOTE: &str = - "\n\n(No evidence was gathered for this answer: the model announced a search it did not perform.)"; - -/// 一轮结束、正文非空、整场没调过工具也没有引用：这个答案什么都不站在上面。 -/// 「问这场对话」的消息也满足这两条，所以追问必须便宜、安静，且留有 DONE 出口。 -fn answer_rests_on_nothing(steps: &[serde_json::Value], sources: &[serde_json::Value]) -> bool { - steps.is_empty() && sources.is_empty() -} - -/// 追问之后模型的回复算哪种 -#[derive(Debug, PartialEq, Eq)] -enum AfterNudge { - /// 调了工具：照常执行，接着走 - Tools, - /// 说上一句已经是答案，或者什么都没说：原样收尾 - Done, - /// 又是一段文字：还在说空话 - Stalled, -} - -fn after_nudge(turn: &utopia_llm::AssistantTurn) -> AfterNudge { - if !turn.tool_calls.is_empty() { - return AfterNudge::Tools; - } - let said = turn - .content - .as_deref() - .map(|t| t.trim().trim_matches(|c: char| !c.is_alphanumeric())) - .unwrap_or_default(); - if said.is_empty() || said.eq_ignore_ascii_case("done") { - AfterNudge::Done - } else { - AfterNudge::Stalled - } -} - /// `remember` 曾整个停用过一段（见 `docs/decisions/0015`）：它那时会把一句话直接 /// 变成图上一条活边，实测里「记住 Acme 把总部搬到了深圳」落成的是一条**空谓词、 /// 0.9 置信**的边，而助手宣称的和图里得到的不是一回事。 @@ -199,7 +146,7 @@ pub(super) fn tools_schema(can_write: bool, data_source_names: &[String]) -> ser /// /// 判据直接取自工具表里的 `required`：加一个必填参数，这里自动跟上， /// 不必记得来改第二处。 -fn check_call( +pub(super) fn check_call( tools: &serde_json::Value, name: &str, raw_args: &str, @@ -623,296 +570,190 @@ pub async fn chat( } } } - let mut msgs: Vec = - vec![json!({ "role": "system", "content": system_prompt })]; - /* **上一轮做过什么，按它当时发生的位置放回去。** - 最后那条助手消息是它的结论；带 `tool_calls` 的消息与 tool 结果 - 发生在它之前，所以插在它前面——顺序就是真实顺序，模型读起来 - 就是「我问了、我查了、我答了」。 - 少了这一段，跨轮之后它只看得见自己写的散文，于是接着说「翻译」 - 时重查一遍（还可能落到另一批同名实体上）。 */ - let last_assistant = history - .turns - .iter() - .rposition(|(role, _)| role == "assistant"); - for (i, (role, content)) in history.turns.iter().enumerate() { - if Some(i) == last_assistant { - for m in &history.last_tool_exchange { - msgs.push(m.clone()); - } - } - msgs.push(json!({ "role": role, "content": content })); - } - // **前几轮已经认下的实体，连 id 一起交回去。** - // - // 少了这一段，模型只看得见上一轮的最终答案文字，不知道自己搜过什么、 - // 拿到过哪些 id，于是从名字重搜一遍。更隐蔽的是同名歧义时两轮可能落到 - // **不同的实体**上，前后两个答案讲的不是同一个节点。 - // - // 贴在历史之后、当前问题之前——位置就是服从性，跟抽取里 known_block - // 紧挨正文是同一条理由。 - if !history.entities.is_empty() { - let lines: Vec = history.entities - .iter() - .take(KNOWN_ENTITY_LIMIT) - .map(|e| { - format!( - "{} | {} | {}", - e["id"].as_str().unwrap_or("?"), - e["name"].as_str().unwrap_or("?"), - e["type"].as_str().unwrap_or("?") - ) - }) - .collect(); - msgs.push(json!({ - "role": "user", - "content": format!( - "Entities already identified earlier in this conversation (id | name | type). Call entity_facts with these ids directly; do not look them up by name again: - {}", - lines.join(" - ") - ) - })); - } // 会话 id 先行下发（新会话由此告知前端） yield Frame::new("conversation", json!({ "id": conversation_id }).to_string()); - // 引用清单与这一轮认下的实体。**攒在工具外面**——`[3]` 里的 3 取决于 - // 之前已经引过几个，各个工具各算各的会让同一个 chunk 拿到两个号 - let mut sink = tools::ToolSink::default(); + // 循环是 rig 的（#546）：工具、策略钩子、历史、实体清单都交给它； + // 这里只把它的事件翻成前端认得的帧，并在结束时落库 + let shared = agent::Shared::new( + state.clone(), + kb_id, + workspace_id, + mounted_sources.clone(), + can_write, + user.id, + tools, + settings.chat_model.clone().unwrap_or_default(), + ); + let policy = agent::Policy { + shared: shared.clone(), + preamble: system_prompt.clone(), + max_rounds: MAX_ROUNDS, + }; + let tool_server = ToolServer::new() + .dynamic_tools(agent::dynamic_tools(&shared)) + .run(); + let rig_agent = AgentBuilder::new(RigModel::new(client.clone())) + .preamble(&system_prompt) + // 工具轮 + 最后那一轮作答；第 MAX_ROUNDS+1 次请求由钩子撤走工具 + .default_max_turns(MAX_ROUNDS + 1) + .add_hook(policy) + .tool_server_handle(tool_server) + .build(); + let mut runner = rig_agent + .runner(Message::user(query.clone())) + .history(agent::history_messages(&history.turns, &history.last_tool_exchange)); + // 贴在历史之后、当前问题之前——位置就是服从性，跟抽取里 known_block + // 紧挨正文是同一条理由（角色与位置由 `rig_model::wire` 定） + if let Some(block) = agent::known_entities_block(&history.entities, KNOWN_ENTITY_LIMIT) { + runner = runner.document(Document { + id: "known_entities".into(), + text: block, + additional_props: Default::default(), + }); + } + let mut run = runner.stream().await; + // 落库累积：assistant 全文与行动轨迹（历史回放用） let mut answer_acc = String::new(); let mut steps_acc: Vec = Vec::new(); - // 这一轮的工具往返，原样留一份落库：下一轮回放它，模型才知道自己做过什么 + // 这一轮的工具往返，按协议原样留一份落库：下一轮回放它，模型才知道自己做过什么 let mut exchange_acc: Vec = Vec::new(); + // 当前模型回合里说的话与发出的调用；回合的结果一到，攒成一条 assistant 消息 + let mut turn_text = String::new(); + let mut turn_calls: Vec = Vec::new(); + let mut finished = false; - let mut rounds = 0usize; - // #509 的追问只给一次 - let mut nudged = false; - loop { - if rounds >= MAX_ROUNDS { - // 弹药耗尽：命令模型就现有证据作答（流式） - msgs.push(json!({ - "role": "user", - "content": "(system) Tool budget exhausted. Answer now from the evidence gathered above.", - })); - match client.chat_stream_raw(&msgs).await { - Ok(deltas) => { - let mut deltas = std::pin::pin!(deltas); - while let Some(item) = deltas.next().await { - match item { - Ok(text) => { answer_acc.push_str(&text); yield delta_event(&text); } - Err(e) => { yield error_event(&e.to_string()); return; } - } - } - let _ = utopia_store::conversations::append_message( - &state.pool, conversation_id, "assistant", &answer_acc, - &utopia_store::conversations::TurnRecord { - steps: serde_json::Value::Array(steps_acc.clone()), - sources: serde_json::Value::Array(sink.sources.clone()), - resolved: serde_json::Value::Array(sink.resolved.clone()), - tool_exchange: serde_json::Value::Array(exchange_acc.clone()), - }, - ).await; - yield done_event(); - } - Err(e) => yield error_event(&e.to_string()), + while let Some(item) = run.next().await { + match item { + Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(t))) => { + answer_acc.push_str(&t.text); + turn_text.push_str(&t.text); + yield delta_event(&t.text); } - return; - } - - // 主链路全程流式：正文增量即时转发，工具调用在流末归并到达 - let deltas = match client.chat_tools_stream(&msgs, &tools).await { - Ok(s) => s, - Err(e) => { - if rounds == 0 { - // 模型可能不支持 tool-calling：降级为一次性 RAG 注入 - tracing::warn!(error = %e, "tool-calling 不可用，降级为一次性 RAG"); - let chunks = - retrieval::hybrid(&state, kb_id, workspace_id, &query, 8, None) - .await - .unwrap_or_default(); - let legacy_sources: Vec = chunks - .iter() - .enumerate() - .map(|(i, c)| source_json(i + 1, c)) - .collect(); - yield Frame::new( - "sources", - serde_json::to_string(&legacy_sources).unwrap_or_else(|_| "[]".into()), - ); - let mut lmsgs = - vec![json!({ "role": "system", "content": legacy_system_prompt(&chunks) })]; - for (role, content) in &history.turns { - lmsgs.push(json!({ "role": role, "content": content })); + Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::ToolCall { + tool_call, .. + })) => { + turn_calls.push(json!({ + "id": tool_call.id.as_str(), + "type": "function", + "function": { + "name": tool_call.function.name, + "arguments": rig_model::args_string(&tool_call.function.arguments), } - match client.chat_stream_raw(&lmsgs).await { - Ok(deltas) => { - let mut deltas = std::pin::pin!(deltas); - while let Some(item) = deltas.next().await { - match item { - Ok(text) => { answer_acc.push_str(&text); yield delta_event(&text); } - Err(e2) => { yield error_event(&e2.to_string()); return; } - } - } - let _ = utopia_store::conversations::append_message( - &state.pool, conversation_id, "assistant", &answer_acc, - &utopia_store::conversations::TurnRecord { - steps: serde_json::Value::Array(steps_acc.clone()), - sources: serde_json::Value::Array(legacy_sources.clone()), - resolved: serde_json::Value::Array(sink.resolved.clone()), - tool_exchange: serde_json::Value::Array(exchange_acc.clone()), - }, - ).await; - yield done_event(); - } - Err(e2) => yield error_event(&e2.to_string()), + })); + } + Ok(MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult { + tool_result, + internal_call_id, + })) => { + if !turn_calls.is_empty() { + // 工具轮带了叙述文本：与后续轮次的正文之间补一个段落分隔 + if !turn_text.is_empty() { + answer_acc.push_str("\n\n"); + yield delta_event("\n\n"); } - return; + exchange_acc.push(json!({ + "role": "assistant", + "content": if turn_text.is_empty() { + serde_json::Value::Null + } else { + serde_json::Value::String(turn_text.clone()) + }, + "tool_calls": std::mem::take(&mut turn_calls), + })); + turn_text.clear(); } - yield error_event(&e.to_string()); - return; - } - }; - let mut turn: Option = None; - { - let mut deltas = std::pin::pin!(deltas); - while let Some(item) = deltas.next().await { - match item { - Ok(utopia_llm::ToolStreamItem::Delta(text)) => { - answer_acc.push_str(&text); - yield delta_event(&text); + let text = rig_model::tool_result_text(&tool_result.content); + // 闸门工具不留轨迹：「你好」下面挂一条「声明不用查」是噪音 + if tool_result.name != agent::NO_EVIDENCE_TOOL { + let mut step = shared.take_step(&internal_call_id).unwrap_or_else(|| { + json!({ "kind": "tool", "label": tool_result.name, "detail": "unknown" }) + }); + // **这一步发生在正文的哪个位置。** + // + // 模型是边说边调的：说一句、查一下、再说一句。SSE 上 `delta` 与 + // `step` 本来就是交替发出去的，顺序不用额外记；而**历史回放没有 + // 那条时间线**——落库的只有拼好的整段正文和一个扁平的 steps 数组， + // 于是重新打开一场对话，所有调用都堆在正文最前面，读起来像是 + // 先查了七次再一口气说完。记下偏移，回放才能把话再断开。 + // + // 单位是 **UTF-16 码元**，因为切分发生在浏览器里，而 JS 的 + // `String.prototype.length` 数的就是它。用字节数或 `chars()` + // 在中文和 emoji 上都会切歪 + if let Some(obj) = step.as_object_mut() { + obj.insert("at".into(), json!(answer_acc.encode_utf16().count())); } - Ok(utopia_llm::ToolStreamItem::Turn(t)) => turn = Some(t), - Err(e) => { - yield error_event(&e.to_string()); - return; + steps_acc.push(step.clone()); + yield Frame::new("step", serde_json::to_string(&step).unwrap_or_default()); + if step["kind"] == "search" || step["kind"] == "docs" { + let sources = shared.sink.lock().await.sources.clone(); + yield Frame::new( + "sources", + serde_json::to_string(&sources).unwrap_or_else(|_| "[]".into()), + ); } } + exchange_acc.push(tool_result_message(tool_result.call.as_str(), &text)); } - } - let Some(mut turn) = turn else { - yield error_event("LLM stream ended unexpectedly"); - return; - }; - - if turn.tool_calls.is_empty() { - if answer_acc.is_empty() { - yield error_event("Model returned an empty answer"); - return; - } - // **没调工具的一轮不一定是答完了，也可能是停住了**（#509）：正文说 - // 「我去查」，然后轮次就结束。循环分不出这两种，靠一次追问让模型自己 - // 表态。追问不流式：模型若只是确认 DONE，用户不该看见那个词 - let mut carry_on_with_tools = false; - if !nudged && answer_rests_on_nothing(&steps_acc, &sink.sources) { - nudged = true; - let model = settings.chat_model.clone().unwrap_or_default(); - msgs.push(turn.to_message()); - msgs.push(json!({ "role": "user", "content": STALL_NUDGE })); - match client.chat_tools(&msgs, &tools).await { - Ok(second) => match after_nudge(&second) { - AfterNudge::Tools => { - tracing::warn!(model, "模型只说了要查没查，追问后调了工具"); - // 追问那轮的叙述没有流过，这里补上，接在承诺后面 - if let Some(text) = second.content.as_deref().filter(|t| !t.trim().is_empty()) { - answer_acc.push_str("\n\n"); - yield delta_event("\n\n"); - answer_acc.push_str(text); - yield delta_event(text); - } - turn = second; - carry_on_with_tools = true; - } - // 也记一笔：DONE 说得对不对没法在这里判，只能靠日志和这条 - // 回答的 sources 为空这个事实，事后一起看 - AfterNudge::Done => tracing::info!(model, "模型追问后说上一句已是答案"), - AfterNudge::Stalled => { - tracing::warn!(model, "模型追问后仍只说不查，答案标注无证据"); - answer_acc.push_str(NO_EVIDENCE_NOTE); - yield delta_event(NO_EVIDENCE_NOTE); - } - }, - // 追问本身失败不能拖垮已经到手的答案 - Err(e) => tracing::warn!(model, error = %e, "追问失败，按原答案收尾"), + // 钩子把一个只说不查的回合退了回去：那段话已经流给用户，收不回来； + // 接下来的正文另起一段 + Ok(MultiTurnStreamItem::ModelTurnRetried { .. }) => { + if !turn_text.is_empty() { + answer_acc.push_str("\n\n"); + yield delta_event("\n\n"); } + turn_text.clear(); + turn_calls.clear(); } - if !carry_on_with_tools { - let _ = utopia_store::conversations::append_message( - &state.pool, conversation_id, "assistant", &answer_acc, - &utopia_store::conversations::TurnRecord { - steps: serde_json::Value::Array(steps_acc.clone()), - sources: serde_json::Value::Array(sink.sources.clone()), - resolved: serde_json::Value::Array(sink.resolved.clone()), - tool_exchange: serde_json::Value::Array(exchange_acc.clone()), - }, - ).await; - yield done_event(); - return; - } - } - - // 工具轮带了叙述文本：与后续轮次的正文之间补一个段落分隔 - if turn.content.is_some() && !answer_acc.is_empty() { - answer_acc.push_str("\n\n"); - yield delta_event("\n\n"); - } - - let call_msg = turn.to_message(); - exchange_acc.push(call_msg.clone()); - msgs.push(call_msg); - for call in &turn.tool_calls { - // **说不清自己要做什么的调用不执行。** 把话回给模型，让它重来 - let args = match check_call(&tools, &call.name, &call.arguments) { - Ok(args) => args, - Err((message, step)) => { - steps_acc.push(step.clone()); - yield Frame::new("step", serde_json::to_string(&step).unwrap_or_default()); - msgs.push(tool_result_message(&call.id, &message)); - continue; + Ok(MultiTurnStreamItem::FinalResponse(_)) => finished = true, + Ok(_) => {} + Err(e) => { + let (message, rejected) = describe(&e); + // **只有「端点拒绝了带工具的请求」才降级**为一次性 RAG。从前首轮 + // 的任何错误都走这条路：一次到 SiliconFlow 的网络抖动被记成 + // 「tool-calling 不可用」，然后 RAG 死在同一个抖动上 + if rejected && answer_acc.is_empty() && steps_acc.is_empty() { + tracing::warn!(error = %message, "端点拒绝工具调用，降级为一次性 RAG"); + let mut legacy = std::pin::pin!(legacy_rag( + state.clone(), + kb_id, + workspace_id, + conversation_id, + query.clone(), + history.turns.clone(), + client.clone(), + )); + while let Some(frame) = legacy.next().await { + yield frame; + } + return; } - }; - let ctx = tools::ToolCtx { - state: &state, - kb_id, - workspace_id, - mounted_sources: &mounted_sources, - can_write, - actor: Some(user.id), - // 网页端对话不经令牌：说话的就是这个人本人 - via_token: None, - }; - let (result, step) = tools::dispatch(&ctx, &mut sink, &call.name, &args).await; - // **这一步发生在正文的哪个位置。** - // - // 模型是边说边调的：说一句、查一下、再说一句。SSE 上 `delta` 与 - // `step` 本来就是交替发出去的，顺序不用额外记；而**历史回放没有 - // 那条时间线**——落库的只有拼好的整段正文和一个扁平的 steps 数组， - // 于是重新打开一场对话，所有调用都堆在正文最前面，读起来像是 - // 先查了七次再一口气说完。记下偏移，回放才能把话再断开。 - // - // 单位是 **UTF-16 码元**，因为切分发生在浏览器里，而 JS 的 - // `String.prototype.length` 数的就是它。用字节数或 `chars()` - // 在中文和 emoji 上都会切歪 - let mut step = step; - if let Some(obj) = step.as_object_mut() { - obj.insert("at".into(), json!(answer_acc.encode_utf16().count())); - } - steps_acc.push(step.clone()); - yield Frame::new("step", serde_json::to_string(&step).unwrap_or_default()); - if step["kind"] == "search" || step["kind"] == "docs" { - yield Frame::new( - "sources", - serde_json::to_string(&sink.sources).unwrap_or_else(|_| "[]".into()), - ); + yield error_event(&message); + return; } - let result_msg = tool_result_message(&call.id, &result); - exchange_acc.push(result_msg.clone()); - msgs.push(result_msg); } - rounds += 1; } + if !finished { + yield error_event("LLM stream ended unexpectedly"); + return; + } + if answer_acc.is_empty() { + yield error_event("Model returned an empty answer"); + return; + } + let sink = shared.sink.lock().await; + let _ = utopia_store::conversations::append_message( + &state.pool, conversation_id, "assistant", &answer_acc, + &utopia_store::conversations::TurnRecord { + steps: serde_json::Value::Array(steps_acc), + sources: serde_json::Value::Array(sink.sources.clone()), + resolved: serde_json::Value::Array(sink.resolved.clone()), + tool_exchange: serde_json::Value::Array(exchange_acc), + }, + ).await; + yield done_event(); }; // 生成登记在案，然后**这条连接也只是去「接上」它**——与刷新之后 @@ -933,6 +774,79 @@ pub async fn chat( Ok(sse_from(attached)) } +/// rig 的错误变成给用户的一句话，外加「是不是端点拒绝了工具调用」。 +/// 我们自己的错误链（限流、欠费、被拒）从 `rig_model` 里取回来，文本与从前一样 +fn describe(err: &StreamingError) -> (String, bool) { + fn completion(ce: &CompletionError) -> (String, bool) { + match rig_model::llm_failure(ce) { + Some(ours) => (ours.to_string(), rig_model::tool_calling_rejected(ce)), + None => (ce.to_string(), false), + } + } + match err { + StreamingError::Completion(ce) => completion(ce), + StreamingError::Prompt(pe) => match pe.as_ref() { + PromptError::CompletionError(ce) => completion(ce), + PromptError::PromptCancelled { reason, .. } => (reason.clone(), false), + other => (other.to_string(), false), + }, + } +} + +/// 降级路径：端点不支持工具调用时的一次性 RAG 注入。 +/// 检索一次、把来源塞进系统提示、流式作答、落库 +fn legacy_rag( + state: AppState, + kb_id: Uuid, + workspace_id: Uuid, + conversation_id: Uuid, + query: String, + turns: Vec<(String, String)>, + client: utopia_llm::LlmClient, +) -> impl Stream { + async_stream::stream! { + let chunks = retrieval::hybrid(&state, kb_id, workspace_id, &query, 8, None) + .await + .unwrap_or_default(); + let legacy_sources: Vec = chunks + .iter() + .enumerate() + .map(|(i, c)| source_json(i + 1, c)) + .collect(); + yield Frame::new( + "sources", + serde_json::to_string(&legacy_sources).unwrap_or_else(|_| "[]".into()), + ); + let mut lmsgs = vec![json!({ "role": "system", "content": legacy_system_prompt(&chunks) })]; + for (role, content) in &turns { + lmsgs.push(json!({ "role": role, "content": content })); + } + let mut answer_acc = String::new(); + match client.chat_stream_raw(&lmsgs).await { + Ok(deltas) => { + let mut deltas = std::pin::pin!(deltas); + while let Some(item) = deltas.next().await { + match item { + Ok(text) => { answer_acc.push_str(&text); yield delta_event(&text); } + Err(e) => { yield error_event(&e.to_string()); return; } + } + } + let _ = utopia_store::conversations::append_message( + &state.pool, conversation_id, "assistant", &answer_acc, + &utopia_store::conversations::TurnRecord { + steps: serde_json::Value::Array(Vec::new()), + sources: serde_json::Value::Array(legacy_sources), + resolved: serde_json::Value::Array(Vec::new()), + tool_exchange: serde_json::Value::Array(Vec::new()), + }, + ).await; + yield done_event(); + } + Err(e) => yield error_event(&e.to_string()), + } + } +} + /// 把一次「接上」变成 SSE：先补一份快照，再照常收增量。 /// /// `None` = 这个会话没有在跑的生成。回一条 `idle` 而不是 404——**客户端 @@ -1254,65 +1168,3 @@ mod tests { .expect("格式错的日期不归这一关管，交给 changes_window"); } } - -#[cfg(test)] -mod stall_tests { - use super::{after_nudge, answer_rests_on_nothing, AfterNudge}; - use utopia_llm::{AssistantTurn, ToolCall}; - - fn says(text: Option<&str>) -> AssistantTurn { - AssistantTurn { - content: text.map(String::from), - tool_calls: Vec::new(), - } - } - - /// 整场没调过工具、没有引用，才算「什么都不站在上面」；任一边有东西都不追问 - #[test] - fn an_answer_with_a_step_or_a_source_is_not_questioned() { - let step = serde_json::json!({ "kind": "search" }); - let source = serde_json::json!({ "n": 1 }); - assert!(answer_rests_on_nothing(&[], &[])); - assert!(!answer_rests_on_nothing(std::slice::from_ref(&step), &[])); - assert!(!answer_rests_on_nothing(&[], std::slice::from_ref(&source))); - } - - /// 追问后调了工具就接着走，不管它顺带说了什么 - #[test] - fn a_tool_call_after_the_nudge_carries_on() { - let turn = AssistantTurn { - content: Some("Searching now.".into()), - tool_calls: vec![ToolCall { - id: "c1".into(), - name: "search_chunks".into(), - arguments: "{}".into(), - }], - }; - assert_eq!(after_nudge(&turn), AfterNudge::Tools); - } - - /// DONE 怎么写都算：大小写、句号、前后空白；什么都没说也算——没有可补的 - #[test] - fn done_in_any_dress_keeps_the_answer() { - for text in ["DONE", "done", " Done. ", "DONE!", ""] { - assert_eq!(after_nudge(&says(Some(text))), AfterNudge::Done, "{text:?}"); - } - assert_eq!(after_nudge(&says(None)), AfterNudge::Done); - } - - /// 再来一段文字，不管哪种语言、说得多客气，都是又停住了 - #[test] - fn more_prose_after_the_nudge_is_a_second_stall() { - for text in [ - "稍等，我正在搜索OpenAI的时间线信息。", - "Let me search for that now.", - "Done searching, here is the timeline: ...", - ] { - assert_eq!( - after_nudge(&says(Some(text))), - AfterNudge::Stalled, - "{text:?}" - ); - } - } -} diff --git a/crates/utopia-server/src/api/mod.rs b/crates/utopia-server/src/api/mod.rs index fd762c37d..afa81d3df 100644 --- a/crates/utopia-server/src/api/mod.rs +++ b/crates/utopia-server/src/api/mod.rs @@ -1,4 +1,5 @@ mod admin_routes; +mod agent; mod alerts_routes; mod auth_routes; mod chat; @@ -14,6 +15,7 @@ mod mcp; mod members_routes; pub(crate) mod ontology_routes; mod review_routes; +mod rig_model; pub(crate) mod rule_routes; mod search_routes; mod settings_routes; diff --git a/crates/utopia-server/src/api/rig_model.rs b/crates/utopia-server/src/api/rig_model.rs new file mode 100644 index 000000000..db4532132 --- /dev/null +++ b/crates/utopia-server/src/api/rig_model.rs @@ -0,0 +1,507 @@ +//! `LlmClient` 站到 rig 的 `CompletionModel` 后面（#546）。 +//! +//! **循环与钩子是 rig 的，线上的字节是我们的。** `LlmClient` 里攒着的东西都是 +//! 事故换来的：读超时（7,459 块的摄取死在第 55 块上没有一条错误）、错误体转发 +//! （#538）、欠费与限流的分类、缓存 token 的日志。换成 rig 自带的 provider +//! client 这些全都会丢。所以这里只做两个方向的翻译：rig 的请求 → OpenAI 协议的 +//! JSON；`LlmClient` 的回合与流 → rig 的回答与流。 + +use futures_util::{stream, Stream, StreamExt}; +use rig_core::completion::{ + CompletionError, CompletionModel, CompletionRequest, CompletionResponse, Usage, +}; +use rig_core::message::{AssistantContent, Message, ToolChoice, ToolResultContent, UserContent}; +use rig_core::streaming::{ + RawStreamingChoice, RawStreamingToolCall, StreamFinal, StreamingCompletionResponse, +}; +use serde_json::{json, Value}; +use utopia_llm::{AssistantTurn, LlmClient, ToolStreamItem}; + +/// rig 要给每个回答挂一个「谁回的」；这里没有厂商之分，全是 `LlmClient` +const PROVIDER: &str = "utopia"; + +#[derive(Clone)] +pub struct RigModel { + client: LlmClient, +} + +impl RigModel { + pub fn new(client: LlmClient) -> Self { + Self { client } + } + + /// 开一条流。**端点不接受 `tool_choice` 时去掉它重发一次**：`required` 不在 + /// 每家的协议里，而循环那头拿到 400 会读成「这家不支持工具调用」整个降级—— + /// 那是两件事。去掉之后模型仍可能只说不查，那一层由 turn-finished 钩子兜。 + async fn open_stream( + &self, + w: &Wire, + ) -> Result> + Send + use<>, CompletionError> + { + match self + .client + .chat_tools_stream_with(&w.messages, w.tools.as_ref(), w.tool_choice.as_ref()) + .await + { + Ok(s) => Ok(s), + Err(e) if w.tool_choice.is_some() && rejected_shape(&e) => { + tracing::warn!(error = %e, "端点不接受 tool_choice，去掉重发"); + self.client + .chat_tools_stream_with(&w.messages, w.tools.as_ref(), None) + .await + .map_err(completion_error) + } + Err(e) => Err(completion_error(e)), + } + } +} + +impl CompletionModel for RigModel { + async fn completion( + &self, + request: CompletionRequest, + ) -> Result { + let w = wire(&request); + let turn = match self + .client + .chat_tools_with(&w.messages, w.tools.as_ref(), w.tool_choice.as_ref()) + .await + { + Ok(t) => t, + Err(e) if w.tool_choice.is_some() && rejected_shape(&e) => { + tracing::warn!(error = %e, "端点不接受 tool_choice，去掉重发"); + self.client + .chat_tools_with(&w.messages, w.tools.as_ref(), None) + .await + .map_err(completion_error)? + } + Err(e) => return Err(completion_error(e)), + }; + Ok(CompletionResponse::new( + choice_of(&turn), + Usage::new(), + PROVIDER, + )) + } + + async fn stream( + &self, + request: CompletionRequest, + ) -> Result { + let w = wire(&request); + let source = self.open_stream(&w).await?; + // `LlmClient` 的流已经把工具调用按 index 归并好、放在流末的 `Turn` 里； + // rig 只要每个调用一条 `ToolCall`，再一条 `FinalResponse` 收尾 + let inner = source.flat_map(|item| { + let items: Vec> = match item { + Ok(ToolStreamItem::Delta(text)) => vec![Ok(RawStreamingChoice::Message(text))], + Ok(ToolStreamItem::Turn(turn)) => { + let mut v: Vec<_> = turn + .tool_calls + .iter() + .map(|c| { + Ok(RawStreamingChoice::ToolCall(RawStreamingToolCall::new( + call_id(&c.id), + c.name.clone(), + args_value(&c.arguments), + ))) + }) + .collect(); + v.push(Ok(RawStreamingChoice::FinalResponse(StreamFinal::new( + PROVIDER, + Usage::new(), + )))); + v + } + Err(e) => vec![Err(completion_error(e))], + }; + stream::iter(items) + }); + Ok(StreamingCompletionResponse::stream( + PROVIDER, + Box::pin(inner), + )) + } +} + +// ---- 错误：整条 anyhow 链穿过 rig ------------------------------------------- + +/// `LlmClient` 的错误穿过 rig 的 `CompletionError`。 +/// +/// rig 的错误枚举里只有 `RequestError(Box)` 能装外来的东西，而循环 +/// 那头要认出限流、欠费、被拒（只在「端点拒绝工具调用」时降级），所以整条 +/// anyhow 链原样装进去，那头 downcast 回来。 +#[derive(Debug)] +pub struct LlmFailure(pub anyhow::Error); + +impl std::fmt::Display for LlmFailure { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +impl std::error::Error for LlmFailure {} + +fn completion_error(err: anyhow::Error) -> CompletionError { + CompletionError::RequestError(Box::new(LlmFailure(err))) +} + +/// rig 错误里我们自己的那条链，拿不到（rig 自己的错误）就是 None +pub fn llm_failure(err: &CompletionError) -> Option<&anyhow::Error> { + match err { + CompletionError::RequestError(inner) => inner.downcast_ref::().map(|f| &f.0), + _ => None, + } +} + +/// 端点拒绝了请求的形状（400/422）。带工具的首个请求撞上它，才是「这家不支持 +/// 工具调用」；网断、密钥错、限流、欠费都不是，降级也救不了它们 +fn rejected_shape(err: &anyhow::Error) -> bool { + utopia_llm::rejected(err).is_some_and(|r| r.status == 400 || r.status == 422) +} + +/// 见 [`rejected_shape`]，从 rig 的错误上判 +pub fn tool_calling_rejected(err: &CompletionError) -> bool { + llm_failure(err).is_some_and(rejected_shape) +} + +// ---- rig → OpenAI 协议 --------------------------------------------------------- + +/// 一次请求在线上的三样：消息、工具清单、`tool_choice` +pub(crate) struct Wire { + pub messages: Vec, + /// None = 不带工具字段（`ToolChoice::None` 或本来就没有工具） + pub tools: Option, + pub tool_choice: Option, +} + +/// rig 的请求摊成 OpenAI 协议。 +/// +/// - `preamble` 是第一条 system。 +/// - `documents`（rig 的「上下文文档」，我们放已认下的实体清单）是一条 system， +/// **插在最后一条 user 消息之前**：位置就是服从性（紧挨当前问题），而角色不再是 +/// user——从前它假扮成用户发言，模型会回复它（「看起来你已经获得了实体 ID…」）。 +/// - `ToolChoice::None` 直接不带工具：`"tool_choice": "none"` 不是每家都认， +/// 没有工具字段则人人都认，效果一样。 +pub(crate) fn wire(req: &CompletionRequest) -> Wire { + let mut messages = Vec::new(); + if let Some(p) = &req.preamble { + messages.push(json!({ "role": "system", "content": p })); + } + for m in &req.chat_history { + push_message(&mut messages, m); + } + if !req.documents.is_empty() { + let text = req + .documents + .iter() + .map(|d| d.text.as_str()) + .collect::>() + .join("\n\n"); + let at = messages + .iter() + .rposition(|m| m["role"] == "user") + .unwrap_or(messages.len()); + messages.insert(at, json!({ "role": "system", "content": text })); + } + let tools_off = req.tools.is_empty() || matches!(req.tool_choice, Some(ToolChoice::None)); + let tools = (!tools_off).then(|| { + Value::Array( + req.tools + .iter() + .map(|t| { + json!({ + "type": "function", + "function": { + "name": t.name, + "description": t.description, + "parameters": t.parameters, + } + }) + }) + .collect(), + ) + }); + let tool_choice = if tools_off { + None + } else { + req.tool_choice.as_ref().and_then(tool_choice_json) + }; + Wire { + messages, + tools, + tool_choice, + } +} + +fn tool_choice_json(choice: &ToolChoice) -> Option { + match choice { + // 默认值不写：写了反而有端点不认 + ToolChoice::Auto => None, + ToolChoice::None => Some(json!("none")), + ToolChoice::Required => Some(json!("required")), + ToolChoice::Specific { function_names } => match function_names.as_slice() { + [one] => Some(json!({ "type": "function", "function": { "name": one } })), + // OpenAI 协议点不了「这几个之一」，退成「必须调一个」 + _ => Some(json!("required")), + }, + } +} + +fn push_message(out: &mut Vec, m: &Message) { + match m { + Message::System { content } => out.push(json!({ "role": "system", "content": content })), + Message::User { content } => { + let mut text = String::new(); + for c in content { + match c { + UserContent::Text(t) => { + if !text.is_empty() { + text.push('\n'); + } + text.push_str(&t.text); + } + // tool 消息要紧跟带 tool_calls 的 assistant 消息，所以先于正文 + UserContent::ToolResult(r) => out.push(json!({ + "role": "tool", + "tool_call_id": r.call.as_str(), + "content": tool_result_text(&r.content), + })), + // 图片、音频、附件：这条线上没有 + _ => {} + } + } + if !text.is_empty() { + out.push(json!({ "role": "user", "content": text })); + } + } + Message::Assistant { content, .. } => { + let mut text = String::new(); + let mut calls = Vec::new(); + 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), + } + })), + // 推理块与图片不回灌 + _ => {} + } + } + let mut msg = json!({ + "role": "assistant", + "content": if text.is_empty() { Value::Null } else { Value::String(text) }, + }); + if !calls.is_empty() { + msg["tool_calls"] = Value::Array(calls); + } + out.push(msg); + } + } +} + +/// 工具结果的正文：文本原样，结构化 JSON 序列化 +pub(crate) fn tool_result_text(content: &[ToolResultContent]) -> String { + let mut out = String::new(); + for c in content { + match c { + ToolResultContent::Text(t) => out.push_str(&t.text), + ToolResultContent::Json { value, .. } => out.push_str(&value.to_string()), + _ => {} + } + } + out +} + +// ---- OpenAI 协议 → rig --------------------------------------------------------- + +/// 参数串 → JSON。**解不开的不补成空对象**：模型撞上 token 上限时参数会在半路 +/// 断掉，补成 `{}` 就变成一次拿用户原话去检索的「正常调用」（见 `check_call`）。 +/// 原串装进一个 JSON 字符串，到了 `check_call` 那里它不是对象，照样拒绝。 +pub(crate) fn args_value(raw: &str) -> Value { + serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.to_string())) +} + +/// [`args_value`] 的反向：原本就是一段残串的，原样还回去 +pub(crate) fn args_string(v: &Value) -> String { + match v { + Value::String(raw) => raw.clone(), + other => other.to_string(), + } +} + +/// 有的端点（本地 Ollama 一类）不给 tool call id。rig 要一个，协议回灌时也要 +/// 一个，所以补一个；端点自己没发过，回去时它也不会对不上 +pub(crate) fn call_id(id: &str) -> String { + if id.is_empty() { + format!("call_{}", uuid::Uuid::now_v7().simple()) + } else { + id.to_string() + } +} + +fn choice_of(turn: &AssistantTurn) -> Vec { + let mut out = Vec::new(); + if let Some(text) = turn.content.as_deref().filter(|t| !t.is_empty()) { + 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), + )); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + use rig_core::completion::{Document, ToolDefinition}; + + fn req(history: Vec) -> CompletionRequest { + CompletionRequest { + model: None, + preamble: Some("be brief".into()), + chat_history: history, + documents: vec![], + tools: vec![ToolDefinition { + name: "search_chunks".into(), + description: "search".into(), + parameters: json!({ "type": "object" }), + }], + temperature: None, + max_tokens: None, + tool_choice: None, + additional_params: None, + output_schema: None, + record_telemetry_content: false, + } + } + + #[test] + fn preamble_leads_and_tools_follow_the_protocol() { + let w = wire(&req(vec![Message::user("hi")])); + assert_eq!(w.messages[0]["role"], "system"); + assert_eq!(w.messages[0]["content"], "be brief"); + assert_eq!(w.messages[1], json!({ "role": "user", "content": "hi" })); + let tools = w.tools.expect("工具在"); + assert_eq!(tools[0]["type"], "function"); + assert_eq!(tools[0]["function"]["name"], "search_chunks"); + assert!(w.tool_choice.is_none(), "没说就不写 tool_choice"); + } + + /// 实体清单从前假扮成 user 发言，模型会回复它。现在是 system，且紧挨当前问题 + #[test] + fn documents_become_a_system_message_right_before_the_question() { + let mut r = req(vec![ + Message::user("first"), + Message::assistant("a1"), + Message::user("second"), + ]); + r.documents.push(Document { + id: "known".into(), + text: "Entities already identified".into(), + additional_props: Default::default(), + }); + let w = wire(&r); + let roles: Vec<_> = w + .messages + .iter() + .map(|m| m["role"].as_str().unwrap()) + .collect(); + assert_eq!(roles, ["system", "user", "assistant", "system", "user"]); + assert_eq!(w.messages[3]["content"], "Entities already identified"); + assert_eq!(w.messages[4]["content"], "second"); + } + + /// 工具往返回灌：assistant 带 tool_calls（参数是字符串），tool 消息紧随其后 + #[test] + fn a_tool_exchange_round_trips_in_protocol_shape() { + let w = wire(&req(vec![ + Message::user("q"), + Message::Assistant { + id: None, + content: vec![ + AssistantContent::text("looking"), + AssistantContent::tool_call("c1", "search_chunks", json!({ "query": "x" })), + ], + }, + Message::tool_result("c1", "search_chunks", "found"), + ])); + assert_eq!(w.messages[2]["role"], "assistant"); + assert_eq!(w.messages[2]["content"], "looking"); + assert_eq!(w.messages[2]["tool_calls"][0]["id"], "c1"); + assert_eq!( + w.messages[2]["tool_calls"][0]["function"]["arguments"], + "{\"query\":\"x\"}" + ); + assert_eq!( + w.messages[3], + json!({ "role": "tool", "tool_call_id": "c1", "content": "found" }) + ); + } + + #[test] + fn tool_choice_maps_and_none_drops_the_tools_entirely() { + let mut r = req(vec![Message::user("q")]); + r.tool_choice = Some(ToolChoice::Required); + assert_eq!(wire(&r).tool_choice, Some(json!("required"))); + r.tool_choice = Some(ToolChoice::Specific { + function_names: vec!["search_chunks".into()], + }); + assert_eq!( + wire(&r).tool_choice, + Some(json!({ "type": "function", "function": { "name": "search_chunks" } })) + ); + r.tool_choice = Some(ToolChoice::None); + let w = wire(&r); + assert!(w.tools.is_none(), "None = 请求里根本没有工具字段"); + assert!(w.tool_choice.is_none()); + } + + /// 半路断掉的参数不能变成空对象（那会变成一次拿原话检索的正常调用） + #[test] + fn cut_off_arguments_stay_a_string_and_come_back_verbatim() { + let v = args_value("{\"query\": \"Acme reven"); + assert_eq!(v, Value::String("{\"query\": \"Acme reven".into())); + assert_eq!(args_string(&v), "{\"query\": \"Acme reven"); + assert_eq!(args_string(&json!({ "a": 1 })), "{\"a\":1}"); + } + + #[test] + fn a_missing_call_id_is_minted_and_a_given_one_kept() { + assert_eq!(call_id("call_abc"), "call_abc"); + let minted = call_id(""); + assert!(minted.starts_with("call_") && minted.len() > 10); + } + + #[test] + fn a_turn_becomes_text_then_tool_calls() { + let turn = AssistantTurn { + content: Some("hm".into()), + tool_calls: vec![utopia_llm::ToolCall { + id: "c9".into(), + name: "find_entities".into(), + arguments: "{\"name\":\"Acme\"}".into(), + }], + }; + let choice = choice_of(&turn); + assert_eq!(choice.len(), 2); + assert!(matches!(&choice[0], AssistantContent::Text(t) if t.text == "hm")); + match &choice[1] { + AssistantContent::ToolCall(tc) => { + assert_eq!(tc.id.as_str(), "c9"); + assert_eq!(tc.function.name, "find_entities"); + assert_eq!(tc.function.arguments["name"], "Acme"); + } + other => panic!("expected a tool call, got {other:?}"), + } + } +} From ce7b5b0ecb01a2a74b597ef004ac27168146820f Mon Sep 17 00:00:00 2001 From: WaylandYang Date: Wed, 9 Sep 2026 20:04:04 +0800 Subject: [PATCH 2/3] rig comes without its provider clients Co-Authored-By: Claude Fable 5.1 Signed-off-by: WaylandYang --- Cargo.lock | 113 -------------------------------- crates/utopia-server/Cargo.toml | 4 +- 2 files changed, 2 insertions(+), 115 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f6bdb5b63..4468f4f3f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -702,15 +702,6 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" -[[package]] -name = "convert_case" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "affbf0190ed2caf063e3def54ff444b449371d55c58e513a95ab98eca50adb49" -dependencies = [ - "unicode-segmentation", -] - [[package]] name = "cookie" version = "0.18.2" @@ -1021,12 +1012,6 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" -[[package]] -name = "data-encoding" -version = "2.11.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4583a4551df46e2792f82ceeac45e850d2e2d5debba0b91f102385cda5b11f06" - [[package]] name = "datasketches" version = "0.2.0" @@ -3325,15 +3310,6 @@ dependencies = [ "elliptic-curve", ] -[[package]] -name = "proc-macro-crate" -version = "3.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" -dependencies = [ - "toml_edit", -] - [[package]] name = "proc-macro-error-attr3" version = "3.1.1" @@ -3760,7 +3736,6 @@ dependencies = [ "http", "indexmap", "rig-core", - "rig-derive", "schemars", "serde", "serde_json", @@ -3792,32 +3767,17 @@ dependencies = [ "ordered-float", "pin-project-lite", "reqwest", - "rig-derive", "schemars", "serde", "serde_json", "sha2 0.10.9", "thiserror 2.0.20", "tokio", - "tokio-tungstenite", "tracing", "tracing-futures", "url", ] -[[package]] -name = "rig-derive" -version = "0.42.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de0a33f1bac45f16e50146c248bcbbfaa44518c7252d274e972c7f4ad71aaba7" -dependencies = [ - "convert_case", - "proc-macro-crate", - "proc-macro2", - "quote", - "syn 2.0.119", -] - [[package]] name = "ring" version = "0.17.14" @@ -5080,22 +5040,6 @@ dependencies = [ "tokio", ] -[[package]] -name = "tokio-tungstenite" -version = "0.29.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f72a05e828585856dacd553fba484c242c46e391fb0e58917c942ee9202915c" -dependencies = [ - "futures-util", - "log", - "rustls", - "rustls-pki-types", - "tokio", - "tokio-rustls", - "tungstenite", - "webpki-roots 0.26.11", -] - [[package]] name = "tokio-util" version = "0.7.19" @@ -5110,36 +5054,6 @@ dependencies = [ "tokio", ] -[[package]] -name = "toml_datetime" -version = "1.1.1+spec-1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" -dependencies = [ - "serde_core", -] - -[[package]] -name = "toml_edit" -version = "0.25.13+spec-1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" -dependencies = [ - "indexmap", - "toml_datetime", - "toml_parser", - "winnow", -] - -[[package]] -name = "toml_parser" -version = "1.1.3+spec-1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" -dependencies = [ - "winnow", -] - [[package]] name = "tower" version = "0.5.3" @@ -5283,24 +5197,6 @@ version = "0.25.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d2df906b07856748fa3f6e0ad0cbaa047052d4a7dd609e231c4f72cee8c36f31" -[[package]] -name = "tungstenite" -version = "0.29.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c01152af293afb9c7c2a57e4b559c5620b421f6d133261c60dd2d0cdb38e6b8" -dependencies = [ - "bytes", - "data-encoding", - "http", - "httparse", - "log", - "rand 0.9.5", - "rustls", - "rustls-pki-types", - "sha1", - "thiserror 2.0.20", -] - [[package]] name = "type1-encoding-parser" version = "0.1.1" @@ -6134,15 +6030,6 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" -[[package]] -name = "winnow" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" -dependencies = [ - "memchr", -] - [[package]] name = "wiremock" version = "0.6.5" diff --git a/crates/utopia-server/Cargo.toml b/crates/utopia-server/Cargo.toml index 133139dc0..1f40c26de 100644 --- a/crates/utopia-server/Cargo.toml +++ b/crates/utopia-server/Cargo.toml @@ -48,8 +48,8 @@ url = "2.5.8" percent-encoding = "2.3.2" base64.workspace = true subtle.workspace = true -rig-agent = "0.42" -rig-core = "0.42" +rig-agent = { version = "0.42", default-features = false } +rig-core = { version = "0.42", default-features = false } [dev-dependencies] wiremock = "0.6.5" From 20a2731510efc9022e789f94dec31ce4fd2f35ac Mon Sep 17 00:00:00 2001 From: WaylandYang Date: Wed, 9 Sep 2026 20:27:08 +0800 Subject: [PATCH 3/3] The prompt says what the base is Co-Authored-By: Claude Fable 5.1 Signed-off-by: WaylandYang --- crates/utopia-server/src/api/chat.rs | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/crates/utopia-server/src/api/chat.rs b/crates/utopia-server/src/api/chat.rs index 4c03985c2..7e0b595a9 100644 --- a/crates/utopia-server/src/api/chat.rs +++ b/crates/utopia-server/src/api/chat.rs @@ -400,6 +400,11 @@ const SYSTEM_PROMPT: &str = "You are the assistant of Utopia, a temporal knowled You have tools: search_chunks (document search) and get_document (the full text of one \ document found by search), find_entities, entity_facts and changes (a bi-temporal \ knowledge graph), and search_docs (Utopia's own manual, the Charter).\n\ + The knowledge base holds whatever its owners ingested: documents, and a graph extracted \ + from them. You do not know what is in it until you look; public companies, well-known \ + people and events are as likely to be there as private material. A question you could \ + answer from memory is still answered from the base, and \"general knowledge\" is never a \ + reason to skip the tools. Never say the base lacks something you have not searched for.\n\ search_chunks returns short excerpts of the best-matching sections only. When a hit is \ clearly the right document but the excerpt does not carry the answer, read the whole \ document with get_document before saying the knowledge base does not have it.\n\ @@ -411,7 +416,7 @@ const SYSTEM_PROMPT: &str = "You are the assistant of Utopia, a temporal knowled Boundary: search_docs answers questions about Utopia itself (features, ingestion, \ permissions, what fields like 'missing' or validity ranges mean); the other tools answer \ questions about the knowledge stored in it. Never mix the manual into answers about the \ - user's data unless they asked about Utopia's behavior.\n\ + base's contents unless they asked about Utopia's behavior.\n\ \n\ Method:\n\ First decide what the message is about. A message about THIS CONVERSATION — translate it, \ @@ -420,8 +425,8 @@ const SYSTEM_PROMPT: &str = "You are the assistant of Utopia, a temporal knowled not merely wasted work — with several entities sharing a name the second pass can land on \ a different one, and the \"translation\" then says something else. Just deliver it — no \ preamble about what you are or are not looking up. Everything below is for messages about \ - the user's data.\n\ - 1. For factual questions — questions about the user's data, never one about this \ + the knowledge base.\n\ + 1. For factual questions — questions about the knowledge base, never one about this \ conversation — ALWAYS gather evidence with tools before answering. Prefer the \ graph tools for questions about people/organizations/projects and time (\"who was X \ when\", \"what changed\"), search_chunks for content and detail questions. Combine both \