Assistant chat: workspace-scoped, persisted, streaming

crates/nvisy-inference/Cargo.toml

@@ -26,6 +26,7 @@ schema = ["dep:schemars"]
 [dependencies]
 # Async runtime
 futures = { workspace = true, features = [] }
+async-stream = { workspace = true, features = [] }
 
 # LLM framework: provider clients + the classic Agent runtime
 rig = { workspace = true, features = [] }

crates/nvisy-inference/src/client/erased_agent.rs

@@ -0,0 +1,76 @@
+//! [`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<String, PromptError>>;
+
+    /// Run one chat turn against `history`, appending the committed messages.
+    fn chat<'a>(
+        &'a self,
+        prompt: String,
+        history: &'a mut Vec<Message>,
+    ) -> BoxFuture<'a, Result<String, PromptError>>;
+
+    /// Stream one chat turn against `history` as text deltas.
+    fn stream_chat<'a>(
+        &'a self,
+        prompt: String,
+        history: Vec<Message>,
+    ) -> BoxFuture<'a, BoxStream<'a, Result<String, Error>>>;
+}
+
+impl<M> ErasedAgent for Agent<M>
+where
+    M: CompletionModel + 'static,
+    M::StreamingResponse: GetTokenUsage,
+{
+    fn prompt(&self, prompt: String) -> BoxFuture<'_, Result<String, PromptError>> {
+        Box::pin(async move { Prompt::prompt(self, prompt).await })
+    }
+
+    fn chat<'a>(
+        &'a self,
+        prompt: String,
+        history: &'a mut Vec<Message>,
+    ) -> BoxFuture<'a, Result<String, PromptError>> {
+        Box::pin(async move { Chat::chat(self, prompt, history).await })
+    }
+
+    fn stream_chat<'a>(
+        &'a self,
+        prompt: String,
+        history: Vec<Message>,
+    ) -> BoxFuture<'a, BoxStream<'a, Result<String, Error>>> {
+        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()
+        })
+    }
+}

crates/nvisy-inference/src/client/mod.rs

@@ -2,55 +2,30 @@
 //!
 //! [`InferenceClient`] is a thin, cloneable wrapper around any provider's rig
 //! [`Agent`] that provides convenience methods for the most common operations —
-//! one-shot [`prompt`](InferenceClient::prompt) and multi-turn
-//! [`chat`](InferenceClient::chat). Every public method is instrumented with
-//! [`tracing`] for observability. It plays the same role for inference that
-//! `ObjectStoreClient` plays for object storage: one runtime handle callers use
-//! regardless of which provider backs it.
+//! one-shot [`prompt`](InferenceClient::prompt), multi-turn
+//! [`chat`](InferenceClient::chat), and streaming
+//! [`stream_chat`](InferenceClient::stream_chat). Every public method is
+//! instrumented with [`tracing`] for observability. It plays the same role for
+//! 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;
 
 use std::sync::Arc;
 
-use futures::future::BoxFuture;
+use async_stream::stream;
+use futures::StreamExt;
 use rig::agent::Agent;
 use rig::client::verify::{VerifyClient, VerifyError};
-/// A single conversation message, re-exported so callers can build the
-/// [`chat`](InferenceClient::chat) history without depending on `rig` directly.
-pub use rig::completion::Message;
-use rig::completion::{Chat, CompletionModel, Prompt, PromptError};
+use rig::completion::{CompletionModel, GetTokenUsage, Message};
 
+use self::erased_agent::ErasedAgent;
+pub use self::token_stream::TokenStream;
+pub use self::turn::{ChatTurn, Role};
 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.
-trait ErasedAgent: Send + Sync {
-    /// Send a single prompt with no prior context.
-    fn prompt(&self, prompt: String) -> BoxFuture<'_, Result<String, PromptError>>;
-
-    /// Run one chat turn against `history`, appending the committed messages.
-    fn chat<'a>(
-        &'a self,
-        prompt: String,
-        history: &'a mut Vec<Message>,
-    ) -> BoxFuture<'a, Result<String, PromptError>>;
-}
-
-impl<M> ErasedAgent for Agent<M>
-where
-    M: CompletionModel + 'static,
-{
-    fn prompt(&self, prompt: String) -> BoxFuture<'_, Result<String, PromptError>> {
-        Box::pin(async move { Prompt::prompt(self, prompt).await })
-    }
-
-    fn chat<'a>(
-        &'a self,
-        prompt: String,
-        history: &'a mut Vec<Message>,
-    ) -> BoxFuture<'a, Result<String, PromptError>> {
-        Box::pin(async move { Chat::chat(self, prompt, history).await })
-    }
-}
-
 /// Cloneable handle to any inference backend (OpenAI, Anthropic, Ollama, ...).
 ///
 /// Wraps a provider's rig agent behind a provider-agnostic interface, so callers
@@ -63,6 +38,7 @@ impl InferenceClient {
     pub(crate) fn new<M>(agent: Agent<M>) -> Self
     where
         M: CompletionModel + 'static,
+        M::StreamingResponse: GetTokenUsage,
     {
         Self(Arc::new(agent))
     }
@@ -79,16 +55,46 @@ impl InferenceClient {
 
     /// Run one chat turn against `history`, returning the model's text response.
     ///
-    /// `history` is caller-owned and updated in place: the prompt and the
-    /// messages the model commits this turn are appended to it, so passing the
-    /// same `Vec` across calls continues the conversation.
+    /// `history` is the prior conversation as [`ChatTurn`]s.
     #[tracing::instrument(name = "inference.chat", skip_all, fields(history_len = history.len()))]
-    pub async fn chat(&self, prompt: &str, history: &mut Vec<Message>) -> Result<String, Error> {
+    pub async fn chat(&self, prompt: &str, history: Vec<ChatTurn>) -> Result<String, Error> {
+        let mut history = to_messages(history);
         self.0
-            .chat(prompt.to_owned(), history)
+            .chat(prompt.to_owned(), &mut history)
             .await
             .map_err(|err| Error::Prompt(err.to_string()))
     }
+
+    /// Stream one chat turn against `history`, yielding the model's response as
+    /// text deltas.
+    ///
+    /// Returns immediately with a [`TokenStream`]; the request opens lazily when
+    /// the stream is first polled. Each item is a token chunk as it arrives, and
+    /// a failure mid-generation ends the stream with an `Err`. `history` is the
+    /// prior conversation as [`ChatTurn`]s; persist the user prompt and the
+    /// assembled reply on the caller side.
+    #[tracing::instrument(name = "inference.stream_chat", skip_all, fields(history_len = history.len()))]
+    pub fn stream_chat(&self, prompt: &str, history: Vec<ChatTurn>) -> TokenStream {
+        // Own the agent handle in the generator so the result is `'static` and
+        // can outlive this `InferenceClient` (moved into a response body). The
+        // provider stream borrows the owned `Arc`, which the coroutine keeps
+        // alive for the whole stream.
+        let agent = Arc::clone(&self.0);
+        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;
+            }
+        };
+        TokenStream::new(inner.boxed())
+    }
+}
+
+/// Converts a provider-agnostic history into rig messages.
+fn to_messages(history: Vec<ChatTurn>) -> Vec<Message> {
+    history.into_iter().map(Message::from).collect()
 }
 
 /// Verifies a built provider client's credentials against the provider.

crates/nvisy-inference/src/client/token_stream.rs

@@ -0,0 +1,37 @@
+//! The [`TokenStream`] type: an assistant response streamed as text deltas.
+
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use futures::Stream;
+use futures::stream::{BoxStream, StreamExt};
+
+use crate::error::Error;
+
+/// A stream of the assistant's response as text deltas.
+///
+/// Each item is a token chunk (`Ok`); a failure mid-generation is the final item
+/// (`Err`) and ends the stream. Owned and `'static`, so it can be moved into a
+/// response body outliving the client that produced it.
+///
+/// Yielded by [`InferenceClient::stream_chat`](crate::InferenceClient::stream_chat).
+/// Poll it with the [`Stream`] API ([`futures::StreamExt`]).
+#[must_use = "a token stream does nothing unless polled"]
+pub struct TokenStream {
+    inner: BoxStream<'static, Result<String, Error>>,
+}
+
+impl TokenStream {
+    /// Wraps an owned delta stream.
+    pub(crate) fn new(inner: BoxStream<'static, Result<String, Error>>) -> Self {
+        Self { inner }
+    }
+}
+
+impl Stream for TokenStream {
+    type Item = Result<String, Error>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.inner.poll_next_unpin(cx)
+    }
+}

crates/nvisy-inference/src/client/turn.rs

@@ -0,0 +1,60 @@
+//! The provider-agnostic conversation types callers build history from.
+//!
+//! These wrap the underlying rig message model so consumers depend on this
+//! crate's own types rather than `rig`'s.
+
+use rig::completion::Message;
+
+/// Who authored a conversation turn.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum Role {
+    /// A system instruction.
+    System,
+    /// A message from the user.
+    User,
+    /// A reply from the assistant.
+    Assistant,
+}
+
+/// One turn of a conversation: a role and its text.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ChatTurn {
+    /// Who authored the turn.
+    pub role: Role,
+    /// The turn's text.
+    pub content: String,
+}
+
+impl ChatTurn {
+    /// A system turn.
+    pub fn system(content: impl Into<String>) -> Self {
+        Self::of(Role::System, content)
+    }
+
+    /// A user turn.
+    pub fn user(content: impl Into<String>) -> Self {
+        Self::of(Role::User, content)
+    }
+
+    /// An assistant turn.
+    pub fn assistant(content: impl Into<String>) -> Self {
+        Self::of(Role::Assistant, content)
+    }
+
+    fn of(role: Role, content: impl Into<String>) -> Self {
+        Self {
+            role,
+            content: content.into(),
+        }
+    }
+}
+
+impl From<ChatTurn> for Message {
+    fn from(turn: ChatTurn) -> Self {
+        match turn.role {
+            Role::System => Message::system(turn.content),
+            Role::User => Message::user(turn.content),
+            Role::Assistant => Message::assistant(turn.content),
+        }
+    }
+}

crates/nvisy-inference/src/lib.rs

@@ -2,8 +2,9 @@
 #![cfg_attr(docsrs, feature(doc_cfg))]
 #![doc = include_str!("../README.md")]
 
-pub mod client;
+mod client;
 mod error;
 pub mod providers;
 
+pub use client::{ChatTurn, InferenceClient, Role, TokenStream, verify};
 pub use error::Error;

crates/nvisy-postgres/src/model/chat_message.rs

@@ -0,0 +1,42 @@
+//! Chat message model for PostgreSQL database operations.
+
+use diesel::prelude::*;
+use jiff_diesel::Timestamp;
+use uuid::Uuid;
+
+use crate::schema::chat_messages;
+use crate::types::ChatRole;
+
+/// One message in a chat session.
+#[derive(Debug, Clone, PartialEq, Queryable, Selectable)]
+#[diesel(table_name = chat_messages)]
+#[diesel(check_for_backend(diesel::pg::Pg))]
+pub struct ChatMessage {
+    /// Unique message identifier.
+    pub id: Uuid,
+    /// Session this message belongs to.
+    pub session_id: Uuid,
+    /// Parent in the conversation tree; `None` is a root.
+    pub parent_id: Option<Uuid>,
+    /// Author of the message.
+    pub role: ChatRole,
+    /// Message text, XChaCha20-Poly1305 encrypted with the workspace key.
+    pub content: Vec<u8>,
+    /// Message creation timestamp.
+    pub created_at: Timestamp,
+}
+
+/// Data for appending a new chat message.
+#[derive(Debug, Clone, Insertable)]
+#[diesel(table_name = chat_messages)]
+#[diesel(check_for_backend(diesel::pg::Pg))]
+pub struct NewChatMessage {
+    /// Session this message belongs to.
+    pub session_id: Uuid,
+    /// Parent in the conversation tree; `None` is a root.
+    pub parent_id: Option<Uuid>,
+    /// Author of the message.
+    pub role: ChatRole,
+    /// Message text, XChaCha20-Poly1305 encrypted with the workspace key.
+    pub content: Vec<u8>,
+}