Source code for evidence_fetcher.models

"""Public data models for normalized evidence results and feedback."""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import asdict, dataclass, replace
from enum import StrEnum
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from evidence_fetcher.ranking import FeedbackAwareEvidenceRanker




@dataclass(frozen=True, slots=True)
class Evidence:
    """A normalized search result returned by evidence-fetcher."""

    title: str
    url: str
    snippet: str
    source: str
    rank: int | None = None

    def to_dict(self) -> dict[str, object]:
        """Return a JSON-serializable representation of the evidence."""
        return asdict(self)





class FeedbackLabel(StrEnum):
    """A relevance judgment for an evidence result."""

    MORE = "more"
    RELEVANT = "relevant"
    LESS = "less"
    IRRELEVANT = "irrelevant"
    UNSURE = "unsure"





@dataclass(frozen=True, slots=True)
class RelevanceFeedback:
    """User feedback for one evidence result."""

    url: str
    label: FeedbackLabel





@dataclass(frozen=True, slots=True)
class SearchSession:
    """State needed for a future retrieve, rank, feedback, rerank workflow."""

    query: str
    candidates: tuple[Evidence, ...]
    ranked_results: tuple[Evidence, ...]
    feedback: tuple[RelevanceFeedback, ...] = ()

    def add_feedback(self, feedback: RelevanceFeedback) -> SearchSession:
        """Return a session with one validated feedback entry appended."""
        resolved_url = _resolve_evidence_url(feedback.url, self.candidates)
        if resolved_url is None:
            raise ValueError("feedback URL is not present in the search session")
        resolved_feedback = replace(feedback, url=resolved_url)
        return replace(self, feedback=(*self.feedback, resolved_feedback))

    def with_feedback(self, url: str, label: FeedbackLabel) -> SearchSession:
        """Return a session with feedback for the evidence identified by URL."""
        return self.add_feedback(RelevanceFeedback(url=url, label=label))

    def rerank(
        self,
        ranker: FeedbackAwareEvidenceRanker | None = None,
    ) -> SearchSession:
        """Return a session with ranked results revised by accumulated feedback."""
        from evidence_fetcher.ranking import TfidfEvidenceRanker

        evidence_ranker = ranker if ranker is not None else TfidfEvidenceRanker()
        ranked = evidence_ranker.rerank(self.query, self.candidates, self.feedback)
        return replace(self, ranked_results=tuple(ranked))

    def with_evidence(self, evidence: Sequence[Evidence]) -> SearchSession:
        """Return a session with new evidence merged by canonical URL."""
        seen_urls = {_canonical_url_key(item.url) for item in self.candidates}
        candidates = list(self.candidates)
        ranked_results = list(self.ranked_results)

        for item in evidence:
            key = _canonical_url_key(item.url)
            if key in seen_urls:
                continue
            seen_urls.add(key)
            candidates.append(item)
            ranked_results.append(item)

        return replace(
            self,
            candidates=tuple(candidates),
            ranked_results=tuple(ranked_results),
        )



def _canonical_url_key(url: str) -> str:
    return url.rstrip("/").casefold()


def _resolve_evidence_url(url: str, evidence: Sequence[Evidence]) -> str | None:
    key = _canonical_url_key(url)
    for item in evidence:
        if _canonical_url_key(item.url) == key:
            return item.url
    return None