Newer
Older
cortex-hub / ai-hub / app / api / schemas.py
from pydantic import BaseModel, Field
from typing import List, Literal, Optional
from datetime import datetime  # <-- Add this import

# --- Chat Schemas ---

class ChatRequest(BaseModel):
    """Defines the shape of a request to the /chat endpoint."""
    prompt: str = Field(..., min_length=1)
    model: Literal["deepseek", "gemini"] = Field("deepseek")

class ChatResponse(BaseModel):
    """Defines the shape of a successful response from the /chat endpoint."""
    answer: str
    model_used: str

# --- Document Schemas ---

class DocumentCreate(BaseModel):
    """Defines the shape for creating a new document."""
    title: str
    text: str
    source_url: Optional[str] = None
    author: Optional[str] = None
    user_id: str = "default_user"

class DocumentResponse(BaseModel):
    """Defines the response after creating a document."""
    message: str

class DocumentInfo(BaseModel):
    id: int
    title: str
    source_url: Optional[str] = None
    status: str
    created_at: datetime

class DocumentListResponse(BaseModel):
    documents: List[DocumentInfo]

class DocumentDeleteResponse(BaseModel):
    message: str
    document_id: int