feat: standardize chat typing (#803)

mhordynski · jakubduda-dsai · web-flow · commit 2ff290d046a2 · 2025-09-10T13:44:36.000+02:00
Co-authored-by: jakubduda-dsai &lt;jakub.duda@deepsense.ai&gt;
diff --git a/README.md b/README.md
@@ -229,8 +229,8 @@ class MyChat(ChatInterface):
     async def chat(
         self,
         message: str,
-        history: ChatFormat | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse]:
         async for result in agent.run_streaming(message):
             match result:
diff --git a/docs/how-to/chatbots/api.md b/docs/how-to/chatbots/api.md
@@ -12,7 +12,8 @@ First, create a chat implementation by subclassing `ChatInterface`. Here's a min
 from collections.abc import AsyncGenerator
 
 from ragbits.chat.interface import ChatInterface
-from ragbits.chat.interface.types import ChatResponse, Message
+from ragbits.chat.interface.types import ChatResponse
+from ragbits.core.prompt import ChatFormat
 from ragbits.core.llms import LiteLLM
 
 class MyChat(ChatInterface):
@@ -22,8 +23,8 @@ class MyChat(ChatInterface):
     async def chat(
         self,
         message: str,
-        history: list[Message] | None = None,
-        context: dict | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse, None]:
         async for chunk in self.llm.generate_streaming([*history, {"role": "user", "content": message}]):
             yield self.create_text_response(chunk)
@@ -58,7 +59,7 @@ Ragbits Chat supports multiple response types that can be yielded from your `cha
 Text responses are the primary way to stream content to users. Use `create_text_response()` to yield text chunks:
 
 ```python
-async def chat(self, message: str, history: list[Message] | None = None, context: dict | None = None) -> AsyncGenerator[ChatResponse, None]:
+async def chat(self, message: str, history: ChatFormat, context: ChatContext) -> AsyncGenerator[ChatResponse, None]:
     # Stream response from LLM
     async for chunk in self.llm.generate_streaming([*history, {"role": "user", "content": message}]):
         yield self.create_text_response(chunk)
@@ -69,7 +70,7 @@ async def chat(self, message: str, history: list[Message] | None = None, context
 References allow you to cite sources, documents, or external links that support your response:
 
 ```python
-async def chat(self, message: str, history: list[Message] | None = None, context: dict | None = None) -> AsyncGenerator[ChatResponse, None]:
+async def chat(self, message: str, history: ChatFormat, context: ChatContext) -> AsyncGenerator[ChatResponse, None]:
     # Add a reference
     yield self.create_reference(
         title="Example Reference",
@@ -85,7 +86,7 @@ You can include images in your responses using `create_image_response()`:
 ```python
 import uuid
 
-async def chat(self, message: str, history: list[Message] | None = None, context: dict | None = None) -> AsyncGenerator[ChatResponse, None]:
+async def chat(self, message: str, history: ChatFormat, context: ChatContext) -> AsyncGenerator[ChatResponse, None]:
     # Add an image to the response
     yield self.create_image_response(
         str(uuid.uuid4()),  # Unique identifier for the image
@@ -98,7 +99,7 @@ async def chat(self, message: str, history: list[Message] | None = None, context
 Provide suggested follow-up questions to guide the conversation:
 
 ```python
-async def chat(self, message: str, history: list[Message] | None = None, context: dict | None = None) -> AsyncGenerator[ChatResponse, None]:
+async def chat(self, message: str, history: ChatFormat, context: ChatContext) -> AsyncGenerator[ChatResponse, None]:
     # Main response...
     async for chunk in self.llm.generate_streaming([*history, {"role": "user", "content": message}]):
         yield self.create_text_response(chunk)
@@ -121,7 +122,7 @@ Live updates show real-time progress for long-running operations (like web searc
 import asyncio
 from ragbits.chat.interface.types import LiveUpdateType
 
-async def chat(self, message: str, history: list[Message] | None = None, context: dict | None = None) -> AsyncGenerator[ChatResponse, None]:
+async def chat(self, message: str, history: ChatFormat, context: ChatContext) -> AsyncGenerator[ChatResponse, None]:
     # Start a live update
     yield self.create_live_update(
         "search_task",  # Unique task ID
@@ -163,12 +164,13 @@ Use `create_state_update()` to store state information that persists across conv
 
 ```python
 from ragbits.chat.interface.types import ChatContext
+from ragbits.core.prompt import ChatFormat
 
 async def chat(
     self,
     message: str,
-    history: list[Message] | None = None,
-    context: ChatContext | None = None
+    history: ChatFormat,
+    context: ChatContext
 ) -> AsyncGenerator[ChatResponse, None]:
     # Access existing state from context
     current_state = context.state if context else {}
@@ -393,7 +395,8 @@ from pydantic import BaseModel, ConfigDict, Field
 
 from ragbits.chat.interface import ChatInterface
 from ragbits.chat.interface.forms import FeedbackConfig, UserSettings
-from ragbits.chat.interface.types import ChatContext, ChatResponse, LiveUpdateType, Message
+from ragbits.chat.interface.types import ChatContext, ChatResponse, LiveUpdateType
+from ragbits.core.prompt import ChatFormat
 from ragbits.chat.interface.ui_customization import HeaderCustomization, PageMetaCustomization, UICustomization
 from ragbits.core.llms import LiteLLM
 
@@ -486,8 +489,8 @@ class MyChat(ChatInterface):
     async def chat(
         self,
         message: str,
-        history: list[Message] | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse, None]:
         """
         Comprehensive chat implementation demonstrating all response types.
diff --git a/docs/index.md b/docs/index.md
@@ -258,8 +258,8 @@ class MyChat(ChatInterface):
     async def chat(
         self,
         message: str,
-        history: ChatFormat | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse]:
         async for result in agent.run_streaming(message):
             match result:
diff --git a/docs/tutorials/chat.md b/docs/tutorials/chat.md
@@ -93,8 +93,8 @@ class BasicMountainChat(ChatInterface):
     async def chat(
         self,
         message: str,
-        history: list[Message] | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse, None]:
         # Basic streaming implementation
         stream = self.agent.run_streaming(
@@ -170,8 +170,8 @@ class BasicMountainChat(ChatInterface):
     async def chat(
         self,
         message: str,
-        history: list[Message] | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse, None]:
         # Enhanced streaming with tool handling
         stream = self.agent.run_streaming(
@@ -402,8 +402,8 @@ class MountainChatWithUI(ChatInterface):
     async def chat(
         self,
         message: str,
-        history: list[Message] | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse, None]:
         # Get user language preference
         language = "English"
@@ -530,11 +530,11 @@ class AuthenticatedMountainChat(ChatInterface):
     async def chat(
         self,
         message: str,
-        history: list[Message] | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse, None]:
         # Check authentication
-        user_info = context.state.get("authenticated_user") if context else None
+        user_info = context.user
 
         if not user_info:
             yield self.create_text_response("⚠️ Authentication information not found.")
@@ -741,8 +741,8 @@ class MyChat(ChatInterface):
             base64_image = base64.b64encode(image_file.read()).decode("utf-8")
             return self.create_image_response(image_filename, f"data:image/png;base64,{base64_image}")
 
-    async def chat(self, message: str, history: list[Message] | None = None, context: ChatContext | None = None) -> AsyncGenerator[ChatResponse, None]:
-        user_info = context.state.get("authenticated_user") if context else None
+    async def chat(self, message: str, history: ChatFormat, context: ChatContext) -> AsyncGenerator[ChatResponse, None]:
+        user_info = context.user
         if not user_info:
             yield self.create_text_response("⚠️ Authentication information not found.")
             return
diff --git a/examples/chat/README_authenticated.md b/examples/chat/README_authenticated.md
@@ -165,7 +165,7 @@ Modify the `chat()` method to customize responses based on user roles:
 
 ```python
 # Get user info from context
-user_info = context.state.get("authenticated_user") if context else None
+user_info = context.user
 user_roles = user_info.roles if user_info else []
 
 if "admin" in user_roles:
diff --git a/examples/chat/authenticated_chat.py b/examples/chat/authenticated_chat.py
@@ -28,9 +28,10 @@
 
 from ragbits.chat.auth import ListAuthenticationBackend
 from ragbits.chat.interface import ChatInterface
-from ragbits.chat.interface.types import ChatContext, ChatResponse, LiveUpdateType, Message
+from ragbits.chat.interface.types import ChatContext, ChatResponse, LiveUpdateType
 from ragbits.chat.interface.ui_customization import HeaderCustomization, UICustomization
 from ragbits.core.llms import LiteLLM
+from ragbits.core.prompt.base import ChatFormat
 
 
 class MyAuthenticatedChat(ChatInterface):
@@ -61,14 +62,14 @@ def __init__(self):
     async def chat(
         self,
         message: str,
-        history: list[Message] | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse, None]:
         """
         Authenticated chat implementation that provides user-specific responses.
 
         This method is called after authentication validation passes.
-        The user information is available in context.state["authenticated_user"].
+        The user information is available in context.user.
 
         Args:
             message: The current user message
@@ -79,7 +80,7 @@ async def chat(
             ChatResponse objects containing different types of content
         """
         # Get authenticated user info
-        user_info = context.state.get("authenticated_user") if context else None
+        user_info = context.user
 
         if not user_info:
             yield self.create_text_response("⚠️ Authentication information not found.")
@@ -100,8 +101,7 @@ async def chat(
         # Create user-specific state update
         yield self.create_state_update(
             {
-                "authenticated_user_id": user_id,
-                "session_context": context.session_id if context and context.session_id else "unknown",
+                "session_context": context.session_id if context.session_id else "unknown",
                 "user_roles": user_roles,
                 "chat_timestamp": asyncio.get_event_loop().time(),
             }
diff --git a/examples/chat/chat.py b/examples/chat/chat.py
@@ -28,9 +28,10 @@
 
 from ragbits.chat.interface import ChatInterface
 from ragbits.chat.interface.forms import FeedbackConfig, UserSettings
-from ragbits.chat.interface.types import ChatContext, ChatResponse, LiveUpdateType, Message
+from ragbits.chat.interface.types import ChatContext, ChatResponse, LiveUpdateType
 from ragbits.chat.interface.ui_customization import HeaderCustomization, PageMetaCustomization, UICustomization
 from ragbits.core.llms import LiteLLM
+from ragbits.core.prompt import ChatFormat
 
 
 class LikeFormExample(BaseModel):
@@ -96,8 +97,8 @@ def __init__(self) -> None:
     async def chat(
         self,
         message: str,
-        history: list[Message] | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse, None]:
         """
         Example implementation of the ChatInterface.
diff --git a/examples/chat/offline_chat.py b/examples/chat/offline_chat.py
@@ -26,9 +26,10 @@
 
 from ragbits.chat.interface import ChatInterface
 from ragbits.chat.interface.forms import FeedbackConfig, UserSettings
-from ragbits.chat.interface.types import ChatContext, ChatResponse, LiveUpdateType, Message
+from ragbits.chat.interface.types import ChatContext, ChatResponse, LiveUpdateType
 from ragbits.chat.interface.ui_customization import HeaderCustomization, UICustomization
 from ragbits.chat.persistence.file import FileHistoryPersistence
+from ragbits.core.prompt import ChatFormat
 
 
 class LikeFormExample(BaseModel):
@@ -100,8 +101,8 @@ async def _generate_response(message: str) -> AsyncGenerator[str, None]:
     async def chat(
         self,
         message: str,
-        history: list[Message] | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse, None]:
         """
         Offline implementation of the ChatInterface.
diff --git a/examples/chat/tutorial.py b/examples/chat/tutorial.py
@@ -44,10 +44,11 @@
 from ragbits.chat.auth import ListAuthenticationBackend
 from ragbits.chat.interface import ChatInterface
 from ragbits.chat.interface.forms import FeedbackConfig, UserSettings
-from ragbits.chat.interface.types import ChatContext, ChatResponse, LiveUpdateType, Message
+from ragbits.chat.interface.types import ChatContext, ChatResponse, LiveUpdateType
 from ragbits.chat.interface.ui_customization import HeaderCustomization, PageMetaCustomization, UICustomization
 from ragbits.core.llms import LiteLLM, ToolCall
 from ragbits.core.prompt import Prompt
+from ragbits.core.prompt.base import ChatFormat
 
 
 class LikeFormExample(BaseModel):
@@ -200,8 +201,8 @@ async def _create_image_response(self, image_path: Path) -> ChatResponse:
     async def chat(
         self,
         message: str,
-        history: list[Message] | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse, None]:
         """
         Example implementation of the ChatInterface.
@@ -218,7 +219,7 @@ async def chat(
             - Live updates for tool execution status
         """
         # Get authenticated user info
-        user_info = context.state.get("authenticated_user") if context else None
+        user_info = context.user
 
         if not user_info:
             yield self.create_text_response("⚠️ Authentication information not found.")
diff --git a/packages/ragbits-chat/CHANGELOG.md b/packages/ragbits-chat/CHANGELOG.md
@@ -1,6 +1,8 @@
 # CHANGELOG
 
 ## Unreleased
+
+- fix: replace authenticated_user state tracking with direct user field in ChatContext
 - Refactor chat handlers in the UI to use registry (#805)
 - Add auth token storage and automatic logout on 401 (#802)
 - Improve user settings storage when history is disabled (#799)
diff --git a/packages/ragbits-chat/src/ragbits/chat/api.py b/packages/ragbits-chat/src/ragbits/chat/api.py
@@ -229,7 +229,7 @@ def _prepare_chat_context(
         # Add session_id to context if authenticated
         if authenticated_user and credentials:
             chat_context.session_id = credentials.credentials
-            chat_context.state["authenticated_user"] = authenticated_user
+            chat_context.user = authenticated_user
 
         # Verify state signature if provided
         if "state" in request.context and "signature" in request.context:
diff --git a/packages/ragbits-chat/src/ragbits/chat/interface/_interface.py b/packages/ragbits-chat/src/ragbits/chat/interface/_interface.py
@@ -37,7 +37,7 @@
 
 
 def with_chat_metadata(
-    func: Callable[["ChatInterface", str, ChatFormat | None, ChatContext | None], AsyncGenerator[ChatResponse, None]],
+    func: Callable[["ChatInterface", str, ChatFormat, ChatContext], AsyncGenerator[ChatResponse, None]],
 ) -> Callable[["ChatInterface", str, ChatFormat | None, ChatContext | None], AsyncGenerator[ChatResponse, None]]:
     """
     Decorator that adds message and conversation metadata to the chat method and handles history persistence.
@@ -51,12 +51,12 @@ async def wrapper(
     ) -> AsyncGenerator[ChatResponse, None]:
         start_time = time.time()
 
-        # Initialize context if None
-        if context is None:
-            context = ChatContext()
+        # Assure history and context are not None
+        history = history or []
+        context = context or ChatContext()
 
         # Track history length
-        history_length = len(history) if history else 0
+        history_length = len(history)
         record_metric(
             ChatHistogramMetric.CHAT_HISTORY_LENGTH,
             history_length,
@@ -307,8 +307,8 @@ async def setup(self) -> None:  # noqa: B027
     async def chat(
         self,
         message: str,
-        history: ChatFormat | None = None,
-        context: ChatContext | None = None,
+        history: ChatFormat,
+        context: ChatContext,
     ) -> AsyncGenerator[ChatResponse, None]:
         """
         Process a chat message and yield responses asynchronously.
diff --git a/packages/ragbits-chat/src/ragbits/chat/interface/types.py b/packages/ragbits-chat/src/ragbits/chat/interface/types.py
@@ -3,6 +3,7 @@
 
 from pydantic import BaseModel, ConfigDict, Field
 
+from ragbits.chat.auth.types import User
 from ragbits.chat.interface.forms import UserSettings
 from ragbits.chat.interface.ui_customization import UICustomization
 from ragbits.core.llms.base import Usage
@@ -129,6 +130,7 @@ class ChatContext(BaseModel):
     conversation_id: str | None = None
     message_id: str | None = None
     state: dict[str, Any] = Field(default_factory=dict)
+    user: User | None = None
     session_id: str | None = None
     model_config = ConfigDict(extra="allow")
 
diff --git a/typescript/@ragbits/api-client/src/autogen.types.ts b/typescript/@ragbits/api-client/src/autogen.types.ts
@@ -76,6 +76,7 @@ export interface ChatContext {
     state: {
         [k: string]: unknown
     }
+    user: User | null
     session_id: string | null
     [k: string]: unknown
 }
