query_engine/models.py

"""
models.py
Pydantic models for the RAG query pipeline.
"""

from typing import List
from pydantic import BaseModel


class SearchResult(BaseModel):
    """A single chunk retrieved from Qdrant."""
    chunk_id:  str
    text:      str
    source:    str
    file_path: str
    score:     float


class RankedResult(BaseModel):
    """A chunk after Cohere reranking."""
    chunk_id:     str
    text:         str
    source:       str
    file_path:    str
    rerank_score: float


class QueryResponse(BaseModel):
    """Full response returned to the Gradio UI."""
    answer:  str
    sources: List[RankedResult]

    def format_for_gradio(self) -> str:
        """
        Format as a single markdown string for Gradio ChatInterface.
        Answer followed by a deduplicated numbered source list.
        """
        lines = [self.answer, "\n**Sources:**"]
        seen = set()
        for i, src in enumerate(self.sources, 1):
            if src.file_path in seen:
                continue
            seen.add(src.file_path)
            lines.append(f"{i}. `{src.source}` — {src.file_path}")
        return "\n".join(lines)