mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-04-26 01:06:23 +02:00
ceb01dc544
- Updated the new chat routes to include handling for mentioned document IDs, allowing users to reference specific documents in their chat. - Modified the NewChatRequest schema to accommodate optional document IDs. - Implemented document mention formatting in the chat streaming service for improved context. - Enhanced the frontend to manage document mentions, including a new atom for state management and UI updates for document selection. - Refactored the DocumentsDataTable component for better integration with the new mention functionality.
165 lines
4.4 KiB
Python
165 lines
4.4 KiB
Python
"""
|
|
Pydantic schemas for the new chat feature with assistant-ui integration.
|
|
|
|
These schemas follow the assistant-ui ThreadHistoryAdapter pattern:
|
|
- ThreadRecord: id, title, archived, createdAt, updatedAt
|
|
- MessageRecord: id, threadId, role, content, createdAt
|
|
"""
|
|
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
from app.db import NewChatMessageRole
|
|
|
|
from .base import IDModel, TimestampModel
|
|
|
|
# =============================================================================
|
|
# Message Schemas
|
|
# =============================================================================
|
|
|
|
|
|
class NewChatMessageBase(BaseModel):
|
|
"""Base schema for new chat messages."""
|
|
|
|
role: NewChatMessageRole
|
|
content: Any # JSONB content - can be text, tool calls, etc.
|
|
|
|
|
|
class NewChatMessageCreate(NewChatMessageBase):
|
|
"""Schema for creating a new message."""
|
|
|
|
thread_id: int
|
|
|
|
|
|
class NewChatMessageRead(NewChatMessageBase, IDModel, TimestampModel):
|
|
"""Schema for reading a message."""
|
|
|
|
thread_id: int
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
|
|
class NewChatMessageAppend(BaseModel):
|
|
"""
|
|
Schema for appending a message via the history adapter.
|
|
This is the format assistant-ui sends when calling append().
|
|
"""
|
|
|
|
role: str # Accept string and validate in route handler
|
|
content: Any
|
|
|
|
|
|
# =============================================================================
|
|
# Thread Schemas
|
|
# =============================================================================
|
|
|
|
|
|
class NewChatThreadBase(BaseModel):
|
|
"""Base schema for new chat threads."""
|
|
|
|
title: str = Field(default="New Chat", max_length=500)
|
|
archived: bool = False
|
|
|
|
|
|
class NewChatThreadCreate(NewChatThreadBase):
|
|
"""Schema for creating a new thread."""
|
|
|
|
search_space_id: int
|
|
|
|
|
|
class NewChatThreadUpdate(BaseModel):
|
|
"""Schema for updating a thread."""
|
|
|
|
title: str | None = None
|
|
archived: bool | None = None
|
|
|
|
|
|
class NewChatThreadRead(NewChatThreadBase, IDModel):
|
|
"""
|
|
Schema for reading a thread (matches assistant-ui ThreadRecord).
|
|
"""
|
|
|
|
search_space_id: int
|
|
created_at: datetime
|
|
updated_at: datetime
|
|
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
|
|
class NewChatThreadWithMessages(NewChatThreadRead):
|
|
"""Schema for reading a thread with its messages."""
|
|
|
|
messages: list[NewChatMessageRead] = []
|
|
|
|
|
|
# =============================================================================
|
|
# History Adapter Response Schemas
|
|
# =============================================================================
|
|
|
|
|
|
class ThreadHistoryLoadResponse(BaseModel):
|
|
"""
|
|
Response format for the ThreadHistoryAdapter.load() method.
|
|
Returns messages array for the current thread.
|
|
"""
|
|
|
|
messages: list[NewChatMessageRead]
|
|
|
|
|
|
class ThreadListItem(BaseModel):
|
|
"""
|
|
Thread list item for sidebar display.
|
|
Matches assistant-ui ThreadListPrimitive expected format.
|
|
"""
|
|
|
|
id: int
|
|
title: str
|
|
archived: bool
|
|
created_at: datetime = Field(alias="createdAt")
|
|
updated_at: datetime = Field(alias="updatedAt")
|
|
|
|
model_config = ConfigDict(from_attributes=True, populate_by_name=True)
|
|
|
|
|
|
class ThreadListResponse(BaseModel):
|
|
"""Response containing list of threads for the sidebar."""
|
|
|
|
threads: list[ThreadListItem]
|
|
archived_threads: list[ThreadListItem]
|
|
|
|
|
|
# =============================================================================
|
|
# Chat Request Schemas (for deep agent)
|
|
# =============================================================================
|
|
|
|
|
|
class ChatMessage(BaseModel):
|
|
"""A single message in the chat history."""
|
|
|
|
role: str # "user" or "assistant"
|
|
content: str
|
|
|
|
|
|
class ChatAttachment(BaseModel):
|
|
"""An attachment with its extracted content for chat context."""
|
|
|
|
id: str # Unique attachment ID
|
|
name: str # Original filename
|
|
type: str # Attachment type: document, image, audio
|
|
content: str # Extracted markdown content from the file
|
|
|
|
|
|
class NewChatRequest(BaseModel):
|
|
"""Request schema for the deep agent chat endpoint."""
|
|
|
|
chat_id: int
|
|
user_query: str
|
|
search_space_id: int
|
|
messages: list[ChatMessage] | None = None # Optional chat history from frontend
|
|
attachments: list[ChatAttachment] | None = (
|
|
None # Optional attachments with extracted content
|
|
)
|
|
mentioned_document_ids: list[int] | None = (
|
|
None # Optional document IDs mentioned with @ in the chat
|
|
)
|