294 lines
11 KiB
Python
294 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
from dataclasses import dataclass, replace
|
|
from difflib import SequenceMatcher
|
|
from enum import StrEnum
|
|
|
|
from .models import EvidenceDecision, QueryIntent, RetrievalResult, SubjectScope
|
|
from .service import RetrievalService
|
|
from .text import WORD_RE, normalize_name
|
|
|
|
__all__ = [
|
|
"WORD_RE",
|
|
"CatalogDrugResolver",
|
|
"DrugResolution",
|
|
"DrugResolutionStatus",
|
|
"QueryRoutingService",
|
|
"normalize_name",
|
|
]
|
|
|
|
|
|
class DrugResolutionStatus(StrEnum):
|
|
RESOLVED = "resolved"
|
|
NOT_FOUND = "not_found"
|
|
AMBIGUOUS = "ambiguous"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DrugResolution:
|
|
status: DrugResolutionStatus
|
|
drug_id: str | None = None
|
|
score: float | None = None
|
|
candidate_drug_ids: tuple[str, ...] = ()
|
|
|
|
|
|
class CatalogDrugResolver:
|
|
def __init__(
|
|
self,
|
|
catalog: dict[str, set[str]],
|
|
fuzzy_threshold: float = 0.84,
|
|
ambiguity_margin: float = 0.04,
|
|
) -> None:
|
|
self._catalog = {
|
|
drug_id: {
|
|
normalized for alias in aliases if (normalized := normalize_name(alias))
|
|
}
|
|
for drug_id, aliases in catalog.items()
|
|
}
|
|
self._aliases = [
|
|
(drug_id, normalized)
|
|
for drug_id, aliases in catalog.items()
|
|
for alias in aliases
|
|
if (normalized := normalize_name(alias))
|
|
]
|
|
self._fuzzy_threshold = fuzzy_threshold
|
|
self._ambiguity_margin = ambiguity_margin
|
|
|
|
def resolve(self, query: str) -> DrugResolution:
|
|
normalized_query = normalize_name(query)
|
|
query_tokens = normalized_query.split()
|
|
exact = [
|
|
(drug_id, alias, match.start(1), match.end(1))
|
|
for drug_id, alias in self._aliases
|
|
for match in [
|
|
re.search(rf"(?:^| )({re.escape(alias)})(?:$| )", normalized_query)
|
|
]
|
|
if match
|
|
]
|
|
if exact:
|
|
maximal = [
|
|
row for row in exact
|
|
if not any(
|
|
other[2] <= row[2] and row[3] <= other[3]
|
|
and (other[2], other[3]) != (row[2], row[3])
|
|
for other in exact
|
|
)
|
|
]
|
|
drug_ids = {drug_id for drug_id, _, _, _ in maximal}
|
|
if len(drug_ids) == 1:
|
|
return DrugResolution(
|
|
DrugResolutionStatus.RESOLVED, next(iter(drug_ids)), 1.0,
|
|
)
|
|
return DrugResolution(
|
|
DrugResolutionStatus.AMBIGUOUS,
|
|
candidate_drug_ids=tuple(sorted(drug_ids)),
|
|
)
|
|
|
|
scores: dict[str, float] = {}
|
|
for drug_id, alias in self._aliases:
|
|
width = len(alias.split())
|
|
if width > len(query_tokens):
|
|
continue
|
|
spans = (
|
|
" ".join(query_tokens[start:start + width])
|
|
for start in range(len(query_tokens) - width + 1)
|
|
)
|
|
score = max(
|
|
(SequenceMatcher(None, alias, span).ratio() for span in spans),
|
|
default=0.0,
|
|
)
|
|
scores[drug_id] = max(scores.get(drug_id, 0.0), score)
|
|
ranked = sorted(scores.items(), key=lambda item: (-item[1], item[0]))
|
|
if not ranked or ranked[0][1] < self._fuzzy_threshold:
|
|
return DrugResolution(DrugResolutionStatus.NOT_FOUND)
|
|
if len(ranked) > 1 and ranked[0][1] - ranked[1][1] < self._ambiguity_margin:
|
|
return DrugResolution(
|
|
DrugResolutionStatus.AMBIGUOUS,
|
|
candidate_drug_ids=(ranked[0][0], ranked[1][0]),
|
|
)
|
|
return DrugResolution(DrugResolutionStatus.RESOLVED, *ranked[0])
|
|
|
|
def aliases_for(self, drug_id: str) -> set[str]:
|
|
return self._catalog.get(drug_id, set())
|
|
|
|
def complete(self, prefix: str, k: int = 8) -> list[str]:
|
|
"""Drug ids whose alias contains `prefix`, for as-you-type autocomplete.
|
|
|
|
Substring match on the normalized alias, ranked prefix-first then by
|
|
alias length, so "para" surfaces "paracetamol" ahead of a drug that only
|
|
contains "para" mid-word. Distinct drug ids, best first.
|
|
"""
|
|
needle = normalize_name(prefix)
|
|
if not needle:
|
|
return []
|
|
matches: list[tuple[tuple[int, int], str]] = []
|
|
for drug_id, alias in self._aliases:
|
|
position = alias.find(needle)
|
|
if position < 0:
|
|
continue
|
|
matches.append(((0 if position == 0 else 1, len(alias)), drug_id))
|
|
matches.sort()
|
|
ordered: list[str] = []
|
|
seen: set[str] = set()
|
|
for _, drug_id in matches:
|
|
if drug_id in seen:
|
|
continue
|
|
seen.add(drug_id)
|
|
ordered.append(drug_id)
|
|
if len(ordered) >= k:
|
|
break
|
|
return ordered
|
|
|
|
def suggest(
|
|
self, query: str, k: int = 3, min_score: float = 0.5
|
|
) -> list[tuple[str, float]]:
|
|
"""Closest drug ids by fuzzy score, for a 'did you mean' on a miss.
|
|
|
|
Uses the same windowed SequenceMatcher scoring as `resolve`, but returns
|
|
the top-k *below* the resolution threshold too, so a typo that does not
|
|
confidently resolve ("metfomin") can still be offered as a suggestion.
|
|
`min_score` keeps a genuinely non-drug query ("cái này thế nào") from
|
|
surfacing spurious suggestions.
|
|
"""
|
|
query_tokens = normalize_name(query).split()
|
|
if not query_tokens:
|
|
return []
|
|
scores: dict[str, float] = {}
|
|
for drug_id, alias in self._aliases:
|
|
width = len(alias.split())
|
|
if width > len(query_tokens):
|
|
continue
|
|
spans = (
|
|
" ".join(query_tokens[start:start + width])
|
|
for start in range(len(query_tokens) - width + 1)
|
|
)
|
|
score = max(
|
|
(SequenceMatcher(None, alias, span).ratio() for span in spans),
|
|
default=0.0,
|
|
)
|
|
scores[drug_id] = max(scores.get(drug_id, 0.0), score)
|
|
ranked = sorted(scores.items(), key=lambda item: (-item[1], item[0]))
|
|
return [(drug_id, score) for drug_id, score in ranked[:k] if score >= min_score]
|
|
|
|
|
|
class QueryRoutingService:
|
|
def __init__(
|
|
self,
|
|
retrieval: RetrievalService,
|
|
resolver: CatalogDrugResolver,
|
|
) -> None:
|
|
self._retrieval = retrieval
|
|
self._resolver = resolver
|
|
|
|
@staticmethod
|
|
def _scope_gate(
|
|
subject_scope: SubjectScope, intent: QueryIntent
|
|
) -> RetrievalResult | None:
|
|
# Scope comes from the API/policy layer. Unknown is deliberately
|
|
# fail-closed; retrieval must not infer clinical scope from keywords.
|
|
if subject_scope == SubjectScope.NON_HUMAN:
|
|
return RetrievalResult(EvidenceDecision.ABSTAIN, "out_of_scope_non_human")
|
|
if subject_scope == SubjectScope.UNKNOWN:
|
|
return RetrievalResult(EvidenceDecision.ABSTAIN, "subject_scope_unknown")
|
|
if intent == QueryIntent.RECOMMENDATION:
|
|
return RetrievalResult(EvidenceDecision.ABSTAIN, "recommendation_out_of_scope")
|
|
if intent == QueryIntent.UNKNOWN:
|
|
return RetrievalResult(EvidenceDecision.ABSTAIN, "query_intent_unknown")
|
|
return None
|
|
|
|
def retrieve_for_drug(
|
|
self,
|
|
query: str,
|
|
drug_id: str,
|
|
subject_scope: SubjectScope = SubjectScope.UNKNOWN,
|
|
intent: QueryIntent = QueryIntent.UNKNOWN,
|
|
) -> RetrievalResult:
|
|
"""Retrieve for an ALREADY-resolved drug, skipping name resolution.
|
|
|
|
The conversational layer has already resolved (and possibly inherited)
|
|
the drug; re-resolving from the rewritten turn text is what produced the
|
|
`drug_resolution_ambiguous` empty answers on follow-ups. The query text
|
|
still drives section routing and the intro/overview decision.
|
|
"""
|
|
gate = self._scope_gate(subject_scope, intent)
|
|
if gate is not None:
|
|
return gate
|
|
result = self._retrieval.retrieve(query, drug_id)
|
|
return replace(
|
|
result,
|
|
resolved_drug_id=drug_id,
|
|
drug_resolution_status=DrugResolutionStatus.RESOLVED,
|
|
)
|
|
|
|
def retrieve(
|
|
self,
|
|
query: str,
|
|
subject_scope: SubjectScope = SubjectScope.UNKNOWN,
|
|
intent: QueryIntent = QueryIntent.UNKNOWN,
|
|
) -> RetrievalResult:
|
|
gate = self._scope_gate(subject_scope, intent)
|
|
if gate is not None:
|
|
return gate
|
|
resolution = self._resolver.resolve(query)
|
|
if resolution.status == DrugResolutionStatus.NOT_FOUND:
|
|
return RetrievalResult(
|
|
EvidenceDecision.ABSTAIN,
|
|
"drug_not_resolved",
|
|
drug_resolution_status=resolution.status,
|
|
)
|
|
if resolution.status == DrugResolutionStatus.AMBIGUOUS:
|
|
disambiguated = self._disambiguate_with_evidence(
|
|
query, resolution.candidate_drug_ids,
|
|
)
|
|
if disambiguated is not None:
|
|
drug_id, result = disambiguated
|
|
return replace(
|
|
result,
|
|
resolved_drug_id=drug_id,
|
|
drug_resolution_status=DrugResolutionStatus.RESOLVED,
|
|
)
|
|
return RetrievalResult(
|
|
EvidenceDecision.ABSTAIN, "drug_resolution_ambiguous",
|
|
drug_resolution_status=resolution.status,
|
|
)
|
|
if resolution.drug_id is None:
|
|
return RetrievalResult(
|
|
EvidenceDecision.ABSTAIN,
|
|
"drug_resolution_invalid_state",
|
|
drug_resolution_status="invalid_state",
|
|
)
|
|
result = self._retrieval.retrieve(query, resolution.drug_id)
|
|
return replace(
|
|
result,
|
|
resolved_drug_id=resolution.drug_id,
|
|
drug_resolution_status=resolution.status,
|
|
)
|
|
|
|
def _disambiguate_with_evidence(
|
|
self,
|
|
query: str,
|
|
candidate_ids: tuple[str, ...],
|
|
) -> tuple[str, RetrievalResult] | None:
|
|
"""Resolve subject-vs-component ambiguity through asymmetric evidence.
|
|
|
|
A candidate wins only when its retrieved evidence explicitly contains
|
|
every other mentioned entity, while the reverse direction does not.
|
|
This keeps genuine multi-drug questions ambiguous.
|
|
"""
|
|
if len(candidate_ids) < 2:
|
|
return None
|
|
winners = []
|
|
for candidate_id in candidate_ids:
|
|
result = self._retrieval.retrieve(query, candidate_id)
|
|
if result.decision == EvidenceDecision.ABSTAIN:
|
|
continue
|
|
evidence_text = normalize_name(" ".join(item.text for item in result.evidence))
|
|
others = [item for item in candidate_ids if item != candidate_id]
|
|
if all(any(
|
|
re.search(rf"(?:^| ){re.escape(alias)}(?:$| )", evidence_text)
|
|
for alias in self._resolver.aliases_for(other)
|
|
) for other in others):
|
|
winners.append((candidate_id, result))
|
|
return winners[0] if len(winners) == 1 else None
|