feat: usage component (#786)

Author: dazy-ds

.github/scripts/check_changelog_update.sh

@@ -13,6 +13,14 @@ fi
 
 CHANGED_PACKAGES=$(echo "$CHANGED_FILES" | grep -oE 'packages/[^/]+/src' | cut -d '/' -f2 | sort -u)
 
+# Treat changes in `typescript` directory as `ragbits-chat` package.
+if echo "$CHANGED_FILES" | grep -q "typescript/"; then
+  CHANGED_PACKAGES="$CHANGED_PACKAGES"$'\n'ragbits-chat
+fi
+
+# Deduplicate
+CHANGED_PACKAGES=$(echo "$CHANGED_PACKAGES" | sort -u)
+
 if [ -z "$CHANGED_PACKAGES" ]; then
   echo "No package changes detected. Skipping changelog check."
   exit 0

examples/chat/chat.py

@@ -88,6 +88,7 @@ class MyChat(ChatInterface):
     )
 
     conversation_history = True
+    show_usage = True
 
     def __init__(self) -> None:
         self.llm = LiteLLM(model_name="gpt-4o-mini")
@@ -143,7 +144,11 @@ async def chat(
             yield live_update
             await asyncio.sleep(2)
 
-        async for chunk in self.llm.generate_streaming([*history, {"role": "user", "content": message}]):
+        streaming_result = self.llm.generate_streaming([*history, {"role": "user", "content": message}])
+        async for chunk in streaming_result:
             yield self.create_text_response(chunk)
 
+        if streaming_result.usage:
+            yield self.create_usage_response(streaming_result.usage)
+
         yield self.create_followup_messages(["Example Response 1", "Example Response 2", "Example Response 3"])

package-lock.json

@@ -2,6 +2,7 @@
 
 ## Unreleased
 - Add `clear_message` event type allowing to reset whole message (#789)
+- Add usage component to UI with backend support (#786)
 - Add authentication handling in the UI (#763)
 - Add backend authentication (#761)
 

packages/ragbits-chat/CHANGELOG.md

@@ -5,6 +5,7 @@
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 from pathlib import Path
+from typing import Any
 
 import uvicorn
 from fastapi import Depends, FastAPI, HTTPException, Request, status
@@ -13,6 +14,7 @@
 from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
 
 from ragbits.chat.auth import AuthenticationBackend, User
 from ragbits.chat.auth.types import LoginRequest, LoginResponse, LogoutRequest
@@ -170,6 +172,7 @@ async def config() -> JSONResponse:
                 user_settings=self.chat_interface.user_settings,
                 debug_mode=self.debug_mode,
                 conversation_history=self.chat_interface.conversation_history,
+                show_usage=self.chat_interface.show_usage,
                 authentication=AuthenticationConfig(
                     enabled=self.auth_backend is not None,
                     auth_types=[AuthType.CREDENTIALS],
@@ -537,12 +540,18 @@ async def _chat_response_to_sse(
         try:
             async for response in responses:
                 chunk_count += 1
+                response_to_send: Any = response.content
+                if isinstance(response.content, dict):
+                    response_to_send = {
+                        key: model.model_dump() if isinstance(model, BaseModel) else model
+                        for key, model in response.content.items()
+                    }
                 data = json.dumps(
                     {
                         "type": response.type.value,
-                        "content": response.content
-                        if isinstance(response.content, str | list | None)
-                        else response.content.model_dump(),
+                        "content": response_to_send.model_dump()
+                        if isinstance(response_to_send, BaseModel)
+                        else response_to_send,
                     }
                 )
                 yield f"data: {data}\n\n"

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

@@ -12,6 +12,7 @@
 from ragbits.chat.interface.ui_customization import UICustomization
 from ragbits.core.audit.metrics import record_metric
 from ragbits.core.audit.metrics.base import MetricType
+from ragbits.core.llms.base import Usage
 from ragbits.core.prompt.base import ChatFormat
 from ragbits.core.utils import get_secret_key
 
@@ -27,6 +28,7 @@
     LiveUpdate,
     LiveUpdateContent,
     LiveUpdateType,
+    MessageUsage,
     Reference,
     StateUpdate,
 )
@@ -179,6 +181,7 @@ class ChatInterface(ABC):
     feedback_config: FeedbackConfig = FeedbackConfig()
     user_settings: UserSettings = UserSettings()
     conversation_history: bool = False
+    show_usage: bool = False
     ui_customization: UICustomization | None = None
     history_persistence: HistoryPersistenceStrategy | None = None
 
@@ -241,6 +244,13 @@ def create_clear_message_response() -> ChatResponse:
         """Helper method to create an clear message response."""
         return ChatResponse(type=ChatResponseType.CLEAR_MESSAGE, content=None)
 
+    @staticmethod
+    def create_usage_response(usage: Usage) -> ChatResponse:
+        return ChatResponse(
+            type=ChatResponseType.USAGE,
+            content={model: MessageUsage.from_usage(usage) for model, usage in usage.model_breakdown.items()},
+        )
+
     @staticmethod
     def _sign_state(state: dict[str, Any]) -> str:
         """

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

@@ -5,6 +5,7 @@
 
 from ragbits.chat.interface.forms import UserSettings
 from ragbits.chat.interface.ui_customization import UICustomization
+from ragbits.core.llms.base import Usage
 
 
 class MessageRole(str, Enum):
@@ -66,6 +67,35 @@ class Image(BaseModel):
     url: str
 
 
+class MessageUsage(BaseModel):
+    """Represents usage for a message. Reflects `Usage` computed properties."""
+
+    n_requests: int
+    estimated_cost: float
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+    @classmethod
+    def from_usage(cls, usage: Usage) -> "MessageUsage":
+        """
+        Create a MessageUsage object from Usage.
+
+        Args:
+            usage: Usage object to be transformed.
+
+        Returns:
+            The corresponding MessageUsage.
+        """
+        return cls(
+            completion_tokens=usage.completion_tokens,
+            estimated_cost=usage.estimated_cost,
+            n_requests=usage.n_requests,
+            prompt_tokens=usage.prompt_tokens,
+            total_tokens=usage.total_tokens,
+        )
+
+
 class ChatResponseType(str, Enum):
     """Types of responses that can be returned by the chat interface."""
 
@@ -78,6 +108,7 @@ class ChatResponseType(str, Enum):
     FOLLOWUP_MESSAGES = "followup_messages"
     IMAGE = "image"
     CLEAR_MESSAGE = "clear_message"
+    USAGE = "usage"
 
 
 class ChatContext(BaseModel):
@@ -94,7 +125,7 @@ class ChatResponse(BaseModel):
     """Container for different types of chat responses."""
 
     type: ChatResponseType
-    content: str | Reference | StateUpdate | LiveUpdate | list[str] | Image | None
+    content: str | Reference | StateUpdate | LiveUpdate | list[str] | Image | dict[str, MessageUsage] | None
 
     def as_text(self) -> str | None:
         """
@@ -164,6 +195,12 @@ def as_clear_message(self) -> None:
         """
         return cast(None, self.content)
 
+    def as_usage(self) -> dict[str, MessageUsage] | None:
+        """
+        Return the content as dict from model name to Usage if this is an usage response, else None
+        """
+        return cast(dict[str, MessageUsage], self.content) if self.type == ChatResponseType.USAGE else None
+
 
 class ChatMessageRequest(BaseModel):
     """Client-side chat request interface."""
@@ -230,5 +267,6 @@ class ConfigResponse(BaseModel):
     customization: UICustomization | None = Field(default=None, description="UI customization")
     user_settings: UserSettings = Field(default_factory=UserSettings, description="User settings")
     debug_mode: bool = Field(default=False, description="Debug mode flag")
-    conversation_history: bool = Field(default=False, description="Debug mode flag")
+    conversation_history: bool = Field(default=False, description="Flag to enable conversation history")
+    show_usage: bool = Field(default=False, description="Flag to enable usage statistics")
     authentication: AuthenticationConfig = Field(..., description="Authentication configuration")

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

@@ -65,6 +65,7 @@ def get_models(self) -> dict[str, type[BaseModel | Enum]]:
                 LiveUpdateType,
                 Message,
                 MessageRole,
+                MessageUsage,
                 Reference,
                 StateUpdate,
             )
@@ -89,6 +90,7 @@ def get_models(self) -> dict[str, type[BaseModel | Enum]]:
                 "ServerState": StateUpdate,
                 "FeedbackItem": FeedbackItem,
                 "Image": Image,
+                "MessageUsage": MessageUsage,
                 # Configuration models
                 "HeaderCustomization": HeaderCustomization,
                 "UICustomization": UICustomization,
@@ -145,6 +147,7 @@ def get_categories(self) -> dict[str, list[str]]:
                 "Image",
                 "JWTToken",
                 "User",
+                "MessageUsage",
             ],
             "configuration": [
                 "HeaderCustomization",

packages/ragbits-chat/src/ragbits/chat/providers/model_provider.py

@@ -200,6 +200,7 @@ def _generate_chat_response_union_type() -> str:
         ("FollowupMessagesChatResponse", "followup_messages", "string[]"),
         ("ImageChatResponse", "image", "Image"),
         ("ClearMessageResponse", "clear_message", "never"),
+        ("MessageUsageChatResponse", "usage", "Record<string, MessageUsage>"),
     ]
 
     for interface_name, response_type, content_type in response_interfaces:

scripts/generate_typescript_from_json_schema.py

@@ -21,6 +21,7 @@ export const ChatResponseType = {
     FollowupMessages: 'followup_messages',
     Image: 'image',
     ClearMessage: 'clear_message',
+    Usage: 'usage',
 } as const
 
 export type ChatResponseType = TypeFrom<typeof ChatResponseType>
@@ -144,6 +145,17 @@ export interface Image {
     url: string
 }
 
+/**
+ * Represents usage for a message. Reflects `Usage` computed properties.
+ */
+export interface MessageUsage {
+    n_requests: number
+    estimated_cost: number
+    prompt_tokens: number
+    completion_tokens: number
+    total_tokens: number
+}
+
 /**
  * Customization for the header section of the UI.
  */
@@ -227,9 +239,13 @@ export interface ConfigResponse {
      */
     debug_mode: boolean
     /**
-     * Debug mode flag
+     * Flag to enable conversation history
      */
     conversation_history: boolean
+    /**
+     * Flag to enable usage statistics
+     */
+    show_usage: boolean
     authentication: AuthenticationConfig
 }
 
@@ -430,6 +446,11 @@ interface ClearMessageResponse {
     content: never
 }
 
+interface MessageUsageChatResponse {
+    type: 'usage'
+    content: Record<string, MessageUsage>
+}
+
 /**
  * Typed chat response union
  */
@@ -443,3 +464,4 @@ export type ChatResponse =
     | FollowupMessagesChatResponse
     | ImageChatResponse
     | ClearMessageResponse
+    | MessageUsageChatResponse