AIRegulation-DocAnalysis/backend/app/schemas/rag.py

"""Define schema models for rag."""

from pydantic import BaseModel
from typing import Optional
# Group related schema definitions so validation rules stay consistent.



class RagChatRequest(BaseModel):
    """Define the Rag Chat Request API model."""
    query: str
    top_k: int = 5
    session_id: Optional[str] = None
    filters: Optional[str] = None
    # Optional document text to inject directly as LLM conversation context.
    # When provided the document content is prepended to the query so the LLM
    # can answer questions about it without requiring vector-store indexing.
    context_text: Optional[str] = None
    context_filename: Optional[str] = None


class RetrievedDoc(BaseModel):
    """Define the Retrieved Doc API model."""
    id: str
    doc_name: str
    clause_id: Optional[str] = None
    score: float
    content: str
    preview: str


class SourceInfo(BaseModel):
    """Define the Source Info API model."""
    name: str
    clause: Optional[str] = None


class QuickQuestion(BaseModel):
    """Define the Quick Question API model."""
    id: str
    question: str
    category: str


class QuickQuestionsResponse(BaseModel):
    """Define the Quick Questions Response API model."""
    questions: list[QuickQuestion]