Skip to content

Test RAG retry telemetry logging #82

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
BigSimmo merged 1 commit into from
Jun 27, 2026
Merged

Test RAG retry telemetry logging #82

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 63 additions & 11 deletions src/lib/rag.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<typeof createAdminClient>,
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(
Expand Down Expand Up @@ -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<string, SearchResult[]>();

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,
Expand Down Expand Up @@ -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],
Comment on lines +5064 to +5066

@chatgpt-codex-connector chatgpt-codex-connector Bot Jun 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve retry counters in fallback logs

When generation falls into this fallback path after a retry has already been scheduled (for example, the fast answer triggers fast_*_retry_strong and the strong model call throws), these counters are present on the returned fallbackAnswer.latencyTimings, but the logRagQuery metadata built in the catch block still only writes context_pack_latency_ms and omits context_pack_cache_hits, answer_retry_count, and answer_retry_reasons. That makes the production rag_queries telemetry lose the retry reason exactly for failed retry attempts, so the new retry logging is incomplete for generation-fallback cases.

Useful? React with 👍 / 👎.

search_latency_ms: searchLatencyMs,
generation_latency_ms: generationLatencyMs,
total_latency_ms: Date.now() - startedAt,
Expand Down Expand Up @@ -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),
Expand All @@ -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;
Expand All @@ -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(
Expand All @@ -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.",
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
3 changes: 3 additions & 0 deletions src/lib/types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
200 changes: 197 additions & 3 deletions tests/rag-answer-fallback.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, unknown> }]>;
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");
Expand Down Expand Up @@ -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 };
Expand Down Expand Up @@ -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");
});

Expand All @@ -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 };
Expand Down
Loading
Loading