diff --git a/Cargo.lock b/Cargo.lock index cdf05b5a..554b9ae1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3058,7 +3058,7 @@ dependencies = [ "reqwest-middleware", "reqwest-retry", "reqwest-tracing", - "rig 0.42.0", + "rig", "schemars 1.2.2", "serde", "serde_json", @@ -6141,7 +6141,7 @@ dependencies = [ "async-stream", "derive_more", "futures", - "rig 0.41.0", + "rig", "schemars 1.2.2", "serde", "serde_json", @@ -7695,52 +7695,19 @@ checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4" [[package]] name = "rig" -version = "0.41.0" +version = "0.42.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ce03971e6115d30ef53fb3244d06718a4e62bbde82c103065600c09459b989a" +checksum = "b989f4060e5f3c1ea4b8b6999f6759949d820704286cc6976b7a688cc0133029" dependencies = [ - "rig-agent 0.41.0", - "rig-core 0.41.0", - "rig-derive 0.41.0", + "rig-agent", + "rig-core", + "rig-derive", "rig-fastembed", "rig-helixdb", "rig-lancedb", "rig-milvus", ] -[[package]] -name = "rig" -version = "0.42.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b989f4060e5f3c1ea4b8b6999f6759949d820704286cc6976b7a688cc0133029" -dependencies = [ - "rig-agent 0.42.0", - "rig-core 0.42.0", - "rig-derive 0.42.0", -] - -[[package]] -name = "rig-agent" -version = "0.41.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b0796bbf47d7b76670401aac975bc619cf7fba3482b22dfe14992edaa9c2e04" -dependencies = [ - "async-stream", - "fastrand", - "futures", - "http", - "indexmap 2.14.0", - "rig-core 0.41.0", - "rig-derive 0.41.0", - "schemars 1.2.2", - "serde", - "serde_json", - "thiserror", - "tokio", - "tracing", - "tracing-futures", -] - [[package]] name = "rig-agent" version = "0.42.0" @@ -7752,49 +7719,15 @@ dependencies = [ "futures", "http", "indexmap 2.14.0", - "rig-core 0.42.0", - "rig-derive 0.42.0", - "schemars 1.2.2", - "serde", - "serde_json", - "thiserror", - "tokio", - "tracing", - "tracing-futures", -] - -[[package]] -name = "rig-core" -version = "0.41.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35f5520515ae8f6851adcbc6fde9eea8e96f657418c062e16c82cd81cce44e8e" -dependencies = [ - "as-any", - "async-stream", - "base64 0.22.1", - "bytes", - "eventsource-stream", - "fastrand", - "futures", - "futures-timer", - "glob", - "http", - "indexmap 2.14.0", - "mime", - "mime_guess", - "ordered-float 5.3.0", - "pin-project-lite", - "reqwest 0.13.4", - "rig-derive 0.41.0", + "rig-core", + "rig-derive", "schemars 1.2.2", "serde", "serde_json", "thiserror", "tokio", - "tokio-tungstenite 0.28.0", "tracing", "tracing-futures", - "url", ] [[package]] @@ -7820,32 +7753,19 @@ dependencies = [ "pin-project-lite", "reqwest 0.13.4", "reqwest-middleware", - "rig-derive 0.42.0", + "rig-derive", "schemars 1.2.2", "serde", "serde_json", "sha2 0.10.9", "thiserror", "tokio", - "tokio-tungstenite 0.29.0", + "tokio-tungstenite", "tracing", "tracing-futures", "url", ] -[[package]] -name = "rig-derive" -version = "0.41.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb868fcebdf3ba425e3afad2e4926bb6d9e1188a856843b00bcee2e15c07424f" -dependencies = [ - "convert_case 0.11.0", - "proc-macro-crate", - "proc-macro2", - "quote", - "syn 2.0.119", -] - [[package]] name = "rig-derive" version = "0.42.0" @@ -7861,12 +7781,12 @@ dependencies = [ [[package]] name = "rig-fastembed" -version = "0.41.0" +version = "0.42.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69ba93356242940918fac76a4b1cfe78b964f757e5cb65bb3feebd597635cfb0" +checksum = "f1dff99f9bea13108bb14df5431b6fc0738beb87f1fe5ebfcf26ae9fe54a188d" dependencies = [ "fastembed", - "rig-core 0.41.0", + "rig-core", "schemars 1.2.2", "serde", "serde_json", @@ -7875,12 +7795,12 @@ dependencies = [ [[package]] name = "rig-helixdb" -version = "0.41.0" +version = "0.42.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1d40ae8fa89af33793626b898ec527e246183ef8b4b5095afee5f6a65cbbb1a" +checksum = "b5fa2b188a1e9e65d29c34738c97465f0822c57540d6d45fd4f04dad906c966a" dependencies = [ "reqwest 0.13.4", - "rig-core 0.41.0", + "rig-core", "serde", "serde_json", "thiserror", @@ -7888,27 +7808,26 @@ dependencies = [ [[package]] name = "rig-lancedb" -version = "0.41.0" +version = "0.42.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f703b9e081ee4d77e60953598a5071a5cff4b8b702a2f06d251310950b0dae7c" +checksum = "c5fe3555202c5bd4e30948414c2e234fddd9260d4520046733fcf0d58718f3a1" dependencies = [ "arrow-array", - "deranged", "futures", "lancedb", - "rig-core 0.41.0", + "rig-core", "serde", "serde_json", ] [[package]] name = "rig-milvus" -version = "0.41.0" +version = "0.42.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76662eb7f75de4e54b5f71e616430dc7f75d96d32f109d9435ad426b511eb565" +checksum = "af9c7b7ff7ed9c16da983a21f1e81f425c9805355427a92f4d2d97ba1fe2aa12" dependencies = [ "reqwest 0.13.4", - "rig-core 0.41.0", + "rig-core", "serde", "serde_json", "uuid", @@ -9213,22 +9132,6 @@ dependencies = [ "tokio-util", ] -[[package]] -name = "tokio-tungstenite" -version = "0.28.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d25a406cddcc431a75d3d9afc6a7c0f7428d4891dd973e4d54c56b46127bf857" -dependencies = [ - "futures-util", - "log", - "rustls", - "rustls-pki-types", - "tokio", - "tokio-rustls", - "tungstenite 0.28.0", - "webpki-roots 0.26.11", -] - [[package]] name = "tokio-tungstenite" version = "0.29.0" @@ -9241,7 +9144,7 @@ dependencies = [ "rustls-pki-types", "tokio", "tokio-rustls", - "tungstenite 0.29.0", + "tungstenite", "webpki-roots 0.26.11", ] @@ -9542,25 +9445,6 @@ dependencies = [ "tokio", ] -[[package]] -name = "tungstenite" -version = "0.28.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8628dcc84e5a09eb3d8423d6cb682965dea9133204e8fb3efee74c2a0c259442" -dependencies = [ - "bytes", - "data-encoding", - "http", - "httparse", - "log", - "rand 0.9.5", - "rustls", - "rustls-pki-types", - "sha1", - "thiserror", - "utf-8", -] - [[package]] name = "tungstenite" version = "0.29.0" @@ -9771,12 +9655,6 @@ dependencies = [ "serde_derive", ] -[[package]] -name = "utf-8" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" - [[package]] name = "utf8-ranges" version = "1.0.5" diff --git a/Cargo.toml b/Cargo.toml index 50bf6ed5..0cdda1e2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -58,7 +58,7 @@ async-stream = { version = "0.3", features = [] } async-trait = { version = "0.1", features = [] } # LLM framework -rig = { version = "0.41", default-features = false, features = ["agent", "derive", "rustls"] } +rig = { version = "0.42", default-features = false, features = ["agent", "derive", "rustls"] } # HTTP client & middleware reqwest = { version = "0.13", default-features = false, features = ["rustls"] } diff --git a/crates/nvisy-inference/src/client/erased_agent.rs b/crates/nvisy-inference/src/client/erased_agent.rs deleted file mode 100644 index ad3061fd..00000000 --- a/crates/nvisy-inference/src/client/erased_agent.rs +++ /dev/null @@ -1,76 +0,0 @@ -//! [`ErasedAgent`]: an object-safe view of a rig [`Agent`], erasing the -//! provider's concrete completion-model type so a single handle can hold any -//! backend. - -use futures::future::BoxFuture; -use futures::stream::{BoxStream, StreamExt}; -use rig::agent::{Agent, MultiTurnStreamItem}; -use rig::completion::message::Text; -use rig::completion::{Chat, CompletionModel, GetTokenUsage, Message, Prompt, PromptError}; -use rig::streaming::{StreamedAssistantContent, StreamingChat}; - -use crate::error::Error; - -/// Object-safe view of a rig [`Agent`], erasing the provider's concrete -/// completion-model type so a single handle can hold any backend. -pub(crate) trait ErasedAgent: Send + Sync { - /// Send a single prompt with no prior context. - fn prompt(&self, prompt: String) -> BoxFuture<'_, Result>; - - /// Run one chat turn against `history`, appending the committed messages. - fn chat<'a>( - &'a self, - prompt: String, - history: &'a mut Vec, - ) -> BoxFuture<'a, Result>; - - /// Stream one chat turn against `history` as text deltas. - fn stream_chat<'a>( - &'a self, - prompt: String, - history: Vec, - ) -> BoxFuture<'a, BoxStream<'a, Result>>; -} - -impl ErasedAgent for Agent -where - M: CompletionModel + 'static, - M::StreamingResponse: GetTokenUsage, -{ - fn prompt(&self, prompt: String) -> BoxFuture<'_, Result> { - Box::pin(async move { Prompt::prompt(self, prompt).await }) - } - - fn chat<'a>( - &'a self, - prompt: String, - history: &'a mut Vec, - ) -> BoxFuture<'a, Result> { - Box::pin(async move { Chat::chat(self, prompt, history).await }) - } - - fn stream_chat<'a>( - &'a self, - prompt: String, - history: Vec, - ) -> BoxFuture<'a, BoxStream<'a, Result>> { - Box::pin(async move { - // Map rig's multi-turn stream down to bare text deltas here, inside - // the concrete-`M` impl, so the boxed stream is provider-agnostic and - // the trait stays object-safe. - let stream = StreamingChat::stream_chat(self, prompt, history).await; - let deltas = stream.filter_map(|item| async move { - match item { - Ok(MultiTurnStreamItem::StreamAssistantItem( - StreamedAssistantContent::Text(Text { text, .. }), - )) => Some(Ok(text)), - // Non-text items (tool calls, reasoning, the final response - // marker) carry no user-visible text: drop them. - Ok(_) => None, - Err(err) => Some(Err(Error::Prompt(err.to_string()))), - } - }); - deltas.boxed() - }) - } -} diff --git a/crates/nvisy-inference/src/client/mod.rs b/crates/nvisy-inference/src/client/mod.rs index 89868f62..8c60b60d 100644 --- a/crates/nvisy-inference/src/client/mod.rs +++ b/crates/nvisy-inference/src/client/mod.rs @@ -9,7 +9,6 @@ //! inference that `ObjectStoreClient` plays for object storage: one runtime //! handle callers use regardless of which provider backs it. -mod erased_agent; mod token_stream; mod turn; @@ -17,29 +16,27 @@ use std::sync::Arc; use async_stream::stream; use futures::StreamExt; -use rig::agent::Agent; +use rig::agent::{Agent, MultiTurnStreamItem}; use rig::client::verify::{VerifyClient, VerifyError}; -use rig::completion::{CompletionModel, GetTokenUsage, Message}; +use rig::completion::message::Text; +use rig::completion::{Chat, Message, Prompt}; +use rig::streaming::{StreamedAssistantContent, StreamingChat}; -use self::erased_agent::ErasedAgent; pub use self::token_stream::TokenStream; pub use self::turn::{ChatTurn, Role}; use crate::error::Error; /// Cloneable handle to any inference backend (OpenAI, Anthropic, Ollama, ...). /// -/// Wraps a provider's rig agent behind a provider-agnostic interface, so callers -/// issue prompts without knowing which provider is configured. +/// Wraps a rig [`Agent`] — which erases its provider's completion-model type +/// internally — behind a provider-agnostic interface, so callers issue prompts +/// without knowing which provider is configured. #[derive(Clone)] -pub struct InferenceClient(Arc); +pub struct InferenceClient(Arc); impl InferenceClient { - /// Wrap a concrete rig [`Agent`]. - pub(crate) fn new(agent: Agent) -> Self - where - M: CompletionModel + 'static, - M::StreamingResponse: GetTokenUsage, - { + /// Wrap a rig [`Agent`]. + pub(crate) fn new(agent: Agent) -> Self { Self(Arc::new(agent)) } @@ -47,8 +44,7 @@ impl InferenceClient { /// text response. #[tracing::instrument(name = "inference.prompt", skip_all)] pub async fn prompt(&self, prompt: &str) -> Result { - self.0 - .prompt(prompt.to_owned()) + Prompt::prompt(&*self.0, prompt.to_owned()) .await .map_err(|err| Error::Prompt(err.to_string())) } @@ -59,8 +55,7 @@ impl InferenceClient { #[tracing::instrument(name = "inference.chat", skip_all, fields(history_len = history.len()))] pub async fn chat(&self, prompt: &str, history: Vec) -> Result { let mut history = to_messages(history); - self.0 - .chat(prompt.to_owned(), &mut history) + Chat::chat(&*self.0, prompt.to_owned(), &mut history) .await .map_err(|err| Error::Prompt(err.to_string())) } @@ -83,9 +78,23 @@ impl InferenceClient { let prompt = prompt.to_owned(); let history = to_messages(history); let inner = stream! { - let mut deltas = agent.stream_chat(prompt, history).await; - while let Some(delta) = deltas.next().await { - yield delta; + let stream = StreamingChat::stream_chat(&*agent, prompt, history).await; + for await item in stream { + // Map rig's multi-turn stream down to bare text deltas: non-text + // items (tool calls, reasoning, the final-response marker) carry + // no user-visible text and are dropped. + match item { + Ok(MultiTurnStreamItem::StreamAssistantItem( + StreamedAssistantContent::Text(Text { text, .. }), + )) => yield Ok(text), + Ok(_) => {} + // An error is terminal: yield it and stop, per TokenStream's + // contract, rather than polling the rig stream again. + Err(err) => { + yield Err(Error::Prompt(err.to_string())); + break; + } + } } }; TokenStream::new(inner.boxed())