Spaces:
Sleeping
Sleeping
File size: 4,450 Bytes
c2b7494 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | """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__)
@dataclass(slots=True)
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
|