Heimgeist/backend/schemas.py

from pydantic import BaseModel, ConfigDict
from typing import List, Optional
from datetime import datetime

class ImageAttachment(BaseModel):
    name: str
    mime_type: Optional[str] = None
    data_url: str

class Message(BaseModel):
    role: str
    content: str
    sources: Optional[List[str]] = None
    attachments: Optional[List[ImageAttachment]] = None

class ChatRequest(BaseModel):
    session_id: str
    model: str
    message: str
    enriched_message: Optional[str] = None
    stream: Optional[bool] = False
    sources: Optional[List[str]] = None
    attachments: Optional[List[ImageAttachment]] = None

class ChatResponse(BaseModel):
    reply: str

class HistoryResponse(BaseModel):
    messages: List[Message]

class GenerateTitleRequest(BaseModel):
    session_id: str
    message: str
    model: str

class GenerateTitleResponse(BaseModel):
    title: str

class CreateSessionRequest(BaseModel):
    session_id: str

class ChatSession(BaseModel):
    id: int
    session_id: str
    name: str
    created_at: datetime

    model_config = ConfigDict(from_attributes=True)

class SessionsResponse(BaseModel):
    sessions: List[ChatSession]

class EditMessageRequest(BaseModel):
    message: str

class RegenerateRequest(BaseModel):
    index: int
    model: Optional[str] = None
    enriched_message: Optional[str] = None
    stream: bool = True
    sources: Optional[List[str]] = None

# Request payload for the web search enrichment endpoint.
class WebSearchRequest(BaseModel):
    prompt: str
    model: str
    messages: Optional[List[Message]] = None
    history_limit: Optional[int] = 8
    searx_url: Optional[str] = None
    engines: Optional[List[str]] = None
    
# Response payload for the web search enrichment endpoint.
class WebSearchResponse(BaseModel):
    enriched_prompt: str
    sources: List[str] = []
    context_block: str = ""


class OllamaPullRequest(BaseModel):
    model: Optional[str] = None


class AudioTranscriptionRequest(BaseModel):
    mime_type: str
    audio_base64: str
    model: Optional[str] = None
    language: Optional[str] = None


class AudioTranscriptionResponse(BaseModel):
    text: str
    language: Optional[str] = None
    model: str