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 {