app.py

from typing import List, Optional

from fastapi import FastAPI, Query
from pydantic import BaseModel, Field

from utils import (
    semantic_search,
    get_schemas,
    get_model_id,
    get_corpus_size,
)


app = FastAPI(title="Semantic Table Search API", version="1.0.0")


class SearchRequest(BaseModel):
    """Client request schema for semantic search.

    Attributes:
    - query: Natural language search text.
    - limit: Optional max number of results to return (fallback to query param
      when not provided). Included here to allow sending within body when
      clients prefer JSON-only interactions.
    """

    query: str = Field(..., min_length=1)
    limit: Optional[int] = Field(None, ge=1, le=50)


class Match(BaseModel):
    """A single search match result."""

    score: float = Field(..., description="Cosine similarity score (-1 to 1)")
    text: str = Field(..., description="Matched table metadata text")
    index: int = Field(..., description="Stable index of the matched corpus entry")


class SearchResponse(BaseModel):
    """Search results with basic service metadata."""

    query: str
    results: List[Match]
    count: int
    limit: int


@app.get("/health")
def health() -> dict:
    """Basic health check including corpus size.

    Returns a JSON indicating the service is up along with a few
    diagnostic fields.
    """

    return {
        "status": "ok",
        "corpus_size": get_corpus_size(),
        "model": get_model_id(),
    }


@app.get("/schemas")
def schemas(
    include_columns: bool = Query(False, description="Include column metadata"),
) -> List[dict]:
    """List available table schemas and optional column metadata.

    Parameters:
    - include_columns: When true, return full schema definitions; otherwise a
        compact view containing table names and descriptions is returned.
    """

    return get_schemas(include_columns=include_columns)


@app.post("/search", response_model=SearchResponse)
def search(
    body: SearchRequest,
    limit: int = Query(5, ge=1, le=50, description="Max number of results"),
) -> SearchResponse:
    """Perform a semantic search over table metadata and return ranked matches.

    The clients provide a natural language
    query and receive the most relevant tables with similarity
    scores and stable corpus indices.
    """

    effective_limit = body.limit or limit
    results = semantic_search(body.query, top_k=effective_limit)
    matches = [
        {"score": score, "text": text, "index": idx} for score, text, idx in results
    ]
    return SearchResponse(
        query=body.query, results=matches, count=len(matches), limit=effective_limit
    )