feat: #729 - images support in chat message, image gallery (#731)

Author: jakubduda-dsai

examples/chat/chat.py

@@ -20,6 +20,7 @@
 #
 
 import asyncio
+import uuid
 from collections.abc import AsyncGenerator
 from typing import Literal
 
@@ -117,6 +118,11 @@ async def chat(
             {"example_state_key": "example_state_value", "example_state_key_2": "example_state_value_2"}
         )
 
+        yield self.create_image_response(
+            str(uuid.uuid4()),
+            "https://media.istockphoto.com/id/1145618475/photo/villefranche-on-sea-in-evening.jpg?s=612x612&w=0&k=20&c=vQGj6uK7UUVt0vQhZc9yhRO_oYBEf8IeeDxGyJKbLKI=",
+        )
+
         example_live_updates = [
             self.create_live_update("0", LiveUpdateType.START, "[EXAMPLE] Searching for examples in the web..."),
             self.create_live_update(

packages/ragbits-chat/CHANGELOG.md

@@ -1,6 +1,7 @@
 # CHANGELOG
 
 ## Unreleased
+- Add images support in chat message, images gallery (#731)
 - Add persistent user settings (#719)
 - Send chat options under `user_settings` key (#721)
 - Add feedback indicator to messages, allow `extensions` in chat messages (#722)

packages/ragbits-chat/src/ragbits/chat/api.py

@@ -234,6 +234,8 @@ async def chat_response() -> AsyncGenerator[str, None]:
                                 outputs.message_id = content
                             case ChatResponseType.CONVERSATION_ID:
                                 outputs.conversation_id = content
+                            case ChatResponseType.IMAGE:
+                                outputs.image_url = content
 
                         yield chunk
 

packages/ragbits-chat/src/ragbits/chat/interface/_interface.py

@@ -22,6 +22,7 @@
     ChatContext,
     ChatResponse,
     ChatResponseType,
+    Image,
     LiveUpdate,
     LiveUpdateContent,
     LiveUpdateType,
@@ -228,6 +229,11 @@ def create_followup_messages(messages: list[str]) -> ChatResponse:
         """Helper method to create a live update response."""
         return ChatResponse(type=ChatResponseType.FOLLOWUP_MESSAGES, content=messages)
 
+    @staticmethod
+    def create_image_response(image_id: str, image_url: str) -> ChatResponse:
+        """Helper method to create an image response."""
+        return ChatResponse(type=ChatResponseType.IMAGE, content=Image(id=image_id, url=image_url))
+
     @staticmethod
     def _sign_state(state: dict[str, Any]) -> str:
         """

packages/ragbits-chat/src/ragbits/chat/interface/types.py

@@ -56,6 +56,13 @@ class LiveUpdate(BaseModel):
     content: LiveUpdateContent
 
 
+class Image(BaseModel):
+    """Represents an image in the conversation."""
+
+    id: str
+    url: str
+
+
 class ChatResponseType(str, Enum):
     """Types of responses that can be returned by the chat interface."""
 
@@ -66,13 +73,14 @@ class ChatResponseType(str, Enum):
     CONVERSATION_ID = "conversation_id"
     LIVE_UPDATE = "live_update"
     FOLLOWUP_MESSAGES = "followup_messages"
+    IMAGE = "image"
 
 
 class ChatResponse(BaseModel):
     """Container for different types of chat responses."""
 
     type: ChatResponseType
-    content: str | Reference | StateUpdate | LiveUpdate | list[str]
+    content: str | Reference | StateUpdate | LiveUpdate | list[str] | Image
 
     def as_text(self) -> str | None:
         """
@@ -130,6 +138,12 @@ def as_followup_messages(self) -> list[str] | None:
         """
         return cast(list[str], self.content) if self.type == ChatResponseType.FOLLOWUP_MESSAGES else None
 
+    def as_image(self) -> Image | None:
+        """
+        Return the content as Image if this is an image response, else None.
+        """
+        return cast(Image, self.content) if self.type == ChatResponseType.IMAGE else None
+
 
 class ChatContext(BaseModel):
     """Represents the context of a chat conversation."""