diff --git a/src/lib/clinical-search.ts b/src/lib/clinical-search.ts index 4f7f03992..e94a2ec8b 100644 --- a/src/lib/clinical-search.ts +++ b/src/lib/clinical-search.ts @@ -1178,11 +1178,19 @@ const genericMedicationDoseQueryTokens = new Set([ "sl", ]); -/** Require dose evidence to carry the medication question's clinical subject. */ +/** Require medication evidence to carry the medication question's clinical subject. */ export function medicationDoseQuerySubjectTokens(query: string) { return normalizedClinicalSearchTokens(query).filter((token) => !genericMedicationDoseQueryTokens.has(token)); } +/** Canonical medication subjects named by a monitoring query. */ +export function medicationMonitoringQuerySubjectTokens(query: string) { + return unique( + medicationTerms(normalizeAnalysisText(query)).flatMap((term) => normalizedClinicalSearchTokens(term)), + 12, + ); +} + const medicationDoseAmountQueryPattern = /\b(?:dose|doses|dosage|dosages|dosing|amount|amounts|maximum|max|min|minimum|mg|mcg|micrograms?|milligrams?|ug)\b|[µμ]g|\bhow\s+much\b/i; const medicationDoseRouteQueryPattern = diff --git a/src/lib/rag.ts b/src/lib/rag.ts index 5c9e068c4..2939b5a67 100644 --- a/src/lib/rag.ts +++ b/src/lib/rag.ts @@ -549,15 +549,21 @@ function recordRetrievalLayer( } } +/** Whether the current result set carries bounded second-stage release scores. */ +function resultsHaveReleaseRankScore(results: SearchResult[]) { + return results.some((result) => result.score_explanation?.releaseRankScore !== undefined); +} + /** - * Keep distinct results in the current ranked order while upgrading duplicate chunks to the - * strongest released-hybrid copy. + * Keep the released result order on the live-eval-proven hybrid and bounded second-stage signals. * * App-layer rank scores remain available to answer evidence ranking and telemetry, but the - * live corpus gate has not validated replacing the final distinct-result order with them. + * live corpus gate has not validated using them as the final retrieval order. Resolve duplicate + * chunks to their strongest released-hybrid copy before sorting the distinct results. */ -export function stabilizeReleasedSearchOrder(results: SearchResult[]) { - const compareReleasedSearchStrength = (left: SearchResult, right: SearchResult) => { +export function stabilizeReleasedSearchOrder(results: SearchResult[], preferSecondStageScore = false) { + const useSecondStageReleaseOrder = preferSecondStageScore && resultsHaveReleaseRankScore(results); + const compareReleasedHybridStrength = (left: SearchResult, right: SearchResult) => { const leftHybrid = left.hybrid_score ?? left.similarity ?? 0; const rightHybrid = right.hybrid_score ?? right.similarity ?? 0; if (rightHybrid !== leftHybrid) return rightHybrid - leftHybrid; @@ -568,18 +574,30 @@ export function stabilizeReleasedSearchOrder(results: SearchResult[]) { return (right.relevance?.score ?? 0) - (left.relevance?.score ?? 0); return left.id.localeCompare(right.id); }; + const compareReleasedSearchOrder = (left: SearchResult, right: SearchResult) => { + if (!useSecondStageReleaseOrder) return compareReleasedHybridStrength(left, right); + const leftReleaseScore = left.score_explanation?.releaseRankScore ?? left.hybrid_score ?? left.similarity ?? 0; + const rightReleaseScore = right.score_explanation?.releaseRankScore ?? right.hybrid_score ?? right.similarity ?? 0; + if (rightReleaseScore !== leftReleaseScore) return rightReleaseScore - leftReleaseScore; + const leftSimilarity = left.similarity ?? 0; + const rightSimilarity = right.similarity ?? 0; + if (rightSimilarity !== leftSimilarity) return rightSimilarity - leftSimilarity; + if (right.relevance?.score !== left.relevance?.score) + return (right.relevance?.score ?? 0) - (left.relevance?.score ?? 0); + return left.id.localeCompare(right.id); + }; const strongestById = new Map(); for (const result of results) { const current = strongestById.get(result.id); - if (!current || compareReleasedSearchStrength(result, current) < 0) strongestById.set(result.id, result); - } - const deduped: SearchResult[] = []; - const seen = new Set(); - for (const result of results) { - if (seen.has(result.id)) continue; - seen.add(result.id); - deduped.push(strongestById.get(result.id) ?? result); + if (!current || compareReleasedHybridStrength(result, current) < 0) strongestById.set(result.id, result); } + const deduped = [...strongestById.values()] + .sort(compareReleasedSearchOrder) + .map((result, index) => + result.score_explanation + ? { ...result, score_explanation: { ...result.score_explanation, finalRank: index + 1 } } + : result, + ); results.length = 0; results.push(...deduped); return results; @@ -601,7 +619,9 @@ function recordSearchScoreTelemetry(telemetry: SearchTelemetry, results: SearchR return; } - stabilizeReleasedSearchOrder(results); + const useSecondStageReleaseOrder = resultsHaveReleaseRankScore(results); + telemetry.second_stage_rerank_used = useSecondStageReleaseOrder; + stabilizeReleasedSearchOrder(results, useSecondStageReleaseOrder); const coverageScores = results .map((result) => Math.max(0, result.hybrid_score ?? result.similarity ?? 0)) .sort((left, right) => right - left); @@ -719,6 +739,12 @@ export function applySecondStageRerankIfNeeded(args: { const secondStage = secondStageScore(result, args.queryClass, index); let rankScore = secondStage.rankScore; let confidenceAdjustment = secondStage.adjustment; + const releasedHybridScore = result.hybrid_score ?? result.similarity ?? 0; + let releaseRankScore = Math.max( + releasedHybridScore, + (result.score_explanation?.finalScore ?? result.hybrid_score ?? result.similarity ?? 0) + + secondStage.adjustment, + ); const priorOccurrences = seenPerDocument.get(result.document_id) ?? 0; seenPerDocument.set(result.document_id, priorOccurrences + 1); if (rankingConfig.documentDiversityPenalty > 0 && priorOccurrences > 0) { @@ -728,6 +754,16 @@ export function applySecondStageRerankIfNeeded(args: { ); rankScore -= diversityPenalty; confidenceAdjustment -= diversityPenalty; + releaseRankScore -= diversityPenalty; + } + const selectionReasons = result.match_explanation?.reasons ?? []; + const clinicalSubjectRequired = selectionReasons.includes("retrieval_required_signal:clinical_subject"); + const clinicalSubjectMatched = selectionReasons.includes("retrieval_signal:clinical_subject"); + if (clinicalSubjectRequired && !clinicalSubjectMatched) { + // A wrong-medication chunk can carry attractive numeric dose/monitoring signals. Keep it + // available at its released hybrid strength, but do not let second-stage evidence boosts + // promote it above chunks that contain the medication subject requested by the query. + releaseRankScore = Math.min(releaseRankScore, releasedHybridScore); } const finalScore = Math.min( 1, @@ -745,6 +781,7 @@ export function applySecondStageRerankIfNeeded(args: { ? { ...result.score_explanation, rankScore: Number(rankScore.toFixed(4)), + releaseRankScore: Number(releaseRankScore.toFixed(4)), preClampFinalScore: Number(rankScore.toFixed(4)), finalScore: Number(finalScore.toFixed(4)), } diff --git a/src/lib/retrieval-selection.ts b/src/lib/retrieval-selection.ts index 57abb4a58..35779ca27 100644 --- a/src/lib/retrieval-selection.ts +++ b/src/lib/retrieval-selection.ts @@ -3,6 +3,7 @@ import { medicationDoseEvidenceQueryIntent, medicationDoseQueryContext, medicationDoseQuerySubjectTokens, + medicationMonitoringQuerySubjectTokens, } from "@/lib/clinical-search"; import type { RagQueryClass, @@ -236,6 +237,18 @@ function signalMatchesText(signal: string, text: string) { } } +function medicationClinicalSubjectMatches(query: string, result: SearchResult, normalizedEvidenceText: string) { + const doseIntent = medicationDoseEvidenceQueryIntent(query); + if (doseIntent.asksAmount || doseIntent.asksRoute || doseIntent.asksFrequency) { + return medicationDoseQueryContext(query, result).matched; + } + const subjectTokens = medicationMonitoringQuerySubjectTokens(query); + if (!subjectTokens.length) return true; + const evidenceTokens = new Set(normalizedEvidenceText.split(" ")); + const hitCount = subjectTokens.filter((token) => evidenceTokens.has(token)).length; + return hitCount >= Math.min(2, subjectTokens.length); +} + function matchedSignalsForResult(args: { query: string; intent: RetrievalIntent; @@ -267,7 +280,7 @@ function matchedSignalsForResult(args: { if (args.intent.needsDoseRouteFrequency && hasRoute(text)) signals.push("route"); if ( args.intent.requiredTermSignals.includes("clinical_subject") && - medicationDoseQueryContext(args.query, args.result).matched + medicationClinicalSubjectMatches(args.query, args.result, text) ) { signals.push("clinical_subject"); } @@ -358,6 +371,11 @@ export function buildRetrievalIntent(query: string, queryClass: RagQueryClass): const asksDoseRoute = medicationEvidenceIntent.asksAmount || medicationEvidenceIntent.asksRoute || medicationEvidenceIntent.asksFrequency; const asksDoseAmount = medicationEvidenceIntent.asksAmount; + const asksMedicationMonitoring = + queryClass === "medication_dose_risk" && /\b(?:monitor\w*|baseline|blood|level)\b/.test(normalizedQuery); + const clinicalSubjectTokens = asksMedicationMonitoring + ? medicationMonitoringQuerySubjectTokens(query) + : medicationDoseQuerySubjectTokens(query); const asksTable = /\b(?:table|chart|matrix|threshold|cutoff|cut off|range|criteria|row)\b/.test(normalizedQuery); const asksSourceImage = /\b(?:source|show|open|view|display|see)\b.*\b(?:image|figure|visual|table|chart|matrix)\b/.test(normalizedQuery) || @@ -400,9 +418,13 @@ export function buildRetrievalIntent(query: string, queryClass: RagQueryClass): if (asksDoseRoute) { if (asksDoseAmount) requiredTermSignals.push("dose_amount"); if (medicationEvidenceIntent.asksRoute) requiredTermSignals.push("route"); - if (queryClass === "medication_dose_risk" && medicationDoseQuerySubjectTokens(query).length > 0) { - requiredTermSignals.push("clinical_subject"); - } + } + if ( + queryClass === "medication_dose_risk" && + (asksDoseRoute || asksMedicationMonitoring) && + clinicalSubjectTokens.length > 0 + ) { + requiredTermSignals.push("clinical_subject"); } if (asksFlowchart) { preferredDocumentSignals.push("flowchart", "pathway", "risk matrix"); @@ -490,9 +512,13 @@ function annotateResultWithSelection( result: SearchResult, candidate: RetrievalCandidate, originalScore: number, + intent: RetrievalIntent, ): SearchResult { const score = Number(Math.max(originalScore, candidate.score).toFixed(4)); const selectionReasons = candidate.matchedSignals.map((signal) => `retrieval_signal:${signal}`); + if (intent.requiredTermSignals.includes("clinical_subject")) { + selectionReasons.push("retrieval_required_signal:clinical_subject"); + } if (candidate.score > originalScore + 0.04) selectionReasons.push("retrieval_intent_rescue"); return { @@ -593,7 +619,7 @@ export function selectRetrievalEvidence(args: { .map((candidate) => { const result = byId.get(candidate.chunkId); if (!result) return null; - return annotateResultWithSelection(result, candidate, originalScoreById.get(candidate.chunkId) ?? 0); + return annotateResultWithSelection(result, candidate, originalScoreById.get(candidate.chunkId) ?? 0, intent); }) .filter((result): result is SearchResult => Boolean(result)); const rescueApplied = selectedCandidates.some( diff --git a/src/lib/types.ts b/src/lib/types.ts index eb4fbe85c..d2441c88b 100644 --- a/src/lib/types.ts +++ b/src/lib/types.ts @@ -421,6 +421,8 @@ export type SearchScoreExplanation = { * this retains the magnitude of boosts and signed penalties after confidence saturates. */ rankScore: number; + /** Live-eval-gated retrieval score used for the final released search order. */ + releaseRankScore?: number; /** Existing public confidence signal, clamped to the inclusive 0-1 range. */ finalScore: number; finalRank?: number; diff --git a/tests/clinical-search.test.ts b/tests/clinical-search.test.ts index 10c14170c..8cddfaff9 100644 --- a/tests/clinical-search.test.ts +++ b/tests/clinical-search.test.ts @@ -9,6 +9,7 @@ import { hasDoseEvidenceSupport, hasNumericOrTableEvidence, hasStructuredThresholdEvidence, + medicationMonitoringQuerySubjectTokens, normalizedClinicalSearchTokens, rankClinicalResults, } from "../src/lib/clinical-search"; @@ -34,6 +35,13 @@ function result(overrides: Partial): SearchResult { } describe("clinical search query normalization", () => { + it("keeps only the medication subject for monitoring questions", () => { + expect(medicationMonitoringQuerySubjectTokens("What monitoring is required for lithium therapy?")).toEqual([ + "lithium", + ]); + expect(medicationMonitoringQuerySubjectTokens("What monitoring escalation is required?")).toEqual([]); + }); + it("classifies common RAG query shapes for routing and observability", () => { expect(classifyRagQuery("Find the NOCC document").queryClass).toBe("document_lookup"); expect(classifyRagQuery("What should a patient safety plan include?").queryClass).toBe("document_lookup"); diff --git a/tests/rag-second-stage-ranking.test.ts b/tests/rag-second-stage-ranking.test.ts index 52c7dfe81..d6979d2e9 100644 --- a/tests/rag-second-stage-ranking.test.ts +++ b/tests/rag-second-stage-ranking.test.ts @@ -106,9 +106,47 @@ describe("second-stage rank score", () => { expect(ranked[0].id).toBe("high-rank"); expect(ranked[0].score_explanation?.rankScore).toBe(1.49); + expect(ranked[0].score_explanation?.releaseRankScore).toBe(0.8); + expect(ranked[1].score_explanation?.releaseRankScore).toBeGreaterThan(0.8); expect(ranked[0].score_explanation?.preClampFinalScore).toBe(1.49); expect(ranked[0].score_explanation?.finalScore).toBe(0.69); expect(ranked.map((item) => item.score_explanation?.finalRank)).toEqual([1, 2]); + expect(stabilizeReleasedSearchOrder(ranked, true).map((item) => item.id)).toEqual(["runner-up", "high-rank"]); + }); + + it("caps wrong-subject monitoring evidence at its released hybrid score", () => { + const wrongSubject = explanation(1.4); + wrongSubject.finalScore = 1; + const rightSubject = explanation(0.79); + const ranked = applySecondStageRerankIfNeeded({ + queryClass: "medication_dose_risk", + results: [ + result({ + id: "cardiac-dose", + document_id: "cardiac-doc", + content: "Cardiac monitoring includes a 5 mg treatment note.", + hybrid_score: 0.8, + score_explanation: wrongSubject, + match_explanation: { reasons: ["retrieval_required_signal:clinical_subject"] }, + }), + result({ + id: "lithium-monitoring", + document_id: "lithium-doc", + content: "Lithium monitoring requires serum levels.", + hybrid_score: 0.79, + score_explanation: rightSubject, + match_explanation: { + reasons: ["retrieval_required_signal:clinical_subject", "retrieval_signal:clinical_subject"], + }, + }), + ], + telemetry: {} as SearchTelemetry, + topK: 2, + }); + + const stabilized = stabilizeReleasedSearchOrder(ranked, true); + expect(stabilized.map((item) => item.id)).toEqual(["lithium-monitoring", "cardiac-dose"]); + expect(stabilized[1].score_explanation?.releaseRankScore).toBe(0.8); }); it("sorts by the computed second-stage score even when score explanations are absent", () => { @@ -137,7 +175,7 @@ describe("second-stage rank score", () => { expect(ranked[1].score_explanation?.finalRank).toBe(2); }); - it("keeps distinct rank order while upgrading duplicate chunks to the strongest released-hybrid copy", () => { + it("sorts distinct chunks by the bounded release score and keeps the strongest duplicate copy", () => { const telemetry = {} as SearchTelemetry; const strongerHybridDuplicate = result({ id: "stronger-hybrid", @@ -180,10 +218,54 @@ describe("second-stage rank score", () => { }); expect(secondStage.map((item) => item.id)).toEqual(["higher-rank-score", "stronger-hybrid", "stronger-hybrid"]); - const stabilized = stabilizeReleasedSearchOrder(secondStage); + const stabilized = stabilizeReleasedSearchOrder(secondStage, true); expect(stabilized.map((item) => item.id)).toEqual(["higher-rank-score", "stronger-hybrid"]); expect(stabilized[1].hybrid_score).toBe(0.8); expect(stabilized[1].content).toBe(strongerHybridDuplicate.content); + expect(stabilized[0].score_explanation?.releaseRankScore).toBeGreaterThan( + stabilized[1].score_explanation?.releaseRankScore ?? 0, + ); + expect(stabilized.map((item) => item.score_explanation?.finalRank)).toEqual([1, 2]); + }); + + it("ignores sticky second-stage preference when the final result set has no releaseRankScore", () => { + const clinicallyOrdered = [ + result({ + id: "lithium-monitoring", + hybrid_score: 0.79, + score_explanation: explanation(0.9), + }), + result({ + id: "cardiac-dose", + hybrid_score: 0.8, + score_explanation: explanation(0.5), + }), + ]; + + expect(stabilizeReleasedSearchOrder([...clinicallyOrdered], true).map((item) => item.id)).toEqual( + stabilizeReleasedSearchOrder([...clinicallyOrdered], false).map((item) => item.id), + ); + }); + + it("only applies bounded release ordering when the current result set carries releaseRankScore", () => { + const withReleaseScore = result({ + id: "release-ranked", + hybrid_score: 0.79, + score_explanation: { ...explanation(0.79), releaseRankScore: 0.95 }, + }); + const withoutReleaseScore = result({ + id: "hybrid-leading", + hybrid_score: 0.84, + score_explanation: explanation(0.84), + }); + + expect(stabilizeReleasedSearchOrder([withoutReleaseScore, withReleaseScore], true).map((item) => item.id)).toEqual([ + "release-ranked", + "hybrid-leading", + ]); + expect(stabilizeReleasedSearchOrder([withoutReleaseScore, withReleaseScore], false).map((item) => item.id)).toEqual( + ["hybrid-leading", "release-ranked"], + ); }); }); diff --git a/tests/retrieval-selection.test.ts b/tests/retrieval-selection.test.ts index 7e77d327f..3889ddf1d 100644 --- a/tests/retrieval-selection.test.ts +++ b/tests/retrieval-selection.test.ts @@ -194,6 +194,36 @@ describe("retrieval source selection", () => { expect(selected.summary.matchedSignals).toContain("route"); }); + it("anchors medication-monitoring selection to the requested clinical subject", () => { + const selection = selectRetrievalEvidence({ + query: "What monitoring is required for lithium therapy?", + queryClass: "medication_dose_risk", + topK: 2, + maxResultsPerDocument: 2, + results: [ + source({ + id: "generic-cardiac-monitoring", + document_id: "cardiac-doc", + title: "Cardiac Monitoring", + content: "Admission monitoring requirements include continuous cardiac observations.", + hybrid_score: 0.95, + }), + source({ + id: "lithium-monitoring", + document_id: "lithium-doc", + title: "Lithium Therapy", + content: "Lithium monitoring requires serum levels and renal and thyroid function tests.", + hybrid_score: 0.55, + }), + ], + }); + + expect(selection.intent.requiredTermSignals).toContain("clinical_subject"); + expect(selection.results[0].id).toBe("lithium-monitoring"); + expect(selection.results[1].match_explanation?.reasons).toContain("retrieval_required_signal:clinical_subject"); + expect(selection.results[1].match_explanation?.reasons).not.toContain("retrieval_signal:clinical_subject"); + }); + it("promotes flowchart next-step evidence for red-zone pathway questions", () => { const selection = selectRetrievalEvidence({ query: "In the clinical flowchart, what is the next step after red-zone risk?",