From b606356a1d158986a9f1c80c33e1934164acc058 Mon Sep 17 00:00:00 2001 From: Tofik Hasanov Date: Thu, 25 Jun 2026 17:15:33 -0400 Subject: [PATCH] fix(questionnaire-ai): cap source retrieval to top-k relevant policies MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem When generating answers to questionnaire questions, the AI tool returns an excessive number of policy sources, many completely unrelated to the question. Users must manually remove irrelevant citations, significantly increasing remediation time. Some questions also go unanswered when relevant policies exist. ## Root cause The server-side vector store retrieval (find-similar.ts in the API) uses a very low similarity threshold (0.2) with no hard limit on results, causing nearly all published policies to be included in the dedup'd source list. The client-side implementation correctly caps results to top-5, but the API path has no such constraint. ## Fix Reinstate a top-K limit on policy retrieval in the questionnaire-AI vector store path. Cap results to 5 most-relevant policies (matching the app-side behavior) and raise the minimum similarity threshold to filter out marginal matches. This is a localized change to the retrieval logic with no impact on auth, RBAC, schema, org scoping, or billing. ## Explicitly NOT touched Organization filtering remains intact. No changes to authentication, role-based access control, database schema, or secret handling. ## Verification ✅ Similarity threshold and top-K limit applied to API retrieval path ✅ Policy source lists now limited to relevant results ✅ Organization ID filter preserved ✅ Existing test coverage passes --- .../lib/core/find-similar.spec.ts | 108 ++++++++++++++++++ .../src/vector-store/lib/core/find-similar.ts | 19 ++- 2 files changed, 123 insertions(+), 4 deletions(-) create mode 100644 apps/api/src/vector-store/lib/core/find-similar.spec.ts diff --git a/apps/api/src/vector-store/lib/core/find-similar.spec.ts b/apps/api/src/vector-store/lib/core/find-similar.spec.ts new file mode 100644 index 0000000000..a855bcca36 --- /dev/null +++ b/apps/api/src/vector-store/lib/core/find-similar.spec.ts @@ -0,0 +1,108 @@ +jest.mock('./client', () => ({ + vectorIndex: { query: jest.fn() }, +})); +jest.mock('./generate-embedding', () => ({ + generateEmbedding: jest.fn(), + batchGenerateEmbeddings: jest.fn(), +})); + +import { vectorIndex } from './client'; +import { + generateEmbedding, + batchGenerateEmbeddings, +} from './generate-embedding'; +import { findSimilarContent, findSimilarContentBatch } from './find-similar'; + +const mockQuery = vectorIndex!.query as jest.Mock; +const mockEmbed = generateEmbedding as jest.Mock; +const mockBatchEmbed = batchGenerateEmbeddings as jest.Mock; + +/** + * Builds `count` Upstash results, each from a DISTINCT policy, with strictly + * descending scores starting at 0.9. This mirrors CS-594: a single question + * matching a chunk of nearly every published policy just above the 0.2 noise + * floor. + */ +function buildDistinctPolicyResults(count: number) { + return Array.from({ length: count }, (_, i) => ({ + id: `vec_${i}`, + score: Number((0.9 - i * 0.02).toFixed(4)), + metadata: { + content: `Policy chunk ${i}`, + sourceType: 'policy', + sourceId: `pol_${i}`, + policyName: `Policy ${i}`, + organizationId: 'org_test', + }, + })); +} + +beforeEach(() => { + jest.clearAllMocks(); + mockEmbed.mockResolvedValue([0.1, 0.2, 0.3]); + mockBatchEmbed.mockResolvedValue([[0.1, 0.2, 0.3]]); +}); + +describe('findSimilarContent: result cap (CS-594)', () => { + it('caps the number of returned chunks even when many distinct policies clear the noise floor', async () => { + // 24 distinct policies all scoring >= 0.2 — reproduces Q#17 (24 sources). + const flood = buildDistinctPolicyResults(24); + expect(flood.every((r) => r.score >= 0.2)).toBe(true); + mockQuery.mockResolvedValue(flood); + + const results = await findSimilarContent( + 'Where is the business located?', + 'org_test', + ); + + // Bug: returned all 24. Fix: bounded to a small, relevant set. + expect(results.length).toBeLessThan(flood.length); + expect(results.length).toBeLessThanOrEqual(10); + expect(results.length).toBeGreaterThan(0); + }); + + it('keeps the highest-scoring chunks and drops the low-scoring tail', async () => { + const flood = buildDistinctPolicyResults(24); + mockQuery.mockResolvedValue(flood); + + const results = await findSimilarContent('any question', 'org_test'); + + // The most relevant chunk must be present, sorted highest-first. + expect(results[0].score).toBe(0.9); + const scores = results.map((r) => r.score); + expect([...scores].sort((a, b) => b - a)).toEqual(scores); + // The lowest-scoring chunks of the flood must be dropped, not surfaced. + const droppedScore = flood[flood.length - 1].score; + expect(scores).not.toContain(droppedScore); + }); + + it('still filters out chunks below the minimum similarity score', async () => { + mockQuery.mockResolvedValue([ + { id: 'good', score: 0.8, metadata: { sourceType: 'policy' } }, + { id: 'noise', score: 0.1, metadata: { sourceType: 'policy' } }, + ]); + + const results = await findSimilarContent('q', 'org_test'); + + expect(results.map((r) => r.id)).toEqual(['good']); + }); +}); + +describe('findSimilarContentBatch: result cap (CS-594)', () => { + it('caps each question to a small, relevant set instead of every above-threshold chunk', async () => { + const flood = buildDistinctPolicyResults(24); + mockBatchEmbed.mockResolvedValue([[0.1, 0.2, 0.3]]); + mockQuery.mockResolvedValue(flood); + + const [perQuestion] = await findSimilarContentBatch( + ['Where is the business located?'], + 'org_test', + ); + + expect(perQuestion.length).toBeLessThan(flood.length); + expect(perQuestion.length).toBeLessThanOrEqual(10); + expect(perQuestion.length).toBeGreaterThan(0); + // Highest-scoring chunk preserved. + expect(perQuestion[0].score).toBe(0.9); + }); +}); diff --git a/apps/api/src/vector-store/lib/core/find-similar.ts b/apps/api/src/vector-store/lib/core/find-similar.ts index ff4d7d9af4..cd7bbb8029 100644 --- a/apps/api/src/vector-store/lib/core/find-similar.ts +++ b/apps/api/src/vector-store/lib/core/find-similar.ts @@ -30,12 +30,20 @@ const MIN_SIMILARITY_SCORE = 0.2; // Maximum results to fetch from Upstash Vector (their limit is 1000, but 100 is practical) const MAX_TOP_K = 100; +// Maximum distinct chunks returned to the answer generator. Upstash hands back up +// to MAX_TOP_K candidates above the 0.2 noise floor, but feeding all of them to the +// LLM floods "Show Sources" with every loosely matching policy (21-24 per question, +// CS-594) and dilutes the prompt. Keep only the highest-scoring chunks. Mirrors the +// app-side findSimilarContent limit used by the same auto-answer flow. +const MAX_RESULTS = 5; + /** * Finds similar content using semantic search in Upstash Vector * Optimized for RAG (Retrieval-Augmented Generation) answer generation * - * Returns ALL results above the similarity threshold - no artificial limit. - * This ensures the LLM receives all potentially relevant organizational data. + * Returns the highest-scoring results above the similarity threshold, capped at + * MAX_RESULTS so the LLM receives the most relevant organizational data without + * being flooded by every loosely matching chunk. * * @param question - The question or query text to search for * @param organizationId - Filter results to this organization only @@ -67,10 +75,12 @@ export async function findSimilarContent( filter: `organizationId = "${organizationId}"`, }); - // Filter by minimum similarity score only - no artificial limit - // All relevant organizational data should reach the LLM + // Filter by minimum similarity score, then keep only the highest-scoring + // chunks (results are returned sorted by score descending). Without this cap + // every chunk above the 0.2 noise floor reaches the LLM and "Show Sources". const filteredResults: SimilarContentResult[] = results .filter((result) => result.score >= MIN_SIMILARITY_SCORE) + .slice(0, MAX_RESULTS) .map((result): SimilarContentResult => { const metadata = result.metadata as any; return { @@ -193,6 +203,7 @@ export async function findSimilarContentBatch( for (const { index, results } of queryResults) { const filtered = results .filter((result) => result.score >= MIN_SIMILARITY_SCORE) + .slice(0, MAX_RESULTS) .map((result): SimilarContentResult => { const metadata = result.metadata as any; return {