models.py

"""
Pydantic models for the Dynamic RAG Bot.
Simplified schema for main bot and document loader.
"""

from pydantic import BaseModel, Field
from typing import List, Optional


class Source(BaseModel):
    """Citation information for document sources."""
    file: str = Field(description="Source markdown file name (e.g., 'france.md')")
    section: str = Field(description="Section heading from the document")
    line_start: int = Field(description="Starting line number of the cited content")
    line_end: int = Field(description="Ending line number of the cited content")
    quote: str = Field(description="Relevant excerpt from the document")
    
    def format_citation(self) -> str:
        """Format source as [file:section:line_start-line_end]"""
        return f"[{self.file}:{self.section}:{self.line_start}-{self.line_end}]"


class BotResponse(BaseModel):
    """Main bot's response with thoughts, action, and content."""
    thoughts: str = Field(description="Bot's step-by-step thinking about the query")
    action: str = Field(description="Action to take: 'answer' | 'ask_specifics' | 'fetch:country:question'")
    content: str = Field(description="Response to show user (empty string if fetch action)")


class DocumentResponse(BaseModel):
    """Document loader's structured response with content and sources."""
    content: str = Field(description="Extracted information from the document")
    sources: List[Source] = Field(default=[], description="Citations from the document")
    
    def format_with_sources(self) -> str:
        """Format content with appended source citations."""
        if not self.sources:
            return self.content
        
        citations = " ".join([source.format_citation() for source in self.sources])
        return f"{self.content}\n\nSources: {citations}"