models.py

"""
GCAS Search Engine – Pydantic request / response models
"""
from __future__ import annotations
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, Field


# ---------------------------------------------------------------------------
# Request
# ---------------------------------------------------------------------------

class SearchRequest(BaseModel):
    """Body for POST /search"""

    query: str = Field(
        ...,
        description="Natural-language search query, e.g. 'engineering colleges in Ahmedabad with hostel'"
    )
    top_k: int = Field(
        default=10, ge=1, le=100,
        description="Number of results to return"
    )
    tables: Optional[List[str]] = Field(
        default=None,
        description=(
            "Restrict search to specific table names (stem of the Excel filename). "
            "null / omit = search all indexed tables."
        )
    )
    use_llm_rerank: bool = Field(
        default=True,
        description="Pass FAISS candidates through an LLM for semantic reranking"
    )
    # Per-request LLM overrides (useful when caller wants to supply its own key/model)
    llm_provider: Optional[Literal["openai", "anthropic"]] = Field(
        default=None, description="Override the server-level LLM provider"
    )
    llm_model: Optional[str] = Field(
        default=None, description="Override the server-level LLM model name"
    )
    api_key: Optional[str] = Field(
        default=None, description="Override API key for the chosen LLM provider"
    )


# ---------------------------------------------------------------------------
# Search response
# ---------------------------------------------------------------------------

class SearchResult(BaseModel):
    table: str = Field(..., description="Source table (Excel filename stem)")
    row_index: int = Field(..., description="Original row position in the Excel file")
    score: float = Field(..., description="Relevance score (higher = more relevant)")
    llm_reason: Optional[str] = Field(
        default=None, description="LLM explanation of why this result matches"
    )
    data: Dict[str, Any] = Field(..., description="Full row data as key-value pairs")


class EntityCorrection(BaseModel):
    original_span: str  = Field(..., description="The token as the user typed/spoke it")
    corrected_to:  str  = Field(..., description="Canonical entity name from the database")
    entity_type:   str  = Field(..., description="college | university | district | taluka | program | subject")
    match_score:   float = Field(..., description="Fuzzy match score 0-100")
    method:        str  = Field(..., description="exact | fuzzy | phonetic")


class SearchResponse(BaseModel):
    query: str = Field(..., description="Original query as submitted")
    total_results: int
    results: List[SearchResult]
    search_time_ms: float
    reranked: bool = Field(default=False, description="True if LLM reranking was applied")

    # ── Multilingual / ASR intelligence ──────────────────────────────────────
    detected_language: str = Field(
        default="en",
        description="Detected query language: en | hi | gu | unknown"
    )
    corrected_query: Optional[str] = Field(
        default=None,
        description="Query after alias resolution, transliteration, and entity correction. "
                    "Null if no corrections were needed."
    )
    entity_corrections: List[EntityCorrection] = Field(
        default_factory=list,
        description="List of entity spelling / ASR corrections applied to the query"
    )
    confidence_level: str = Field(
        default="high",
        description="Result confidence: high | medium | low"
    )
    detected_intent: str = Field(
        default="general",
        description=(
            "Inferred query intent: fees | hostel | cutoff | facilities | "
            "courses | naac | contact | general. "
            "Determines which fields are returned in each result."
        )
    )
    did_you_mean: List[str] = Field(
        default_factory=list,
        description="Suggestions shown when confidence is low, e.g. 'Did you mean GTU?'"
    )


# ---------------------------------------------------------------------------
# System endpoints
# ---------------------------------------------------------------------------

class ReindexResponse(BaseModel):
    status: str
    tables_indexed: List[str]
    total_rows_indexed: int
    time_taken_ms: float


class TableSchema(BaseModel):
    columns: List[str]
    row_count: int
    file: str


class SchemaResponse(BaseModel):
    tables: Dict[str, TableSchema]


class HealthResponse(BaseModel):
    status: str
    indexed_tables: List[str]
    total_indexed_rows: int
    embedding_provider: str
    llm_provider: str