Fix streaming issues and add citations support in chat messages

Author: joulei

Cargo.lock

@@ -5,6 +5,9 @@ edition = "2021"
 
 [dependencies]
 futures = "0.3.31"
+url = "2.4.10"
+link-preview = { version = "0.1.1", features = ["fetch"] }
+robius-open = "0.1.1"
 
 makepad-widgets = { git = "https://github.com/makepad/makepad", branch = "rik" }
 makepad-code-editor = { git = "https://github.com/makepad/makepad", branch = "rik" }

moly-kit/Cargo.toml

@@ -77,15 +77,18 @@ enum Role {
     Assistant,
 }
 
+/// The Choice object as part of a streaming response.
 #[derive(Clone, Debug, Deserialize)]
 struct Choice {
     pub delta: IncomingMessage,
 }
 
-/// Response from the completions endpoint.
+/// Response from the completions endpoint
 #[derive(Clone, Debug, Deserialize)]
-struct Completation {
+struct Completion {
     pub choices: Vec<Choice>,
+    #[serde(default)]
+    pub citations: Option<Vec<String>>,
 }
 
 #[derive(Clone, Debug, Default)]
@@ -186,11 +189,12 @@ impl BotClient for MolyClient {
         Box::new(self.clone())
     }
 
+    /// Stream pieces of content back as a ChatDelta instead of just a String.
     fn send_stream(
         &mut self,
         bot: &BotId,
         messages: &[Message],
-    ) -> MolyStream<'static, Result<String, ()>> {
+    ) -> MolyStream<'static, Result<ChatDelta, ()>> {
         let moly_messages: Vec<OutcomingMessage> = messages
             .iter()
             .filter_map(|m| m.clone().try_into().ok())
@@ -255,22 +259,28 @@ impl BotClient for MolyClient {
                     .filter(|m| m.trim() != "[DONE]");
 
                 for m in messages {
-                    let completition: Completation = match serde_json::from_str(m) {
-                        Ok(completition) => completition,
+                    let completion: Completion = match serde_json::from_str(m) {
+                        Ok(c) => c,
                         Err(error) => {
                             log!("Error: {:?}", error);
                             yield Err(());
                             return;
                         }
                     };
 
-                    let text = completition
+                    // Combine all partial choices content
+                    let content_delta = completion
                         .choices
                         .iter()
                         .map(|c| c.delta.content.as_str())
                         .collect::<String>();
 
-                    yield Ok(text);
+                    let citations = completion.citations.clone();
+
+                    yield Ok(ChatDelta {
+                        content_delta,
+                        citations,
+                    });
                 }
 
                 buffer = incomplete_message.to_vec();

moly-kit/src/clients/moly.rs

@@ -28,7 +28,7 @@ impl BotClient for MultiClient {
         &mut self,
         bot: &BotId,
         messages: &[Message],
-    ) -> MolyStream<'static, Result<String, ()>> {
+    ) -> MolyStream<'static, Result<ChatDelta, ()>> {
         let mut client = self
             .clients_with_bots
             .lock()
@@ -43,7 +43,7 @@ impl BotClient for MultiClient {
             })
             .expect("no client for bot");
 
-        client.send_stream(&bot, messages)
+        client.send_stream(bot, messages)
     }
 
     fn clone_box(&self) -> Box<dyn BotClient> {

moly-kit/src/clients/multi.rs

@@ -76,6 +76,15 @@ pub struct Message {
     ///
     /// If `false`, it means the message will not change anymore.
     pub is_writing: bool,
+    /// Citations for the message.
+    pub citations: Vec<String>,
+}
+
+/// A new structure to hold both text delta and optional metadata like citations.
+#[derive(Clone, Debug)]
+pub struct ChatDelta {
+    pub content_delta: String,
+    pub citations: Option<Vec<String>>,
 }
 
 /// A standard interface to fetch bots information and send messages to them.
@@ -92,7 +101,7 @@ pub trait BotClient: Send {
         &mut self,
         bot: &BotId,
         messages: &[Message],
-    ) -> MolyStream<'static, Result<String, ()>>;
+    ) -> MolyStream<'static, Result<ChatDelta, ()>>;
 
     /// Interrupt the bot's current operation.
     // TODO: There may be many chats with the same bot/model/agent so maybe this
@@ -109,24 +118,37 @@ pub trait BotClient: Send {
     fn clone_box(&self) -> Box<dyn BotClient>;
 
     /// Send a message to a bot expecting a full response at once.
-    // TODO: messages may end up being a little bit more complex, using string while thinking.
-    // TODO: Should support a way of passing, unknown, backend-specific, inference parameters.
     fn send(
         &mut self,
         bot: &BotId,
         messages: &[Message],
-    ) -> MolyFuture<'static, Result<String, ()>> {
+    ) -> MolyFuture<'static, Result<Message, ()>> {
         let stream = self.send_stream(bot, messages);
+        let bot = bot.clone();
 
         let future = async move {
-            let parts = stream.collect::<Vec<_>>().await;
-
-            if parts.contains(&Err(())) {
-                return Err(());
+            let mut content = String::new();
+            let mut citations = Vec::new();
+
+            let mut stream = stream;
+            while let Some(delta) = stream.next().await {
+                match delta {
+                    Ok(chat_delta) => {
+                        content.push_str(&chat_delta.content_delta);
+                        if let Some(cits) = chat_delta.citations {
+                            citations = cits;
+                        }
+                    }
+                    Err(()) => return Err(()),
+                }
             }
 
-            let message = parts.into_iter().filter_map(Result::ok).collect::<String>();
-            Ok(message)
+            Ok(Message {
+                from: EntityId::Bot(bot),
+                body: content,
+                is_writing: false,
+                citations,
+            })
         };
 
         moly_future(future)

moly-kit/src/protocol.rs

@@ -88,8 +88,8 @@ pub enum ChatTask {
     /// When received back, it will re-write the message history with the given messages.
     SetMessages(Vec<Message>),
 
-    /// When received back, it will insert a message at the given index with the given text and entity.
-    InsertMessage(usize, EntityId, String),
+    /// When received back, it will insert a message at the given index.
+    InsertMessage(usize, Message),
 
     /// When received back, it will delete the message at the given index.
     DeleteMessage(usize),
@@ -292,8 +292,12 @@ impl Chat {
             if !text.is_empty() {
                 composition.push(ChatTask::InsertMessage(
                     next_index,
-                    EntityId::User,
-                    text.clone(),
+                    Message {
+                        from: EntityId::User,
+                        body: text.clone(),
+                        is_writing: false,
+                        citations: vec![],
+                    },
                 ));
             }
 
@@ -322,6 +326,7 @@ impl Chat {
                 from: EntityId::Bot(bot_id.clone()),
                 body: String::new(),
                 is_writing: true,
+                citations: vec![],
             });
 
             self.dispatch(cx, ChatTask::ScrollToBottom.into());
@@ -343,24 +348,28 @@ impl Chat {
             let mut message_stream = client.send_stream(&bot_id, &context);
 
             while let Some(delta) = message_stream.next().await {
-                let delta = delta.unwrap_or_else(|_| "An error occurred".to_string());
+                let delta = match delta {
+                    Ok(chat_delta) => chat_delta,
+                    Err(_) => ChatDelta {
+                        content_delta: "An error occurred".to_string(),
+                        citations: None,
+                    },
+                };
 
                 ui.defer_with_redraw(move |me, cx, _scope| {
                     me.messages_ref().write_with(|messages| {
                         let last_index = messages.messages.len() - 1;
-                        let message = messages
-                            .messages
-                            .last_mut()
-                            .expect("no message where to put delta");
+                        let message = messages.messages.last_mut().expect("no message where to put delta");
 
-                        message.body.push_str(&delta);
-                        let updated_body = message.body.clone();
+                        // Append new text
+                        message.body.push_str(&delta.content_delta);
 
-                        me.dispatch(
-                            cx,
-                            ChatTask::UpdateMessage(last_index, updated_body).into(),
-                        );
+                        // If the chunk contains citations, store them
+                        if let Some(cits) = &delta.citations {
+                            message.citations = cits.clone();
+                        }
 
+                        me.dispatch(cx, ChatTask::UpdateMessage(last_index, message.body.clone()).into());
                         if messages.is_at_bottom() {
                             me.dispatch(cx, ChatTask::ScrollToBottom.into());
                         }
@@ -458,14 +467,10 @@ impl Chat {
                 self.messages_ref().write().messages.remove(*index);
                 self.redraw(cx);
             }
-            ChatTask::InsertMessage(index, entity, text) => {
+            ChatTask::InsertMessage(index, message) => {
                 self.messages_ref().write().messages.insert(
                     *index,
-                    Message {
-                        from: entity.clone(),
-                        body: text.clone(),
-                        is_writing: false,
-                    },
+                    message.clone(),
                 );
                 self.redraw(cx);
             }