Files
duocthu/apps/ai-service/rag/models.py
T

91 lines
2.1 KiB
Python

from __future__ import annotations
from dataclasses import dataclass, field
from enum import StrEnum
class EvidenceDecision(StrEnum):
ANSWERABLE = "answerable"
VERIFY_PDF = "verify_pdf"
ABSTAIN = "abstain"
class SubjectScope(StrEnum):
HUMAN = "human"
NON_HUMAN = "non_human"
UNKNOWN = "unknown"
class QueryIntent(StrEnum):
FACT_LOOKUP = "fact_lookup"
RECOMMENDATION = "recommendation"
UNKNOWN = "unknown"
@dataclass(frozen=True)
class SourceRef:
physical_page: int
precision: str
block_id: str | None = None
bbox: tuple[float, float, float, float] | None = None
source_crop: str | None = None
page_range: tuple[int, int] | None = None
printed_page: int | None = None
printed_page_range: tuple[int, int] | None = None
@dataclass(frozen=True)
class RetrievalDocument:
doc_id: str
drug_id: str
kind: str
text: str
section_key: str
source_refs: tuple[SourceRef, ...]
parent_id: str | None = None
requires_visual_check: bool = False
drug_name: str | None = None
part_index: int | None = None
part_count: int | None = None
context_labels: tuple[str, ...] = field(default_factory=tuple)
@dataclass(frozen=True)
class ParentDocument:
parent_id: str
kind: str
text: str
source_refs: tuple[SourceRef, ...]
requires_visual_check: bool = False
@dataclass(frozen=True)
class SearchHit:
document: RetrievalDocument
score: float
@dataclass(frozen=True)
class Evidence:
evidence_id: str
matched_doc_id: str
kind: str
text: str
score: float
source_refs: tuple[SourceRef, ...]
hydrated_from_parent: bool
requires_visual_check: bool
@dataclass(frozen=True)
class RetrievalResult:
decision: EvidenceDecision
reason: str
evidence: tuple[Evidence, ...] = field(default_factory=tuple)
resolved_drug_id: str | None = None
drug_resolution_status: str = "not_attempted"
# True when the user typed only a drug name (no attribute): the answer layer
# should introduce the drug (what it is + what it treats), not restate a
# section verbatim.
is_drug_overview: bool = False