feat: rerank what a learner is shown, with Cohere through the proxy
Retrieval fused a bi-encoder and BM25 by reciprocal rank. A bi-encoder embeds a document long before the question exists, so the two never meet: it is good at "same topic" and mediocre at "answers this". A cross-encoder reads the pair. The proxy already serves three — `cohere-rerank-v4.0-pro` is the default and measurably better than the fast variant. Query text goes exactly where the embeddings already go, and nothing new was signed up for. It found a defect nobody was looking for. In AI Mode each finder scored `1/(1+rank)` *within its own corpus*, so the best article, section, question and card all scored 1.0 and the shortlist was a meaningless round-robin. A cross-encoder is the first thing in this system that can compare a question with a section. Candidates per kind widened so it can select rather than merely reorder. Measured against labels neither ranker produced. Questions, 60 disease tags: precision@3 0.394 → 0.483. Sections, 60 article titles: 0.772 → 0.833. "Management of bronchiolitis" led with influenza transmission and a pregnancy question; "when do you image a first febrile seizure" returned the definition rather than the sentence saying imaging is unnecessary. And the honest negative, in docs/reranking.md: board vignettes are written *not* to name their diagnosis, so on "what causes croup" it prefers a question that says the word in passing over the barking-cough vignette that never says it. Some of the bi-encoder's strength is traded away. Not on the typeahead. A page of results is a choice being made and worth a third of a second; a typeahead is a word being finished, runs on every keystroke, and has nothing to judge yet. The three-state thresholds stay on cosine, argued at the constant: a reranker only ever sees a shortlist and structurally cannot answer the corpus-wide question those numbers ask, and whether an answer claims to come from the library is a promise that must not depend on a network hop. Every failure returns None and leaves the order alone — unconfigured, no proxy, connect error, bare 502, timeout, non-JSON, a duplicate or out-of-range index, a non-numeric score, a list the wrong length. Verified against the running site with a bogus model name: same results, fused order, no error to the reader. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TqXevQJhxFrM7jJg82cgZN
This commit is contained in:
parent
4ee808c200
commit
031de53034
9 changed files with 852 additions and 8 deletions
|
|
@ -15,6 +15,10 @@ class Settings(BaseSettings):
|
||||||
LITELLM_API_KEY: str = ""
|
LITELLM_API_KEY: str = ""
|
||||||
LITELLM_API_BASE: str = ""
|
LITELLM_API_BASE: str = ""
|
||||||
LITELLM_EMBEDDING_MODEL: str = ""
|
LITELLM_EMBEDDING_MODEL: str = ""
|
||||||
|
# Cross-encoder reranker, named as the proxy serves it. Blank turns
|
||||||
|
# reranking off and leaves every result list in the order rank fusion
|
||||||
|
# produced, which is what a deployment without this proxy gets.
|
||||||
|
LITELLM_RERANK_MODEL: str = "cohere-rerank-v4.0-pro"
|
||||||
OPENAI_API_KEY: str = ""
|
OPENAI_API_KEY: str = ""
|
||||||
ELEVENLABS_API_KEY: str = ""
|
ELEVENLABS_API_KEY: str = ""
|
||||||
GOOGLE_TTS_API_KEY: str = ""
|
GOOGLE_TTS_API_KEY: str = ""
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,5 @@
|
||||||
import logging
|
import logging
|
||||||
|
import time
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||||
|
|
@ -606,10 +607,14 @@ def get_settings(admin: User = Depends(require_admin)):
|
||||||
r = redis_lib.from_url(settings.REDIS_URL, decode_responses=True)
|
r = redis_lib.from_url(settings.REDIS_URL, decode_responses=True)
|
||||||
registration_enabled = r.get("settings:registration_enabled")
|
registration_enabled = r.get("settings:registration_enabled")
|
||||||
embedding_model = r.get("settings:embedding_model")
|
embedding_model = r.get("settings:embedding_model")
|
||||||
|
rerank_model = r.get("settings:rerank_model")
|
||||||
sso_only = r.get("settings:sso_only")
|
sso_only = r.get("settings:sso_only")
|
||||||
return {
|
return {
|
||||||
"registration_enabled": registration_enabled != "false",
|
"registration_enabled": registration_enabled != "false",
|
||||||
"embedding_model": embedding_model or settings.LITELLM_EMBEDDING_MODEL or "",
|
"embedding_model": embedding_model or settings.LITELLM_EMBEDDING_MODEL or "",
|
||||||
|
# Blank is a valid answer and means result lists keep the order rank
|
||||||
|
# fusion gave them, so it is stored and read as written, not defaulted.
|
||||||
|
"rerank_model": rerank_model if rerank_model is not None else (settings.LITELLM_RERANK_MODEL or ""),
|
||||||
"sso_only": sso_only == "true",
|
"sso_only": sso_only == "true",
|
||||||
"sso_configured": bool(settings.OIDC_PROVIDER_URL and settings.OIDC_CLIENT_ID),
|
"sso_configured": bool(settings.OIDC_PROVIDER_URL and settings.OIDC_CLIENT_ID),
|
||||||
"sso_provider_name": settings.OIDC_PROVIDER_NAME,
|
"sso_provider_name": settings.OIDC_PROVIDER_NAME,
|
||||||
|
|
@ -619,6 +624,7 @@ def get_settings(admin: User = Depends(require_admin)):
|
||||||
return {
|
return {
|
||||||
"registration_enabled": True,
|
"registration_enabled": True,
|
||||||
"embedding_model": settings.LITELLM_EMBEDDING_MODEL or "",
|
"embedding_model": settings.LITELLM_EMBEDDING_MODEL or "",
|
||||||
|
"rerank_model": settings.LITELLM_RERANK_MODEL or "",
|
||||||
"sso_only": False,
|
"sso_only": False,
|
||||||
"sso_configured": bool(settings.OIDC_PROVIDER_URL and settings.OIDC_CLIENT_ID),
|
"sso_configured": bool(settings.OIDC_PROVIDER_URL and settings.OIDC_CLIENT_ID),
|
||||||
"sso_provider_name": settings.OIDC_PROVIDER_NAME,
|
"sso_provider_name": settings.OIDC_PROVIDER_NAME,
|
||||||
|
|
@ -647,6 +653,9 @@ def update_settings(
|
||||||
if "embedding_model" in settings_data:
|
if "embedding_model" in settings_data:
|
||||||
r.set("settings:embedding_model", settings_data["embedding_model"])
|
r.set("settings:embedding_model", settings_data["embedding_model"])
|
||||||
|
|
||||||
|
if "rerank_model" in settings_data:
|
||||||
|
r.set("settings:rerank_model", (settings_data["rerank_model"] or "").strip())
|
||||||
|
|
||||||
if "sso_only" in settings_data:
|
if "sso_only" in settings_data:
|
||||||
value = "true" if settings_data["sso_only"] else "false"
|
value = "true" if settings_data["sso_only"] else "false"
|
||||||
r.set("settings:sso_only", value)
|
r.set("settings:sso_only", value)
|
||||||
|
|
@ -669,6 +678,38 @@ def test_embedding(admin: User = Depends(require_admin)):
|
||||||
return {"model": model, "dimensions": len(result), "status": "ok"}
|
return {"model": model, "dimensions": len(result), "status": "ok"}
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/rerank/test")
|
||||||
|
def test_rerank(admin: User = Depends(require_admin)):
|
||||||
|
"""Check that the configured reranker answers, and that it answers sensibly.
|
||||||
|
|
||||||
|
A reranker that returns 200 and ranks the decoy first is worse than one that
|
||||||
|
is switched off, and nothing else on the site would ever tell you: its whole
|
||||||
|
output is an order somebody has to already know the right answer to judge.
|
||||||
|
"""
|
||||||
|
from app.services.rerank_service import rerank, rerank_model
|
||||||
|
|
||||||
|
model = rerank_model()
|
||||||
|
if not model:
|
||||||
|
raise HTTPException(status_code=400, detail="No rerank model configured")
|
||||||
|
documents = [
|
||||||
|
"Sourdough bread needs a starter culture and a long, cool proof.",
|
||||||
|
"Croup is a viral laryngotracheitis, usually parainfluenza, and presents "
|
||||||
|
"with a barking cough and inspiratory stridor.",
|
||||||
|
]
|
||||||
|
started = time.perf_counter()
|
||||||
|
scores = rerank("what causes croup in a toddler", documents)
|
||||||
|
elapsed_ms = int((time.perf_counter() - started) * 1000)
|
||||||
|
if scores is None:
|
||||||
|
raise HTTPException(status_code=500, detail=f"Rerank failed for model: {model}")
|
||||||
|
return {
|
||||||
|
"model": model,
|
||||||
|
"elapsed_ms": elapsed_ms,
|
||||||
|
"scores": [round(score, 4) for score in scores],
|
||||||
|
"ordered_correctly": scores[1] > scores[0],
|
||||||
|
"status": "ok" if scores[1] > scores[0] else "suspect",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
@router.get("/classification-snapshots")
|
@router.get("/classification-snapshots")
|
||||||
def list_classification_snapshots(
|
def list_classification_snapshots(
|
||||||
limit: int = Query(10, ge=1, le=50),
|
limit: int = Query(10, ge=1, le=50),
|
||||||
|
|
|
||||||
|
|
@ -13,6 +13,11 @@ Two decisions worth stating:
|
||||||
* A section match is reported under its article, not beside it. Ten sections of
|
* A section match is reported under its article, not beside it. Ten sections of
|
||||||
one article are one result with ten places to start reading, not ten results
|
one article are one result with ten places to start reading, not ten results
|
||||||
that bury everything else.
|
that bury everything else.
|
||||||
|
* The results page is reranked by a cross-encoder; the typeahead below it is
|
||||||
|
not. A page is a choice being made, and worth a third of a second to get
|
||||||
|
right. A typeahead is a word being finished, runs on every keystroke, and has
|
||||||
|
nothing to judge yet — it is prefix matching on titles, which is exactly what
|
||||||
|
finishing a word wants.
|
||||||
"""
|
"""
|
||||||
import logging
|
import logging
|
||||||
import re
|
import re
|
||||||
|
|
@ -29,7 +34,7 @@ from app.models.question import Question
|
||||||
from app.models.user import User
|
from app.models.user import User
|
||||||
from app.routers.media import readable_libraries
|
from app.routers.media import readable_libraries
|
||||||
from app.services.quiz_builder import bank_query, exam_scope_predicate
|
from app.services.quiz_builder import bank_query, exam_scope_predicate
|
||||||
from app.services.search_service import article_ids_with_sections, hybrid_ids
|
from app.services.search_service import article_ids_with_sections, hybrid_ids, rerank_ids
|
||||||
from app.utils.auth import get_current_user
|
from app.utils.auth import get_current_user
|
||||||
|
|
||||||
router = APIRouter()
|
router = APIRouter()
|
||||||
|
|
@ -107,7 +112,12 @@ def _questions(db, user, q, limit):
|
||||||
scope = exam_scope_predicate(db, user)
|
scope = exam_scope_predicate(db, user)
|
||||||
if scope is not None:
|
if scope is not None:
|
||||||
query = query.filter(scope)
|
query = query.filter(scope)
|
||||||
rows = _ordered(query.all(), ranked)[:limit]
|
# Reranked after the visibility rules, not before: spending a cross-encoder's
|
||||||
|
# candidate slots on rows this learner will never be shown is how a shortlist
|
||||||
|
# of fifty turns into a page of three.
|
||||||
|
visible = {row.id: row for row in query.all()}
|
||||||
|
order = rerank_ids(db, q, "question", [row_id for row_id in ranked if row_id in visible])
|
||||||
|
rows = [visible[row_id] for row_id in order][:limit]
|
||||||
return [{
|
return [{
|
||||||
"id": row.id,
|
"id": row.id,
|
||||||
"snippet": _snippet(row.question_text, q),
|
"snippet": _snippet(row.question_text, q),
|
||||||
|
|
@ -185,7 +195,9 @@ def suggest(
|
||||||
"""Titles for a typeahead — cheap enough to run on every keystroke.
|
"""Titles for a typeahead — cheap enough to run on every keystroke.
|
||||||
|
|
||||||
Deliberately lexical and prefix-based: a typeahead is finishing the word you
|
Deliberately lexical and prefix-based: a typeahead is finishing the word you
|
||||||
are typing, and a semantic neighbour of half a word is noise.
|
are typing, and a semantic neighbour of half a word is noise. No embedding
|
||||||
|
and no reranker for the same reason — both would be asked to judge relevance
|
||||||
|
to half a word, at a round trip each, several times a second.
|
||||||
"""
|
"""
|
||||||
query_text = (q or "").strip()
|
query_text = (q or "").strip()
|
||||||
if len(query_text) < 2:
|
if len(query_text) < 2:
|
||||||
|
|
|
||||||
|
|
@ -15,6 +15,15 @@ They are assertions, so they are trusted, and a retrieved row that carries one t
|
||||||
another retrieved row is boosted: two things an educator already tied together
|
another retrieved row is boosted: two things an educator already tied together
|
||||||
answering the same query is evidence, not coincidence. *Retrieved* links are
|
answering the same query is evidence, not coincidence. *Retrieved* links are
|
||||||
ranked guesses computed per query and stored nowhere.
|
ranked guesses computed per query and stored nowhere.
|
||||||
|
|
||||||
|
The shortlist is chosen twice. Each corpus offers twice as many candidates as it
|
||||||
|
can have places for, and a cross-encoder reads all of them against the query and
|
||||||
|
decides which survive and in what order. Unlike everywhere else on the site,
|
||||||
|
where reranking is a strict permutation of a page of results, here it selects —
|
||||||
|
because here the shortlist is not a page somebody skims past but the entire
|
||||||
|
evidence an answer may be built from, and a source that does not make it is one
|
||||||
|
the model cannot lean on. If no reranker answers, each corpus keeps the six it
|
||||||
|
elected, in the order it elected them, exactly as before there was one.
|
||||||
"""
|
"""
|
||||||
import logging
|
import logging
|
||||||
import re
|
import re
|
||||||
|
|
@ -26,6 +35,7 @@ from app.models.article import Article, ArticleSectionIndex, QuestionArticleLink
|
||||||
from app.models.flashcard import Flashcard, FlashcardDeck
|
from app.models.flashcard import Flashcard, FlashcardDeck
|
||||||
from app.models.question import Question
|
from app.models.question import Question
|
||||||
from app.models.user import User
|
from app.models.user import User
|
||||||
|
from app.services import rerank_service
|
||||||
from app.services.quiz_builder import bank_query, exam_scope_predicate
|
from app.services.quiz_builder import bank_query, exam_scope_predicate
|
||||||
from app.services.search_service import hybrid_ids, top_similarity
|
from app.services.search_service import hybrid_ids, top_similarity
|
||||||
|
|
||||||
|
|
@ -34,6 +44,12 @@ logger = logging.getLogger(__name__)
|
||||||
# How many of each kind retrieval offers the model. Small on purpose: a
|
# How many of each kind retrieval offers the model. Small on purpose: a
|
||||||
# shortlist the model can hold is worth more than a corpus it skims.
|
# shortlist the model can hold is worth more than a corpus it skims.
|
||||||
PER_KIND = 6
|
PER_KIND = 6
|
||||||
|
# How many of each kind the cross-encoder gets to choose from. Handing it only
|
||||||
|
# the six a corpus already elected would leave the membership of the shortlist
|
||||||
|
# entirely to rankers that never read the query against the document; twelve
|
||||||
|
# lets a rank-nine section reach the answer, and keeps the union of four corpora
|
||||||
|
# inside a single rerank call.
|
||||||
|
CANDIDATES_PER_KIND = PER_KIND * 2
|
||||||
MAX_SOURCES = 14
|
MAX_SOURCES = 14
|
||||||
# An excerpt long enough to answer from, short enough that fourteen of them fit.
|
# An excerpt long enough to answer from, short enough that fourteen of them fit.
|
||||||
EXCERPT_CHARS = 700
|
EXCERPT_CHARS = 700
|
||||||
|
|
@ -65,7 +81,7 @@ def _articles(db: Session, user: User, query: str) -> list[dict]:
|
||||||
"title": a.title,
|
"title": a.title,
|
||||||
"text": _clean(a.summary or a.content),
|
"text": _clean(a.summary or a.content),
|
||||||
"score": 1.0 / (1 + order.get(a.id, 0)),
|
"score": 1.0 / (1 + order.get(a.id, 0)),
|
||||||
} for a in rows[:PER_KIND]]
|
} for a in rows[:CANDIDATES_PER_KIND]]
|
||||||
|
|
||||||
|
|
||||||
def _sections(db: Session, user: User, query: str) -> list[dict]:
|
def _sections(db: Session, user: User, query: str) -> list[dict]:
|
||||||
|
|
@ -80,7 +96,7 @@ def _sections(db: Session, user: User, query: str) -> list[dict]:
|
||||||
order = {rid: i for i, rid in enumerate(ranked)}
|
order = {rid: i for i, rid in enumerate(ranked)}
|
||||||
rows.sort(key=lambda r: order.get(r.id, len(order)))
|
rows.sort(key=lambda r: order.get(r.id, len(order)))
|
||||||
out = []
|
out = []
|
||||||
for row in rows[:PER_KIND]:
|
for row in rows[:CANDIDATES_PER_KIND]:
|
||||||
article = articles.get(row.article_id)
|
article = articles.get(row.article_id)
|
||||||
if not article:
|
if not article:
|
||||||
continue
|
continue
|
||||||
|
|
@ -116,7 +132,7 @@ def _questions(db: Session, user: User, query: str) -> list[dict]:
|
||||||
# printed it would hand away the practice it is meant to prepare for.
|
# printed it would hand away the practice it is meant to prepare for.
|
||||||
"text": _clean(row.question_text, 320),
|
"text": _clean(row.question_text, 320),
|
||||||
"score": 1.0 / (1 + order.get(row.id, 0)),
|
"score": 1.0 / (1 + order.get(row.id, 0)),
|
||||||
} for row in rows[:PER_KIND]]
|
} for row in rows[:CANDIDATES_PER_KIND]]
|
||||||
|
|
||||||
|
|
||||||
def _cards(db: Session, user: User, query: str) -> list[dict]:
|
def _cards(db: Session, user: User, query: str) -> list[dict]:
|
||||||
|
|
@ -136,7 +152,7 @@ def _cards(db: Session, user: User, query: str) -> list[dict]:
|
||||||
"title": _clean(row.front, 90),
|
"title": _clean(row.front, 90),
|
||||||
"text": _clean(row.back, 260),
|
"text": _clean(row.back, 260),
|
||||||
"score": 1.0 / (1 + order.get(row.id, 0)),
|
"score": 1.0 / (1 + order.get(row.id, 0)),
|
||||||
} for row in rows[:PER_KIND]]
|
} for row in rows[:CANDIDATES_PER_KIND]]
|
||||||
|
|
||||||
|
|
||||||
def _apply_curated_boost(db: Session, sources: list[dict]) -> None:
|
def _apply_curated_boost(db: Session, sources: list[dict]) -> None:
|
||||||
|
|
@ -166,6 +182,32 @@ def _apply_curated_boost(db: Session, sources: list[dict]) -> None:
|
||||||
source["curated"] = True
|
source["curated"] = True
|
||||||
|
|
||||||
|
|
||||||
|
def _cross_encode(query: str, sources: list[dict]) -> None:
|
||||||
|
"""Rescore the whole shortlist on one scale, in place, if a reranker answers.
|
||||||
|
|
||||||
|
The scores the finders assign are `1/(1+rank)` *within their own corpus*, so
|
||||||
|
the best article, the best section, the best question and the best card all
|
||||||
|
score 1.0 and the order between them is whatever the sort happened to do.
|
||||||
|
A cross-encoder is the first thing in this pipeline that can compare a
|
||||||
|
section against a question, because it reads both against the same query.
|
||||||
|
|
||||||
|
The new score is the position, not the raw relevance: 0.75 from one reranker
|
||||||
|
and 0.78 from another mean nothing to `CURATED_BOOST`, whereas `1/(1+rank)`
|
||||||
|
keeps a curated pair worth the same few places it has always been worth.
|
||||||
|
|
||||||
|
Silence leaves the scores alone. That is the whole degradation story here:
|
||||||
|
without a reranker the shortlist is exactly what it was before there was one.
|
||||||
|
"""
|
||||||
|
if len(sources) < 2:
|
||||||
|
return
|
||||||
|
scores = rerank_service.rerank(query, [f"{s['title']}. {s['text']}" for s in sources])
|
||||||
|
if scores is None:
|
||||||
|
return
|
||||||
|
ranking = sorted(range(len(sources)), key=lambda i: (-scores[i], i))
|
||||||
|
for position, index in enumerate(ranking):
|
||||||
|
sources[index]["score"] = 1.0 / (1 + position)
|
||||||
|
|
||||||
|
|
||||||
def retrieve(db: Session, user: User, query: str) -> list[dict]:
|
def retrieve(db: Session, user: User, query: str) -> list[dict]:
|
||||||
"""The only things the model will be allowed to cite for this message."""
|
"""The only things the model will be allowed to cite for this message."""
|
||||||
query = (query or "").strip()
|
query = (query or "").strip()
|
||||||
|
|
@ -178,9 +220,24 @@ def retrieve(db: Session, user: User, query: str) -> list[dict]:
|
||||||
except Exception:
|
except Exception:
|
||||||
# One corpus failing narrows the answer; it does not end it.
|
# One corpus failing narrows the answer; it does not end it.
|
||||||
logger.warning("AI Mode retrieval failed for %s", finder.__name__, exc_info=True)
|
logger.warning("AI Mode retrieval failed for %s", finder.__name__, exc_info=True)
|
||||||
|
_cross_encode(query, sources)
|
||||||
_apply_curated_boost(db, sources)
|
_apply_curated_boost(db, sources)
|
||||||
sources.sort(key=lambda s: -s["score"])
|
sources.sort(key=lambda s: -s["score"])
|
||||||
return sources[:MAX_SOURCES]
|
|
||||||
|
# Still no more than `PER_KIND` of anything. A shortlist the cross-encoder
|
||||||
|
# filled with fourteen sections of one article would score well and read
|
||||||
|
# like a single source quoted fourteen times, and it would leave the chat
|
||||||
|
# with no question to send the learner to practise.
|
||||||
|
kept: list[dict] = []
|
||||||
|
seen: dict[str, int] = {}
|
||||||
|
for source in sources:
|
||||||
|
if seen.get(source["kind"], 0) >= PER_KIND:
|
||||||
|
continue
|
||||||
|
seen[source["kind"]] = seen.get(source["kind"], 0) + 1
|
||||||
|
kept.append(source)
|
||||||
|
if len(kept) >= MAX_SOURCES:
|
||||||
|
break
|
||||||
|
return kept
|
||||||
|
|
||||||
|
|
||||||
def sources_block(sources: list[dict]) -> str:
|
def sources_block(sources: list[dict]) -> str:
|
||||||
|
|
@ -209,6 +266,16 @@ def sources_block(sources: list[dict]) -> str:
|
||||||
# Worth re-measuring when the corpus changes size or subject; the method and
|
# Worth re-measuring when the corpus changes size or subject; the method and
|
||||||
# the full measurement are in docs/retrieval-thresholds.md. Anything else is
|
# the full measurement are in docs/retrieval-thresholds.md. Anything else is
|
||||||
# tuning by feel against numbers nobody wrote down.
|
# tuning by feel against numbers nobody wrote down.
|
||||||
|
#
|
||||||
|
# Deliberately not the cross-encoder's score, though it is the better judge of a
|
||||||
|
# pair. The question here is "is there anything in this library about this at
|
||||||
|
# all", and `top_similarity` answers it by scanning every embedded row in two
|
||||||
|
# corpora through the vector index. A reranker can only score candidates that
|
||||||
|
# were shortlisted first, so a reranked closeness cannot tell "the library does
|
||||||
|
# not cover this" from "retrieval had a bad day", and it would put a network hop
|
||||||
|
# in the path of a decision about what the answer *claims* — the one place where
|
||||||
|
# a service being down must not change the output. Order is a preference;
|
||||||
|
# whether an answer says it came from the library is a promise.
|
||||||
STRONG_MATCH = 0.55
|
STRONG_MATCH = 0.55
|
||||||
ADJACENT_MATCH = 0.50
|
ADJACENT_MATCH = 0.50
|
||||||
|
|
||||||
|
|
|
||||||
225
backend/app/services/rerank_service.py
Normal file
225
backend/app/services/rerank_service.py
Normal file
|
|
@ -0,0 +1,225 @@
|
||||||
|
"""Cross-encoder reranking: the model reads the query and the document together.
|
||||||
|
|
||||||
|
Everything upstream of here ranks without ever putting the two side by side. A
|
||||||
|
bi-encoder embedded each row long before anybody typed anything, so cosine
|
||||||
|
distance measures "same topic", which is not the same as "answers this". BM25
|
||||||
|
measures which words happen to coincide. Reciprocal rank fusion combines those
|
||||||
|
two opinions but cannot add a third.
|
||||||
|
|
||||||
|
A cross-encoder is that third opinion: one forward pass over the pair, so it can
|
||||||
|
tell that a section titled "Croup" which spends its whole body on epiglottitis
|
||||||
|
is not the thing to put first. The price is that it cannot be precomputed — the
|
||||||
|
work is per pair, per query — which is why nothing here ever sees the corpus.
|
||||||
|
It sees a shortlist the cheap rankers already produced, and reorders it.
|
||||||
|
|
||||||
|
**Reranking is an opinion about order and nothing else.** Every caller keeps the
|
||||||
|
rows retrieval found; a reranker that is unconfigured, down, slow or talking
|
||||||
|
nonsense returns `None` from `rerank()` and the caller ships today's ordering.
|
||||||
|
A search page with no results because a scoring service is unwell is a much
|
||||||
|
worse failure than a search page in a slightly worse order, and the second
|
||||||
|
failure is invisible to a user while the first ends their session.
|
||||||
|
|
||||||
|
Served through the same proxy as everything else (`LITELLM_API_BASE`), so a
|
||||||
|
learner's query goes where the embeddings already go and nowhere new. Nothing
|
||||||
|
here knows which reranker is behind that name: `POST /v1/rerank` with
|
||||||
|
`{model, query, documents, top_n}`, back `{results: [{index, relevance_score}]}`,
|
||||||
|
is the shape Cohere defined and Jina, BGE and LiteLLM's own rerank route all
|
||||||
|
speak. This deployment's proxy serves `cohere-rerank-v4.0-pro` and
|
||||||
|
`-fast` alongside `jina-reranker-v2-base-multilingual`, which is open-weight —
|
||||||
|
moving to a self-hosted one is a change to `LITELLM_RERANK_MODEL` and no change
|
||||||
|
to any code.
|
||||||
|
"""
|
||||||
|
import hashlib
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from app.config import settings
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
#: Per-document budget. Rerankers truncate at roughly 1k tokens per document
|
||||||
|
#: anyway, and the part of a clinical vignette that decides relevance is its
|
||||||
|
#: opening, not the option list; sending more buys nothing and costs latency
|
||||||
|
#: linear in total tokens.
|
||||||
|
DOC_CHARS = 1200
|
||||||
|
|
||||||
|
#: A cross-encoder is trained on query-length queries. Handing it a whole
|
||||||
|
#: uploaded document (which `search_service._query_terms` exists to cope with)
|
||||||
|
#: is out of distribution as well as slow, so the query is clamped too.
|
||||||
|
QUERY_CHARS = 2000
|
||||||
|
|
||||||
|
#: Measured against this proxy: 10 candidates ≈ 250 ms, 30 ≈ 380 ms, 50 ≈ 430 ms,
|
||||||
|
#: 100 ≈ 580 ms — mostly the round trip, then roughly 2 ms a pair. The fused pool
|
||||||
|
#: is 60 wide and a page shows ten, so past about fifty the reranker is being
|
||||||
|
#: paid to confirm that rows neither cheap ranker wanted are still not wanted.
|
||||||
|
#: The numbers and the method are in docs/reranking.md.
|
||||||
|
MAX_CANDIDATES = 50
|
||||||
|
|
||||||
|
#: Deterministic in (model, query, documents), and the document text is part of
|
||||||
|
#: the key, so an edited row invalidates its own entry. That leaves nothing a
|
||||||
|
#: day can make wrong — the same reasoning as the query-embedding cache.
|
||||||
|
CACHE_TTL = 24 * 3600
|
||||||
|
|
||||||
|
#: Long enough for a cold model to answer, short enough that a wedged proxy
|
||||||
|
#: costs one page load rather than the request. On timeout the caller keeps the
|
||||||
|
#: order it already had.
|
||||||
|
TIMEOUT_SECONDS = 8
|
||||||
|
|
||||||
|
|
||||||
|
def rerank_model() -> str:
|
||||||
|
"""The active reranker, Redis override first — same shape as the embedder.
|
||||||
|
|
||||||
|
An administrator can move between served rerankers without a redeploy;
|
||||||
|
unset everywhere means the feature is simply off.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
import redis as redis_lib
|
||||||
|
|
||||||
|
client = redis_lib.from_url(settings.REDIS_URL, decode_responses=True,
|
||||||
|
socket_connect_timeout=1)
|
||||||
|
chosen = client.get("settings:rerank_model")
|
||||||
|
if chosen is not None:
|
||||||
|
return chosen.strip()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return (settings.LITELLM_RERANK_MODEL or "").strip()
|
||||||
|
|
||||||
|
|
||||||
|
def is_configured() -> bool:
|
||||||
|
"""Whether reranking can be attempted at all."""
|
||||||
|
return bool(rerank_model() and settings.LITELLM_API_KEY and settings.LITELLM_API_BASE)
|
||||||
|
|
||||||
|
|
||||||
|
def _cache_key(model: str, query: str, documents: list[str]) -> str:
|
||||||
|
digest = hashlib.sha256()
|
||||||
|
digest.update(query.encode("utf-8", "ignore"))
|
||||||
|
for document in documents:
|
||||||
|
digest.update(b"\x00")
|
||||||
|
digest.update(document.encode("utf-8", "ignore"))
|
||||||
|
return f"rerank:{model}:{digest.hexdigest()[:40]}"
|
||||||
|
|
||||||
|
|
||||||
|
def _redis():
|
||||||
|
try:
|
||||||
|
import redis as redis_lib
|
||||||
|
|
||||||
|
return redis_lib.from_url(settings.REDIS_URL, decode_responses=True,
|
||||||
|
socket_connect_timeout=1)
|
||||||
|
except Exception:
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def _parse(payload, count: int) -> list[float] | None:
|
||||||
|
"""Scores aligned to the documents that went in, or None if they are not.
|
||||||
|
|
||||||
|
Positional trust is the bug waiting to happen here: the response is sorted
|
||||||
|
by score, so reading it in order attaches the best document's score to the
|
||||||
|
first document. Every score is placed by the `index` the server echoes, and
|
||||||
|
a response that does not cover every input is refused outright rather than
|
||||||
|
half-applied — a partial reordering is indistinguishable from a bad one.
|
||||||
|
"""
|
||||||
|
if not isinstance(payload, dict):
|
||||||
|
return None
|
||||||
|
results = payload.get("results")
|
||||||
|
if not isinstance(results, list) or len(results) != count:
|
||||||
|
return None
|
||||||
|
scores: list[float | None] = [None] * count
|
||||||
|
for item in results:
|
||||||
|
if not isinstance(item, dict):
|
||||||
|
return None
|
||||||
|
index, score = item.get("index"), item.get("relevance_score")
|
||||||
|
if not isinstance(index, int) or not isinstance(score, (int, float)):
|
||||||
|
return None
|
||||||
|
if not 0 <= index < count or scores[index] is not None:
|
||||||
|
return None
|
||||||
|
scores[index] = float(score)
|
||||||
|
if any(score is None for score in scores):
|
||||||
|
return None
|
||||||
|
return scores # type: ignore[return-value]
|
||||||
|
|
||||||
|
|
||||||
|
def rerank(query: str, documents: list[str]) -> list[float] | None:
|
||||||
|
"""Relevance of each document to the query, in the order they were given.
|
||||||
|
|
||||||
|
`None` means "no opinion" and is the answer to every kind of trouble: not
|
||||||
|
configured, unreachable, slow, or a reply that does not line up with the
|
||||||
|
request. Callers must read it as "keep what you had".
|
||||||
|
"""
|
||||||
|
query = " ".join((query or "").split())[:QUERY_CHARS]
|
||||||
|
if not query or not documents:
|
||||||
|
return None
|
||||||
|
model = rerank_model()
|
||||||
|
if not is_configured():
|
||||||
|
return None
|
||||||
|
|
||||||
|
clipped = [" ".join((document or "").split())[:DOC_CHARS] or " " for document in documents]
|
||||||
|
key = _cache_key(model, query, clipped)
|
||||||
|
cache = _redis()
|
||||||
|
if cache is not None:
|
||||||
|
try:
|
||||||
|
hit = cache.get(key)
|
||||||
|
if hit:
|
||||||
|
cached = json.loads(hit)
|
||||||
|
if isinstance(cached, list) and len(cached) == len(clipped):
|
||||||
|
return [float(score) for score in cached]
|
||||||
|
except Exception:
|
||||||
|
logger.debug("Rerank cache unreadable", exc_info=True)
|
||||||
|
|
||||||
|
api_base = (settings.LITELLM_API_BASE or "").rstrip("/").removesuffix("/v1")
|
||||||
|
try:
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
response = httpx.post(
|
||||||
|
f"{api_base}/v1/rerank",
|
||||||
|
headers={"Authorization": f"Bearer {settings.LITELLM_API_KEY}",
|
||||||
|
"Content-Type": "application/json"},
|
||||||
|
# `top_n` is the whole list on purpose. Asking for fewer would save
|
||||||
|
# nothing — the pairs are scored either way — and would leave the
|
||||||
|
# tail unscored, so it could only be appended in its old order.
|
||||||
|
json={"model": model, "query": query, "documents": clipped,
|
||||||
|
"top_n": len(clipped)},
|
||||||
|
timeout=TIMEOUT_SECONDS,
|
||||||
|
)
|
||||||
|
response.raise_for_status()
|
||||||
|
scores = _parse(response.json(), len(clipped))
|
||||||
|
except Exception as error:
|
||||||
|
logger.warning("Rerank unavailable (%s); keeping fused order", error)
|
||||||
|
return None
|
||||||
|
|
||||||
|
if scores is None:
|
||||||
|
logger.warning("Rerank returned a reply that did not match the request; "
|
||||||
|
"keeping fused order")
|
||||||
|
return None
|
||||||
|
|
||||||
|
if cache is not None:
|
||||||
|
try:
|
||||||
|
cache.setex(key, CACHE_TTL, json.dumps(scores))
|
||||||
|
except Exception:
|
||||||
|
logger.debug("Could not cache rerank scores", exc_info=True)
|
||||||
|
return scores
|
||||||
|
|
||||||
|
|
||||||
|
def reorder(query: str, items: list, text_of) -> list:
|
||||||
|
"""`items` best-first by cross-encoder score, unchanged if there is no score.
|
||||||
|
|
||||||
|
Only the head is reordered. Beyond `MAX_CANDIDATES` the list is passed
|
||||||
|
through untouched, which keeps the cost of a query independent of how many
|
||||||
|
rows matched it and keeps pagination past the first page stable.
|
||||||
|
|
||||||
|
The return value is always a permutation of the input: nothing is dropped
|
||||||
|
and nothing is added, whatever the reranker does or fails to do. That is the
|
||||||
|
property that makes this safe to put in front of a search page.
|
||||||
|
"""
|
||||||
|
if len(items) < 2:
|
||||||
|
return list(items)
|
||||||
|
head = list(items[:MAX_CANDIDATES])
|
||||||
|
tail = list(items[len(head):])
|
||||||
|
scores = rerank(query, [text_of(item) for item in head])
|
||||||
|
# The length is `rerank`'s contract, and a caller that stubs or replaces it
|
||||||
|
# is exactly when a reordering half the list long would go unnoticed.
|
||||||
|
if scores is None or len(scores) != len(head):
|
||||||
|
return list(items)
|
||||||
|
# Ties keep the fused order, so agreement between the cheap rankers still
|
||||||
|
# settles anything the cross-encoder is indifferent about.
|
||||||
|
order = sorted(range(len(head)), key=lambda index: (-scores[index], index))
|
||||||
|
return [head[index] for index in order] + tail
|
||||||
|
|
@ -14,12 +14,18 @@ scales; RRF only needs each ranker's ordering.
|
||||||
Retrieval is always hybrid. A keyword-only mode looks precise but silently drops
|
Retrieval is always hybrid. A keyword-only mode looks precise but silently drops
|
||||||
the question that asks the same thing in different words, which is exactly the
|
the question that asks the same thing in different words, which is exactly the
|
||||||
question a learner searching a concept wants.
|
question a learner searching a concept wants.
|
||||||
|
|
||||||
|
Fusion decides *which* rows are candidates. `rerank_ids` may then reorder the
|
||||||
|
head of that list with a cross-encoder (see `rerank_service`), which is a strict
|
||||||
|
permutation: recall belongs to the two rankers here, and a scoring service being
|
||||||
|
unwell must never subtract a result.
|
||||||
"""
|
"""
|
||||||
import hashlib
|
import hashlib
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import re
|
import re
|
||||||
|
|
||||||
|
from sqlalchemy import bindparam as sa_bindparam
|
||||||
from sqlalchemy import text as sa_text
|
from sqlalchemy import text as sa_text
|
||||||
from sqlalchemy.orm import Session
|
from sqlalchemy.orm import Session
|
||||||
|
|
||||||
|
|
@ -263,6 +269,13 @@ def article_ids_with_sections(db: Session, query_text: str,
|
||||||
|
|
||||||
ranked, _ = hybrid_ids(db, query_text, "article", limit=limit)
|
ranked, _ = hybrid_ids(db, query_text, "article", limit=limit)
|
||||||
section_ranked, _ = hybrid_ids(db, query_text, "article_section", limit=limit)
|
section_ranked, _ = hybrid_ids(db, query_text, "article_section", limit=limit)
|
||||||
|
# Only the sections are reranked, and it is the article ordering that this
|
||||||
|
# buys. An article row carries a title, a summary and a topical vector; the
|
||||||
|
# prose a query is actually about is in its sections, so handing a
|
||||||
|
# cross-encoder the article row means handing it a stub and asking about a
|
||||||
|
# document it cannot see. Its best section stands for it instead — which is
|
||||||
|
# already how the fusion below decides an article's place.
|
||||||
|
section_ranked = rerank_ids(db, query_text, "article_section", section_ranked)
|
||||||
|
|
||||||
rows = {row.id: row for row in db.query(ArticleSectionIndex).filter(
|
rows = {row.id: row for row in db.query(ArticleSectionIndex).filter(
|
||||||
ArticleSectionIndex.id.in_(section_ranked)).all()} if section_ranked else {}
|
ArticleSectionIndex.id.in_(section_ranked)).all()} if section_ranked else {}
|
||||||
|
|
@ -289,3 +302,53 @@ def article_ids_with_sections(db: Session, query_text: str,
|
||||||
|
|
||||||
ordered = sorted(scores, key=lambda article_id: (-scores[article_id], article_id))[:limit]
|
ordered = sorted(scores, key=lambda article_id: (-scores[article_id], article_id))[:limit]
|
||||||
return ordered, by_article
|
return ordered, by_article
|
||||||
|
|
||||||
|
|
||||||
|
# What a cross-encoder is shown of each row. Not the same as the columns the
|
||||||
|
# lexical fallback scans: `CAST(options AS TEXT)` helps a substring match find a
|
||||||
|
# drug name buried in an option, but as prose it is a JSON array, and paying a
|
||||||
|
# cross-encoder to read punctuation makes its judgement worse rather than
|
||||||
|
# better. Title first everywhere, because a truncated document keeps its head.
|
||||||
|
RERANK_TEXT = {
|
||||||
|
"question": "question_text",
|
||||||
|
"article": "coalesce(title, '') || '. ' || coalesce(summary, '')",
|
||||||
|
"article_section": "coalesce(title, '') || '. ' || coalesce(content, '')",
|
||||||
|
"flashcard": "coalesce(front, '') || ' — ' || coalesce(back, '')",
|
||||||
|
"media": "coalesce(title, '') || '. ' || coalesce(caption, '') || ' ' || coalesce(alt_text, '')",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def rerank_ids(db: Session, query_text: str, kind: str, ranked: list[int]) -> list[int]:
|
||||||
|
"""`ranked` reordered by a cross-encoder, or exactly `ranked` if there isn't one.
|
||||||
|
|
||||||
|
The contract is narrow on purpose: this is a permutation. Fusion decides
|
||||||
|
which rows exist and every visibility rule downstream still runs on the same
|
||||||
|
set, so a reranker that is off, down or wrong costs a worse order and
|
||||||
|
nothing else.
|
||||||
|
|
||||||
|
Only the head is fetched and scored. The rows past it are ones neither cheap
|
||||||
|
ranker put near the front, and reading a 200-row pool out of the database to
|
||||||
|
ask a model about it would cost more than the answer is worth.
|
||||||
|
"""
|
||||||
|
from app.services import rerank_service
|
||||||
|
|
||||||
|
query_text = (query_text or "").strip()
|
||||||
|
if len(ranked) < 2 or kind not in RERANK_TEXT or not query_text:
|
||||||
|
return ranked
|
||||||
|
if not rerank_service.is_configured():
|
||||||
|
return ranked
|
||||||
|
|
||||||
|
head = ranked[:rerank_service.MAX_CANDIDATES]
|
||||||
|
table, _ = CORPORA[kind]
|
||||||
|
try:
|
||||||
|
statement = sa_text(
|
||||||
|
f"SELECT id, {RERANK_TEXT[kind]} AS body FROM {table} WHERE id IN :ids"
|
||||||
|
).bindparams(sa_bindparam("ids", expanding=True))
|
||||||
|
bodies = {row.id: row.body or "" for row in db.execute(statement, {"ids": head}).fetchall()}
|
||||||
|
except Exception:
|
||||||
|
logger.warning("Could not read %s text for reranking; keeping fused order",
|
||||||
|
kind, exc_info=True)
|
||||||
|
return ranked
|
||||||
|
if not bodies:
|
||||||
|
return ranked
|
||||||
|
return rerank_service.reorder(query_text, ranked, lambda row_id: bodies.get(row_id, ""))
|
||||||
|
|
|
||||||
226
backend/tests/test_rerank.py
Normal file
226
backend/tests/test_rerank.py
Normal file
|
|
@ -0,0 +1,226 @@
|
||||||
|
"""What the cross-encoder is allowed to change, and what happens when it cannot.
|
||||||
|
|
||||||
|
No network: the point of these is not whether a reranker ranks well — that is
|
||||||
|
measured against the real corpus, in docs/reranking.md — but that every way it
|
||||||
|
can fail leaves the same rows on the page, in some order, with no error.
|
||||||
|
|
||||||
|
Run: DATABASE_URL=sqlite:// PYTHONPATH=backend python -m unittest discover -s backend/tests
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import unittest
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
os.environ.setdefault("DATABASE_URL", "sqlite://")
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
from sqlalchemy import create_engine
|
||||||
|
from sqlalchemy.orm import Session
|
||||||
|
from sqlalchemy.pool import StaticPool
|
||||||
|
|
||||||
|
from app.database import Base
|
||||||
|
from app.models.article import Article, ArticleSectionIndex # noqa — mapper registration.
|
||||||
|
from app.models.course import Course # noqa — Quiz.course_id FK needs the table in metadata.
|
||||||
|
from app.models.media import MediaAsset # noqa — media is an embeddable corpus.
|
||||||
|
from app.models.question import Question
|
||||||
|
from app.models.user import User
|
||||||
|
from app.services import rerank_service, search_service
|
||||||
|
|
||||||
|
|
||||||
|
def response(payload, status=200):
|
||||||
|
"""An httpx reply in the shape the proxy sends one."""
|
||||||
|
return httpx.Response(status, json=payload, request=httpx.Request("POST", "http://proxy/v1/rerank"))
|
||||||
|
|
||||||
|
|
||||||
|
def scored(*pairs):
|
||||||
|
"""A rerank body: (index, score) pairs, deliberately not in index order."""
|
||||||
|
return {"id": "test", "results": [{"index": i, "relevance_score": s} for i, s in pairs]}
|
||||||
|
|
||||||
|
|
||||||
|
class ConfigurationTests(unittest.TestCase):
|
||||||
|
def test_no_model_means_no_call_and_no_opinion(self):
|
||||||
|
with patch.object(rerank_service, "rerank_model", return_value=""), \
|
||||||
|
patch("httpx.post", side_effect=AssertionError("must not reach the proxy")):
|
||||||
|
self.assertIsNone(rerank_service.rerank("croup", ["a", "b"]))
|
||||||
|
|
||||||
|
def test_a_model_without_a_proxy_is_not_configured(self):
|
||||||
|
with patch.object(rerank_service, "rerank_model", return_value="a-reranker"), \
|
||||||
|
patch.object(rerank_service.settings, "LITELLM_API_BASE", ""):
|
||||||
|
self.assertFalse(rerank_service.is_configured())
|
||||||
|
|
||||||
|
|
||||||
|
class DegradationTests(unittest.TestCase):
|
||||||
|
"""Every one of these must return None, which every caller reads as "keep it"."""
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
self.configured = patch.multiple(
|
||||||
|
rerank_service.settings, LITELLM_API_BASE="http://proxy",
|
||||||
|
LITELLM_API_KEY="k", LITELLM_RERANK_MODEL="a-reranker")
|
||||||
|
self.configured.start()
|
||||||
|
self.addCleanup(self.configured.stop)
|
||||||
|
# Redis is absent in tests; the cache must be optional, not required.
|
||||||
|
self.no_cache = patch.object(rerank_service, "_redis", return_value=None)
|
||||||
|
self.no_cache.start()
|
||||||
|
self.addCleanup(self.no_cache.stop)
|
||||||
|
|
||||||
|
def none_for(self, **post):
|
||||||
|
with patch("httpx.post", **post):
|
||||||
|
return rerank_service.rerank("croup", ["a", "b", "c"])
|
||||||
|
|
||||||
|
def test_unreachable_proxy(self):
|
||||||
|
self.assertIsNone(self.none_for(side_effect=httpx.ConnectError("refused")))
|
||||||
|
|
||||||
|
def test_a_bare_502_while_the_proxy_restarts(self):
|
||||||
|
self.assertIsNone(self.none_for(return_value=response({"error": "bad gateway"}, status=502)))
|
||||||
|
|
||||||
|
def test_a_timeout(self):
|
||||||
|
self.assertIsNone(self.none_for(side_effect=httpx.ReadTimeout("slow")))
|
||||||
|
|
||||||
|
def test_a_body_that_is_not_json(self):
|
||||||
|
self.assertIsNone(self.none_for(return_value=httpx.Response(
|
||||||
|
200, text="<html>gateway</html>",
|
||||||
|
request=httpx.Request("POST", "http://proxy/v1/rerank"))))
|
||||||
|
|
||||||
|
def test_a_reply_missing_a_document(self):
|
||||||
|
# Two scores for three documents: the third would have to be guessed,
|
||||||
|
# and a guess here is a reordering nobody asked for.
|
||||||
|
self.assertIsNone(self.none_for(return_value=response(scored((0, 0.1), (1, 0.9)))))
|
||||||
|
|
||||||
|
def test_a_reply_with_an_index_out_of_range(self):
|
||||||
|
self.assertIsNone(self.none_for(return_value=response(
|
||||||
|
scored((0, 0.1), (1, 0.9), (7, 0.5)))))
|
||||||
|
|
||||||
|
def test_a_reply_with_a_duplicated_index(self):
|
||||||
|
self.assertIsNone(self.none_for(return_value=response(
|
||||||
|
scored((0, 0.1), (0, 0.9), (2, 0.5)))))
|
||||||
|
|
||||||
|
def test_a_score_that_is_not_a_number(self):
|
||||||
|
self.assertIsNone(self.none_for(return_value=response(
|
||||||
|
{"results": [{"index": 0, "relevance_score": "high"},
|
||||||
|
{"index": 1, "relevance_score": 0.2},
|
||||||
|
{"index": 2, "relevance_score": 0.3}]})))
|
||||||
|
|
||||||
|
def test_scores_are_placed_by_the_index_the_server_echoed(self):
|
||||||
|
# The proxy answers best-first, so reading positionally would give the
|
||||||
|
# first document the best document's score.
|
||||||
|
with patch("httpx.post", return_value=response(scored((2, 0.9), (0, 0.1), (1, 0.4)))):
|
||||||
|
self.assertEqual(rerank_service.rerank("croup", ["a", "b", "c"]), [0.1, 0.4, 0.9])
|
||||||
|
|
||||||
|
|
||||||
|
class ReorderTests(unittest.TestCase):
|
||||||
|
def test_the_result_is_always_a_permutation(self):
|
||||||
|
items = list(range(8))
|
||||||
|
with patch.object(rerank_service, "rerank", return_value=[0.5] * 8):
|
||||||
|
self.assertCountEqual(rerank_service.reorder("q", items, str), items)
|
||||||
|
|
||||||
|
def test_no_opinion_leaves_the_order_alone(self):
|
||||||
|
items = [3, 1, 2]
|
||||||
|
with patch.object(rerank_service, "rerank", return_value=None):
|
||||||
|
self.assertEqual(rerank_service.reorder("q", items, str), items)
|
||||||
|
|
||||||
|
def test_ties_keep_the_order_fusion_gave_them(self):
|
||||||
|
with patch.object(rerank_service, "rerank", return_value=[0.4, 0.4, 0.9]):
|
||||||
|
self.assertEqual(rerank_service.reorder("q", [10, 20, 30], str), [30, 10, 20])
|
||||||
|
|
||||||
|
def test_only_the_head_is_scored_and_the_tail_is_kept(self):
|
||||||
|
items = list(range(rerank_service.MAX_CANDIDATES + 5))
|
||||||
|
seen = {}
|
||||||
|
|
||||||
|
def fake(query, documents):
|
||||||
|
seen["count"] = len(documents)
|
||||||
|
return list(range(len(documents))) # exactly reverses the head
|
||||||
|
|
||||||
|
with patch.object(rerank_service, "rerank", side_effect=fake):
|
||||||
|
out = rerank_service.reorder("q", items, str)
|
||||||
|
self.assertEqual(seen["count"], rerank_service.MAX_CANDIDATES)
|
||||||
|
self.assertEqual(out[0], rerank_service.MAX_CANDIDATES - 1)
|
||||||
|
self.assertEqual(out[-5:], items[-5:])
|
||||||
|
self.assertCountEqual(out, items)
|
||||||
|
|
||||||
|
|
||||||
|
class CorpusReorderTests(unittest.TestCase):
|
||||||
|
"""`rerank_ids` against a real (SQLite) corpus, where the text comes from."""
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
self.engine = create_engine("sqlite://", connect_args={"check_same_thread": False},
|
||||||
|
poolclass=StaticPool)
|
||||||
|
Base.metadata.create_all(self.engine)
|
||||||
|
self.db = Session(self.engine)
|
||||||
|
self.db.add(User(id=1, name="Mod", email="mod@example.test",
|
||||||
|
hashed_password="unused", role="moderator"))
|
||||||
|
for question_id, text in [
|
||||||
|
(1, "A child with fever and a rash"),
|
||||||
|
(2, "A toddler with a barking cough and stridor"),
|
||||||
|
(3, "An infant with jaundice on day three"),
|
||||||
|
]:
|
||||||
|
self.db.add(Question(id=question_id, user_id=1, question_text=text,
|
||||||
|
question_type="mcq", options=["yes", "no"], correct_answer="yes"))
|
||||||
|
self.db.commit()
|
||||||
|
self.addCleanup(self.engine.dispose)
|
||||||
|
self.addCleanup(self.db.close)
|
||||||
|
|
||||||
|
def test_the_reranker_sees_the_row_text_and_reorders_by_it(self):
|
||||||
|
seen = {}
|
||||||
|
|
||||||
|
def fake(query, documents):
|
||||||
|
seen["documents"] = documents
|
||||||
|
return [0.1, 0.9, 0.2]
|
||||||
|
|
||||||
|
with patch.object(rerank_service, "is_configured", return_value=True), \
|
||||||
|
patch.object(rerank_service, "rerank", side_effect=fake):
|
||||||
|
out = search_service.rerank_ids(self.db, "croup", "question", [1, 2, 3])
|
||||||
|
self.assertEqual(out, [2, 3, 1])
|
||||||
|
self.assertIn("barking cough", " ".join(seen["documents"]))
|
||||||
|
|
||||||
|
def test_an_unconfigured_reranker_costs_neither_a_query_nor_the_order(self):
|
||||||
|
with patch.object(rerank_service, "is_configured", return_value=False), \
|
||||||
|
patch.object(rerank_service, "rerank", side_effect=AssertionError("no call")):
|
||||||
|
self.assertEqual(search_service.rerank_ids(self.db, "croup", "question", [3, 1, 2]),
|
||||||
|
[3, 1, 2])
|
||||||
|
|
||||||
|
def test_ids_that_are_no_longer_in_the_table_are_still_returned(self):
|
||||||
|
# Retrieval decided these exist; a row deleted between the two queries
|
||||||
|
# is scored on empty text, and is still on the page afterwards.
|
||||||
|
with patch.object(rerank_service, "is_configured", return_value=True), \
|
||||||
|
patch.object(rerank_service, "rerank", return_value=[0.2, 0.8, 0.5]):
|
||||||
|
out = search_service.rerank_ids(self.db, "croup", "question", [1, 2, 999])
|
||||||
|
self.assertEqual(out, [2, 999, 1])
|
||||||
|
|
||||||
|
def test_a_score_list_that_does_not_match_the_shortlist_is_refused(self):
|
||||||
|
with patch.object(rerank_service, "is_configured", return_value=True), \
|
||||||
|
patch.object(rerank_service, "rerank", return_value=[0.2, 0.8]):
|
||||||
|
self.assertEqual(search_service.rerank_ids(self.db, "croup", "question", [1, 2, 3]),
|
||||||
|
[1, 2, 3])
|
||||||
|
|
||||||
|
|
||||||
|
class AiModeShortlistTests(unittest.TestCase):
|
||||||
|
"""The shortlist is the answer's evidence, so membership matters most."""
|
||||||
|
|
||||||
|
def sources(self):
|
||||||
|
return [
|
||||||
|
{"kind": "section", "id": 1, "title": "Croup", "text": "barking cough", "score": 1.0},
|
||||||
|
{"kind": "section", "id": 2, "title": "Asthma", "text": "wheeze", "score": 0.5},
|
||||||
|
{"kind": "question", "id": 9, "title": "A toddler", "text": "stridor", "score": 1.0},
|
||||||
|
]
|
||||||
|
|
||||||
|
def test_one_scale_replaces_four_incomparable_ones(self):
|
||||||
|
from app.services import ai_mode_service
|
||||||
|
|
||||||
|
sources = self.sources()
|
||||||
|
with patch.object(ai_mode_service.rerank_service, "rerank",
|
||||||
|
return_value=[0.2, 0.1, 0.9]):
|
||||||
|
ai_mode_service._cross_encode("croup", sources)
|
||||||
|
# The best section and the best question both scored 1.0 before; now
|
||||||
|
# every source sits on one ranking, so the sort can mean something.
|
||||||
|
self.assertEqual([s["score"] for s in sources], [0.5, 1 / 3, 1.0])
|
||||||
|
|
||||||
|
def test_silence_leaves_every_score_exactly_as_it_was(self):
|
||||||
|
from app.services import ai_mode_service
|
||||||
|
|
||||||
|
sources = self.sources()
|
||||||
|
with patch.object(ai_mode_service.rerank_service, "rerank", return_value=None):
|
||||||
|
ai_mode_service._cross_encode("croup", sources)
|
||||||
|
self.assertEqual([s["score"] for s in sources], [1.0, 0.5, 1.0])
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
181
docs/reranking.md
Normal file
181
docs/reranking.md
Normal file
|
|
@ -0,0 +1,181 @@
|
||||||
|
# Reranking: the second opinion on order
|
||||||
|
|
||||||
|
Retrieval fuses two rankers that never see the query and the document together.
|
||||||
|
A bi-encoder embedded every row long before anybody typed anything, so cosine
|
||||||
|
distance answers "is this the same topic"; `ts_rank_cd` answers "do these words
|
||||||
|
coincide". Neither answers "does this document answer this question", and that
|
||||||
|
is the question a learner is actually asking.
|
||||||
|
|
||||||
|
A cross-encoder answers it, at the cost of one forward pass per pair, per query,
|
||||||
|
which is why it can only ever be given a shortlist the cheap rankers produced.
|
||||||
|
It reorders. It never decides what exists.
|
||||||
|
|
||||||
|
## What is served, and where the query goes
|
||||||
|
|
||||||
|
The proxy at `LITELLM_API_BASE` already serves three rerankers — `/model/info`
|
||||||
|
reports `"mode": "rerank"` for each:
|
||||||
|
|
||||||
|
| model | notes |
|
||||||
|
|---|---|
|
||||||
|
| `cohere-rerank-v4.0-pro` | **in use** — the best of the three here, see the numbers below |
|
||||||
|
| `cohere-rerank-v4.0-fast` | roughly 150 ms quicker, and measurably worse on this corpus |
|
||||||
|
| `jina-reranker-v2-base-multilingual` | open-weight, but rate-limited to the point of unusability on this proxy: a 60-query sweep got 429s on nearly every call |
|
||||||
|
|
||||||
|
Both `POST /rerank` and `POST /v1/rerank` answer, with the body Cohere defined —
|
||||||
|
`{model, query, documents, top_n}` in, `{results: [{index, relevance_score}]}`
|
||||||
|
out. Jina, BGE and LiteLLM's own rerank route all speak it, so replacing the
|
||||||
|
model behind `LITELLM_RERANK_MODEL` with a self-hosted open-weight one is a
|
||||||
|
configuration change and nothing else.
|
||||||
|
|
||||||
|
## Where it is applied
|
||||||
|
|
||||||
|
| place | reranked | why |
|
||||||
|
|---|---|---|
|
||||||
|
| AI Mode shortlist (`ai_mode_service.retrieve`) | yes | The shortlist *is* the evidence the answer may use, and the turn is already waiting on a chat model. It is also the only place where four corpora have to be compared with each other — see below. |
|
||||||
|
| Question search (`/questions?q=`, `/search`) | yes | A page of results is a choice being made. |
|
||||||
|
| Article search (`/articles?q=`, `/search`) | yes, through its sections | An article row is a title, a summary and a topical vector; the prose is in its sections. Handing a cross-encoder the article row is handing it a stub and asking about a document it cannot see. Its best section already decides its place in the fusion, so reranking the sections reranks the articles. |
|
||||||
|
| Test builder from a description | yes | A sentence about what somebody wants to study is exactly the query shape a cross-encoder is trained on, and the top of that list becomes the test rather than a page to scroll past. |
|
||||||
|
| Test builder from an uploaded file | **no** | The "query" is a whole handout. Clipped to 2000 characters it judges every candidate against whichever part that happened to be — confidently, and about the wrong thing. |
|
||||||
|
| Flashcard and media search | **no** | Short fields (front/back, title/caption) where the lexical ranker is already doing the whole job, and both are small personal collections. Latency for nothing. |
|
||||||
|
| Typeahead (`/search/suggest`) | **no** | It runs on every keystroke and has half a word to judge. It is prefix matching on titles, which is what finishing a word wants. Adding a 400 ms round trip per character would make the fastest thing on the site the slowest. |
|
||||||
|
|
||||||
|
## Latency
|
||||||
|
|
||||||
|
Measured from the backend container against `cohere-rerank-v4.0-pro`, median of
|
||||||
|
three, real question stems as documents:
|
||||||
|
|
||||||
|
| candidates | round trip |
|
||||||
|
|---|---|
|
||||||
|
| 10 | 250 ms |
|
||||||
|
| 30 | 380 ms |
|
||||||
|
| 50 | 430 ms |
|
||||||
|
| 100 | 580 ms |
|
||||||
|
|
||||||
|
Mostly fixed cost, then roughly 2 ms a pair. `MAX_CANDIDATES = 50`: the fused
|
||||||
|
pool is 60 wide and a page shows ten, so past fifty the model is being paid to
|
||||||
|
confirm that rows neither cheap ranker wanted are still not wanted.
|
||||||
|
|
||||||
|
In place, including reading the candidate text out of Postgres, over 120 live
|
||||||
|
queries never asked before: **median 480 ms, p90 550 ms**. AI Mode's whole
|
||||||
|
retrieval step goes from about 85 ms to about 500 ms — one call over roughly
|
||||||
|
forty candidates across all four corpora.
|
||||||
|
|
||||||
|
Two caches sit in front of that, and both matter more than the raw figure. Ours
|
||||||
|
is Redis, keyed on model, query and the document text, for a day: **3–5 ms** on
|
||||||
|
a hit, which is most repeat queries and every reload of a results page. The
|
||||||
|
proxy keeps its own cache of identical rerank requests and answers those in
|
||||||
|
about 30 ms, so a query one learner has already asked is cheap for the next even
|
||||||
|
after our key expires. A full four-corpus search page measured 94 ms warm,
|
||||||
|
236 ms with our cache cleared but the proxy's still warm, and about a second on
|
||||||
|
a genuinely new query, which is two corpora reranked one after the other.
|
||||||
|
Running the per-corpus finders concurrently would recover half of that and has
|
||||||
|
not been done.
|
||||||
|
|
||||||
|
## Does it help
|
||||||
|
|
||||||
|
Two measurements, both against labels neither ranker produced.
|
||||||
|
|
||||||
|
**Questions**, query = a disease tag's name, relevant = the questions carrying
|
||||||
|
that tag, 60 tags sampled at random from those with 3–40 questions:
|
||||||
|
|
||||||
|
| | precision@3 | MRR@10 |
|
||||||
|
|---|---|---|
|
||||||
|
| fused | 0.394 | 0.595 |
|
||||||
|
| + `cohere-rerank-v4.0-pro` | **0.483** | **0.641** |
|
||||||
|
| + `cohere-rerank-v4.0-fast` | 0.439 | 0.642 |
|
||||||
|
|
||||||
|
Top result changed on 38 of 60; precision@3 improved on 21 and fell on 9.
|
||||||
|
|
||||||
|
**Article sections**, query = an article's title, relevant = that article's
|
||||||
|
sections, 60 articles with three or more sections:
|
||||||
|
|
||||||
|
| | precision@3 | MRR@10 |
|
||||||
|
|---|---|---|
|
||||||
|
| fused | 0.772 | 0.847 |
|
||||||
|
| + `cohere-rerank-v4.0-pro` | **0.833** | **0.908** |
|
||||||
|
| + `cohere-rerank-v4.0-fast` | 0.811 | 0.917 |
|
||||||
|
|
||||||
|
Top result changed on 41 of 60; improved on 12, fell on 4.
|
||||||
|
|
||||||
|
### Queries where the old top result was wrong
|
||||||
|
|
||||||
|
* *"management of bronchiolitis in an infant"* — fused first three sections were
|
||||||
|
influenza transmission, foreign body ingestion, and the management of a
|
||||||
|
pregnancy. Reranked: three supportive-management sections, the third of them
|
||||||
|
bronchiolitis's own and the first two transient tachypnoea of the newborn.
|
||||||
|
Better, and not right: on this corpus the reranker will take a *Management*
|
||||||
|
section from a neighbouring topic over a differently-titled one from the right
|
||||||
|
topic.
|
||||||
|
* *"when do you image a first febrile seizure"* — fused returned the definition
|
||||||
|
of a febrile seizure. Reranked returned "routine laboratory studies and acute
|
||||||
|
neuroimaging are unnecessary…", which is the answer.
|
||||||
|
* *"teenager with knee pain worse after sport"* — fused led with a 6-year-old
|
||||||
|
girl's case and, in the articles, with bone pain that wakes a child at night.
|
||||||
|
Reranked led with a 15-year-old's sports knee pain, and with patellofemoral
|
||||||
|
pain syndrome.
|
||||||
|
* *"delayed passage of meconium"* — fused returned functional constipation.
|
||||||
|
Reranked returned "obstruction presents with failure to pass meconium within
|
||||||
|
24–48 hours of birth".
|
||||||
|
* *"child limping with a fever and refusing to bear weight"* — fused led with a
|
||||||
|
urinary tract infection's presentation and acute gastroenteritis. Reranked led
|
||||||
|
with osteomyelitis, septic arthritis, and the differential for a limping
|
||||||
|
toddler.
|
||||||
|
* *"why does my patient with milk in the bottle have anemia"* (AI Mode) — fused
|
||||||
|
led with the general anemia article; reranked led with iron deficiency anemia.
|
||||||
|
|
||||||
|
### Where it does not help, and one place it is worse
|
||||||
|
|
||||||
|
The gain is uneven across corpora, and the reason is worth keeping in mind: a
|
||||||
|
bank question is a vignette written so as *not* to name its diagnosis. Asked
|
||||||
|
"what causes croup", the reranker prefers a question that says the word croup
|
||||||
|
in passing (a post-influenza bacterial tracheitis case) over the vignette with
|
||||||
|
the barking cough that never says it. That is defensible — someone searching
|
||||||
|
"croup" is not badly served by either — but it is the bi-encoder's strength being
|
||||||
|
partly traded away, and it is why the question-corpus gain (+0.09 precision@3)
|
||||||
|
is smaller than the prose-corpus one relative to how wrong the old order was.
|
||||||
|
|
||||||
|
The failure to watch for is a query the library does not cover. "How is croup
|
||||||
|
treated at home", against a corpus with no croup treatment section, is answered
|
||||||
|
with *Heat-Related Illness › Treatment*, *Tetanus › Treatment*, *Patellofemoral
|
||||||
|
Pain Syndrome › Treatment* — the cross-encoder matching the shape of the
|
||||||
|
question when it cannot match the subject. Fusion's answer was no better, but it
|
||||||
|
was wrong in a way that looked wrong. Two things contain this: AI Mode's
|
||||||
|
closeness gate is measured on cosine over the whole corpus and still says
|
||||||
|
"nothing here covers this" regardless of what the reranker thought (see
|
||||||
|
retrieval-thresholds.md), and search results carry their own snippet. A minimum
|
||||||
|
relevance score would be the third, and it is deliberately not implemented: it
|
||||||
|
would need the same measured calibration the cosine thresholds have, on a scale
|
||||||
|
that changes with the model, and getting it wrong empties a page.
|
||||||
|
|
||||||
|
## When it is not there
|
||||||
|
|
||||||
|
Unset `LITELLM_RERANK_MODEL` (or the `settings:rerank_model` override in Redis,
|
||||||
|
which wins) and every list is in fusion order, with no error anywhere. The same
|
||||||
|
is true of a proxy that is down, slow, rate-limiting, or answering with
|
||||||
|
something that does not line up with the request: `rerank()` returns `None` and
|
||||||
|
every caller reads that as "keep what you had". The 429 storm from the Jina
|
||||||
|
model above is what that looks like in practice — a warning line per call and
|
||||||
|
results in their previous order.
|
||||||
|
|
||||||
|
`POST /admin/rerank/test` checks the configured model end to end, including
|
||||||
|
whether it puts an obviously relevant document above an obviously irrelevant
|
||||||
|
one. A reranker that answers 200 and ranks the decoy first is worse than one
|
||||||
|
that is switched off, and nothing else on the site would ever tell you.
|
||||||
|
|
||||||
|
## Re-measuring
|
||||||
|
|
||||||
|
```
|
||||||
|
docker compose exec -T backend python - <<'EOF'
|
||||||
|
from app.database import SessionLocal
|
||||||
|
from app.services import search_service as ss
|
||||||
|
db = SessionLocal()
|
||||||
|
for q in ["management of bronchiolitis in an infant", "delayed passage of meconium"]:
|
||||||
|
ranked, _ = ss.hybrid_ids(db, q, "article_section", limit=200)
|
||||||
|
print(q, "\n before", ranked[:3], "\n after ", ss.rerank_ids(db, q, "article_section", ranked)[:3])
|
||||||
|
EOF
|
||||||
|
```
|
||||||
|
|
||||||
|
The precision figures above come from labels that already exist in the database
|
||||||
|
— disease tags on questions, and an article owning its sections. Re-run that
|
||||||
|
comparison rather than judging by eye after a corpus change; ten queries read by
|
||||||
|
hand will agree with whichever ordering was looked at second.
|
||||||
|
|
@ -54,6 +54,31 @@ It feels adjacent to a paediatrics library — attachment, behaviour — and it
|
||||||
not: 0.49 is where anything written in English lands against any corpus. That is
|
not: 0.49 is where anything written in English lands against any corpus. That is
|
||||||
the reading to keep in mind. A number in the 0.4s is noise, not a weak signal.
|
the reading to keep in mind. A number in the 0.4s is noise, not a weak signal.
|
||||||
|
|
||||||
|
## Why these are not the reranker's score
|
||||||
|
|
||||||
|
Retrieval now has a cross-encoder in it (docs/reranking.md), which is a better
|
||||||
|
judge of a query-document pair than cosine distance is by a wide margin. The
|
||||||
|
three-state decision still does not go through it, on purpose.
|
||||||
|
|
||||||
|
The question here is *"is there anything in this library about this at all"*,
|
||||||
|
and that is a question about the corpus, not about a shortlist. `top_similarity`
|
||||||
|
answers it by scanning every embedded row in two corpora through the vector
|
||||||
|
index, in about 25 ms. A reranker can only score the candidates something else
|
||||||
|
already shortlisted, so a reranked closeness could not tell "the library does
|
||||||
|
not cover this" from "retrieval had a bad day" — and on this corpus it is
|
||||||
|
exactly the uncovered query where the cross-encoder is least trustworthy: asked
|
||||||
|
how croup is treated at home, with no croup treatment section in the library, it
|
||||||
|
promotes the *Treatment* section of whatever else is lying around.
|
||||||
|
|
||||||
|
There is also a failure argument. Order is a preference, so a reranker being
|
||||||
|
down costs a worse-ordered page and nothing else. Whether the answer *claims to
|
||||||
|
come from the library* is a promise, and putting a network hop in the path of a
|
||||||
|
promise means a proxy restart changes what the assistant asserts.
|
||||||
|
|
||||||
|
So: the cross-encoder decides the order of the shortlist; cosine decides what
|
||||||
|
the answer is allowed to say about it. The numbers below are unchanged and did
|
||||||
|
not need re-measuring, because nothing that feeds them changed.
|
||||||
|
|
||||||
## Why these are not `SEMANTIC_FLOOR`
|
## Why these are not `SEMANTIC_FLOOR`
|
||||||
|
|
||||||
`search_service.SEMANTIC_FLOOR` (0.45) decides what is worth putting in a list,
|
`search_service.SEMANTIC_FLOOR` (0.45) decides what is worth putting in a list,
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue