Spaces:
Sleeping
Sleeping
| """LLM-based re-ranking of candidate codes using full EDF semantic context.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| from src.data.loaders import EventDefinitionForm | |
| from src.llm.client import call_llm_json | |
| from src.llm.prompts import RANK_SYSTEM_PROMPT, build_rank_prompt | |
| from src.retrieval.hybrid import HybridCandidate | |
| from src.utils.logging import get_logger | |
| logger = get_logger(__name__) | |
| class RankedCandidate: | |
| rank: int | |
| code: str | |
| description: str | |
| vocabulary: str | |
| relevance_score: float | |
| retrieval_score: float | |
| def __lt__(self, other: object) -> bool: | |
| if not isinstance(other, RankedCandidate): | |
| return NotImplemented | |
| return self.relevance_score > other.relevance_score | |
| def skip_rank(candidates: list[HybridCandidate]) -> list[RankedCandidate]: | |
| """Wrap hybrid retrieval candidates into ``RankedCandidate`` objects without | |
| calling the LLM, preserving their original hybrid retrieval order. | |
| The LLM ranking step (:func:`llm_rank`) is currently disabled in the | |
| pipeline/UI (an extra LLM call that duplicates work also done by | |
| classification) but is kept available for future use. This helper lets | |
| downstream code (``llm_classify``) keep consuming ``RankedCandidate`` | |
| objects uniformly regardless of whether ranking actually ran. | |
| """ | |
| results = [ | |
| RankedCandidate( | |
| rank=c.rank, | |
| code=c.code, | |
| description=c.description, | |
| vocabulary=c.vocabulary, | |
| relevance_score=0.0, | |
| retrieval_score=c.score, | |
| ) | |
| for c in candidates | |
| ] | |
| return results | |
| def llm_rank( | |
| candidates: list[HybridCandidate], | |
| edf: EventDefinitionForm, | |
| provider: str = "auto", | |
| model: str | None = None, | |
| max_retries: int = 3, | |
| ) -> list[RankedCandidate]: | |
| """Re-rank hybrid retrieval candidates using an LLM given the full EDF context. | |
| Parameters | |
| ---------- | |
| candidates : list[HybridCandidate] | |
| Candidates produced by :func:`src.retrieval.hybrid.hybrid_retrieval`. | |
| edf : EventDefinitionForm | |
| Structured phenotype definition providing semantic context. | |
| provider : str | |
| "auto", "anthropic", or "openai". | |
| model : str, optional | |
| Model identifier. Defaults to a provider-specific default. | |
| max_retries : int | |
| Number of LLM call attempts before failing. | |
| Returns | |
| ------- | |
| list[RankedCandidate] | |
| Candidates ordered by descending LLM-assigned relevance score. | |
| """ | |
| if not candidates: | |
| return [] | |
| lookup = {(c.code, c.vocabulary): c for c in candidates} | |
| payload = [ | |
| {"code": c.code, "vocabulary": c.vocabulary, "description": c.description} | |
| for c in candidates | |
| ] | |
| prompt = build_rank_prompt(edf, payload) | |
| response = call_llm_json( | |
| system_prompt=RANK_SYSTEM_PROMPT, | |
| user_prompt=prompt, | |
| provider=provider, | |
| model=model, | |
| max_retries=max_retries, | |
| ) | |
| ranked_raw = response.get("ranked_codes", []) | |
| results: list[RankedCandidate] = [] | |
| seen: set[tuple[str, str]] = set() | |
| for item in ranked_raw: | |
| key = (item.get("code", ""), item.get("vocabulary", "")) | |
| source = lookup.get(key) | |
| if source is None: | |
| continue | |
| seen.add(key) | |
| results.append( | |
| RankedCandidate( | |
| rank=0, | |
| code=source.code, | |
| description=source.description, | |
| vocabulary=source.vocabulary, | |
| relevance_score=float(item.get("relevance_score", 0.0)), | |
| retrieval_score=source.score, | |
| ) | |
| ) | |
| # Append any candidates the LLM omitted, ordered by original retrieval score, | |
| # so no candidate is silently dropped from downstream classification. | |
| for c in candidates: | |
| key = (c.code, c.vocabulary) | |
| if key not in seen: | |
| results.append( | |
| RankedCandidate( | |
| rank=0, | |
| code=c.code, | |
| description=c.description, | |
| vocabulary=c.vocabulary, | |
| relevance_score=0.0, | |
| retrieval_score=c.score, | |
| ) | |
| ) | |
| results.sort(key=lambda r: (-r.relevance_score, -r.retrieval_score)) | |
| for r, item in enumerate(results): | |
| item.rank = r + 1 | |
| return results | |