SurfSense/surfsense_backend/app/schemas/new_chat.py

"""
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 uuid import UUID

from pydantic import BaseModel, ConfigDict, Field

from app.db import ChatVisibility, 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
    author_id: UUID | None = None
    author_display_name: str | None = None
    author_avatar_url: str | None = None
    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
    # Visibility defaults to PRIVATE, but can be set on creation
    visibility: ChatVisibility = ChatVisibility.PRIVATE


class NewChatThreadUpdate(BaseModel):
    """Schema for updating a thread."""

    title: str | None = None
    archived: bool | None = None


class NewChatThreadVisibilityUpdate(BaseModel):
    """Schema for updating thread visibility/sharing settings."""

    visibility: ChatVisibility


class NewChatThreadRead(NewChatThreadBase, IDModel):
    """
    Schema for reading a thread (matches assistant-ui ThreadRecord).
    """

    search_space_id: int
    visibility: ChatVisibility
    created_by_id: UUID | None = None
    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] = []
    has_comments: bool = False


# =============================================================================
# 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
    visibility: ChatVisibility
    created_by_id: UUID | None = None
    is_own_thread: bool = False
    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
    )
    mentioned_surfsense_doc_ids: list[int] | None = (
        None  # Optional SurfSense documentation IDs mentioned with @ in the chat
    )


class RegenerateRequest(BaseModel):
    """
    Request schema for regenerating an AI response.

    This supports two operations:
    1. Edit: Provide a new user_query to replace the last user message and regenerate
    2. Reload: Leave user_query empty to regenerate the last AI response with the same query

    Both operations rewind the LangGraph checkpointer to the appropriate state.
    """

    search_space_id: int
    user_query: str | None = (
        None  # New user query (for edit). None = reload with same query
    )
    attachments: list[ChatAttachment] | None = None
    mentioned_document_ids: list[int] | None = None
    mentioned_surfsense_doc_ids: list[int] | None = None


# =============================================================================
# Public Chat Snapshot Schemas
# =============================================================================


class SnapshotCreateResponse(BaseModel):
    """Response after creating a public snapshot."""

    snapshot_id: int
    share_token: str
    public_url: str
    is_new: bool  # False if existing snapshot returned (same content)


class SnapshotInfo(BaseModel):
    """Info about a single snapshot."""

    id: int
    share_token: str
    public_url: str
    created_at: datetime
    message_count: int


class SnapshotListResponse(BaseModel):
    """List of snapshots for a thread."""

    snapshots: list[SnapshotInfo]


# =============================================================================
# Public Chat View Schemas (for unauthenticated access)
# =============================================================================


class PublicAuthor(BaseModel):
    display_name: str | None = None
    avatar_url: str | None = None


class PublicChatMessage(BaseModel):
    role: NewChatMessageRole
    content: Any
    author: PublicAuthor | None = None
    created_at: datetime


class PublicChatThread(BaseModel):
    title: str
    created_at: datetime


class PublicChatResponse(BaseModel):
    thread: PublicChatThread
    messages: list[PublicChatMessage]


class CloneResponse(BaseModel):
    """Response after cloning a public snapshot."""

    thread_id: int
    search_space_id: int