openrag/sdks/python/openrag_sdk/models.py

"""OpenRAG SDK data models."""

from datetime import datetime
from typing import Any, Literal

from pydantic import BaseModel, Field


# Chat models
class Source(BaseModel):
    """A source document returned in chat/search results."""

    filename: str
    text: str
    score: float
    page: int | None = None
    mimetype: str | None = None


class ChatResponse(BaseModel):
    """Response from a non-streaming chat request."""

    response: str
    chat_id: str | None = None
    sources: list[Source] = Field(default_factory=list)


class StreamEvent(BaseModel):
    """Base class for streaming events."""

    type: Literal["content", "sources", "done"]


class ContentEvent(StreamEvent):
    """A content delta event during streaming."""

    type: Literal["content"] = "content"
    delta: str


class SourcesEvent(StreamEvent):
    """A sources event containing retrieved documents."""

    type: Literal["sources"] = "sources"
    sources: list[Source]


class DoneEvent(StreamEvent):
    """Indicates the stream is complete."""

    type: Literal["done"] = "done"
    chat_id: str | None = None


# Search models
class SearchResult(BaseModel):
    """A single search result."""

    filename: str
    text: str
    score: float
    page: int | None = None
    mimetype: str | None = None


class SearchResponse(BaseModel):
    """Response from a search request."""

    results: list[SearchResult]


# Document models
class IngestResponse(BaseModel):
    """Response from document ingestion (async task-based)."""

    task_id: str
    status: str | None = None  # Optional - we poll for actual status
    filename: str | None = None


class IngestTaskStatus(BaseModel):
    """Status of an ingestion task."""

    task_id: str
    status: str  # "pending", "running", "completed", "failed"
    total_files: int = 0
    processed_files: int = 0
    successful_files: int = 0
    failed_files: int = 0
    files: dict = {}  # Detailed per-file status


class DeleteDocumentResponse(BaseModel):
    """Response from document deletion."""

    success: bool
    deleted_chunks: int = 0


# Chat history models
class Message(BaseModel):
    """A message in a conversation."""

    role: str
    content: str
    timestamp: str | None = None


class Conversation(BaseModel):
    """A conversation summary."""

    chat_id: str
    title: str = ""
    created_at: str | None = None
    last_activity: str | None = None
    message_count: int = 0


class ConversationDetail(Conversation):
    """A conversation with full message history."""

    messages: list[Message] = Field(default_factory=list)


class ConversationListResponse(BaseModel):
    """Response from listing conversations."""

    conversations: list[Conversation]


# Settings models
class AgentSettings(BaseModel):
    """Agent configuration settings."""

    llm_provider: str | None = None
    llm_model: str | None = None


class KnowledgeSettings(BaseModel):
    """Knowledge base configuration settings."""

    embedding_provider: str | None = None
    embedding_model: str | None = None
    chunk_size: int | None = None
    chunk_overlap: int | None = None


class SettingsResponse(BaseModel):
    """Response from settings endpoint."""

    agent: AgentSettings = Field(default_factory=AgentSettings)
    knowledge: KnowledgeSettings = Field(default_factory=KnowledgeSettings)


# Request models
class SearchFilters(BaseModel):
    """Filters for search requests."""

    data_sources: list[str] | None = None
    document_types: list[str] | None = None


# Settings update models
class SettingsUpdateOptions(BaseModel):
    """Options for updating settings."""

    llm_model: str | None = None
    llm_provider: str | None = None
    system_prompt: str | None = None
    embedding_model: str | None = None
    embedding_provider: str | None = None
    chunk_size: int | None = None
    chunk_overlap: int | None = None
    table_structure: bool | None = None
    ocr: bool | None = None
    picture_descriptions: bool | None = None


class SettingsUpdateResponse(BaseModel):
    """Response from settings update."""

    message: str


# Knowledge filter models
class KnowledgeFilterQueryData(BaseModel):
    """Query configuration stored in a knowledge filter."""

    query: str | None = None
    filters: dict[str, list[str]] | None = None
    limit: int | None = None
    score_threshold: float | None = Field(default=None, alias="scoreThreshold")
    color: str | None = None
    icon: str | None = None

    model_config = {"populate_by_name": True}


class KnowledgeFilter(BaseModel):
    """A knowledge filter definition."""

    id: str
    name: str
    description: str | None = None
    query_data: KnowledgeFilterQueryData | None = Field(default=None, alias="queryData")
    owner: str | None = None
    created_at: str | None = Field(default=None, alias="createdAt")
    updated_at: str | None = Field(default=None, alias="updatedAt")

    model_config = {"populate_by_name": True}


class CreateKnowledgeFilterOptions(BaseModel):
    """Options for creating a knowledge filter."""

    name: str
    description: str | None = None
    query_data: KnowledgeFilterQueryData = Field(alias="queryData")

    model_config = {"populate_by_name": True}


class UpdateKnowledgeFilterOptions(BaseModel):
    """Options for updating a knowledge filter."""

    name: str | None = None
    description: str | None = None
    query_data: KnowledgeFilterQueryData | None = Field(default=None, alias="queryData")

    model_config = {"populate_by_name": True}


class CreateKnowledgeFilterResponse(BaseModel):
    """Response from creating a knowledge filter."""

    success: bool
    id: str | None = None
    error: str | None = None


class KnowledgeFilterSearchResponse(BaseModel):
    """Response from searching knowledge filters."""

    success: bool
    filters: list[KnowledgeFilter] = Field(default_factory=list)


class GetKnowledgeFilterResponse(BaseModel):
    """Response from getting a knowledge filter."""

    success: bool
    filter: KnowledgeFilter | None = None
    error: str | None = None


class DeleteKnowledgeFilterResponse(BaseModel):
    """Response from deleting a knowledge filter."""

    success: bool
    error: str | None = None