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, Literal
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 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
    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
    )
    disabled_tools: list[str] | None = (
        None  # Optional list of tool names the user has disabled from the UI
    )
    # Cloud mode: override the search space's agent_llm_id with a system model
    # (negative ID from global_llm_config.yaml, selected via SystemModelSelector).
    # Self-hosted mode: leave None and the search space config is used as before.
    model_id: int | None = None


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
    )
    mentioned_document_ids: list[int] | None = None
    mentioned_surfsense_doc_ids: list[int] | None = None
    disabled_tools: list[str] | None = None
    model_id: int | None = None  # Cloud mode: override with system model ID


# =============================================================================
# Agent Tools Schemas
# =============================================================================


class AgentToolInfo(BaseModel):
    """Schema for a single agent tool's public metadata."""

    name: str
    description: str
    enabled_by_default: bool


class ResumeDecision(BaseModel):
    type: Literal["approve", "edit", "reject"]
    edited_action: dict[str, Any] | None = None


class ResumeRequest(BaseModel):
    search_space_id: int
    decisions: list[ResumeDecision]
    model_id: int | None = None  # Cloud mode: override with system model ID


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


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

    snapshot_id: int
    share_token: str
    public_url: str
    is_new: bool


class PublicChatSnapshotInfo(BaseModel):
    """Info about a single public chat snapshot."""

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


class PublicChatSnapshotListResponse(BaseModel):
    """List of public chat snapshots for a thread."""

    snapshots: list[PublicChatSnapshotInfo]


class PublicChatSnapshotDetail(BaseModel):
    """Public chat snapshot with thread context."""

    id: int
    share_token: str
    public_url: str
    created_at: datetime
    message_count: int
    thread_id: int
    thread_title: str
    created_by_user_id: str | None = None


class PublicChatSnapshotsBySpaceResponse(BaseModel):
    """List of public chat snapshots for a search space."""

    snapshots: list[PublicChatSnapshotDetail]


# =============================================================================
# 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