AIRegulation-DocAnalysis/backend/app/domain/retrieval/models.py

"""Define domain models for retrieval."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any
# Keep module behavior explicit so the backend flow stays easy to audit.



@dataclass
class RetrievalQuery:
    """Represent the Retrieval Query type."""
    query: str
    top_k: int
    filters: str | None = None


@dataclass(init=False)
class RetrievedChunk:
    """Represent the retrieved chunk payload with legacy aliases."""

    chunk_id: str
    doc_id: str
    doc_title: str
    text: str
    score: float
    chunk_type: str = ""
    section_title: str = ""
    page_start: int = 0
    page_end: int = 0
    section_level: int = 0
    chunk_index: int = 0
    piece_index: int = 0
    metadata: dict[str, Any] = field(default_factory=dict)

    def __init__(
        self,
        *,
        chunk_id: str,
        doc_id: str,
        doc_title: str | None = None,
        text: str | None = None,
        score: float = 0.0,
        chunk_type: str = "",
        section_title: str = "",
        page_start: int = 0,
        page_end: int = 0,
        section_level: int = 0,
        chunk_index: int = 0,
        piece_index: int = 0,
        metadata: dict[str, Any] | None = None,
        doc_name: str | None = None,
        content: str | None = None,
        page_number: int | None = None,
        block_type: str | None = None,
        **_: Any,
    ) -> None:
        """Initialize the retrieved chunk while accepting legacy field names."""
        self.chunk_id = chunk_id
        self.doc_id = doc_id
        self.doc_title = doc_title if doc_title is not None else (doc_name or "")
        self.text = text if text is not None else (content or "")
        self.score = float(score)
        self.chunk_type = chunk_type or (block_type or "")
        self.section_title = section_title
        self.page_start = int(page_start or page_number or 0)
        self.page_end = int(page_end or self.page_start)
        self.section_level = int(section_level or 0)
        self.chunk_index = int(chunk_index or 0)
        self.piece_index = int(piece_index or 0)
        self.metadata = dict(metadata or {})

    @property
    def doc_name(self) -> str:
        """Return the legacy document name alias."""
        return self.doc_title

    @doc_name.setter
    def doc_name(self, value: str) -> None:
        """Update the legacy document name alias."""
        self.doc_title = value

    @property
    def content(self) -> str:
        """Return the legacy content alias."""
        return self.text

    @content.setter
    def content(self, value: str) -> None:
        """Update the legacy content alias."""
        self.text = value

    @property
    def page_number(self) -> int:
        """Return the legacy page number alias."""
        return self.page_start

    @page_number.setter
    def page_number(self, value: int) -> None:
        """Update the legacy page number alias."""
        self.page_start = value
        self.page_end = max(self.page_end, value)