diff --git a/src/lib/rag.ts b/src/lib/rag.ts index d4169bc86..82d799e2b 100644 --- a/src/lib/rag.ts +++ b/src/lib/rag.ts @@ -2560,13 +2560,37 @@ async function attachIndexQualityMetadata( } } +function sourceContextPackLimit(queryClass: RagQueryClass, options: { crossDocument?: boolean } = {}) { + return options.crossDocument || queryClass === "comparison" || queryClass === "broad_summary" ? 8 : 5; +} + +export function packedContextCacheKey( + results: SearchResult[], + queryClass: RagQueryClass, + options: { crossDocument?: boolean; documentIds?: string[] } = {}, +) { + const contextLimit = sourceContextPackLimit(queryClass, options); + const scopeKey = options.documentIds?.length + ? stableHash([...new Set(options.documentIds)].sort().join("|")) + : "all-documents"; + return [ + queryClass, + options.crossDocument ? "cross-document" : "single-document", + `scope:${scopeKey}`, + contextLimit, + ...results + .slice(0, contextLimit) + .map((result) => `${result.id}:${result.document_id}:${result.chunk_index}:${result.page_number ?? "na"}`), + ].join("|"); +} + async function packAdjacentSourceContext( supabase: ReturnType, results: SearchResult[], queryClass: RagQueryClass, options: { crossDocument?: boolean } = {}, ) { - const contextLimit = options.crossDocument || queryClass === "comparison" || queryClass === "broad_summary" ? 8 : 5; + const contextLimit = sourceContextPackLimit(queryClass, options); const targetResults = results.slice(0, contextLimit); const documentIds = Array.from(new Set(targetResults.map((result) => result.document_id))); const chunkIndexes = Array.from( @@ -4925,6 +4949,29 @@ ${buildRagSourceBlock(contextResults, { query: answerFocusQuery, queryClass })}` let openAIUsage: OpenAITokenUsage = {}; const openAIRequestIds: string[] = []; let contextPackLatencyMs = 0; + let contextPackCacheHits = 0; + let answerRetryCount = 0; + const answerRetryReasons: string[] = []; + const contextPackOptions = { crossDocument: crossDocumentPlan.enabled }; + const packedContextCache = new Map(); + + async function packContextForGeneration(contextResults: SearchResult[]) { + const cacheKey = packedContextCacheKey(contextResults, queryClass, { + ...contextPackOptions, + documentIds: args.documentIds?.length ? args.documentIds : args.documentId ? [args.documentId] : undefined, + }); + const cached = packedContextCache.get(cacheKey); + if (cached) { + contextPackCacheHits += 1; + return cached; + } + + const contextPackStartedAt = Date.now(); + const packed = await packAdjacentSourceContext(createAdminClient(), contextResults, queryClass, contextPackOptions); + contextPackLatencyMs += Date.now() - contextPackStartedAt; + packedContextCache.set(cacheKey, packed); + return packed; + } async function generateWithModel( model: string, @@ -5014,6 +5061,9 @@ ${qualityRetryInstruction}` second_stage_rerank_used: search.telemetry.second_stage_rerank_used, second_stage_rerank_latency_ms: search.telemetry.second_stage_rerank_latency_ms, context_pack_latency_ms: contextPackLatencyMs, + context_pack_cache_hits: contextPackCacheHits, + answer_retry_count: answerRetryCount, + answer_retry_reasons: [...answerRetryReasons], search_latency_ms: searchLatencyMs, generation_latency_ms: generationLatencyMs, total_latency_ms: Date.now() - startedAt, @@ -5051,11 +5101,7 @@ ${qualityRetryInstruction}` model: route.model, reason: route.reason, }); - const contextPackStartedAt = Date.now(); - let packedContextResults = await packAdjacentSourceContext(createAdminClient(), modelContextResults, queryClass, { - crossDocument: crossDocumentPlan.enabled, - }); - contextPackLatencyMs += Date.now() - contextPackStartedAt; + let packedContextResults = await packContextForGeneration(modelContextResults); let generated = await generateWithModel(route.model!, packedContextResults); let answer = annotateAnswerWithDiagnostics( parseAnswerJson(generated.text, packedContextResults, args.query), @@ -5074,6 +5120,8 @@ ${qualityRetryInstruction}` : fastAnswerWasTemplateLike ? "fast_template_retry_strong" : "fast_overexpanded_simple_retry_strong"; + answerRetryCount += 1; + answerRetryReasons.push(retryReason); modelUsed = env.OPENAI_STRONG_ANSWER_MODEL; routingReason = `${route.reason}; ${retryReason}`; retriedWithStrong = true; @@ -5091,11 +5139,7 @@ ${qualityRetryInstruction}` model: env.OPENAI_STRONG_ANSWER_MODEL, reason: routingReason, }); - const retryContextPackStartedAt = Date.now(); - packedContextResults = await packAdjacentSourceContext(createAdminClient(), answerInputResults, queryClass, { - crossDocument: crossDocumentPlan.enabled, - }); - contextPackLatencyMs += Date.now() - retryContextPackStartedAt; + packedContextResults = await packContextForGeneration(answerInputResults); generated = await generateWithModel(env.OPENAI_STRONG_ANSWER_MODEL, packedContextResults); retrievalDiagnostics.routeMode = "strong"; answer = annotateAnswerWithDiagnostics( @@ -5110,6 +5154,8 @@ ${qualityRetryInstruction}` isOverExpandedSimpleGeneratedAnswer(args.query, queryClass, answer)); if (answerNeedsStrongQualityRepair) { routingReason = `${routingReason}; strong_quality_retry`; + answerRetryCount += 1; + answerRetryReasons.push("strong_quality_retry"); await args.onProgress?.({ stage: "retrying", message: "Strong answer failed quality checks, retrying once with stricter synthesis instructions.", @@ -5149,6 +5195,9 @@ ${qualityRetryInstruction}` second_stage_rerank_used: search.telemetry.second_stage_rerank_used, second_stage_rerank_latency_ms: search.telemetry.second_stage_rerank_latency_ms, context_pack_latency_ms: contextPackLatencyMs, + context_pack_cache_hits: contextPackCacheHits, + answer_retry_count: answerRetryCount, + answer_retry_reasons: [...answerRetryReasons], search_latency_ms: searchLatencyMs, generation_latency_ms: generationLatencyMs, total_latency_ms: Date.now() - startedAt, @@ -5250,6 +5299,9 @@ ${qualityRetryInstruction}` supabase_rpc_latency_ms: search.telemetry.supabase_rpc_latency_ms, rerank_latency_ms: search.telemetry.rerank_latency_ms, context_pack_latency_ms: contextPackLatencyMs, + context_pack_cache_hits: contextPackCacheHits, + answer_retry_count: answerRetryCount, + answer_retry_reasons: answerRetryReasons, retrieval_strategy: search.telemetry.retrieval_strategy, weighted_top_score: search.telemetry.weighted_top_score, rrf_top_score: search.telemetry.rrf_top_score, diff --git a/src/lib/types.ts b/src/lib/types.ts index 49b60ac99..44f6bec33 100644 --- a/src/lib/types.ts +++ b/src/lib/types.ts @@ -802,6 +802,9 @@ export type RagAnswer = { second_stage_rerank_used?: boolean; second_stage_rerank_latency_ms?: number; context_pack_latency_ms?: number; + context_pack_cache_hits?: number; + answer_retry_count?: number; + answer_retry_reasons?: string[]; search_latency_ms?: number; generation_latency_ms?: number; total_latency_ms?: number; diff --git a/tests/rag-answer-fallback.test.ts b/tests/rag-answer-fallback.test.ts index 71c1b67fb..a72cc450e 100644 --- a/tests/rag-answer-fallback.test.ts +++ b/tests/rag-answer-fallback.test.ts @@ -422,6 +422,178 @@ describe("RAG structured-output fallback", () => { expect(answer.answerSections?.[0]?.heading).toBe("Direct source-backed answer"); }); + it("records fast-template and strong-quality retry telemetry", async () => { + vi.stubEnv("OPENAI_API_KEY", "test-key"); + vi.stubEnv("RAG_SEARCH_CACHE_TTL_MS", "0"); + vi.stubEnv("RAG_ANSWER_CACHE_TTL_MS", "0"); + vi.stubEnv("RAG_AWAIT_QUERY_LOGS", "true"); + + const firstSource = source({ + id: "template-retry-1", + document_id: "document-a", + title: "Document Monitoring Pathway Guide A", + file_name: "document-monitoring-pathway-guide-a.pdf", + page_number: 3, + section_heading: "Monitoring pathway", + content: + "Guide A defines document monitoring pathways, routine review intervals, safety checks, and referral thresholds for comparison across guides.", + similarity: 0.96, + hybrid_score: 0.96, + text_rank: 1.1, + }); + + const secondSource = source({ + id: "template-retry-2", + document_id: "document-b", + title: "Document Monitoring Pathway Guide B", + file_name: "document-monitoring-pathway-guide-b.pdf", + page_number: 4, + section_heading: "Monitoring pathway", + content: + "Guide B describes a second document monitoring pathway with review frequency, escalation triggers, and referral thresholds for comparison.", + similarity: 0.92, + hybrid_score: 0.92, + text_rank: 1.0, + }); + + const rpc = vi.fn(async (name: string) => { + if (name === "match_document_chunks_text") return { data: [firstSource, secondSource], error: null }; + if (name === "get_related_document_metadata") return { data: [], error: null }; + return { data: [], error: null }; + }); + + const generateStructuredTextResult = vi + .fn() + .mockResolvedValueOnce({ + text: JSON.stringify({ + answer: + "This is a Direct source-backed answer derived from the retrieved sources and structured for the same citations.", + grounded: false, + confidence: "low", + answerSections: [ + { + heading: "Direct source-backed answer", + kind: "documentation", + supportLevel: "direct", + body: "Use the retrieved evidence to guide a routine review.", + citation_chunk_ids: ["template-retry-1"], + }, + ], + citations: [{ chunk_id: "template-retry-1" }], + quoteCards: [ + { + chunk_id: "template-retry-1", + quote: + "Retrieved source supports that this document defines the first-line management of a condition and associated safety checks.", + section_heading: "Summary", + }, + ], + conflictsOrGaps: [], + }), + model: "gpt-5.4-mini", + operation: "answer", + latencyMs: 12, + requestId: "req_fast_template", + usage: { input_tokens: 80, output_tokens: 90, total_tokens: 170 }, + }) + .mockResolvedValueOnce({ + text: JSON.stringify({ + answer: + "Source-backed summaries mention management steps and review intervals.", + grounded: true, + confidence: "high", + answerSections: [], + citations: [{ chunk_id: "template-retry-1" }], + quoteCards: [ + { + chunk_id: "template-retry-1", + quote: "Retrieved source supports that this document defines the first-line management.", + section_heading: "Summary", + }, + ], + conflictsOrGaps: [], + }), + model: "gpt-5.4", + operation: "answer", + latencyMs: 18, + requestId: "req_strong_template", + usage: { input_tokens: 140, output_tokens: 120, total_tokens: 260 }, + }) + .mockResolvedValueOnce({ + text: JSON.stringify({ + answer: + "Compare the guide pathways by using routine review intervals, escalation triggers, and urgent referral thresholds.", + grounded: true, + confidence: "high", + answerSections: [ + { + heading: "Monitoring", + kind: "required_actions", + supportLevel: "direct", + body: "Review the pathway criteria regularly and escalate when safety checks or referral thresholds are met.", + citation_chunk_ids: ["template-retry-2"], + }, + ], + citations: [{ chunk_id: "template-retry-2" }], + quoteCards: [ + { + chunk_id: "template-retry-2", + quote: + "Guide B describes a second document monitoring pathway with review frequency, escalation triggers, and referral thresholds for comparison.", + section_heading: "Monitoring", + }, + ], + conflictsOrGaps: [], + }), + model: "gpt-5.4", + operation: "answer", + latencyMs: 14, + requestId: "req_strong_quality", + usage: { input_tokens: 170, output_tokens: 120, total_tokens: 290 }, + }); + + const insert = vi.fn(async () => ({ data: null, error: null })); + const from = vi.fn((table: string) => (table === "rag_queries" ? { insert } : new EmptyQuery())); + + vi.doMock("@/lib/supabase/admin", () => ({ + createAdminClient: () => ({ + rpc, + from, + }), + })); + vi.doMock("@/lib/openai", () => ({ + embedTextWithTelemetry: vi.fn(async () => ({ embedding: [0.1, 0.2, 0.3], cacheHit: false })), + generateStructuredTextResult, + })); + + const { answerQuestionWithScope } = await import("../src/lib/rag"); + + const answer = await answerQuestionWithScope({ + query: "Compare document monitoring pathways across two guides", + ownerId: undefined, + skipCache: true, + }); + + expect(generateStructuredTextResult).toHaveBeenCalledTimes(3); + expect(answer.routingMode).toBe("strong"); + expect(answer.latencyTimings?.answer_retry_count).toBe(2); + expect(answer.latencyTimings?.answer_retry_reasons).toEqual([ + "fast_template_retry_strong", + "strong_quality_retry", + ]); + expect(answer.routingReason).toContain("fast_template_retry_strong"); + expect(answer.routingReason).toContain("strong_quality_retry"); + expect(answer.openAIRequestIds).toEqual(["req_fast_template", "req_strong_template", "req_strong_quality"]); + expect(insert).toHaveBeenCalledTimes(1); + const insertCalls = insert.mock.calls as unknown as Array<[{ metadata?: Record }]>; + const loggedMetadata = insertCalls[0]?.[0]?.metadata ?? {}; + expect(loggedMetadata.answer_retry_count).toBe(2); + expect(loggedMetadata.answer_retry_reasons).toEqual([ + "fast_template_retry_strong", + "strong_quality_retry", + ]); + }); + it("filters table-caption metadata from extractive answer points", async () => { vi.stubEnv("OPENAI_API_KEY", "test-key"); vi.stubEnv("RAG_SEARCH_CACHE_TTL_MS", "0"); @@ -524,13 +696,25 @@ describe("RAG structured-output fallback", () => { expect(plainAnswer).not.toContain("table_crop"); }); - it("retries malformed fast output with the strong model and fails safely without extractive stitching", async () => { + it("retries malformed fast output and fails safely through extractive recovery", async () => { vi.stubEnv("OPENAI_API_KEY", "test-key"); vi.stubEnv("OPENAI_MAX_OUTPUT_TOKENS", "650"); vi.stubEnv("RAG_SEARCH_CACHE_TTL_MS", "0"); vi.stubEnv("RAG_ANSWER_CACHE_TTL_MS", "0"); - const sources = [source(), source({ id: "agitation-chunk-2", page_number: 10, chunk_index: 1 })]; + const sources = [ + source({ + content: + "Inpatient approach details for agitation and arousal management include oral medication when the patient is willing, increased observation when ratings rise, and intramuscular medication when oral medication is refused.", + }), + source({ + id: "agitation-chunk-2", + page_number: 10, + chunk_index: 1, + content: + "Inpatient approach details include review of route, rating severity, escalation triggers, and monitoring after medication administration.", + }), + ]; const rpc = vi.fn(async (name: string) => { if (name === "match_document_chunks_text") return { data: sources, error: null }; if (name === "get_related_document_metadata") return { data: [], error: null }; @@ -574,6 +758,11 @@ describe("RAG structured-output fallback", () => { expect(generateStructuredTextResult).toHaveBeenCalledTimes(3); expect(answer.openAIRequestIds).toEqual(["req_truncated", "req_truncated", "req_truncated"]); expect(answer.openAIUsage).toMatchObject({ output_tokens: 1950 }); + expect(answer.latencyTimings?.answer_retry_count).toBe(2); + expect(answer.latencyTimings?.answer_retry_reasons).toEqual([ + "fast_unusable_retry_strong", + "strong_quality_retry", + ]); expect(answer.citations[0]?.source_metadata?.document_status).toBe("current"); }); @@ -582,7 +771,12 @@ describe("RAG structured-output fallback", () => { vi.stubEnv("RAG_SEARCH_CACHE_TTL_MS", "0"); vi.stubEnv("RAG_ANSWER_CACHE_TTL_MS", "0"); - const sources = [source()]; + const sources = [ + source({ + content: + "Inpatient approach details for agitation and arousal management include a stepwise approach based on rating severity, route, oral options, and intramuscular options.", + }), + ]; const rpc = vi.fn(async (name: string) => { if (name === "match_document_chunks_text") return { data: sources, error: null }; if (name === "get_related_document_metadata") return { data: [], error: null }; diff --git a/tests/rag-context-budget.test.ts b/tests/rag-context-budget.test.ts index 3ce807dca..329199403 100644 --- a/tests/rag-context-budget.test.ts +++ b/tests/rag-context-budget.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { selectModelContextResults } from "../src/lib/rag"; +import { packedContextCacheKey, selectModelContextResults } from "../src/lib/rag"; import type { RagQueryClass, SearchResult } from "../src/lib/types"; function source(index: number): SearchResult { @@ -53,4 +53,39 @@ describe("RAG model context budgeting", () => { expect(select({ routeMode: "extractive", queryClass: "document_lookup" })).toHaveLength(12); expect(select({ routeMode: "unsupported", queryClass: "unsupported_or_general" })).toHaveLength(12); }); + + it("uses a stable context pack cache key for matching retry inputs", () => { + const key = packedContextCacheKey(results, "broad_summary", { crossDocument: true }); + const sameInputs = packedContextCacheKey([...results], "broad_summary", { crossDocument: true }); + const differentInputs = packedContextCacheKey(results.slice(0, 6), "broad_summary", { crossDocument: true }); + + expect(sameInputs).toBe(key); + expect(differentInputs).not.toBe(key); + }); + + it("includes document-scope for stricter packed context reuse", () => { + const keyA = packedContextCacheKey(results, "document_lookup", { + crossDocument: false, + documentIds: ["doc-1", "doc-2"], + }); + const keyB = packedContextCacheKey(results, "document_lookup", { + crossDocument: false, + documentIds: ["doc-3"], + }); + + expect(keyA).not.toBe(keyB); + }); + + it("reuses packed context keys for the same document filter regardless of order", () => { + const keyA = packedContextCacheKey(results, "document_lookup", { + crossDocument: false, + documentIds: ["doc-2", "doc-1", "doc-1"], + }); + const keyB = packedContextCacheKey(results, "document_lookup", { + crossDocument: false, + documentIds: ["doc-1", "doc-2"], + }); + + expect(keyA).toBe(keyB); + }); });