app/backend/metadata_schema.py

"""
Metadata schema for FinBot RAG system.
Defines data structures for chunks, users, and retrieval results.
"""

from dataclasses import dataclass, field, asdict
from typing import Optional, List
from enum import Enum


class ChunkType(str, Enum):
    """Type of content in a chunk."""
    TEXT = "text"
    TABLE = "table"
    HEADING = "heading"
    CODE = "code"


@dataclass
class Chunk:
    """
    Represents a hierarchically-chunked document segment.
    
    This is the fundamental unit stored in the vector database.
    Each chunk carries metadata about its source, hierarchy, and access controls.
    """
    # Content
    id: str  # Unique identifier (e.g., "doc_name_chunk_0")
    text: str  # The actual text content of this chunk
    
    # Document source metadata (REQUIRED)
    source_document: str  # Filename (e.g., "system_architecture.md")
    collection: str  # Collection name (general, finance, engineering, marketing, hr)
    access_roles: List[str]  # Roles that can access this chunk (e.g., ["engineering", "c_level"])
    
    # Hierarchical structure metadata
    section_title: Optional[str] = None  # Parent section heading
    subsection_title: Optional[str] = None  # Sub-heading if applicable
    page_number: Optional[int] = None  # Page number in source document
    chunk_type: ChunkType = ChunkType.TEXT  # Type of content (text, table, heading, code)
    parent_chunk_id: Optional[str] = None  # ID of parent section chunk for hierarchy
    parent_summary: Optional[str] = None  # Summary of parent section
    
    # For tracking hierarchy depth
    depth: int = 0  # Depth in document tree (0 = root)
    
    # Embedding (populated after vectorization)
    embedding: Optional[List[float]] = field(default_factory=list)
    
    def to_qdrant_payload(self) -> dict:
        """
        Convert chunk to Qdrant payload format.
        Used when storing in vector database.
        """
        return {
            "source_document": self.source_document,
            "collection": self.collection,
            "access_roles": self.access_roles,
            "section_title": self.section_title or "",
            "subsection_title": self.subsection_title or "",
            "page_number": self.page_number or 0,
            "chunk_type": self.chunk_type.value,
            "parent_chunk_id": self.parent_chunk_id or "",
            "parent_summary": self.parent_summary or "",
            "depth": self.depth,
            "text": self.text,
        }
    
    def to_dict(self) -> dict:
        """Convert chunk to dictionary (excludes embedding)."""
        return asdict(self)


@dataclass
class User:
    """
    Represents a FinSolve employee with role and permissions.
    """
    username: str
    name: str
    role: str  # UserRole enum value (employee, finance, engineering, marketing, c_level)
    department: str
    
    def to_dict(self) -> dict:
        """Convert user to dictionary."""
        return asdict(self)


@dataclass
class QueryMetadata:
    """
    Metadata captured for every query for auditing and logging.
    """
    user_role: str
    user_department: str
    query_text: str
    route_selected: str
    collections_queried: List[str]
    chunks_retrieved: int
    guardrail_flags: List[str] = field(default_factory=list)  # e.g., ["prompt_injection_detected"]
    rbac_denied: bool = False
    answer: Optional[str] = None
    sources: List[str] = field(default_factory=list)  # List of source doc names
    
    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return asdict(self)


@dataclass
class RetrievalResult:
    """
    Result from a RBAC-checked retrieval operation.
    """
    chunks: List[Chunk]
    rbac_passed: bool
    reason: Optional[str] = None  # If RBAC failed, explain why
    
    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "chunks": [c.to_dict() for c in self.chunks],
            "rbac_passed": self.rbac_passed,
            "reason": self.reason,
        }


@dataclass
class RAGResponse:
    """
    Final response from the RAG pipeline.
    Contains answer, sources, metadata, and any warnings.
    """
    answer: str
    sources: List[dict]  # List of {document, page_number, section_title}
    route: str
    user_role: str
    accessible_collections: List[str]
    guardrail_flags: List[str] = field(default_factory=list)
    guardrail_warnings: List[str] = field(default_factory=list)
    rbac_denied: bool = False
    rbac_reason: Optional[str] = None
    
    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return asdict(self)


# Validation helpers

def validate_chunk_metadata(chunk: Chunk) -> tuple[bool, str]:
    """
    Validate that a chunk has all required metadata.
    Returns (is_valid, error_message).
    """
    errors = []
    
    if not chunk.id:
        errors.append("Chunk id is required")
    if not chunk.text:
        errors.append("Chunk text is required")
    if not chunk.source_document:
        errors.append("source_document is required")
    if not chunk.collection:
        errors.append("collection is required")
    if not chunk.access_roles or len(chunk.access_roles) == 0:
        errors.append("access_roles must not be empty")
    
    if errors:
        return False, "; ".join(errors)
    return True, ""


def validate_user(user: User) -> tuple[bool, str]:
    """
    Validate that a user has required fields.
    Returns (is_valid, error_message).
    """
    errors = []
    
    if not user.username:
        errors.append("username is required")
    if not user.name:
        errors.append("name is required")
    if not user.role:
        errors.append("role is required")
    if not user.department:
        errors.append("department is required")
    
    if errors:
        return False, "; ".join(errors)
    return True, ""