feat: add support for Tasks in API and integrate UI

Author: dazy-ds

examples/chat/chat.py

@@ -26,6 +26,7 @@
 
 from pydantic import BaseModel, ConfigDict, Field
 
+from ragbits.agents.tools.todo import Task, TaskStatus
 from ragbits.chat.interface import ChatInterface
 from ragbits.chat.interface.forms import FeedbackConfig, UserSettings
 from ragbits.chat.interface.types import ChatContext, ChatResponse, LiveUpdateType
@@ -141,10 +142,25 @@ async def chat(
             ),
         ]
 
+        parentTask = Task(id="task_id_1", description="Example task with a subtask")
+        subtaskTask = Task(
+            id="task_id_2", description="Example subtask", status=TaskStatus.IN_PROGRESS, parent_id="task_id_1"
+        )
+
         for live_update in example_live_updates:
             yield live_update
             await asyncio.sleep(2)
 
+        yield self.create_todo_item_response(parentTask)
+        yield self.create_todo_item_response(subtaskTask)
+
+        await asyncio.sleep(1)
+        subtaskTask.status = TaskStatus.COMPLETED
+        parentTask.status = TaskStatus.COMPLETED
+
+        yield self.create_todo_item_response(parentTask)
+        yield self.create_todo_item_response(subtaskTask)
+
         streaming_result = self.llm.generate_streaming([*history, {"role": "user", "content": message}])
         async for chunk in streaming_result:
             yield self.create_text_response(chunk)

packages/ragbits-agents/src/ragbits/agents/tools/todo.py

@@ -2,31 +2,35 @@
 
 import uuid
 from dataclasses import dataclass, field
-from datetime import datetime
 from enum import Enum
 from typing import Any, Literal
 
+from pydantic import BaseModel
+
 
 class TaskStatus(str, Enum):
     """Task status options."""
+
     PENDING = "pending"
     IN_PROGRESS = "in_progress"
     COMPLETED = "completed"
 
 
-@dataclass
-class Task:
+class Task(BaseModel):
     """Simple task representation."""
+
     id: str
     description: str
     status: TaskStatus = TaskStatus.PENDING
     order: int = 0
     summary: str | None = None
+    parent_id: str | None = None
 
 
 @dataclass
 class TodoList:
     """Simple todo list for one agent run."""
+
     tasks: list[Task] = field(default_factory=list)
     current_index: int = 0
 
@@ -36,14 +40,15 @@ def get_current_task(self) -> Task | None:
             return self.tasks[self.current_index]
         return None
 
-    def advance_to_next(self):
+    def advance_to_next(self) -> None:
         """Move to next task."""
         self.current_index += 1
 
 
 # Storage - just one todo list per agent run
 _current_todo: TodoList | None = None
 
+
 def todo_manager(
     action: Literal["create", "get_current", "start_task", "complete_task", "get_final_summary"],
     tasks: list[str] | None = None,
@@ -67,18 +72,14 @@ def todo_manager(
 
         _current_todo = TodoList()
         for i, desc in enumerate(tasks):
-            task = Task(
-                id=str(uuid.uuid4()),
-                description=desc.strip(),
-                order=i
-            )
+            task = Task(id=str(uuid.uuid4()), description=desc.strip(), order=i)
             _current_todo.tasks.append(task)
 
         return {
             "action": "create",
             "tasks": [{"id": t.id, "description": t.description, "order": t.order} for t in _current_todo.tasks],
             "total_count": len(_current_todo.tasks),
-            "message": f"Created {len(tasks)} tasks"
+            "message": f"Created {len(tasks)} tasks",
         }
 
     if not _current_todo:
@@ -91,14 +92,14 @@ def todo_manager(
                 "action": "get_current",
                 "current_task": None,
                 "all_completed": True,
-                "message": "All tasks completed!"
+                "message": "All tasks completed!",
             }
 
         return {
             "action": "get_current",
             "current_task": {"id": current.id, "description": current.description, "status": current.status.value},
             "progress": f"{_current_todo.current_index + 1}/{len(_current_todo.tasks)}",
-            "message": f"Current task: {current.description}"
+            "message": f"Current task: {current.description}",
         }
 
     elif action == "start_task":
@@ -110,7 +111,7 @@ def todo_manager(
         return {
             "action": "start_task",
             "task": {"id": current.id, "description": current.description, "status": current.status.value},
-            "message": f"Started task: {current.description}"
+            "message": f"Started task: {current.description}",
         }
 
     elif action == "complete_task":
@@ -137,18 +138,14 @@ def todo_manager(
             "next_task": {"id": next_task.id, "description": next_task.description} if next_task else None,
             "progress": f"{completed_count}/{len(_current_todo.tasks)}",
             "all_completed": next_task is None,
-            "message": f"Completed: {current.description}"
+            "message": f"Completed: {current.description}",
         }
 
     elif action == "get_final_summary":
         completed_tasks = [t for t in _current_todo.tasks if t.status == TaskStatus.COMPLETED]
 
         if not completed_tasks:
-            return {
-                "action": "get_final_summary",
-                "final_summary": "",
-                "message": "No completed tasks found."
-            }
+            return {"action": "get_final_summary", "final_summary": "", "message": "No completed tasks found."}
 
         # Create comprehensive final summary
         final_content = []
@@ -167,7 +164,7 @@ def todo_manager(
             "action": "get_final_summary",
             "final_summary": final_summary,
             "total_completed": len(completed_tasks),
-            "message": f"Final summary with {len(completed_tasks)} completed tasks."
+            "message": f"Final summary with {len(completed_tasks)} completed tasks.",
         }
 
     else:

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

@@ -9,6 +9,7 @@
 from collections.abc import AsyncGenerator, Callable
 from typing import Any
 
+from ragbits.agents.tools.todo import Task
 from ragbits.chat.interface.ui_customization import UICustomization
 from ragbits.core.audit.metrics import record_metric
 from ragbits.core.audit.metrics.base import MetricType
@@ -251,6 +252,10 @@ def create_usage_response(usage: Usage) -> ChatResponse:
             content={model: MessageUsage.from_usage(usage) for model, usage in usage.model_breakdown.items()},
         )
 
+    @staticmethod
+    def create_todo_item_response(task: Task) -> ChatResponse:
+        return ChatResponse(type=ChatResponseType.TODO_ITEM, content=task)
+
     @staticmethod
     def _sign_state(state: dict[str, Any]) -> str:
         """

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

@@ -3,6 +3,7 @@
 
 from pydantic import BaseModel, ConfigDict, Field
 
+from ragbits.agents.tools.todo import Task
 from ragbits.chat.auth.types import User
 from ragbits.chat.interface.forms import UserSettings
 from ragbits.chat.interface.ui_customization import UICustomization
@@ -122,6 +123,7 @@ class ChatResponseType(str, Enum):
     CHUNKED_CONTENT = "chunked_content"
     CLEAR_MESSAGE = "clear_message"
     USAGE = "usage"
+    TODO_ITEM = "todo_item"
 
 
 class ChatContext(BaseModel):
@@ -140,7 +142,16 @@ class ChatResponse(BaseModel):
 
     type: ChatResponseType
     content: (
-        str | Reference | StateUpdate | LiveUpdate | list[str] | Image | dict[str, MessageUsage] | ChunkedContent | None
+        str
+        | Reference
+        | StateUpdate
+        | LiveUpdate
+        | list[str]
+        | Image
+        | dict[str, MessageUsage]
+        | ChunkedContent
+        | None
+        | Task
     )
 
     def as_text(self) -> str | None:
@@ -217,6 +228,12 @@ def as_usage(self) -> dict[str, MessageUsage] | None:
         """
         return cast(dict[str, MessageUsage], self.content) if self.type == ChatResponseType.USAGE else None
 
+    def as_task(self) -> Task | None:
+        """
+        Return the content as Task if this is an todo_item response, else None.
+        """
+        return cast(Task, self.content) if self.type == ChatResponseType.TODO_ITEM else None
+
 
 class ChatMessageRequest(BaseModel):
     """Client-side chat request interface."""

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

@@ -10,6 +10,7 @@
 
 from pydantic import BaseModel
 
+from ragbits.agents.tools.todo import Task, TaskStatus
 from ragbits.chat.interface.types import AuthType
 
 
@@ -82,6 +83,7 @@ def get_models(self) -> dict[str, type[BaseModel | Enum]]:
                 "FeedbackType": FeedbackType,
                 "LiveUpdateType": LiveUpdateType,
                 "MessageRole": MessageRole,
+                "TaskStatus": TaskStatus,
                 # Core data models
                 "ChatContext": ChatContext,
                 "ChunkedContent": ChunkedContent,
@@ -93,6 +95,7 @@ def get_models(self) -> dict[str, type[BaseModel | Enum]]:
                 "FeedbackItem": FeedbackItem,
                 "Image": Image,
                 "MessageUsage": MessageUsage,
+                "Task": Task,
                 # Configuration models
                 "HeaderCustomization": HeaderCustomization,
                 "UICustomization": UICustomization,
@@ -151,6 +154,8 @@ def get_categories(self) -> dict[str, list[str]]:
                 "JWTToken",
                 "User",
                 "MessageUsage",
+                "Task",
+                "TaskStatus",
             ],
             "configuration": [
                 "HeaderCustomization",

scripts/generate_typescript_from_json_schema.py

@@ -201,6 +201,7 @@ def _generate_chat_response_union_type() -> str:
         ("ImageChatResponse", "image", "Image"),
         ("ClearMessageResponse", "clear_message", "never"),
         ("MessageUsageChatResponse", "usage", "Record<string, MessageUsage>"),
+        ("TodoItemChatResonse", "todo_item", "Task"),
     ]
 
     internal_response_interfaces = [

typescript/@ragbits/api-client/package.json

@@ -5,7 +5,7 @@
     "repository": {
         "type": "git",
         "url": "https://github.com/deepsense-ai/ragbits"
-     },
+    },
     "main": "dist/index.cjs",
     "module": "dist/index.js",
     "types": "dist/index.d.ts",

typescript/@ragbits/api-client/src/autogen.types.ts

@@ -23,6 +23,7 @@ export const ChatResponseType = {
     ChunkedContent: 'chunked_content',
     ClearMessage: 'clear_message',
     Usage: 'usage',
+    TodoItem: 'todo_item',
 } as const
 
 export type ChatResponseType = TypeFrom<typeof ChatResponseType>
@@ -58,6 +59,17 @@ export const MessageRole = {
 
 export type MessageRole = TypeFrom<typeof MessageRole>
 
+/**
+ * Represents the TaskStatus enum
+ */
+export const TaskStatus = {
+    Pending: 'pending',
+    InProgress: 'in_progress',
+    Completed: 'completed',
+} as const
+
+export type TaskStatus = TypeFrom<typeof TaskStatus>
+
 /**
  * Represents the AuthType enum
  */
@@ -170,6 +182,21 @@ export interface MessageUsage {
     total_tokens: number
 }
 
+/**
+ * Simple task representation.
+ */
+export interface Task {
+    id: string
+    description: string
+    /**
+     * Task status options.
+     */
+    status: 'pending' | 'in_progress' | 'completed'
+    order: number
+    summary: string | null
+    parent_id: string | null
+}
+
 /**
  * Customization for the header section of the UI.
  */
@@ -465,6 +492,11 @@ export interface MessageUsageChatResponse {
     content: Record<string, MessageUsage>
 }
 
+export interface TodoItemChatResonse {
+    type: 'todo_item'
+    content: Task
+}
+
 export interface ChunkedChatResponse {
     type: 'chunked_content'
     content: ChunkedContent
@@ -484,3 +516,4 @@ export type ChatResponse =
     | ImageChatResponse
     | ClearMessageResponse
     | MessageUsageChatResponse
+    | TodoItemChatResonse