From 681818f823b0d88323468d4703984e1d637ee2e5 Mon Sep 17 00:00:00 2001 From: heybeaux Date: Sun, 6 Sep 2026 18:18:21 -0700 Subject: [PATCH 1/6] fix(ensemble): enforce ownership and reject unsafe embedding results --- .../ensemble-audit-regression.spec.ts | 230 ++++++++++++++ src/ensemble/ensemble-http-isolation.spec.ts | 290 ++++++++++++++++++ src/ensemble/ensemble-model.types.ts | 2 + src/ensemble/ensemble.controller.spec.ts | 50 ++- src/ensemble/ensemble.controller.ts | 76 ++++- src/ensemble/ensemble.service.spec.ts | 4 +- src/ensemble/ensemble.service.ts | 158 +++++++--- .../pgvector-ensemble.provider.spec.ts | 51 ++- src/ensemble/pgvector-ensemble.provider.ts | 58 +++- 9 files changed, 825 insertions(+), 94 deletions(-) create mode 100644 src/ensemble/ensemble-audit-regression.spec.ts create mode 100644 src/ensemble/ensemble-http-isolation.spec.ts diff --git a/src/ensemble/ensemble-audit-regression.spec.ts b/src/ensemble/ensemble-audit-regression.spec.ts new file mode 100644 index 00000000..35a7f0e5 --- /dev/null +++ b/src/ensemble/ensemble-audit-regression.spec.ts @@ -0,0 +1,230 @@ +import { EnsembleService } from './ensemble.service'; +import { PgVectorEnsembleProvider } from './pgvector-ensemble.provider'; + +const vector = (value = 0.1) => Array(768).fill(value); +const result = (model = 'bge-base', value = 0.1) => ({ + model, + embedding: vector(value), + dimensions: 768, + latencyMs: 0, +}); + +describe('Ensemble audit regressions', () => { + let service: EnsembleService; + let storage: any; + const originalFetch = global.fetch; + beforeEach(() => { + storage = { + upsertEmbeddings: jest.fn(), + queryWithModelEmbeddings: jest.fn(), + }; + service = new EnsembleService( + { + get: (key: string, fallback: any) => + key === 'ENSEMBLE_ENABLED' + ? true + : key === 'ENSEMBLE_MODELS' + ? 'bge-base,nomic' + : fallback, + } as any, + storage, + { + memory: { + findMany: async () => + ['a', 'b', 'c'].map((id) => ({ id, raw: id, userId: 'user' })), + }, + } as any, + {} as any, + {} as any, + ); + }); + afterEach(() => { + global.fetch = originalFetch; + }); + + it.each([ + [{ embedding: vector(0.1) }, {}, { embedding: vector(0.3) }], + [ + { index: 0, embedding: vector() }, + { index: 0, embedding: vector() }, + { index: 2, embedding: vector() }, + ], + [ + { index: 0, embedding: vector() }, + { index: 2, embedding: vector() }, + ], + [ + { index: 0, embedding: vector() }, + { embedding: vector() }, + { index: 2, embedding: vector() }, + ], + ])( + 'rejects ambiguous/missing batch vectors before any write (%#)', + async (...entries) => { + global.fetch = jest.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + embeddings: [{ model: 'bge-base', data: entries }], + }), + }); + await expect( + service.embedBatchForMemories(['a', 'b', 'c'], ['bge-base']), + ).rejects.toThrow(); + expect(storage.upsertEmbeddings).not.toHaveBeenCalled(); + }, + ); + + it('reorders explicit indices to preserve memory identity', async () => { + global.fetch = jest.fn().mockResolvedValue({ + ok: true, + json: async () => ({ + embeddings: [ + { + model: 'bge-base', + data: [2, 0, 1].map((index) => ({ + index, + embedding: vector(index + 1), + })), + }, + ], + }), + }); + await service.embedBatchForMemories(['a', 'b', 'c'], ['bge-base']); + expect( + storage.upsertEmbeddings.mock.calls[0][0].map((row: any) => [ + row.memoryId, + row.embedding[0], + ]), + ).toEqual([ + ['a', 1], + ['b', 2], + ['c', 3], + ]); + }); + + it('rejects compressed cloud batches before positional assignment', async () => { + (service as any).useCloud = true; + (service as any).cloudEnsemble = { + embedBatch: async () => ({ + embeddings: [result(), result('bge-base', 0.3)], + }), + }; + await expect( + service.embedBatchForMemories(['a', 'b', 'c'], ['bge-base']), + ).rejects.toThrow('Incomplete batch'); + expect(storage.upsertEmbeddings).not.toHaveBeenCalled(); + }); + + it('returns retryable failure for no query or write embeddings', async () => { + jest + .spyOn(service, 'embedAll') + .mockResolvedValue({ embeddings: [], totalMs: 1 }); + await expect( + service.query({ query: 'fact', userId: 'user' }), + ).rejects.toMatchObject({ status: 503 }); + await expect( + service.upsert({ memoryId: 'a', content: 'fact', userId: 'user' }), + ).rejects.toMatchObject({ status: 503 }); + expect(storage.upsertEmbeddings).not.toHaveBeenCalled(); + expect(storage.queryWithModelEmbeddings).not.toHaveBeenCalled(); + }); + + it('distinguishes degraded successful empty search from total backend failure', async () => { + jest + .spyOn(service, 'embedAll') + .mockResolvedValue({ embeddings: [result() as any], totalMs: 1 }); + storage.queryWithModelEmbeddings.mockResolvedValue( + new Map([['bge-base', []]]), + ); + const response = await service.query({ query: 'fact', userId: 'user' }); + expect(response.results).toEqual([]); + expect(response.metadata).toMatchObject({ + degraded: true, + failedModels: ['nomic'], + modelsQueried: ['bge-base'], + }); + storage.queryWithModelEmbeddings.mockResolvedValue(new Map()); + await expect( + service.query({ query: 'fact', userId: 'user' }), + ).rejects.toMatchObject({ status: 503 }); + }); + + it('honors an explicit model subset', async () => { + jest.spyOn(service, 'embedAll').mockResolvedValue({ + embeddings: [result() as any, result('nomic') as any], + totalMs: 1, + }); + storage.queryWithModelEmbeddings.mockResolvedValue( + new Map([['bge-base', []]]), + ); + const response = await service.query({ + query: 'fact', + userId: 'user', + models: ['bge-base'], + }); + expect([ + ...storage.queryWithModelEmbeddings.mock.calls[0][0].keys(), + ]).toEqual(['bge-base']); + expect(response.metadata.degraded).toBe(false); + }); +}); + +describe('Ensemble storage validation', () => { + let db: any; + let provider: PgVectorEnsembleProvider; + beforeEach(() => { + db = { + $executeRawUnsafe: jest.fn().mockResolvedValue(1), + $transaction: jest.fn(async (fn) => fn(db)), + }; + provider = new PgVectorEnsembleProvider(db); + }); + it.each( + [ + new Array(768), + vector().map((v, i) => (i === 10 ? NaN : v)), + vector().map((v, i) => (i === 10 ? Infinity : v)), + [0.1], + [], + ].map((embedding) => [embedding]), + )( + 'rejects malformed vectors at single and batch storage boundaries (%#)', + async (embedding) => { + const record = { + memoryId: 'm', + modelId: 'bge-base' as const, + dimensions: 768, + embedding, + }; + await expect(provider.upsertEmbedding(record)).rejects.toThrow(); + await expect(provider.upsertEmbeddings([record])).rejects.toThrow(); + expect(db.$executeRawUnsafe).not.toHaveBeenCalled(); + expect(db.$transaction).not.toHaveBeenCalled(); + }, + ); + it('rejects metadata dimensions inconsistent with the model', async () => { + await expect( + provider.upsertEmbedding({ + memoryId: 'm', + modelId: 'bge-base', + dimensions: 384, + embedding: vector(), + }), + ).rejects.toThrow('Invalid dimensions'); + }); + it('does not call a zero-row write successful', async () => { + db.$executeRawUnsafe.mockResolvedValue(0); + const record = { + memoryId: 'missing', + modelId: 'bge-base' as const, + dimensions: 768, + embedding: vector(), + }; + await expect(provider.upsertEmbedding(record)).rejects.toMatchObject({ + status: 503, + }); + await expect(provider.upsertEmbeddings([record])).rejects.toMatchObject({ + status: 503, + }); + }); +}); diff --git a/src/ensemble/ensemble-http-isolation.spec.ts b/src/ensemble/ensemble-http-isolation.spec.ts new file mode 100644 index 00000000..3b623843 --- /dev/null +++ b/src/ensemble/ensemble-http-isolation.spec.ts @@ -0,0 +1,290 @@ +/** Opt-in integration: ENSEMBLE_AUDIT_PG=1. Uses only a new local scratch DB. */ +import { Test } from '@nestjs/testing'; +import { ConfigService } from '@nestjs/config'; +import { JwtService } from '@nestjs/jwt'; +import { Client } from 'pg'; +import request from 'supertest'; +import { EnsembleController } from './ensemble.controller'; +import { EnsembleService } from './ensemble.service'; +import { PgVectorEnsembleProvider } from './pgvector-ensemble.provider'; +import { PrismaService } from '../prisma/prisma.service'; +import { NightlyReembedService } from './nightly-reembed.service'; +import { DriftDetectionService } from './drift-detection.service'; + +const run = process.env.ENSEMBLE_AUDIT_PG === '1' ? describe : describe.skip; +run( + 'Ensemble two-account HTTP isolation with a non-superuser BYPASSRLS store', + () => { + const database = `engram_audit_ensemble_${process.pid}`; + const role = `engram_audit_vectors_${process.pid}`; + const local = { host: '/tmp', database: 'postgres' }; + const admin = new Client(local); + let db: Client; + let app: any; + let service: EnsembleService; + let tokenA: string; + let tokenB: string; + let createdDatabase = false; + let createdRole = false; + const vector = Array(768).fill(0.1); + + beforeAll(async () => { + await admin.connect(); + await admin.query(`CREATE DATABASE ${database}`); + createdDatabase = true; + await admin.query(`CREATE ROLE ${role} NOLOGIN NOSUPERUSER BYPASSRLS`); + createdRole = true; + db = new Client({ ...local, database }); + await db.connect(); + await db.query(`CREATE EXTENSION vector; + CREATE TABLE users(id text PRIMARY KEY, account_id text, external_id text, deleted_at timestamptz); + CREATE TABLE agents(id text PRIMARY KEY, account_id text); + CREATE TABLE accounts(id text PRIMARY KEY, is_admin boolean); + CREATE TABLE memories(id text PRIMARY KEY, user_id text REFERENCES users(id), raw text, deleted_at timestamptz); + CREATE TABLE memory_embeddings(id text PRIMARY KEY, memory_id text REFERENCES memories(id), model_id text, dimensions int, embedding vector, created_at timestamptz, updated_at timestamptz, UNIQUE(memory_id, model_id)); + INSERT INTO accounts VALUES ('account-a', false), ('account-b', false); + INSERT INTO agents VALUES ('agent-a','account-a'), ('agent-b','account-b'); + INSERT INTO users VALUES ('ua','account-a','a',NULL), ('ub','account-b','b',NULL); + INSERT INTO memories VALUES ('ma','ua','alpha',NULL), ('mb','ub','beta',NULL); + ALTER TABLE memories ENABLE ROW LEVEL SECURITY; + ALTER TABLE memories FORCE ROW LEVEL SECURITY; + CREATE POLICY account_scope ON memories USING (user_id IN (SELECT id FROM users WHERE account_id = current_setting('app.account_id', true))); + GRANT USAGE ON SCHEMA public TO ${role}; + GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA public TO ${role};`); + await db.query(`SET ROLE ${role}`); + const facade: any = { + $executeRawUnsafe: async (sql: string, ...args: any[]) => + (await db.query(sql, args)).rowCount, + $queryRawUnsafe: async (sql: string, ...args: any[]) => + (await db.query(sql, args)).rows, + $transaction: async (fn: any) => { + await db.query('BEGIN'); + try { + const result = await fn(facade); + await db.query('COMMIT'); + return result; + } catch (error) { + await db.query('ROLLBACK'); + throw error; + } + }, + user: { + findFirst: async ({ where }: any) => + ( + await db.query( + 'SELECT id FROM users WHERE id=$1 AND account_id=$2 AND deleted_at IS NULL', + [where.id, where.accountId], + ) + ).rows[0], + findUnique: async ({ where }: any) => + ( + await db.query( + 'SELECT id, account_id AS "accountId" FROM users WHERE account_id=$1 AND external_id=$2', + [ + where.accountId_externalId.accountId, + where.accountId_externalId.externalId, + ], + ) + ).rows[0], + }, + agent: { + findFirst: async ({ where }: any) => + ( + await db.query( + 'SELECT id, account_id AS "accountId" FROM agents WHERE account_id=$1', + [where.accountId], + ) + ).rows[0], + }, + account: { + findUnique: async ({ where }: any) => + ( + await db.query( + 'SELECT is_admin AS "isAdmin" FROM accounts WHERE id=$1', + [where.id], + ) + ).rows[0], + }, + memory: { + findFirst: async ({ where }: any) => + ( + await db.query( + 'SELECT id FROM memories WHERE id=$1 AND user_id=$2 AND deleted_at IS NULL', + [where.id, where.userId], + ) + ).rows[0], + }, + }; + const config = new ConfigService({ + ENSEMBLE_ENABLED: true, + ENSEMBLE_MODELS: 'bge-base', + EDITION: 'cloud', + NODE_ENV: 'production', + TRUST_LOCAL_NETWORK: 'false', + }); + const jwt = new JwtService({ secret: 'isolated-audit-fixture-only' }); + service = new EnsembleService( + config, + new PgVectorEnsembleProvider(facade), + facade, + {} as any, + {} as any, + ); + // Deterministic embedder; all auth/controller/service/storage code is real. + jest.spyOn(service, 'embedAll').mockImplementation(async () => ({ + embeddings: [ + { + model: 'bge-base', + dimensions: 768, + embedding: vector, + latencyMs: 0, + }, + ], + totalMs: 0, + })); + // Avoid onModuleInit's network model-health discovery in this focused HTTP fixture. + jest.spyOn(service, 'onModuleInit').mockResolvedValue(); + await service.upsert({ memoryId: 'ma', userId: 'ua', content: 'alpha' }); + await service.upsert({ memoryId: 'mb', userId: 'ub', content: 'beta' }); + const module = await Test.createTestingModule({ + controllers: [EnsembleController], + providers: [ + { provide: EnsembleService, useValue: service }, + { provide: ConfigService, useValue: config }, + { provide: JwtService, useValue: jwt }, + { provide: PrismaService, useValue: facade }, + { provide: NightlyReembedService, useValue: {} }, + { provide: DriftDetectionService, useValue: {} }, + ], + }).compile(); + app = module.createNestApplication(); + await app.init(); + tokenA = jwt.sign({ sub: 'account-a', email: 'a' }); + tokenB = jwt.sign({ sub: 'account-b', email: 'b' }); + }, 20000); + + afterAll(async () => { + if (app) await app.close(); + if (db) await db.end(); + if (createdDatabase) await admin.query(`DROP DATABASE ${database}`); + if (createdRole) await admin.query(`DROP ROLE ${role}`); + await admin.end(); + }); + + it('runs the vector store with BYPASSRLS but without superuser privileges', async () => { + const row = ( + await db.query( + 'SELECT rolbypassrls, rolsuper FROM pg_roles WHERE rolname=current_user', + ) + ).rows[0]; + expect(row).toEqual({ rolbypassrls: true, rolsuper: false }); + }); + it('authenticates two accounts and retrieves only each account’s own memory', async () => { + for (const [token, userId, memoryId] of [ + [tokenA, 'ua', 'ma'], + [tokenB, 'ub', 'mb'], + ]) { + const response = await request(app.getHttpServer()) + .post('/v1/ensemble/query') + .auth(token, { type: 'bearer' }) + .send({ query: 'fact', userId }) + .expect(200); + expect(response.body.results.map((r: any) => r.memoryId)).toEqual([ + memoryId, + ]); + } + }); + it.each(['query', 'compare'])( + 'rejects attacker-selected user scope on %s', + async (endpoint) => { + await request(app.getHttpServer()) + .post(`/v1/ensemble/${endpoint}`) + .auth(tokenA, { type: 'bearer' }) + .send({ query: 'victim', userId: 'ub' }) + .expect(403); + }, + ); + it('rejects victim embedding writes/status before embedding or reading', async () => { + const embed = jest.spyOn(service, 'embedAll'); + embed.mockClear(); + await request(app.getHttpServer()) + .post('/v1/ensemble/upsert') + .auth(tokenA, { type: 'bearer' }) + .send({ memoryId: 'mb', userId: 'ua', content: 'poison' }) + .expect(404); + await request(app.getHttpServer()) + .get('/v1/ensemble/memories/mb/embeddings') + .auth(tokenA, { type: 'bearer' }) + .expect(404); + expect(embed).not.toHaveBeenCalled(); + }); + it('allows an authorized write and verifies the persisted vector', async () => { + await request(app.getHttpServer()) + .post('/v1/ensemble/upsert') + .auth(tokenA, { type: 'bearer' }) + .send({ memoryId: 'ma', userId: 'ua', content: 'alpha' }) + .expect(200); + expect( + ( + await db.query( + "SELECT count(*)::int AS n FROM memory_embeddings WHERE memory_id='ma' AND vector_dims(embedding)=768", + ) + ).rows[0].n, + ).toBe(1); + }); + it('enforces ownership at SQL write boundary even when called without controller', async () => { + await expect( + service.upsert({ memoryId: 'mb', userId: 'ua', content: 'poison' }), + ).rejects.toMatchObject({ status: 503 }); + const value = ( + await db.query( + "SELECT embedding::text FROM memory_embeddings WHERE memory_id='mb'", + ) + ).rows[0].embedding; + expect(JSON.parse(value)).toEqual(vector); + }); + it.each(['reembed', 'reembed/targeted', 'drift/analyze'])( + 'rejects a non-admin global mutation: %s', + async (endpoint) => { + await request(app.getHttpServer()) + .post(`/v1/ensemble/${endpoint}`) + .auth(tokenA, { type: 'bearer' }) + .send({ memoryIds: ['mb'], models: ['bge-base'] }) + .expect(403); + }, + ); + it.each([ + 'models', + 'coverage', + 'ab-results', + 'drift', + 'drift/history', + 'reembed/status', + ])('rejects non-admin global diagnostics: %s', async (endpoint) => { + await request(app.getHttpServer()) + .get(`/v1/ensemble/${endpoint}`) + .auth(tokenA, { type: 'bearer' }) + .expect(403); + }); + it('allows an explicitly authorized administrator to read global diagnostics', async () => { + jest.spyOn(service, 'getModels').mockResolvedValue([]); + await db.query("UPDATE accounts SET is_admin=true WHERE id='account-a'"); + try { + await request(app.getHttpServer()) + .get('/v1/ensemble/models') + .auth(tokenA, { type: 'bearer' }) + .expect(200); + } finally { + await db.query( + "UPDATE accounts SET is_admin=false WHERE id='account-a'", + ); + } + }); + it('rejects missing authentication', async () => { + await request(app.getHttpServer()) + .post('/v1/ensemble/query') + .send({ query: 'fact', userId: 'ua' }) + .expect(401); + }); + }, +); diff --git a/src/ensemble/ensemble-model.types.ts b/src/ensemble/ensemble-model.types.ts index 05b27c29..a40c03d0 100644 --- a/src/ensemble/ensemble-model.types.ts +++ b/src/ensemble/ensemble-model.types.ts @@ -205,6 +205,8 @@ export interface EnsembleQueryResult { modelsQueried: ModelId[]; candidatesEvaluated: number; fusionAlgorithm: string; + degraded?: boolean; + failedModels?: ModelId[]; }; } diff --git a/src/ensemble/ensemble.controller.spec.ts b/src/ensemble/ensemble.controller.spec.ts index b55eab0f..ca9c9c6c 100644 --- a/src/ensemble/ensemble.controller.spec.ts +++ b/src/ensemble/ensemble.controller.spec.ts @@ -49,7 +49,13 @@ describe('EnsembleController', () => { findMany: jest.fn(), create: jest.fn(), }, + user: { + findFirst: jest + .fn() + .mockImplementation(async ({ where }) => ({ id: where.id })), + }, memory: { + findFirst: jest.fn().mockResolvedValue({ id: 'm1' }), findMany: jest.fn(), }, } as any; @@ -99,7 +105,10 @@ describe('EnsembleController', () => { modelsUsed: ['text-embedding-3-small'], } as any); - const result = await controller.query(dto as any); + const result = await controller.query(dto as any, { + accountId: 'a1', + user: { id: dto.userId }, + }); expect(result.results[0].modelScores).toEqual({ 'text-embedding-3-small': { rank: 1, score: 0.95 }, }); @@ -107,9 +116,12 @@ describe('EnsembleController', () => { it('should throw BadRequestException when ensemble is disabled', async () => { ensembleService.isEnabled.mockReturnValue(false); - await expect(controller.query(dto as any)).rejects.toThrow( - BadRequestException, - ); + await expect( + controller.query(dto as any, { + accountId: 'a1', + user: { id: dto.userId }, + }), + ).rejects.toThrow(BadRequestException); }); it('should throw BadRequestException when query is missing', async () => { @@ -133,7 +145,10 @@ describe('EnsembleController', () => { it('should upsert and return success', async () => { ensembleService.upsert.mockResolvedValue(undefined); - const result = await controller.upsert(dto as any); + const result = await controller.upsert(dto as any, { + accountId: 'a1', + user: { id: dto.userId }, + }); expect(result).toEqual({ success: true }); expect(ensembleService.upsert).toHaveBeenCalledWith({ memoryId: 'm1', @@ -145,9 +160,12 @@ describe('EnsembleController', () => { it('should throw when ensemble is disabled', async () => { ensembleService.isEnabled.mockReturnValue(false); - await expect(controller.upsert(dto as any)).rejects.toThrow( - BadRequestException, - ); + await expect( + controller.upsert(dto as any, { + accountId: 'a1', + user: { id: dto.userId }, + }), + ).rejects.toThrow(BadRequestException); }); it('should throw when required fields are missing', async () => { @@ -190,10 +208,13 @@ describe('EnsembleController', () => { singleModel, } as any); - const result = await controller.compare({ - query: 'q', - userId: 'u', - } as any); + const result = await controller.compare( + { + query: 'q', + userId: 'u', + } as any, + { accountId: 'a1', user: { id: 'u' } }, + ); expect(result.ensemble.results).toHaveLength(1); expect(result.singleModel['text-embedding-3-small']).toHaveLength(1); }); @@ -447,7 +468,10 @@ describe('EnsembleController', () => { { model: 'model-a', exists: true, dimensions: 1536 }, ] as any); - const result = await controller.getMemoryEmbeddings('m1'); + const result = await controller.getMemoryEmbeddings('m1', { + accountId: 'a1', + user: { id: 'u' }, + }); expect(result.memoryId).toBe('m1'); expect(result.embeddings).toHaveLength(1); }); diff --git a/src/ensemble/ensemble.controller.ts b/src/ensemble/ensemble.controller.ts index f6d9f833..8d494f4b 100644 --- a/src/ensemble/ensemble.controller.ts +++ b/src/ensemble/ensemble.controller.ts @@ -17,6 +17,9 @@ import { NotFoundException, UseGuards, Logger, + Req, + ForbiddenException, + UnauthorizedException, } from '@nestjs/common'; import { ApiTags, ApiOperation, ApiResponse, ApiQuery } from '@nestjs/swagger'; import { @@ -41,6 +44,7 @@ import { MemoryEmbeddingStatus, ABTestResult, } from './ensemble.types'; +import { AdminGuard } from '../common/guards/admin.guard'; import { ApiKeyOrJwtGuard } from '../common/guards/api-key-or-jwt.guard'; // ============================================================================ @@ -169,7 +173,10 @@ export class EnsembleController { @HttpCode(HttpStatus.OK) @ApiOperation({ summary: 'Query memories with multi-model ensemble' }) @ApiResponse({ status: 200, description: 'Fused query results' }) - async query(@Body() dto: EnsembleQueryDto): Promise { + async query( + @Body() dto: EnsembleQueryDto, + @Req() request?: any, + ): Promise { if (!this.ensembleService.isEnabled()) { throw new BadRequestException('Ensemble retrieval is not enabled'); } @@ -178,9 +185,10 @@ export class EnsembleController { throw new BadRequestException('query and userId are required'); } + const userId = await this.authorizedUser(request, dto.userId); const result = await this.ensembleService.query({ query: dto.query, - userId: dto.userId, + userId, limit: dto.limit ? Math.min(dto.limit, 1000) : undefined, k: dto.k ? Math.min(dto.k, 1000) : undefined, models: dto.models, @@ -207,7 +215,10 @@ export class EnsembleController { @HttpCode(HttpStatus.OK) @ApiOperation({ summary: 'Upsert memory with ensemble embeddings' }) @ApiResponse({ status: 200, description: 'Memory upserted' }) - async upsert(@Body() dto: EnsembleUpsertDto): Promise<{ success: boolean }> { + async upsert( + @Body() dto: EnsembleUpsertDto, + @Req() request?: any, + ): Promise<{ success: boolean }> { if (!this.ensembleService.isEnabled()) { throw new BadRequestException('Ensemble retrieval is not enabled'); } @@ -218,10 +229,12 @@ export class EnsembleController { ); } + const userId = await this.authorizedUser(request, dto.userId); + await this.authorizedMemory(dto.memoryId, userId); await this.ensembleService.upsert({ memoryId: dto.memoryId, content: dto.content, - userId: dto.userId, + userId, metadata: dto.metadata, }); @@ -235,7 +248,10 @@ export class EnsembleController { @HttpCode(HttpStatus.OK) @ApiOperation({ summary: 'Compare ensemble vs single-model retrieval' }) @ApiResponse({ status: 200, description: 'Comparison results' }) - async compare(@Body() dto: CompareQueryDto): Promise<{ + async compare( + @Body() dto: CompareQueryDto, + @Req() request?: any, + ): Promise<{ ensemble: EnsembleQueryResponse; singleModel: Record< ModelId, @@ -250,9 +266,10 @@ export class EnsembleController { throw new BadRequestException('query and userId are required'); } + const userId = await this.authorizedUser(request, dto.userId); const result = await this.ensembleService.compare( dto.query, - dto.userId, + userId, dto.limit ?? 10, ); @@ -318,6 +335,7 @@ export class EnsembleController { /** * Get all registered models with status and configuration */ + @UseGuards(AdminGuard) @Get('models') @ApiOperation({ summary: 'List all registered ensemble models' }) @ApiResponse({ @@ -332,6 +350,7 @@ export class EnsembleController { * Get embedding coverage statistics * Transforms perModel object to array format for dashboard compatibility */ + @UseGuards(AdminGuard) @Get('coverage') @ApiOperation({ summary: 'Get embedding coverage statistics' }) @ApiResponse({ status: 200, description: 'Coverage stats per model' }) @@ -369,11 +388,14 @@ export class EnsembleController { }) async getMemoryEmbeddings( @Param('id') memoryId: string, + @Req() request?: any, ): Promise<{ memoryId: string; embeddings: MemoryEmbeddingStatus[] }> { if (!memoryId) { throw new BadRequestException('Memory ID is required'); } + const userId = await this.authorizedUser(request); + await this.authorizedMemory(memoryId, userId); const embeddings = await this.ensembleService.getMemoryEmbeddings(memoryId); return { memoryId, embeddings }; } @@ -381,6 +403,7 @@ export class EnsembleController { /** * Get A/B test results */ + @UseGuards(AdminGuard) @Get('ab-results') @ApiOperation({ summary: 'Get A/B test results' }) @ApiQuery({ @@ -409,6 +432,7 @@ export class EnsembleController { /** * Trigger re-embedding for specified models */ + @UseGuards(AdminGuard) @Post('reembed') @HttpCode(HttpStatus.OK) @ApiOperation({ summary: 'Trigger batch re-embedding for models' }) @@ -432,6 +456,7 @@ export class EnsembleController { /** * Get active re-embed job status */ + @UseGuards(AdminGuard) @Get('reembed/status') @ApiOperation({ summary: 'Get active re-embed job status' }) @ApiResponse({ @@ -445,6 +470,7 @@ export class EnsembleController { /** * Re-embed specific memories with specific models (direct endpoint) */ + @UseGuards(AdminGuard) @Post('reembed/targeted') @HttpCode(HttpStatus.OK) @ApiOperation({ summary: 'Re-embed specific memories with specific models' }) @@ -481,6 +507,7 @@ export class EnsembleController { /** * Get latest drift analysis per model */ + @UseGuards(AdminGuard) @Get('drift') @ApiOperation({ summary: 'Get latest drift snapshot per model' }) @ApiResponse({ status: 200, description: 'Latest drift per model' }) @@ -534,6 +561,7 @@ export class EnsembleController { /** * Get drift snapshots over time (for charting) */ + @UseGuards(AdminGuard) @Get('drift/history') @ApiOperation({ summary: 'Get drift history for charting' }) @ApiQuery({ @@ -584,6 +612,7 @@ export class EnsembleController { /** * Trigger a new drift analysis and persist snapshots */ + @UseGuards(AdminGuard) @Post('drift/analyze') @HttpCode(HttpStatus.OK) @ApiOperation({ summary: 'Trigger drift analysis and persist results' }) @@ -699,6 +728,41 @@ export class EnsembleController { return embeddings; } + /** Request-facing routes must scope the privileged vector store explicitly. */ + private async authorizedUser( + request: any, + suppliedUserId?: string, + ): Promise { + const userId = request?.user?.id ?? request?.userId; + const accountId = request?.accountId; + if (!userId || !accountId) { + throw new UnauthorizedException( + 'Authenticated user and account are required', + ); + } + const user = await this.prisma.user.findFirst({ + where: { id: userId, accountId, deletedAt: null }, + select: { id: true }, + }); + if (!user || (suppliedUserId && suppliedUserId !== user.id)) { + throw new ForbiddenException( + 'User scope does not match authenticated identity', + ); + } + return user.id; + } + + private async authorizedMemory( + memoryId: string, + userId: string, + ): Promise { + const memory = await this.prisma.memory.findFirst({ + where: { id: memoryId, userId, deletedAt: null }, + select: { id: true }, + }); + if (!memory) throw new NotFoundException('Memory not found'); + } + private async processTargetedReembed( jobId: string, memoryIds: string[], diff --git a/src/ensemble/ensemble.service.spec.ts b/src/ensemble/ensemble.service.spec.ts index d5eea4fb..1313e8f2 100644 --- a/src/ensemble/ensemble.service.spec.ts +++ b/src/ensemble/ensemble.service.spec.ts @@ -1194,7 +1194,7 @@ describe('PgVectorEnsembleProvider', () => { await provider.upsertEmbedding({ memoryId: 'mem-1', modelId: 'bge-base', - embedding: [0.1, 0.2, 0.3], + embedding: Array(768).fill(0.1), dimensions: 768, }); @@ -1216,7 +1216,7 @@ describe('PgVectorEnsembleProvider', () => { const results = await provider.queryByModel({ userId: 'user-1', modelId: 'bge-base', - embedding: [0.1, 0.2, 0.3], + embedding: Array(768).fill(0.1), limit: 10, }); diff --git a/src/ensemble/ensemble.service.ts b/src/ensemble/ensemble.service.ts index d46b0201..971c75ee 100644 --- a/src/ensemble/ensemble.service.ts +++ b/src/ensemble/ensemble.service.ts @@ -12,7 +12,13 @@ * - Support for batch embedding (for nightly re-embed) */ -import { Injectable, Logger, OnModuleInit } from '@nestjs/common'; +import { + Injectable, + Logger, + OnModuleInit, + ServiceUnavailableException, +} from '@nestjs/common'; +import { assertValidEmbedding } from '../embedding/embedding-validation.util'; import { ConfigService } from '@nestjs/config'; import { OnEvent } from '@nestjs/event-emitter'; import { PrismaService } from '../prisma/prisma.service'; @@ -378,12 +384,11 @@ export class EnsembleService implements OnModuleInit { */ async upsert(options: EnsembleUpsertOptions): Promise { if (!this.config.enabled) { - this.logger.debug('Ensemble disabled, skipping upsert'); - return; + throw new ServiceUnavailableException('Ensemble is not enabled'); } // Generate embeddings from all models (plan-limited for cloud) - const { embeddings, errors } = await this.embedAll( + const { embeddings } = await this.embedAll( options.content, 'document', options.userId, @@ -393,12 +398,15 @@ export class EnsembleService implements OnModuleInit { this.logger.warn( `No embeddings generated for memory ${options.memoryId} — embed service may be down`, ); - return; + throw new ServiceUnavailableException( + 'No embeddings generated; retry the write', + ); } // Prepare records for batch upsert const records: EnsembleEmbeddingRecord[] = embeddings.map((embResult) => ({ memoryId: options.memoryId, + userId: options.userId, modelId: embResult.model, embedding: embResult.embedding, dimensions: embResult.dimensions, @@ -417,7 +425,7 @@ export class EnsembleService implements OnModuleInit { */ async upsertEmbeddings(records: EnsembleEmbeddingRecord[]): Promise { if (!this.config.enabled) { - return; + throw new ServiceUnavailableException('Ensemble is not enabled'); } await this.pgvectorProvider.upsertEmbeddings(records); } @@ -430,7 +438,7 @@ export class EnsembleService implements OnModuleInit { models: ModelId[], ): Promise { if (!this.config.enabled) { - return; + throw new ServiceUnavailableException('Ensemble is not enabled'); } // Fetch memories @@ -439,11 +447,17 @@ export class EnsembleService implements OnModuleInit { select: { id: true, raw: true, userId: true }, }); - if (memories.length === 0) return; + if (memories.length === 0) + throw new ServiceUnavailableException('No memories found to embed'); // Generate embeddings - embedBatch returns embeddings grouped by model, then by text index const texts = memories.map((m) => m.raw); const batchResult = await this.embedBatch(texts, models); + if (batchResult.errors?.length) { + throw new ServiceUnavailableException( + 'Incomplete embedding batch; retry the batch', + ); + } // Group embeddings by model to track text index const embeddingsByModel = new Map(); @@ -457,6 +471,11 @@ export class EnsembleService implements OnModuleInit { // Convert to records for storage const records: EnsembleEmbeddingRecord[] = []; for (const [modelId, modelEmbeddings] of embeddingsByModel) { + if (modelEmbeddings.length !== memories.length) { + throw new ServiceUnavailableException( + 'Embedding batch cardinality mismatch', + ); + } for (let i = 0; i < memories.length && i < modelEmbeddings.length; i++) { const memory = memories[i]; const embedding = modelEmbeddings[i]; @@ -472,9 +491,11 @@ export class EnsembleService implements OnModuleInit { } // Store embeddings - if (records.length > 0) { - await this.pgvectorProvider.upsertEmbeddings(records); - } + if (records.length === 0) + throw new ServiceUnavailableException( + 'No embeddings generated; retry the batch', + ); + await this.pgvectorProvider.upsertEmbeddings(records); } /** @@ -498,7 +519,16 @@ export class EnsembleService implements OnModuleInit { 'query', options.userId, ); - const embeddingMap = new Map(embeddings.map((e) => [e.model, e.embedding])); + const embeddingMap = new Map( + embeddings + .filter((e) => models.includes(e.model)) + .map((e) => [e.model, e.embedding]), + ); + if (embeddingMap.size === 0) { + throw new ServiceUnavailableException( + 'No query embeddings generated; retry retrieval', + ); + } // Query each model using pgvector const modelResults = await this.pgvectorProvider.queryWithModelEmbeddings( @@ -507,6 +537,12 @@ export class EnsembleService implements OnModuleInit { topKPerModel, ); + if (modelResults.size === 0) + throw new ServiceUnavailableException( + 'No ensemble search backend succeeded', + ); + const failedModels = models.filter((model) => !modelResults.has(model)); + // Apply RRF fusion const weights = { ...this.config.weights, ...options.weights }; const fusedResults = this.reciprocalRankFusion(modelResults, k, weights); @@ -518,6 +554,8 @@ export class EnsembleService implements OnModuleInit { modelsQueried: Array.from(modelResults.keys()), candidatesEvaluated: fusedResults.length, fusionAlgorithm: 'rrf', + degraded: failedModels.length > 0, + failedModels, }, }; } @@ -608,7 +646,9 @@ export class EnsembleService implements OnModuleInit { ): Promise { // Route to cloud providers if active if (this.useCloud) { - return this.cloudEnsemble.embedBatch(texts, models); + const result = await this.cloudEnsemble.embedBatch(texts, models); + this.validateBatch(result.embeddings, texts.length, models); + return result; } const start = Date.now(); @@ -636,38 +676,47 @@ export class EnsembleService implements OnModuleInit { const data = await response.json(); - // Handle multi-model response - if (data.embeddings) { - for (const modelEmbed of data.embeddings) { - const modelId = modelEmbed.model as ModelId; - if (models.includes(modelId) && modelEmbed.data) { - // For batch embeddings, we get an array of embeddings - for (let i = 0; i < modelEmbed.data.length; i++) { - if (modelEmbed.data[i]?.embedding) { - embeddings.push({ - model: modelId, - dimensions: modelEmbed.dimensions, - embedding: modelEmbed.data[i].embedding, - latencyMs: data.timing?.per_model?.[modelId] ?? 0, - }); - } - } - } + // A model is accepted only as a complete, unambiguous input-aligned batch. + // Positional responses require exact cardinality; indexed responses are reordered. + const responses = data.embeddings ?? [ + { model: models[0], data: data.data }, + ]; + for (const modelEmbed of responses) { + const modelId = modelEmbed.model as ModelId; + if (!models.includes(modelId)) continue; + const entries = modelEmbed.data; + if (!Array.isArray(entries) || entries.length !== texts.length) { + throw new Error(`Incomplete batch for ${modelId}`); } - } else if (data.data) { - // Single model response format - for (let i = 0; i < data.data.length; i++) { - if (data.data[i]?.embedding) { - embeddings.push({ - model: models[0], - dimensions: data.data[i].embedding.length, - embedding: data.data[i].embedding, - latencyMs: Date.now() - start, - }); + const ordered: EmbeddingResult[] = new Array(texts.length); + const indexed = entries.some((entry) => entry?.index !== undefined); + for (let i = 0; i < entries.length; i++) { + const entry = entries[i]; + const index = indexed ? entry?.index : i; + if ( + !Number.isInteger(index) || + index < 0 || + index >= texts.length || + ordered[index] + ) { + throw new Error(`Invalid or duplicate batch index for ${modelId}`); } + const embedding = assertValidEmbedding(entry?.embedding, { + expectedDimensions: MODEL_CONFIGS[modelId]?.dimensions, + context: `batch ${modelId} input ${index}`, + }); + ordered[index] = { + model: modelId, + dimensions: embedding.length, + embedding, + latencyMs: data.timing?.per_model?.[modelId] ?? Date.now() - start, + }; } + embeddings.push(...ordered); } + this.validateBatch(embeddings, texts.length, models); } catch (error) { + embeddings.length = 0; // Never expose a partially accepted positional batch. this.logger.error('Failed to get batch embeddings', error); for (const model of models) { errors.push({ @@ -685,6 +734,24 @@ export class EnsembleService implements OnModuleInit { }; } + private validateBatch( + embeddings: EmbeddingResult[], + count: number, + models: ModelId[], + ): void { + for (const model of models) { + const batch = embeddings.filter((entry) => entry.model === model); + if (batch.length !== count) + throw new ServiceUnavailableException(`Incomplete batch for ${model}`); + for (const entry of batch) { + assertValidEmbedding(entry.embedding, { + expectedDimensions: MODEL_CONFIGS[model]?.dimensions, + context: `batch ${model}`, + }); + } + } + } + /** * Compare ensemble vs single-model retrieval (for debugging/benchmarking) */ @@ -701,7 +768,16 @@ export class EnsembleService implements OnModuleInit { // Get individual model results const { embeddings } = await this.embedAll(query); - const embeddingMap = new Map(embeddings.map((e) => [e.model, e.embedding])); + const embeddingMap = new Map( + embeddings + .filter((e) => this.config.models.includes(e.model)) + .map((e) => [e.model, e.embedding]), + ); + if (embeddingMap.size === 0) { + throw new ServiceUnavailableException( + 'No query embeddings generated; retry retrieval', + ); + } const singleModel = await this.pgvectorProvider.queryWithModelEmbeddings( embeddingMap, diff --git a/src/ensemble/pgvector-ensemble.provider.spec.ts b/src/ensemble/pgvector-ensemble.provider.spec.ts index 744d9a5a..bb50b4d4 100644 --- a/src/ensemble/pgvector-ensemble.provider.spec.ts +++ b/src/ensemble/pgvector-ensemble.provider.spec.ts @@ -42,17 +42,18 @@ describe('PgVectorEnsembleProvider', () => { await provider.upsertEmbedding({ memoryId: 'mem-1', modelId: 'bge-base' as any, - embedding: [0.1, 0.2, 0.3], - dimensions: 3, + embedding: Array(768).fill(0.1), + dimensions: 768, }); expect(mockPrisma.$executeRawUnsafe).toHaveBeenCalledWith( expect.stringContaining('INSERT INTO memory_embeddings'), 'mem-1', 'bge-base', - 3, - '[0.1,0.2,0.3]', + 768, + JSON.stringify(Array(768).fill(0.1)), expect.any(Date), + null, ); }); }); @@ -69,14 +70,14 @@ describe('PgVectorEnsembleProvider', () => { { memoryId: 'mem-1', modelId: 'bge-base' as any, - embedding: [0.1], - dimensions: 1, + embedding: Array(768).fill(0.1), + dimensions: 768, }, { memoryId: 'mem-2', modelId: 'bge-base' as any, - embedding: [0.2], - dimensions: 1, + embedding: Array(768).fill(0.2), + dimensions: 768, }, ]); @@ -94,7 +95,7 @@ describe('PgVectorEnsembleProvider', () => { const results = await provider.queryByModel({ userId: 'user1', modelId: 'bge-base' as any, - embedding: [0.1, 0.2, 0.3], + embedding: Array(768).fill(0.1), limit: 5, }); @@ -104,9 +105,9 @@ describe('PgVectorEnsembleProvider', () => { ]); expect(mockPrisma.$queryRawUnsafe).toHaveBeenCalledWith( expect.stringContaining('FROM memory_embeddings'), - '[0.1,0.2,0.3]', + JSON.stringify(Array(768).fill(0.1)), 'bge-base', - 3, + 768, 'user1', 5, ); @@ -118,7 +119,7 @@ describe('PgVectorEnsembleProvider', () => { const results = await provider.queryByModel({ userId: 'user1', modelId: 'bge-base' as any, - embedding: [0.1], + embedding: Array(768).fill(0.1), limit: 5, }); @@ -148,12 +149,14 @@ describe('PgVectorEnsembleProvider', () => { it('should skip models with mismatched dimensions', async () => { // 3-dim embedding won't match any model config (768, 384, etc.) - const results = await provider.queryAllModels( - [0.1, 0.2, 0.3], - 'user1', - ['bge-base' as any], - 5, - ); + await expect( + provider.queryAllModels( + [0.1, 0.2, 0.3], + 'user1', + ['bge-base' as any], + 5, + ), + ).rejects.toThrow('No ensemble search backend succeeded'); // Should skip bge-m3 since dims don't match expect(mockPrisma.$queryRawUnsafe).not.toHaveBeenCalled(); @@ -163,15 +166,9 @@ describe('PgVectorEnsembleProvider', () => { mockPrisma.$queryRawUnsafe.mockRejectedValue(new Error('DB error')); const embedding768 = new Array(768).fill(0.1); - const results = await provider.queryAllModels( - embedding768, - 'user1', - ['bge-base' as any], - 5, - ); - - // Should not throw, just return empty results for failed model - expect(results.get('bge-base' as any)).toBeUndefined(); + await expect( + provider.queryAllModels(embedding768, 'user1', ['bge-base'], 5), + ).rejects.toThrow('No ensemble search backend succeeded'); }); }); diff --git a/src/ensemble/pgvector-ensemble.provider.ts b/src/ensemble/pgvector-ensemble.provider.ts index d5e0ade4..ea1b5eee 100644 --- a/src/ensemble/pgvector-ensemble.provider.ts +++ b/src/ensemble/pgvector-ensemble.provider.ts @@ -11,7 +11,12 @@ * - Supports upsert (update on conflict) */ -import { Injectable, Logger } from '@nestjs/common'; +import { + Injectable, + Logger, + ServiceUnavailableException, +} from '@nestjs/common'; +import { assertValidEmbedding } from '../embedding/embedding-validation.util'; import { ServicePrismaService } from '../prisma/service-prisma.service'; import { ModelId, @@ -27,6 +32,8 @@ import { } from './ensemble.types'; export interface EnsembleEmbeddingRecord { + /** Required for request-facing writes; omitted only by trusted background jobs. */ + userId?: string; memoryId: string; modelId: ModelId; embedding: number[]; @@ -52,17 +59,33 @@ export class PgVectorEnsembleProvider { constructor(private prisma: ServicePrismaService) {} + private validateRecord(record: EnsembleEmbeddingRecord): void { + const model = MODEL_CONFIGS[record.modelId]; + if (!model || record.dimensions !== model.dimensions) { + throw new Error( + `Invalid dimensions for ensemble model ${record.modelId}`, + ); + } + assertValidEmbedding(record.embedding, { + expectedDimensions: model.dimensions, + context: `ensemble store ${record.memoryId}/${record.modelId}`, + }); + } + /** * Upsert a single embedding for a memory/model pair */ async upsertEmbedding(record: EnsembleEmbeddingRecord): Promise { + this.validateRecord(record); const embeddingStr = `[${record.embedding.join(',')}]`; const now = new Date(); - await this.prisma.$executeRawUnsafe( + const written = await this.prisma.$executeRawUnsafe( ` INSERT INTO memory_embeddings (id, memory_id, model_id, dimensions, embedding, created_at, updated_at) - VALUES (gen_random_uuid()::text, $1, $2, $3, $4::vector, $5, $5) + SELECT gen_random_uuid()::text, m.id, $2, $3, $4::vector, $5, $5 + FROM memories m WHERE m.id = $1 AND m.deleted_at IS NULL + AND ($6::text IS NULL OR m.user_id = $6) ON CONFLICT (memory_id, model_id) DO UPDATE SET embedding = $4::vector, @@ -74,13 +97,19 @@ export class PgVectorEnsembleProvider { record.dimensions, embeddingStr, now, + record.userId ?? null, ); + if (written !== 1) + throw new ServiceUnavailableException('Embedding was not stored'); } /** * Upsert multiple embeddings in a batch */ async upsertEmbeddings(records: EnsembleEmbeddingRecord[]): Promise { + if (records.length === 0) + throw new ServiceUnavailableException('No embeddings to store'); + records.forEach((record) => this.validateRecord(record)); // Use a transaction for atomicity await this.prisma.$transaction( async (tx) => { @@ -88,10 +117,12 @@ export class PgVectorEnsembleProvider { const embeddingStr = `[${record.embedding.join(',')}]`; const now = new Date(); - await tx.$executeRawUnsafe( + const written = await tx.$executeRawUnsafe( ` INSERT INTO memory_embeddings (id, memory_id, model_id, dimensions, embedding, created_at, updated_at) - VALUES (gen_random_uuid()::text, $1, $2, $3, $4::vector, $5, $5) + SELECT gen_random_uuid()::text, m.id, $2, $3, $4::vector, $5, $5 + FROM memories m WHERE m.id = $1 AND m.deleted_at IS NULL + AND ($6::text IS NULL OR m.user_id = $6) ON CONFLICT (memory_id, model_id) DO UPDATE SET embedding = $4::vector, @@ -103,7 +134,10 @@ export class PgVectorEnsembleProvider { record.dimensions, embeddingStr, now, + record.userId ?? null, ); + if (written !== 1) + throw new ServiceUnavailableException('Embedding was not stored'); } }, { timeout: 120000 }, @@ -127,6 +161,12 @@ export class PgVectorEnsembleProvider { async queryByModel( options: EnsembleSearchOptions, ): Promise { + if (!MODEL_CONFIGS[options.modelId]) + throw new Error(`Unknown ensemble model ${options.modelId}`); + assertValidEmbedding(options.embedding, { + expectedDimensions: MODEL_CONFIGS[options.modelId]?.dimensions, + context: `ensemble query ${options.modelId}`, + }); const embeddingStr = `[${options.embedding.join(',')}]`; const dimensions = options.embedding.length; @@ -208,6 +248,10 @@ export class PgVectorEnsembleProvider { }), ); + if (results.size === 0) + throw new ServiceUnavailableException( + 'No ensemble search backend succeeded', + ); return results; } @@ -245,6 +289,10 @@ export class PgVectorEnsembleProvider { }), ); + if (results.size === 0) + throw new ServiceUnavailableException( + 'No ensemble search backend succeeded', + ); return results; } From 8028582d04ebb707b13428442ad3ada5c469de38 Mon Sep 17 00:00:00 2001 From: heybeaux Date: Sun, 6 Sep 2026 18:17:22 -0700 Subject: [PATCH 2/6] fix: preserve source memories across consolidation and semantic dedup --- .../audit/consolidation-reconciliation.sql | 26 +++ scripts/audit/preservation-postgres.cjs | 89 ++++++++++ .../stages/dream-cycle-archival.stage.spec.ts | 20 ++- .../stages/dream-cycle-archival.stage.ts | 51 ++++-- .../dream-cycle-consolidation.stage.spec.ts | 8 +- .../stages/dream-cycle-consolidation.stage.ts | 42 +++-- .../stages/dream-cycle-tiering.stage.ts | 101 ++++++----- .../stages/preservation-regression.spec.ts | 168 ++++++++++++++++++ src/memory/memory-dedup.service.spec.ts | 19 +- src/memory/memory-dedup.service.ts | 29 ++- src/vector/embedding-write.service.spec.ts | 18 +- src/vector/embedding-write.service.ts | 48 +++-- 12 files changed, 502 insertions(+), 117 deletions(-) create mode 100644 scripts/audit/consolidation-reconciliation.sql create mode 100644 scripts/audit/preservation-postgres.cjs create mode 100644 src/consolidation/stages/preservation-regression.spec.ts diff --git a/scripts/audit/consolidation-reconciliation.sql b/scripts/audit/consolidation-reconciliation.sql new file mode 100644 index 00000000..29a0f05c --- /dev/null +++ b/scripts/audit/consolidation-reconciliation.sql @@ -0,0 +1,26 @@ +-- READ ONLY inventory. Supply the deployment's actual retrieval model and dims: +-- psql ... -v model_id=bge-base -v expected_dims=768 -f scripts/audit/consolidation-reconciliation.sql +-- No repairs are performed. Snapshot/back up before planning any restoration. +BEGIN READ ONLY; +WITH replacements AS ( + SELECT s.user_id, s.consolidated_into AS replacement_id, + count(*) AS linked_sources, + count(*) FILTER (WHERE s.deleted_at IS NULL AND s.superseded_by_id IS NULL AND s.searchable) AS active_sources, + array_agg(s.id ORDER BY s.id) AS source_ids + FROM memories s + WHERE s.consolidated_into IS NOT NULL + GROUP BY s.user_id, s.consolidated_into +) +SELECT r.*, m.id IS NULL AS missing_replacement, + m.user_id IS DISTINCT FROM r.user_id AS ownership_mismatch, + m.searchable, m.deleted_at IS NOT NULL AS replacement_deleted, + me.model_id, me.dimensions, vector_dims(me.embedding) AS actual_dimensions +FROM replacements r +LEFT JOIN memories m ON m.id = r.replacement_id +LEFT JOIN memory_embeddings me ON me.memory_id = m.id AND me.model_id = :'model_id' +WHERE m.id IS NULL OR m.user_id IS DISTINCT FROM r.user_id + OR m.deleted_at IS NOT NULL OR NOT m.searchable OR m.superseded_by_id IS NOT NULL + OR me.embedding IS NULL OR me.dimensions <> :expected_dims + OR vector_dims(me.embedding) <> :expected_dims +ORDER BY r.user_id, r.replacement_id; +ROLLBACK; diff --git a/scripts/audit/preservation-postgres.cjs b/scripts/audit/preservation-postgres.cjs new file mode 100644 index 00000000..7af1cae6 --- /dev/null +++ b/scripts/audit/preservation-postgres.cjs @@ -0,0 +1,89 @@ +// Isolated H1 regression: never reads application DATABASE_URL or touches its DB. +// Run: node scripts/audit/preservation-postgres.cjs (local socket /tmp, pgvector). +const path = require('node:path'); +const assert = require('node:assert/strict'); +const repo = path.resolve(__dirname, '../..'); +require('ts-node').register({ transpileOnly: true, project: path.join(repo, 'tsconfig.json') }); +const { Client } = require('pg'); +const { DreamCycleConsolidationStage } = require(path.join(repo, 'src/consolidation/stages/dream-cycle-consolidation.stage.ts')); +const { EmbeddingWriteService } = require(path.join(repo, 'src/vector/embedding-write.service.ts')); +const db = `engram_preservation_test_${process.pid}`; +const config = { host: '/tmp', database: 'postgres' }; +(async () => { + const admin = new Client(config); await admin.connect(); + await admin.query(`CREATE DATABASE ${db}`); + const tx = new Client({ ...config, database: db }); + try { + await tx.connect(); + await tx.query('CREATE EXTENSION vector'); + await tx.query(`CREATE TABLE memories(id text PRIMARY KEY, user_id text DEFAULT 'u', raw text, + embedding vector(768), searchable boolean, superseded_by_id text, tier text, + consolidated boolean DEFAULT false, consolidated_into text, user_pinned boolean DEFAULT false, + deleted_at timestamp, created_at timestamp DEFAULT now())`); + await tx.query(`CREATE TABLE memory_embeddings(id text PRIMARY KEY, memory_id text REFERENCES memories(id), + model_id text, dimensions int, embedding vector, created_at timestamp, updated_at timestamp, + UNIQUE(memory_id, model_id))`); + const sql = { + $queryRawUnsafe: async (q, ...args) => (await tx.query(q, args)).rows, + $executeRawUnsafe: async (q, ...args) => (await tx.query(q, args)).rowCount, + }; + let failPublication = false; + const client = { ...sql, memory: { + create: async ({ data }) => { + await tx.query('INSERT INTO memories(id,raw,searchable,tier) VALUES ($1,$2,$3,$4)', ['new', data.raw, data.searchable, data.tier]); + return { id: 'new' }; + }, + update: async ({ data }) => tx.query("UPDATE memories SET searchable=$1 WHERE id='new'", [data.searchable]), + updateMany: async ({ data }) => { + if (failPublication) await tx.query('SELECT 1/0'); + return tx.query(`UPDATE memories SET consolidated=$1,consolidated_into=$2, + superseded_by_id=$3,tier=COALESCE($4,tier) WHERE id IN ('a','b','c')`, + [data.consolidated, data.consolidatedInto, data.supersededById ?? null, data.tier ?? null]); + }, + } }; + const prisma = { $queryRaw: async (strings, ...args) => (await tx.query(strings.reduce((s, part, i) => s + (i ? '$' + i : '') + part, ''), args)).rows, $transaction: async (fn) => { + await tx.query('BEGIN'); + try { const result = await fn(client); await tx.query('COMMIT'); return result; } + catch (e) { await tx.query('ROLLBACK'); throw e; } + } }; + const embeddingWriter = new EmbeddingWriteService({ + $queryRawUnsafe: () => { throw new Error('Used external connection'); }, + $executeRawUnsafe: () => { throw new Error('Used external connection'); }, + }); + for (const [name, model, dims, vector, fail] of [ + ['768 success', 'bge-base', 768, Array(768).fill(0.1), false], + ['1536 success', 'openai-small', 1536, Array(1536).fill(0.1), false], + ['publication DB failure', 'bge-base', 768, Array(768).fill(0.1), true], + ['missing', 'bge-base', 768, undefined, true], + ['sparse', 'bge-base', 768, Array(768), true], + ['wrong dimensions', 'bge-base', 768, Array(1536).fill(0.1), true], + ]) { + failPublication = name === 'publication DB failure'; + process.env.EMBEDDING_MODEL = model; + await tx.query('TRUNCATE memories, memory_embeddings'); + await tx.query("INSERT INTO memories(id,raw,searchable,tier) VALUES ('a','date: September 12',true,'COLD'),('b','exception: weekends',true,'COLD'),('c','name: Ada',true,'COLD')"); + const vec = `[${Array(dims).fill(0.1).join(',')}]`; + await tx.query('INSERT INTO memory_embeddings SELECT id,id,$1,$2,$3::vector,now(),now() FROM memories', [model, dims, vec]); + const stage = new DreamCycleConsolidationStage(prisma, { get: () => undefined }, + { embed: async () => [vector] }, { chat: async () => ({ content: 'Lossy summary' }) }, embeddingWriter); + assert.equal((await stage.fetchColdMemories('u')).length, 3); + const call = () => stage.consolidateCluster([{ id: 'a', content: 'date: September 12' }, { id: 'b', content: 'exception: weekends' }, { id: 'c', content: 'name: Ada' }], 'u'); + if (fail) await assert.rejects(call); else await call(); + const eligible = (await tx.query(`SELECT m.id, me.dimensions FROM memories m JOIN memory_embeddings me ON me.memory_id=m.id + WHERE m.searchable AND m.superseded_by_id IS NULL AND me.model_id=$1 ORDER BY m.id`, [model])).rows; + assert.equal(eligible.length, fail ? 3 : 4); + assert.deepEqual(eligible.filter(r => r.id !== 'new').map(r => r.id), ['a','b','c']); + assert.ok(eligible.every(r => r.dimensions === dims)); + console.log(JSON.stringify({ name, eligible: eligible.length, sourcesPreserved: 3, rollback: fail })); + } + await tx.query("UPDATE memories SET consolidated_into='missing-replacement' WHERE id IN ('a','b','c')"); + const inventory = require('node:fs').readFileSync(path.join(__dirname, 'consolidation-reconciliation.sql'), 'utf8') + .replace(/:'model_id'/g, "'bge-base'").replace(/:expected_dims/g, '768'); + const inventoryResult = await tx.query(inventory); + assert.equal(inventoryResult[1].rows.length, 1); + assert.equal(inventoryResult[1].rows[0].missing_replacement, true); + console.log(JSON.stringify({ reconciliationReadOnly: true, missingReplacementDetected: true })); + } finally { + await tx.end(); await admin.query(`DROP DATABASE ${db}`); await admin.end(); + } +})().catch(e => { console.error(e); process.exitCode = 1; }); diff --git a/src/consolidation/stages/dream-cycle-archival.stage.spec.ts b/src/consolidation/stages/dream-cycle-archival.stage.spec.ts index 8075e667..6f209de0 100644 --- a/src/consolidation/stages/dream-cycle-archival.stage.spec.ts +++ b/src/consolidation/stages/dream-cycle-archival.stage.spec.ts @@ -13,7 +13,11 @@ describe('DreamCycleArchivalStage', () => { prisma = { memory: { findMany: jest.fn().mockResolvedValue([]), - updateMany: jest.fn().mockResolvedValue({ count: 0 }), + updateMany: jest + .fn() + .mockImplementation(({ where }) => + Promise.resolve({ count: where.id.in.length }), + ), }, }; @@ -59,7 +63,12 @@ describe('DreamCycleArchivalStage', () => { expect(result.byLayer).toEqual({ SESSION: 1 }); expect(result.byType).toEqual({ EVENT: 1 }); expect(prisma.memory.updateMany).toHaveBeenCalledWith({ - where: { id: { in: ['m1'] }, userId: 'user1' }, + where: expect.objectContaining({ + id: { in: ['m1'] }, + userId: 'user1', + userPinned: false, + createdAt: { lt: expect.any(Date) }, + }), data: { archivedReason: 'low_importance', searchable: false, @@ -272,7 +281,12 @@ describe('DreamCycleArchivalStage', () => { expect(result.skippedRecentlyRetrieved).toBe(1); expect(result.skippedFrequentlyUsed).toBe(1); expect(prisma.memory.updateMany).toHaveBeenCalledWith({ - where: { id: { in: ['m1'] }, userId: 'user1' }, + where: expect.objectContaining({ + id: { in: ['m1'] }, + userId: 'user1', + userPinned: false, + createdAt: { lt: expect.any(Date) }, + }), data: expect.objectContaining({ archivedReason: 'low_importance' }), }); }); diff --git a/src/consolidation/stages/dream-cycle-archival.stage.ts b/src/consolidation/stages/dream-cycle-archival.stage.ts index b1c28fd5..ec567056 100644 --- a/src/consolidation/stages/dream-cycle-archival.stage.ts +++ b/src/consolidation/stages/dream-cycle-archival.stage.ts @@ -1,3 +1,4 @@ +import { Prisma } from '@prisma/client'; import { Injectable, Logger } from '@nestjs/common'; import { ConfigService } from '@nestjs/config'; import { ServicePrismaService } from '../../prisma/service-prisma.service'; @@ -53,14 +54,36 @@ export class DreamCycleArchivalStage { retrievalCutoff.getDate() - this.retrievalWindowDays, ); + const minAgeDays = Math.max( + 1, + Number(this.config.get('DREAM_ARCHIVAL_MIN_AGE_DAYS') ?? 90) || 90, + ); + const creationCutoff = new Date(Date.now() - minAgeDays * 86_400_000); + const eligibility: Prisma.MemoryWhereInput = { + userId, + deletedAt: null, + searchable: true, + userPinned: false, + supersededById: null, + OR: [{ tier: null }, { tier: { not: 'ARCHIVED' } }], + createdAt: { lt: creationCutoff }, + importanceScore: { lt: this.importanceThreshold }, + layer: { notIn: ['IDENTITY', 'PROJECT'] }, + usedCount: { lte: this.maxUsedCount }, + AND: [ + { + OR: [ + { lastRetrievedAt: null }, + { lastRetrievedAt: { lt: retrievalCutoff } }, + ], + }, + { OR: [{ lastUsedAt: null }, { lastUsedAt: { lt: retrievalCutoff } }] }, + ], + }; + // Fetch candidate memories: low importance, active, not already archived const candidates = await this.prisma.memory.findMany({ - where: { - userId, - deletedAt: null, - searchable: true, - importanceScore: { lt: this.importanceThreshold }, - }, + where: eligibility, select: { id: true, layer: true, @@ -124,9 +147,10 @@ export class DreamCycleArchivalStage { byType[type] = (byType[type] ?? 0) + 1; } + let archived = toArchive.length; if (toArchive.length > 0 && !dryRun) { - await this.prisma.memory.updateMany({ - where: { id: { in: toArchive }, userId }, + const updated = await this.prisma.memory.updateMany({ + where: { ...eligibility, id: { in: toArchive } }, data: { archivedReason: 'low_importance', searchable: false, @@ -134,13 +158,20 @@ export class DreamCycleArchivalStage { }, }); + archived = updated.count; + // updateMany returns no identities. Do not claim per-layer counts when + // concurrent pin/access changes prevented some writes. + if (archived !== toArchive.length) { + for (const key of Object.keys(byLayer)) delete byLayer[key]; + for (const key of Object.keys(byType)) delete byType[key]; + } this.logger.log( - `Archived ${toArchive.length} memories for user ${userId}: ${JSON.stringify({ byLayer, byType })}`, + `Archived ${archived} memories for user ${userId}: ${JSON.stringify({ byLayer, byType })}`, ); } const result: ArchivalStageResult = { - archived: toArchive.length, + archived, skippedProtectedLayer, skippedRecentlyRetrieved, skippedFrequentlyUsed, diff --git a/src/consolidation/stages/dream-cycle-consolidation.stage.spec.ts b/src/consolidation/stages/dream-cycle-consolidation.stage.spec.ts index 1fb0ade1..84bb080e 100644 --- a/src/consolidation/stages/dream-cycle-consolidation.stage.spec.ts +++ b/src/consolidation/stages/dream-cycle-consolidation.stage.spec.ts @@ -57,7 +57,7 @@ describe('DreamCycleConsolidationStage', () => { { provide: EmbeddingWriteService, useValue: { - writeLegacyInlineEmbedding: jest.fn().mockResolvedValue(undefined), + getCurrentModelId: jest.fn().mockReturnValue('bge-base'), writeMemoryEmbedding: jest.fn().mockResolvedValue(undefined), }, }, @@ -173,7 +173,7 @@ describe('DreamCycleConsolidationStage', () => { const result = await stage.run('user1', false); expect(result.clustersFound).toBe(1); expect(result.consolidated).toBe(1); - expect(result.archived).toBe(3); + expect(result.archived).toBe(0); expect(result.llmCalls).toBe(1); expect(llmService.chat).toHaveBeenCalledTimes(1); expect(embeddingService.embed).toHaveBeenCalledTimes(1); @@ -255,7 +255,7 @@ describe('DreamCycleConsolidationStage', () => { { provide: EmbeddingWriteService, useValue: { - writeLegacyInlineEmbedding: jest.fn().mockResolvedValue(undefined), + getCurrentModelId: jest.fn().mockReturnValue('bge-base'), writeMemoryEmbedding: jest.fn().mockResolvedValue(undefined), }, }, @@ -326,7 +326,7 @@ describe('DreamCycleConsolidationStage', () => { { provide: EmbeddingWriteService, useValue: { - writeLegacyInlineEmbedding: jest.fn().mockResolvedValue(undefined), + getCurrentModelId: jest.fn().mockReturnValue('bge-base'), writeMemoryEmbedding: jest.fn().mockResolvedValue(undefined), }, }, diff --git a/src/consolidation/stages/dream-cycle-consolidation.stage.ts b/src/consolidation/stages/dream-cycle-consolidation.stage.ts index 2c09d013..09dd0464 100644 --- a/src/consolidation/stages/dream-cycle-consolidation.stage.ts +++ b/src/consolidation/stages/dream-cycle-consolidation.stage.ts @@ -88,7 +88,7 @@ export class DreamCycleConsolidationStage { result.llmCalls++; } result.consolidated++; - result.archived += cluster.length; + // Summaries are not proof of factual equivalence: keep sources searchable. } catch (err) { const msg = err instanceof Error ? err.message : String(err); this.logger.error(`Failed to consolidate cluster: ${msg}`); @@ -105,13 +105,18 @@ export class DreamCycleConsolidationStage { const memories = await this.prisma.$queryRaw< Array<{ id: string; content: string; embedding: string | null }> >` - SELECT id, raw AS content, embedding::text - FROM memories + SELECT m.id, m.raw AS content, me.embedding::text + FROM memories m + JOIN memory_embeddings me ON me.memory_id = m.id + AND me.model_id = ${this.embeddingWrite.getCurrentModelId()} WHERE user_id = ${userId} AND deleted_at IS NULL AND tier = 'COLD' AND consolidated = false - ORDER BY created_at ASC + AND searchable = true + AND superseded_by_id IS NULL + AND user_pinned = false + ORDER BY m.created_at ASC `; return memories.map((m) => ({ @@ -208,7 +213,9 @@ Write a single consolidated memory that captures all the information above.`; consolidatedContent, ]); - // Create consolidated memory and archive originals in a transaction + if (!embedding) throw new Error('Consolidation embedding is missing'); + + // Atomically publish the replacement in the same model store as retrieval. await this.prisma.$transaction(async (tx) => { // Create the new consolidated memory; searchable is set after embedding is written const newMemory = await tx.memory.create({ @@ -224,25 +231,26 @@ Write a single consolidated memory that captures all the information above.`; }, }); - if (embedding) { - await this.embeddingWrite.writeLegacyInlineEmbedding(newMemory.id, embedding); - // Mark searchable now that embedding exists - await tx.memory.update({ - where: { id: newMemory.id }, - data: { searchable: true }, - }); - } + await this.embeddingWrite.writeMemoryEmbedding( + newMemory.id, + this.embeddingWrite.getCurrentModelId(), + embedding, + false, + tx, + ); + await tx.memory.update({ + where: { id: newMemory.id }, + data: { searchable: true, embeddingStatus: 'COMPLETE' }, + }); - // Link originals to the consolidated memory and archive them. - // Set supersededById so resolveSuperseded() can follow the chain in benchmarks. + // An LLM summary cannot establish equivalence. Link processed sources, + // but preserve their independent retrieval until equivalence is verified. const originalIds = cluster.map((m) => m.id); await tx.memory.updateMany({ where: { id: { in: originalIds }, userId }, data: { consolidatedInto: newMemory.id, consolidated: true, - tier: 'ARCHIVED', - supersededById: newMemory.id, }, }); }); diff --git a/src/consolidation/stages/dream-cycle-tiering.stage.ts b/src/consolidation/stages/dream-cycle-tiering.stage.ts index 95173381..b9e24f96 100644 --- a/src/consolidation/stages/dream-cycle-tiering.stage.ts +++ b/src/consolidation/stages/dream-cycle-tiering.stage.ts @@ -34,57 +34,62 @@ export class DreamCycleTieringStage { `Starting memory tiering for user ${userId} (dryRun: ${dryRun})`, ); - const memories = await this.prisma.memory.findMany({ - where: { - userId, - deletedAt: null, - }, - select: { - id: true, - tier: true, - userPinned: true, - createdAt: true, - lastRetrievedAt: true, - retrievalCount: true, - }, - take: this.batchSize, - orderBy: { createdAt: 'asc' }, - }); - - if (memories.length === 0) { - this.logger.log(`No memories found for user ${userId}`); - return { promoted: 0, demoted: 0, unchanged: 0 }; - } - - this.logger.log(`Processing ${memories.length} memories for tiering`); - const now = new Date(); - - for (const memory of memories) { - const newTier = this.calculateTier(memory, now); - const currentTier = memory.tier ?? 'WARM'; - - if (newTier === currentTier) { - unchanged++; - continue; + let afterId: string | undefined; + while (true) { + const memories = await this.prisma.memory.findMany({ + where: { + userId, + deletedAt: null, + searchable: true, + supersededById: null, + OR: [{ tier: null }, { tier: { not: 'ARCHIVED' } }], + createdAt: { lte: now }, + ...(afterId ? { id: { gt: afterId } } : {}), + }, + select: { + id: true, + tier: true, + userPinned: true, + createdAt: true, + lastRetrievedAt: true, + retrievalCount: true, + }, + take: this.batchSize, + orderBy: { id: 'asc' }, + }); + + if (memories.length === 0) break; + + for (const memory of memories) { + const newTier = this.calculateTier(memory, now); + const currentTier = memory.tier ?? 'WARM'; + + if (newTier === currentTier) { + unchanged++; + continue; + } + + const currentOrder = + TIER_ORDER[currentTier as keyof typeof TIER_ORDER] ?? 1; + const newOrder = TIER_ORDER[newTier as keyof typeof TIER_ORDER]; + + if (newOrder < currentOrder) { + promoted++; + } else { + demoted++; + } + + if (!dryRun) { + await this.prisma.memory.update({ + where: { id: memory.id }, + data: { tier: newTier }, + }); + } } - const currentOrder = - TIER_ORDER[currentTier as keyof typeof TIER_ORDER] ?? 1; - const newOrder = TIER_ORDER[newTier as keyof typeof TIER_ORDER]; - - if (newOrder < currentOrder) { - promoted++; - } else { - demoted++; - } - - if (!dryRun) { - await this.prisma.memory.update({ - where: { id: memory.id }, - data: { tier: newTier }, - }); - } + afterId = memories[memories.length - 1].id; + if (memories.length < this.batchSize) break; } const result = { promoted, demoted, unchanged }; diff --git a/src/consolidation/stages/preservation-regression.spec.ts b/src/consolidation/stages/preservation-regression.spec.ts new file mode 100644 index 00000000..12e786b5 --- /dev/null +++ b/src/consolidation/stages/preservation-regression.spec.ts @@ -0,0 +1,168 @@ +import { DreamCycleConsolidationStage } from './dream-cycle-consolidation.stage'; +import { DreamCycleTieringStage } from './dream-cycle-tiering.stage'; +import { DreamCycleArchivalStage } from './dream-cycle-archival.stage'; +import { EmbeddingWriteService } from '../../vector/embedding-write.service'; +import { MemoryDedupService } from '../../memory/memory-dedup.service'; + +const config = { get: () => undefined }; + +describe('audit preservation regressions', () => { + it.each([0.99, 0.88])( + 'preserves corrected facts at similarity %s', + async (score) => { + const prisma = { + memory: { + findUnique: jest.fn().mockResolvedValue({ + id: 'old', + userId: 'u', + raw: 'Deadline September 12', + searchable: true, + }), + }, + mergeCandidate: { create: jest.fn() }, + }; + const service = new MemoryDedupService( + prisma as any, + { + generate: async () => [1], + search: async () => [{ id: 'old', score }], + } as any, + ); + expect( + await service.findDuplicateV2('u', 'Deadline September 13'), + ).toEqual({ action: 'create' }); + }, + ); + + it('processes all 501 eligible records, even when first 500 are unchanged', async () => { + const rows = Array.from({ length: 501 }, (_, i) => ({ + id: `m${String(i).padStart(4, '0')}`, + tier: 'HOT', + userPinned: true, + createdAt: new Date(), + lastRetrievedAt: null, + retrievalCount: 0, + })); + const findMany = jest + .fn() + .mockResolvedValueOnce(rows.slice(0, 500)) + .mockResolvedValueOnce(rows.slice(500)); + const stage = new DreamCycleTieringStage( + { memory: { findMany } } as any, + config as any, + ); + expect((await stage.run('u', true)).unchanged).toBe(501); + expect(findMany.mock.calls[1][0].where).toMatchObject({ + id: { gt: 'm0499' }, + searchable: true, + supersededById: null, + OR: [{ tier: null }, { tier: { not: 'ARCHIVED' } }], + }); + }); + + it('rechecks pin, age and access guards at archival update and reports actual writes', async () => { + const findMany = jest + .fn() + .mockResolvedValue([{ id: 'm', layer: 'SESSION', usedCount: 0 }]); + const updateMany = jest.fn().mockResolvedValue({ count: 0 }); // pinned after selection + const stage = new DreamCycleArchivalStage( + { memory: { findMany, updateMany } } as any, + config as any, + ); + expect((await stage.run('u', false)).archived).toBe(0); + const selection = findMany.mock.calls[0][0].where; + expect(selection).toMatchObject({ + userPinned: false, + createdAt: { lt: expect.any(Date) }, + AND: [ + { + OR: [ + { lastRetrievedAt: null }, + { lastRetrievedAt: { lt: expect.any(Date) } }, + ], + }, + { + OR: [{ lastUsedAt: null }, { lastUsedAt: { lt: expect.any(Date) } }], + }, + ], + }); + expect(updateMany.mock.calls[0][0].where).toEqual({ + ...selection, + id: { in: ['m'] }, + }); + }); + + it('rejects zero-row model writes and never uses a separate connection when given a transaction', async () => { + const external = { + $executeRawUnsafe: jest.fn(), + $queryRawUnsafe: jest.fn(), + }; + const tx = { + $queryRawUnsafe: jest.fn().mockResolvedValue([{ exists: 1 }]), + $executeRawUnsafe: jest.fn().mockResolvedValue(0), + }; + const writer = new EmbeddingWriteService(external as any); + await expect( + writer.writeMemoryEmbedding( + 'new', + 'openai-small', + Array(1536).fill(0.1), + false, + tx as any, + ), + ).rejects.toThrow('Expected one model row'); + expect(external.$executeRawUnsafe).not.toHaveBeenCalled(); + expect(external.$queryRawUnsafe).not.toHaveBeenCalled(); + }); + + it('keeps source retrieval and provenance intact after publishing a lossy summary', async () => { + const tx = { + memory: { + create: jest.fn().mockResolvedValue({ id: 'new' }), + update: jest.fn(), + updateMany: jest.fn(), + }, + }; + const writeMemoryEmbedding = jest.fn(); + const stage = new DreamCycleConsolidationStage( + { $transaction: (fn: any) => fn(tx) } as any, + config as any, + { embed: async () => [[0.1]] } as any, + { chat: async () => ({ content: 'Incomplete summary' }) } as any, + { + writeMemoryEmbedding, + getCurrentModelId: () => 'configured-model', + } as any, + ); + await (stage as any).consolidateCluster( + [{ id: 'source', content: 'Original precise date' }], + 'u', + ); + expect(writeMemoryEmbedding).toHaveBeenCalledWith( + 'new', + 'configured-model', + [0.1], + false, + tx, + ); + expect(tx.memory.updateMany.mock.calls[0][0].data).toEqual({ + consolidatedInto: 'new', + consolidated: true, + }); + }); + + it('does not retire sources on missing embedding', async () => { + const transaction = jest.fn(); + const stage = new DreamCycleConsolidationStage( + { $transaction: transaction } as any, + config as any, + { embed: async () => [] } as any, + { chat: async () => ({ content: 'summary' }) } as any, + {} as any, + ); + await expect( + (stage as any).consolidateCluster([{ id: 's', content: 'fact' }], 'u'), + ).rejects.toThrow('missing'); + expect(transaction).not.toHaveBeenCalled(); + }); +}); diff --git a/src/memory/memory-dedup.service.spec.ts b/src/memory/memory-dedup.service.spec.ts index 55697c47..a6827792 100644 --- a/src/memory/memory-dedup.service.spec.ts +++ b/src/memory/memory-dedup.service.spec.ts @@ -42,6 +42,8 @@ describe('MemoryDedupService', () => { describe('findDuplicateV2', () => { const mockMemory = { id: 'mem-1', + userId: 'user1', + raw: 'test text', deletedAt: null, confidence: 0.7, importanceScore: 0.5, @@ -217,7 +219,12 @@ describe('MemoryDedupService', () => { describe('findDuplicate (legacy)', () => { it('should return existing memory when duplicate found', async () => { - const mockMemory = { id: 'mem-1', deletedAt: null }; + const mockMemory = { + id: 'mem-1', + userId: 'user1', + raw: 'test', + deletedAt: null, + }; mockEmbedding.generate.mockResolvedValue([0.1]); mockEmbedding.search.mockResolvedValue([{ id: 'mem-1', score: 0.95 }]); mockPrisma.memory.findUnique.mockResolvedValue(mockMemory); @@ -239,6 +246,7 @@ describe('MemoryDedupService', () => { it('should boost confidence and update counters', async () => { mockPrisma.memory.findUnique.mockResolvedValue({ id: 'mem-1', + raw: 'new content', confidence: 0.7, }); @@ -254,11 +262,9 @@ describe('MemoryDedupService', () => { it('should not update when memory not found', async () => { mockPrisma.memory.findUnique.mockResolvedValue(null); - await service.autoMergeMemory( - 'nonexistent', - 'new content', - 'SYSTEM' as any, - ); + await expect( + service.autoMergeMemory('nonexistent', 'new content', 'SYSTEM' as any), + ).rejects.toThrow('non-equivalent'); expect(mockPrisma.$executeRaw).not.toHaveBeenCalled(); }); @@ -266,6 +272,7 @@ describe('MemoryDedupService', () => { it('should cap boosted confidence at 1.0', async () => { mockPrisma.memory.findUnique.mockResolvedValue({ id: 'mem-1', + raw: 'content', confidence: 0.98, }); diff --git a/src/memory/memory-dedup.service.ts b/src/memory/memory-dedup.service.ts index 4dd5069d..833618b1 100644 --- a/src/memory/memory-dedup.service.ts +++ b/src/memory/memory-dedup.service.ts @@ -53,10 +53,9 @@ export class MemoryDedupService { } /** - * Three-tier semantic deduplication (v2) - * - ≥0.93: auto-merge (combine content, boost confidence) - * - ≥0.85: reinforce (increment accessCount, update lastAccessedAt) - * - ≥0.78: flag for review (add to MergeCandidate table) + * Search for exact content duplicates; similarity alone only queues review. + * Identical content above the merge/reinforcement thresholds can boost + * confidence/counters. All non-identical observations remain independently stored. */ async findDuplicateV2( userId: string, @@ -107,7 +106,14 @@ export class MemoryDedupService { where: { id: match.id }, }); - if (!candidate || candidate.deletedAt) { + if ( + !candidate || + candidate.deletedAt || + candidate.userId !== userId || + candidate.searchable === false || + candidate.supersededById || + candidate.embeddingStatus === 'DUPLICATE' + ) { continue; } @@ -127,7 +133,10 @@ export class MemoryDedupService { let result: DedupResult; - if (bestMatch.score >= threshold) { + // Similarity is not equivalence: changed numbers, negation and added facts + // must remain retrievable. Until a verifier exists only identical text merges. + const equivalent = existingMemory.raw === text; + if (equivalent && bestMatch.score >= threshold) { this.logger.log( `[Dedup] Auto-merge: score=${bestMatch.score.toFixed(3)} memory=${bestMatch.id}`, ); @@ -136,7 +145,7 @@ export class MemoryDedupService { existingMemory, similarityScore: bestMatch.score, }; - } else if (bestMatch.score >= DEDUP_REINFORCE_THRESHOLD) { + } else if (equivalent && bestMatch.score >= DEDUP_REINFORCE_THRESHOLD) { this.logger.log( `[Dedup] Reinforce: score=${bestMatch.score.toFixed(3)} memory=${bestMatch.id}`, ); @@ -199,7 +208,7 @@ export class MemoryDedupService { } /** - * Auto-merge: combine content from new memory into existing, boost confidence + * Reinforce only identical content; semantic similarity cannot prove equivalence. */ async autoMergeMemory( existingId: string, @@ -209,7 +218,9 @@ export class MemoryDedupService { const existing = await this.prisma.memory.findUnique({ where: { id: existingId }, }); - if (!existing) return; + if (!existing || existing.raw !== newContent) { + throw new Error('Cannot reinforce non-equivalent memory content'); + } const newConfidence = SOURCE_CONFIDENCE[newSource] ?? 1.0; const boostedConfidence = Math.min( diff --git a/src/vector/embedding-write.service.spec.ts b/src/vector/embedding-write.service.spec.ts index b6cab7b7..156fc7ba 100644 --- a/src/vector/embedding-write.service.spec.ts +++ b/src/vector/embedding-write.service.spec.ts @@ -42,9 +42,15 @@ describe('EmbeddingWriteService', () => { }); it('passes for correct dimensions', () => { - expect(() => service.validateDimensions('bge-base', make(768))).not.toThrow(); - expect(() => service.validateDimensions('openai-small', make(1536))).not.toThrow(); - expect(() => service.validateDimensions('openai-large', make(3072))).not.toThrow(); + expect(() => + service.validateDimensions('bge-base', make(768)), + ).not.toThrow(); + expect(() => + service.validateDimensions('openai-small', make(1536)), + ).not.toThrow(); + expect(() => + service.validateDimensions('openai-large', make(3072)), + ).not.toThrow(); }); it('passes for unknown model (no guard applied)', () => { @@ -98,11 +104,13 @@ describe('EmbeddingWriteService', () => { expect(calls[0][4]).toBe(1536); }); - it('no-ops when memory does not exist (existence check returns empty)', async () => { + it('rejects when memory does not exist', async () => { (prisma.$queryRawUnsafe as jest.Mock).mockResolvedValueOnce([]); const vector = make(768); - await service.writeMemoryEmbedding('missing-mem', 'bge-base', vector); + await expect( + service.writeMemoryEmbedding('missing-mem', 'bge-base', vector), + ).rejects.toThrow('not found'); expect(prisma.$executeRawUnsafe).not.toHaveBeenCalled(); }); diff --git a/src/vector/embedding-write.service.ts b/src/vector/embedding-write.service.ts index fff9797f..62ea7c81 100644 --- a/src/vector/embedding-write.service.ts +++ b/src/vector/embedding-write.service.ts @@ -1,3 +1,4 @@ +import { Prisma } from '@prisma/client'; import { Injectable, Logger } from '@nestjs/common'; import { PrismaService } from '../prisma/prisma.service'; import { @@ -41,6 +42,10 @@ export class EmbeddingWriteService { modelId: string, vector: number[], skipLegacyInline = false, + client: Pick< + Prisma.TransactionClient, + '$executeRawUnsafe' | '$queryRawUnsafe' + > = this.prisma, ): Promise { this.validateDimensions(modelId, vector); @@ -50,38 +55,36 @@ export class EmbeddingWriteService { if (!skipLegacyInline) { if (vector.length === LEGACY_INLINE_DIMS) { // Write to legacy memories.embedding column — confirm memory exists - const rows = await this.prisma.$queryRawUnsafe>( + const rows = await client.$queryRawUnsafe>( `SELECT 1 AS exists FROM memories WHERE id = $1`, memoryId, ); if (rows.length === 0) { - this.logger.warn( - `[EmbeddingWrite] Memory ${memoryId} not found — skipping write`, - ); - return; + throw new Error(`[EmbeddingWrite] Memory ${memoryId} not found`); } - await this.prisma.$executeRawUnsafe( + const updated = await client.$executeRawUnsafe( `UPDATE memories SET embedding = $1::vector WHERE id = $2`, embeddingStr, memoryId, ); + if (updated !== 1) + throw new Error( + `[EmbeddingWrite] Expected one legacy row, got ${updated}`, + ); } else { // Non-768 dim: verify memory exists without touching the typed column - const rows = await this.prisma.$queryRawUnsafe>( + const rows = await client.$queryRawUnsafe>( `SELECT 1 AS exists FROM memories WHERE id = $1`, memoryId, ); if (rows.length === 0) { - this.logger.warn( - `[EmbeddingWrite] Memory ${memoryId} not found — skipping write`, - ); - return; + throw new Error(`[EmbeddingWrite] Memory ${memoryId} not found`); } } } // Always write to memory_embeddings (the authoritative multi-model store) - await this.prisma.$executeRawUnsafe( + const written = await client.$executeRawUnsafe( ` INSERT INTO memory_embeddings (id, memory_id, model_id, dimensions, embedding, created_at, updated_at) VALUES ( @@ -101,6 +104,10 @@ export class EmbeddingWriteService { modelId, vector.length, ); + if (written !== 1) + throw new Error( + `[EmbeddingWrite] Expected one model row, got ${written}`, + ); } /** @@ -124,11 +131,15 @@ export class EmbeddingWriteService { return; } const embeddingStr = this.serializeVector(vector, memoryId); - await this.prisma.$executeRawUnsafe( + const updated = await this.prisma.$executeRawUnsafe( `UPDATE memories SET embedding = $1::vector WHERE id = $2`, embeddingStr, memoryId, ); + if (updated !== 1) + throw new Error( + `[EmbeddingWrite] Expected one legacy row, got ${updated}`, + ); } /** @@ -149,7 +160,9 @@ export class EmbeddingWriteService { private serializeVector(vector: number[], context: string): string { if (!Array.isArray(vector) || vector.length === 0) { - throw new Error(`[EmbeddingWrite] Empty or invalid vector for ${context}`); + throw new Error( + `[EmbeddingWrite] Empty or invalid vector for ${context}`, + ); } if (Object.keys(vector).length !== vector.length) { throw new Error( @@ -173,7 +186,12 @@ export class EmbeddingWriteService { skipLegacyInline = false, ): Promise { const modelId = resolveEmbeddingModelId(); - return this.writeMemoryEmbedding(memoryId, modelId, vector, skipLegacyInline); + return this.writeMemoryEmbedding( + memoryId, + modelId, + vector, + skipLegacyInline, + ); } /** Expose current model/dimension config for callers that need it */ From 18d16899a5636ffae24aa6e50c9d364f6f4390fc Mon Sep 17 00:00:00 2001 From: heybeaux Date: Sun, 6 Sep 2026 18:19:53 -0700 Subject: [PATCH 3/6] fix(recall): preserve ranking and retry correctness with honest quality gates --- .github/workflows/ci-cloud.yml | 22 +- README.md | 45 +- src/memory/contextual-recall.service.spec.ts | 2 +- src/memory/contextual-recall.service.ts | 111 +++-- src/memory/memory-query-ranking.service.ts | 1 + src/memory/memory-query.service.ts | 40 +- src/memory/memory-recall-audit.spec.ts | 147 +++++++ test/benchmark/harness/autoresearch-sweep.ts | 115 ++--- test/benchmark/harness/simulate.ts | 35 +- test/benchmark/harness/sweep.ts | 92 +--- test/benchmark/history.ts | 22 +- .../recall-benchmark-dream.e2e-spec.ts | 28 +- test/benchmark/recall-benchmark.e2e-spec.ts | 64 +-- test/benchmark/scoring.spec.ts | 71 ++++ test/benchmark/scoring.ts | 392 ++++++------------ test/fixtures/types.ts | 2 + 16 files changed, 560 insertions(+), 629 deletions(-) create mode 100644 src/memory/memory-recall-audit.spec.ts create mode 100644 test/benchmark/scoring.spec.ts diff --git a/.github/workflows/ci-cloud.yml b/.github/workflows/ci-cloud.yml index 28fcbfd8..deaf7dc3 100644 --- a/.github/workflows/ci-cloud.yml +++ b/.github/workflows/ci-cloud.yml @@ -93,7 +93,7 @@ jobs: # ── 4a. Test shard — memory (large module, needs own slot) ─────────────── test-memory: - name: Test / memory + name: Test / shard 1 of 5 runs-on: ubuntu-latest timeout-minutes: 25 needs: install @@ -133,11 +133,13 @@ jobs: - name: Test — memory module run: | pnpm test -- --no-coverage --forceExit --ci \ - --testPathPatterns='src/memory' + --shard=1/5 + - name: Verify benchmark scoring definitions + run: pnpm exec jest --config test/jest-e2e.json --testRegex='benchmark/scoring.spec.ts$' --runInBand --ci # ── 4b. Test shard — core (identity, account, common) ──────────────────── test-core: - name: Test / core + name: Test / shard 2 of 5 runs-on: ubuntu-latest timeout-minutes: 25 needs: install @@ -177,11 +179,11 @@ jobs: - name: Test — core modules run: | pnpm test -- --no-coverage --forceExit --ci \ - --testPathPatterns='src/(identity|account|common|app\.controller)' + --shard=2/5 # ── 4b. Test shard — intelligence ──────────────────────────────────────── test-intelligence: - name: Test / intelligence + name: Test / shard 3 of 5 runs-on: ubuntu-latest timeout-minutes: 12 needs: install @@ -221,11 +223,11 @@ jobs: - name: Test — intelligence modules run: | pnpm test -- --no-coverage --forceExit --ci \ - --testPathPatterns='src/(deduplication|awareness|multi-query|consolidation|prefetch|retrieval-signals|graph|clustering|hierarchy|ensemble)' + --shard=3/5 # ── 4c. Test shard — cloud-edition features ─────────────────────────────── test-cloud: - name: Test / cloud features + name: Test / shard 4 of 5 features runs-on: ubuntu-latest timeout-minutes: 12 needs: install @@ -265,11 +267,11 @@ jobs: - name: Test — cloud-edition modules run: | pnpm test -- --no-coverage --forceExit --ci \ - --testPathPatterns='src/(cloud-link|cloud-sync|instance|webhooks|analytics|monitoring|stripe|reembedding)' + --shard=4/5 # ── 4d. Test shard — remaining features ────────────────────────────────── test-features: - name: Test / features + name: Test / shard 5 of 5 runs-on: ubuntu-latest timeout-minutes: 12 needs: install @@ -309,7 +311,7 @@ jobs: - name: Test — remaining feature modules run: | pnpm test -- --no-coverage --forceExit --ci \ - --testPathPatterns='src/(llm|import|import-v2|inbound-email|agent|agent-recall|agent-session|anticipatory|auto|correction|entity-profile|events|feedback|fog-index|health|memory-access-log|memory-pool|queue|scoped-context|scripts|session-indexing|storage|summarization|teams|utils|eval|dashboard)' + --shard=5/5 # ── 5. Gate ─────────────────────────────────────────────────────────────── tests-passed: diff --git a/README.md b/README.md index dcc21394..a0111282 100644 --- a/README.md +++ b/README.md @@ -227,46 +227,11 @@ pnpm longmemeval --subset full --resume results/full-.jsonl # Resume a ## Recall Benchmark -Engram includes a comprehensive recall benchmark suite that tests semantic retrieval quality across 81 queries in 7 categories. Every PR runs the benchmark in CI with real embeddings (bge-base-en-v1.5) and ensemble reranking. - -### Latest Results (March 2026) - -| Metric | Pre-Dream Cycle | Post-Dream Cycle | -|--------|:-:|:-:| -| **Precision@5** | 95.1% ✅ | 96.9% ✅ | -| **Recall@20** | 95.7% | 96.9% | -| **MRR** | 0.836 | 0.874 | -| **Isolation** | 100% ✅ | 100% ✅ | -| **Queries Passed** | 78/81 | 79/81 | - -
-Category Breakdown (Post-Dream Cycle) - -| Category | Queries | Passed | P@5 | R@20 | MRR | Isolation | -|----------|:-------:|:------:|:---:|:----:|:---:|:---------:| -| Adversarial | 10 | 10 | 100% | 100% | 1.00 | 100% | -| Cross-feature | 10 | 10 | 100% | 100% | 0.88 | 100% | -| Edge case | 16 | 16 | 100% | 100% | 0.95 | 100% | -| Emotional | 10 | 9 | 85% | 95% | 0.73 | 100% | -| RLS isolation | 10 | 10 | 100% | 100% | 0.88 | 100% | -| Semantic | 14 | 13 | 93% | 100% | 0.79 | 100% | -| Temporal | 11 | 11 | 100% | 82% | 0.87 | 100% | - -
- -### Benchmark Progression - -The recall pipeline evolved through several iterations: - -| Phase | P@5 | Notes | -|-------|:---:|-------| -| Baseline (cosine only) | ~45% | Single-model vector search, no reranking | -| + BM25 hybrid | ~65% | Full-text search fusion for keyword recall | -| + Sentiment polarity | ~75% | Penalizes opposite-emotion memories | -| + Ensemble reranking | ~85% | 2 cross-encoder models via RRF | -| + Fixture enrichment | ~90% | Realistic gold memories, noise calibration | -| + Temporal patterns | 95.1% | Month/year/week recognition, expanded BM25 pool | -| + Dream Cycle | **96.9%** | Post-consolidation: cleaner corpus, higher signal | +The recall suite exercises 81 fixture queries in seven categories. Metric version 2 reports **required-item coverage@5** on the 57 queries with top-five labels, **expected-item coverage@20** on the ten queries with top-20 labels, and reciprocal rank of the first required hit. Missing labels are excluded, not scored as perfect. An empty expected list is not an explicit no-answer judgment. + +The previous published 95.1%/96.9% “precision” and associated recall/MRR tables used a different, inflated scoring definition. Those historical numbers are withdrawn as precision claims and must not be compared directly with version 2. Partial required-item labels do not establish conventional precision or exhaustive recall; fully judged held-out queries are still needed. + +The real-embedding suite always enforces the unchanged 95% required-coverage gate, forbidden-ID isolation, and any explicitly labelled no-answer cases. Reports include label denominators and a metric version. A passing forbidden-ID fixture does not prove production database-role isolation. ### Running the Benchmark diff --git a/src/memory/contextual-recall.service.spec.ts b/src/memory/contextual-recall.service.spec.ts index 312fee53..4c8388b6 100644 --- a/src/memory/contextual-recall.service.spec.ts +++ b/src/memory/contextual-recall.service.spec.ts @@ -279,7 +279,7 @@ describe('ContextualRecallService', () => { sessionKey: 'sess-clear', } as any); - service.clearSession('sess-clear'); + service.clearSession(userId, 'sess-clear'); // Next recall should treat as first message (topic shift) embedding.search.mockResolvedValue([]); diff --git a/src/memory/contextual-recall.service.ts b/src/memory/contextual-recall.service.ts index ef8b8012..8f027b48 100644 --- a/src/memory/contextual-recall.service.ts +++ b/src/memory/contextual-recall.service.ts @@ -47,7 +47,9 @@ export class ContextualRecallService { dto: ContextualRecallDto, ): Promise { const startTime = Date.now(); - const session = this.getOrCreateSession(dto.sessionKey); + const session = this.getOrCreateSession( + this.sessionScope(userId, dto.sessionKey), + ); // 1. Generate embedding for the incoming text const queryEmbedding = await this.embedding.generate(dto.text); @@ -55,12 +57,6 @@ export class ContextualRecallService { // 2. Detect topic shift const topicShift = this.detectTopicShift(queryEmbedding, session); - // Store this embedding for future comparisons - session.recentEmbeddings.push(queryEmbedding); - if (session.recentEmbeddings.length > this.MAX_RECENT_EMBEDDINGS) { - session.recentEmbeddings.shift(); - } - // 3. If no topic shift, return empty (no recall needed) if (!topicShift) { return { @@ -118,14 +114,47 @@ export class ContextualRecallService { ? await this.embedding.generate(semanticText) : queryEmbedding; - const vectorResults = await this.embedding.search( - userId ?? 'default', - searchEmbedding, - limit + excludeSet.size, // over-fetch to account for filtering - undefined, - undefined, - poolIds, - ); + // Providers do not all support date predicates. Refill the candidate window + // until enough eligible records survive or the provider is exhausted. Dates + // here deliberately mean creation time, matching the contextual API contract. + let window = Math.max(1, limit + excludeSet.size); + let vectorResults: Awaited>; + let memories: Awaited>; + let previousIds = ''; + for (;;) { + const found = await this.embedding.search( + userId ?? 'default', + searchEmbedding, + window, + undefined, + undefined, + poolIds, + ); + const ids = found + .filter((r) => r.score >= minScore && !excludeSet.has(r.id)) + .map((r) => r.id); + memories = await this.prisma.memory.findMany({ + where: { + id: { in: ids }, + deletedAt: null, + supersededById: null, + searchable: { not: false }, + ...(temporalCreatedAt ? { createdAt: temporalCreatedAt } : {}), + }, + include: { extraction: true }, + }); + const eligible = new Set(memories.map((m) => m.id)); + vectorResults = found.filter((r) => eligible.has(r.id)); + const signature = JSON.stringify(found.map((r) => r.id)); + if ( + vectorResults.length >= limit || + found.length < window || + signature === previousIds + ) + break; + previousIds = signature; + window *= 2; + } // HEY-189: If delegation context provided, fetch delegator's memories and merge let delegatorMemoryIds: Set | undefined; @@ -167,6 +196,7 @@ export class ContextualRecallService { .map((r) => ({ id: r.id, score: r.score })); if (filteredIds.length === 0) { + this.recordSuccessfulTopic(session, queryEmbedding); return { memories: [], topicShift: true, @@ -177,18 +207,7 @@ export class ContextualRecallService { // 6. Fetch full memory records (apply temporal filter if present) const scoreMap = new Map(filteredIds.map((r) => [r.id, r.score])); - const memories = await this.prisma.memory.findMany({ - where: { - id: { in: filteredIds.map((r) => r.id) }, - deletedAt: null, - supersededById: null, - searchable: { not: false }, - ...(temporalCreatedAt ? { createdAt: temporalCreatedAt } : {}), - }, - include: { - extraction: true, - }, - }); + memories = memories.filter((memory) => scoreMap.has(memory.id)); // 7. Build response, respecting token budget const maxTokens = dto.maxTokens ?? 500; @@ -210,16 +229,13 @@ export class ContextualRecallService { raw: rawText, layer: memory.layer, score: scoreMap.get(memory.id) ?? 0, - topics: memory.extraction?.topics ?? [], + topics: + (memory as typeof memory & { extraction?: { topics: string[] } }) + .extraction?.topics ?? [], }); tokenCount += approxTokens; - session.recalledIds.add(memory.id); } - // Update session state - session.lastRecallAt = Date.now(); - session.recallCount++; - // Update retrieval counts const resultIds = result.map((m) => m.id); if (resultIds.length > 0) { @@ -239,6 +255,10 @@ export class ContextualRecallService { } } + // Commit suppression state only after the entire retrieval succeeds. + this.recordSuccessfulTopic(session, queryEmbedding); + for (const id of resultIds) session.recalledIds.add(id); + return { memories: result, topicShift: true, @@ -282,6 +302,27 @@ export class ContextualRecallService { return magnitude === 0 ? 0 : dotProduct / magnitude; } + private sessionScope( + userId: string | string[] | null, + sessionKey: string, + ): string { + const users = [ + ...new Set(Array.isArray(userId) ? userId : [userId ?? 'default']), + ].sort(); + return JSON.stringify([users, sessionKey]); + } + + private recordSuccessfulTopic( + session: SessionState, + embedding: number[], + ): void { + session.recentEmbeddings.push(embedding); + if (session.recentEmbeddings.length > this.MAX_RECENT_EMBEDDINGS) + session.recentEmbeddings.shift(); + session.lastRecallAt = Date.now(); + session.recallCount++; + } + private getOrCreateSession(sessionKey: string): SessionState { // HEY-361: Periodic sweep of stale sessions this.sweepStaleSessions(); @@ -327,8 +368,8 @@ export class ContextualRecallService { /** * Clean up session state (call on session end) */ - clearSession(sessionKey: string): void { - this.sessions.delete(sessionKey); + clearSession(userId: string | string[] | null, sessionKey: string): void { + this.sessions.delete(this.sessionScope(userId, sessionKey)); } /** diff --git a/src/memory/memory-query-ranking.service.ts b/src/memory/memory-query-ranking.service.ts index a13b00b9..55c0820d 100644 --- a/src/memory/memory-query-ranking.service.ts +++ b/src/memory/memory-query-ranking.service.ts @@ -288,6 +288,7 @@ export class MemoryQueryRankingService { const finalScore = (r.score * 0.85 + importanceScore * 0.15) * sp; return { ...mem, score: finalScore }; }) + .sort((a, b) => (b.score ?? 0) - (a.score ?? 0)) .slice(0, limit); this.logger.debug( diff --git a/src/memory/memory-query.service.ts b/src/memory/memory-query.service.ts index ea5952c4..cb2150a9 100644 --- a/src/memory/memory-query.service.ts +++ b/src/memory/memory-query.service.ts @@ -150,7 +150,6 @@ export class MemoryQueryService { const temporalRangeFilter = this.buildTemporalRangeFilter(dto); let scoredMemories: MemoryWithScore[]; - const keywordRescueMap = new Map(); if (hasTemporalIntent) { // TEMPORAL PATH — HEY-575: Adaptive window expansion @@ -431,10 +430,10 @@ export class MemoryQueryService { LIMIT 100`, searchQuery, ); - // RRF fusion (k=60): BM25 rank contributes 1/(k+rank) so rank-1 - // BM25 hit scores ≈0.016, rank-100 ≈0.006. This prevents a flat - // 0.75 override from promoting low-quality exact-keyword matches - // above high-quality semantic matches. + // Lexical rank admits new candidates with a conservative prior. Keep + // measured semantic scores for existing candidates; never overwrite them + // with an out-of-range lexical boost. The final reranker/blend evaluates + // all admitted candidates, including lexical-only matches. const RRF_K = 60; let ftsAdded = 0; for (let ftsRank = 0; ftsRank < ftsResults.length; ftsRank++) { @@ -443,11 +442,9 @@ export class MemoryQueryService { ftsResultIds.add(row.id); keywordRescueIds.add(row.id); if (!scoreMap.has(row.id)) { - scoreMap.set(row.id, 1.25); + scoreMap.set(row.id, bm25Score); memoryIds.push(row.id); ftsAdded++; - } else { - scoreMap.set(row.id, Math.max(scoreMap.get(row.id)!, 1.25)); } } if (ftsAdded > 0) { @@ -521,11 +518,9 @@ export class MemoryQueryService { ftsResultIds.add(row.id); keywordRescueIds.add(row.id); if (!scoreMap.has(row.id)) { - scoreMap.set(row.id, 1.1); + scoreMap.set(row.id, 1 / 61); memoryIds.push(row.id); ilikeAdded++; - } else { - scoreMap.set(row.id, Math.max(scoreMap.get(row.id)!, 1.1)); } } if (ilikeAdded > 0) { @@ -572,11 +567,9 @@ export class MemoryQueryService { for (const row of identityResults) { keywordRescueIds.add(row.id); if (!scoreMap.has(row.id)) { - scoreMap.set(row.id, 1.15); + scoreMap.set(row.id, 1 / 61); memoryIds.push(row.id); identityAdded++; - } else { - scoreMap.set(row.id, Math.max(scoreMap.get(row.id)!, 1.15)); } } if (identityAdded > 0) { @@ -629,10 +622,6 @@ export class MemoryQueryService { const topIds = new Set(sorted.map((m) => m.id)); const memoryMap = new Map(sorted.map((m) => [m.id, m])); - for (const id of keywordRescueIds) { - const mem = memoryMap.get(id); - if (mem) keywordRescueMap.set(id, mem); - } const forcedFts: MemoryWithScore[] = []; for (const id of ftsResultIds) { if (!topIds.has(id)) { @@ -717,21 +706,6 @@ export class MemoryQueryService { limit, ); - // Exact keyword/ILIKE rescued memories are deterministic high-signal hits. - // Keep them sticky after reranking so the cross-encoder cannot drop fresh - // exact-match writes from the final top-N. - const missingKeywordHits = [...keywordRescueMap.entries()] - .filter(([id]) => !scoredMemories.some((m) => m.id === id)) - .map( - ([, mem]) => - ({ ...mem, score: Math.max(mem.score ?? 0, 1.1) }) as MemoryWithScore, - ); - if (missingKeywordHits.length > 0) { - scoredMemories = [...missingKeywordHits, ...scoredMemories] - .sort((a, b) => (b.score ?? 0) - (a.score ?? 0)) - .slice(0, limit); - } - // v1.7: Agent-scoped filter if (dto.filterAgentId) { scoredMemories = scoredMemories.filter( diff --git a/src/memory/memory-recall-audit.spec.ts b/src/memory/memory-recall-audit.spec.ts new file mode 100644 index 00000000..5aada2ad --- /dev/null +++ b/src/memory/memory-recall-audit.spec.ts @@ -0,0 +1,147 @@ +import { ContextualRecallService } from './contextual-recall.service'; +import { MemoryQueryService } from './memory-query.service'; +import { MemoryQueryRankingService } from './memory-query-ranking.service'; + +describe('Recall audit regressions', () => { + const dto = { text: 'fact', sessionKey: 'shared', maxResults: 1 }; + function context() { + const prisma = { + memory: { + findMany: jest.fn().mockResolvedValue([]), + updateMany: jest.fn().mockResolvedValue({ count: 0 }), + }, + }; + const embedding = { + generate: jest.fn().mockResolvedValue([1, 0]), + search: jest.fn().mockResolvedValue([]), + }; + return { + prisma, + embedding, + service: new ContextualRecallService(prisma as any, embedding as any), + }; + } + it('retries after failed vector retrieval instead of suppressing the topic', async () => { + const { embedding, service } = context(); + embedding.search.mockRejectedValueOnce(new Error('unavailable')); + await expect(service.recall('a', dto)).rejects.toThrow('unavailable'); + expect((await service.recall('a', dto)).topicShift).toBe(true); + expect(embedding.search).toHaveBeenCalledTimes(2); + expect((await service.recall('a', dto)).topicShift).toBe(false); + }); + it('isolates identical session keys and clears only the selected user set', async () => { + const { embedding, service } = context(); + await service.recall(['b', 'a'], dto); + expect((await service.recall(['a', 'b', 'a'], dto)).topicShift).toBe(false); + expect((await service.recall('c', dto)).topicShift).toBe(true); + service.clearSession(['a', 'b'], dto.sessionKey); + expect((await service.recall('c', dto)).topicShift).toBe(false); + expect((await service.recall(['b', 'a'], dto)).topicShift).toBe(true); + expect(embedding.search).toHaveBeenCalledTimes(3); + }); + it('refills dated search before applying the score gap and result limit', async () => { + const { prisma, embedding } = context(); + const records = [ + { id: 'old', score: 1, raw: 'old', createdAt: new Date('2020-01-01') }, + { + id: 'eligible', + score: 0.6, + raw: 'fact', + createdAt: new Date('2026-09-01'), + }, + ]; + embedding.search.mockImplementation(async (_u, _q, limit) => + records.slice(0, limit), + ); + prisma.memory.findMany.mockImplementation(async ({ where }) => + records.filter( + (m) => where.id.in.includes(m.id) && m.createdAt >= where.createdAt.gte, + ), + ); + const parser = { + parse: () => ({ + semanticQuery: 'fact', + temporalFilter: { + start: new Date('2026-09-01'), + end: new Date('2026-09-07'), + expression: 'this week', + }, + }), + }; + const service = new ContextualRecallService( + prisma as any, + embedding as any, + undefined, + undefined, + parser as any, + ); + expect((await service.recall('a', dto)).memories.map((m) => m.id)).toEqual([ + 'eligible', + ]); + expect(embedding.search).toHaveBeenCalledTimes(2); + }); + it('does not commit recalled IDs when bookkeeping fails', async () => { + const { prisma, embedding, service } = context(); + embedding.search.mockResolvedValue([{ id: 'm', score: 0.9 }]); + prisma.memory.findMany.mockResolvedValue([{ id: 'm', raw: 'fact' }]); + prisma.memory.updateMany.mockRejectedValueOnce(new Error('db unavailable')); + await expect(service.recall('a', dto)).rejects.toThrow('db unavailable'); + expect((await service.recall('a', dto)).memories.map((m) => m.id)).toEqual([ + 'm', + ]); + }); + it('does not reinsert lexical noise above a stronger reranked match at limit 1', async () => { + const good = { + id: 'semantic', + raw: 'Relevant fact', + importanceScore: 0.5, + extraction: {}, + layer: 'SESSION', + }; + const noise = { ...good, id: 'noise', raw: 'Unrelated keyword' }; + const prisma = { + memory: { + findMany: jest.fn().mockResolvedValue([good, noise]), + updateMany: jest.fn(), + }, + $queryRawUnsafe: jest.fn().mockResolvedValue([{ id: noise.id }]), + }; + const embedding = { + generateForRecall: jest.fn().mockResolvedValue([1, 0]), + search: jest.fn().mockResolvedValue([ + { id: good.id, score: 0.99 }, + { id: noise.id, score: 0.1 }, + ]), + }; + const weights = { applyUsageWeighting: async (m: any) => m }; + const reranker = { + rerank: async (_q: any, texts: string[]) => + texts + .map((t, index) => ({ index, score: t === good.raw ? 0.9 : 0.01 })) + .sort((a, b) => b.score - a.score), + }; + const ranking = new MemoryQueryRankingService( + prisma as any, + embedding as any, + weights as any, + reranker as any, + ); + ranking.surfaceInsights = async (m) => m; + const service = new MemoryQueryService( + prisma as any, + embedding as any, + { + parse: (q: string) => ({ semanticQuery: q, temporalFilter: null }), + } as any, + weights as any, + ranking, + {} as any, + ); + const result = await service.recall('a', { + query: 'relevant fact', + limit: 1, + layers: ['SESSION'], + } as any); + expect(result.memories.map((m) => m.id)).toEqual(['semantic']); + }); +}); diff --git a/test/benchmark/harness/autoresearch-sweep.ts b/test/benchmark/harness/autoresearch-sweep.ts index 00d58c03..8729806a 100644 --- a/test/benchmark/harness/autoresearch-sweep.ts +++ b/test/benchmark/harness/autoresearch-sweep.ts @@ -3,7 +3,7 @@ * * Extends the standard benchmark sweep with durability multipliers to find * optimal scoring parameters that fix the 3 known failing queries (daily_gen - * noise memories beating durable memories) without regressing overall P@5. + * noise memories beating durable memories) without regressing overall Coverage@5. * * The key problem: alice_daily_gen_* noise memories (importanceScore 0.3–0.5) * appear in top 5 for queries where durable memories (health, coffee, identity) @@ -15,7 +15,7 @@ import * as fs from 'fs'; import * as path from 'path'; import { GOLD_QUERIES } from '../../fixtures/queries/gold-queries'; -import { scoreQuery } from '../scoring'; +import { scoreQuery, checkThresholds } from '../scoring'; import type { QueryScore } from '../scoring'; import type { ScoringConfig } from './simulate'; @@ -55,12 +55,12 @@ export interface DurabilityAwareScoringConfig extends ScoringConfig { /** Result for a single swept configuration. */ export interface AutoresearchResult { config: DurabilityAwareScoringConfig; - overallPrecisionAt5: number; + overallRequiredCoverageAt5: number; zeroHits: number; isolationScore: number; passed: boolean; - /** P@5 specifically on the 3 known failing queries */ - focusPrecisionAt5: number; + /** Coverage@5 specifically on the 3 known failing queries */ + focusRequiredCoverageAt5: number; /** How many of the 3 focus queries have their must_top5 in the actual top 5 */ focusHits: number; /** Per-query detail for the focus queries */ @@ -315,68 +315,24 @@ function evaluateConfig( .slice(0, 20) .map((r) => r.id); - const top5Hits = goldQuery.must_top5.filter((id) => topIds.includes(id)); - const precisionAt5 = - goldQuery.must_top5.length > 0 - ? top5Hits.length / goldQuery.must_top5.length - : 1.0; - - const mustAbsentViolations = goldQuery.must_absent.filter((id) => - [...topIds, ...top20].includes(id), + allScores.push( + scoreQuery( + goldQuery, + [...topIds, ...top20.filter((id) => !topIds.includes(id))].slice(0, 20), + ), ); - const isolationPassed = mustAbsentViolations.length === 0; - - const shouldTop20 = goldQuery.should_top20 ?? []; - const top20Hits = shouldTop20.filter((id) => top20.includes(id)); - const recallAt20 = - shouldTop20.length > 0 ? top20Hits.length / shouldTop20.length : 1.0; - - let mrr: number; - if (goldQuery.must_top5.length > 0) { - const allIds = [...new Set([...topIds, ...top20])]; - const reciprocalRanks = goldQuery.must_top5.map((id) => { - const rank = allIds.indexOf(id); - return rank >= 0 ? 1 / (rank + 1) : 0; - }); - mrr = - reciprocalRanks.reduce((sum, rr) => sum + rr, 0) / - goldQuery.must_top5.length; - } else { - mrr = 1.0; - } - - const passed = - isolationPassed && - (goldQuery.must_top5.length === 0 || top5Hits.length > 0); - - allScores.push({ - queryId: goldQuery.id, - category: goldQuery.category, - passed, - precisionAt5, - recallAt20, - mrr, - isolationPassed, - details: { - query: goldQuery.query, - user: goldQuery.user, - expectedTop5: goldQuery.must_top5, - expectedTop20: shouldTop20, - actualIds: [ - ...topIds, - ...top20.filter((id) => !topIds.includes(id)), - ].slice(0, 20), - mustAbsentViolations, - top5Hits, - top20Hits, - }, - }); } - const avg = (vals: number[]) => - vals.length === 0 ? 0 : vals.reduce((s, v) => s + v, 0) / vals.length; + const avg = (vals: (number | null)[]) => { + const labelled = vals.filter((v): v is number => v !== null); + return labelled.length + ? labelled.reduce((s, v) => s + v, 0) / labelled.length + : 0; + }; - const overallPrecisionAt5 = avg(allScores.map((s) => s.precisionAt5)); + const overallRequiredCoverageAt5 = avg( + allScores.map((s) => s.requiredCoverageAt5), + ); const zeroHits = allScores.filter( (s) => s.details.expectedTop5.length > 0 && s.details.top5Hits.length === 0, ).length; @@ -387,7 +343,9 @@ function evaluateConfig( const focusScores = allScores.filter((s) => FOCUS_QUERY_IDS.includes(s.queryId), ); - const focusPrecisionAt5 = avg(focusScores.map((s) => s.precisionAt5)); + const focusRequiredCoverageAt5 = avg( + focusScores.map((s) => s.requiredCoverageAt5), + ); const focusHits = focusScores.filter( (s) => s.details.expectedTop5.length > 0 && s.details.top5Hits.length > 0, ).length; @@ -399,16 +357,15 @@ function evaluateConfig( expected: s.details.expectedTop5, })); - const passed = - overallPrecisionAt5 >= 0.7 && zeroHits === 0 && isolationScore >= 1.0; + const passed = checkThresholds(allScores); return { config, - overallPrecisionAt5, + overallRequiredCoverageAt5, zeroHits, isolationScore, passed, - focusPrecisionAt5, + focusRequiredCoverageAt5, focusHits, focusDetails, }; @@ -530,10 +487,12 @@ function main() { // ── Results analysis ──────────────────────────────────────── - // Primary sort: fixes all 3 focus queries, then by overall P@5 + // Primary sort: fixes all 3 focus queries, then by overall Coverage@5 const fixesAll = allResults .filter((r) => r.focusHits === FOCUS_QUERY_IDS.length && r.passed) - .sort((a, b) => b.overallPrecisionAt5 - a.overallPrecisionAt5); + .sort( + (a, b) => b.overallRequiredCoverageAt5 - a.overallRequiredCoverageAt5, + ); // Secondary: fixes at least some focus queries while passing overall const fixesSome = allResults @@ -544,12 +503,12 @@ function main() { .sort( (a, b) => b.focusHits - a.focusHits || - b.overallPrecisionAt5 - a.overallPrecisionAt5, + b.overallRequiredCoverageAt5 - a.overallRequiredCoverageAt5, ); - // Fallback: best overall P@5 regardless + // Fallback: best overall Coverage@5 regardless const bestOverall = [...allResults].sort( - (a, b) => b.overallPrecisionAt5 - a.overallPrecisionAt5, + (a, b) => b.overallRequiredCoverageAt5 - a.overallRequiredCoverageAt5, ); // ── Print results ───────────────────────────────────────── @@ -577,7 +536,7 @@ function main() { console.log( `\n❌ No config fixes any focus query while passing overall thresholds.`, ); - console.log('\nTop 10 by overall P@5:\n'); + console.log('\nTop 10 by overall Coverage@5:\n'); printResultTable(bestOverall.slice(0, 10)); if (bestOverall.length > 0) printBestConfig(bestOverall[0]); } @@ -607,7 +566,7 @@ function main() { } function printResultTable(results: AutoresearchResult[]) { - const header = `${'Rank'.padEnd(5)} ${'P@5'.padEnd(7)} ${'Focus'.padEnd(7)} ${'ZH'.padEnd(4)} ${'Iso'.padEnd(5)} ${'dB'.padEnd(5)} ${'eP'.padEnd(6)} ${'cW'.padEnd(5)} ${'iW'.padEnd(5)} Focus Detail`; + const header = `${'Rank'.padEnd(5)} ${'Coverage@5'.padEnd(7)} ${'Focus'.padEnd(7)} ${'ZH'.padEnd(4)} ${'Iso'.padEnd(5)} ${'dB'.padEnd(5)} ${'eP'.padEnd(6)} ${'cW'.padEnd(5)} ${'iW'.padEnd(5)} Focus Detail`; console.log(header); console.log('─'.repeat(100)); @@ -621,7 +580,7 @@ function printResultTable(results: AutoresearchResult[]) { console.log( `${String(i + 1).padEnd(5)} ` + - `${(r.overallPrecisionAt5 * 100).toFixed(1).padEnd(6)}% ` + + `${(r.overallRequiredCoverageAt5 * 100).toFixed(1).padEnd(6)}% ` + `${r.focusHits}/${FOCUS_QUERY_IDS.length}`.padEnd(7) + ` ${String(r.zeroHits).padEnd(4)}` + `${(r.isolationScore * 100).toFixed(0).padEnd(5)}% ` + @@ -642,10 +601,10 @@ function printBestConfig(best: AutoresearchResult) { console.log(` importanceFinalWeight: ${best.config.importanceFinalWeight}`); console.log(` preRerankK: ${best.config.preRerankK}`); console.log( - ` Overall P@5: ${(best.overallPrecisionAt5 * 100).toFixed(1)}%`, + ` Overall Coverage@5: ${(best.overallRequiredCoverageAt5 * 100).toFixed(1)}%`, ); console.log( - ` Focus P@5: ${(best.focusPrecisionAt5 * 100).toFixed(1)}%`, + ` Focus Coverage@5: ${(best.focusRequiredCoverageAt5 * 100).toFixed(1)}%`, ); console.log( ` Focus hits: ${best.focusHits}/${FOCUS_QUERY_IDS.length}`, diff --git a/test/benchmark/harness/simulate.ts b/test/benchmark/harness/simulate.ts index 49c2ad5b..db925467 100644 --- a/test/benchmark/harness/simulate.ts +++ b/test/benchmark/harness/simulate.ts @@ -218,35 +218,12 @@ function main() { .slice(0, 20) .map((r) => r.id); - const score = scoreQuery(goldQuery, top20); - // Override actualIds to use top5 from resultMap for precision@5 - const adjustedScore = { - ...score, - precisionAt5: (() => { - const top5Hits = goldQuery.must_top5.filter((id) => - topIds.includes(id), - ); - return goldQuery.must_top5.length > 0 - ? top5Hits.length / goldQuery.must_top5.length - : 1.0; - })(), - details: { - ...score.details, - top5Hits: goldQuery.must_top5.filter((id) => topIds.includes(id)), - actualIds: [ - ...topIds, - ...top20.filter((id) => !topIds.includes(id)), - ].slice(0, 20), - }, - }; - - // Recalculate passed with corrected precisionAt5 - const passed = - adjustedScore.isolationPassed && - (goldQuery.must_top5.length === 0 || - adjustedScore.details.top5Hits.length > 0); - - allScores.push({ ...adjustedScore, passed }); + allScores.push( + scoreQuery( + goldQuery, + [...topIds, ...top20.filter((id) => !topIds.includes(id))].slice(0, 20), + ), + ); } const report = buildReport(allScores, 'offline', 'harness'); diff --git a/test/benchmark/harness/sweep.ts b/test/benchmark/harness/sweep.ts index 6cd97464..9b90e39f 100644 --- a/test/benchmark/harness/sweep.ts +++ b/test/benchmark/harness/sweep.ts @@ -2,7 +2,7 @@ * Benchmark Harness — Parameter Sweep * * Grid searches over scoring config parameters and prints all passing configs. - * Thresholds: P@5 >= 0.70, zero-hits = 0, isolation = 1.0 + * Thresholds: Required coverage@5 >= 0.95; shared benchmark gates * * Run: pnpm benchmark:sweep */ @@ -52,7 +52,7 @@ function loadJson(filename: string): T { interface SweepResult { config: ScoringConfig; - precisionAt5: number; + requiredCoverageAt5: number; zeroHits: number; isolationScore: number; passed: boolean; @@ -89,77 +89,31 @@ function evaluateConfig( .slice(0, 20) .map((r) => r.id); - const top5Hits = goldQuery.must_top5.filter((id) => topIds.includes(id)); - const precisionAt5 = - goldQuery.must_top5.length > 0 - ? top5Hits.length / goldQuery.must_top5.length - : 1.0; - - const mustAbsentViolations = goldQuery.must_absent.filter((id) => - [...topIds, ...top20].includes(id), + allScores.push( + scoreQuery( + goldQuery, + [...topIds, ...top20.filter((id) => !topIds.includes(id))].slice(0, 20), + ), ); - const isolationPassed = mustAbsentViolations.length === 0; - - const shouldTop20 = goldQuery.should_top20 ?? []; - const top20Hits = shouldTop20.filter((id) => top20.includes(id)); - const recallAt20 = - shouldTop20.length > 0 ? top20Hits.length / shouldTop20.length : 1.0; - - let mrr: number; - if (goldQuery.must_top5.length > 0) { - const allIds = [...new Set([...topIds, ...top20])]; - const reciprocalRanks = goldQuery.must_top5.map((id) => { - const rank = allIds.indexOf(id); - return rank >= 0 ? 1 / (rank + 1) : 0; - }); - mrr = - reciprocalRanks.reduce((sum, rr) => sum + rr, 0) / - goldQuery.must_top5.length; - } else { - mrr = 1.0; - } - - const passed = - isolationPassed && - (goldQuery.must_top5.length === 0 || top5Hits.length > 0); - - allScores.push({ - queryId: goldQuery.id, - category: goldQuery.category, - passed, - precisionAt5, - recallAt20, - mrr, - isolationPassed, - details: { - query: goldQuery.query, - user: goldQuery.user, - expectedTop5: goldQuery.must_top5, - expectedTop20: shouldTop20, - actualIds: [ - ...topIds, - ...top20.filter((id) => !topIds.includes(id)), - ].slice(0, 20), - mustAbsentViolations, - top5Hits, - top20Hits, - }, - }); } - const avg = (vals: number[]) => - vals.length === 0 ? 0 : vals.reduce((s, v) => s + v, 0) / vals.length; + const avg = (vals: (number | null)[]) => { + const labelled = vals.filter((v): v is number => v !== null); + return labelled.length + ? labelled.reduce((s, v) => s + v, 0) / labelled.length + : 0; + }; - const precisionAt5 = avg(allScores.map((s) => s.precisionAt5)); + const requiredCoverageAt5 = avg(allScores.map((s) => s.requiredCoverageAt5)); const zeroHits = allScores.filter( (s) => s.details.expectedTop5.length > 0 && s.details.top5Hits.length === 0, ).length; const isolationScore = allScores.filter((s) => s.isolationPassed).length / allScores.length; - const passed = precisionAt5 >= 0.7 && zeroHits === 0 && isolationScore >= 1.0; + const passed = checkThresholds(allScores); - return { config, precisionAt5, zeroHits, isolationScore, passed }; + return { config, requiredCoverageAt5, zeroHits, isolationScore, passed }; } function main() { @@ -211,7 +165,7 @@ function main() { // Filter passing configs const passing = allResults .filter((r) => r.passed) - .sort((a, b) => b.precisionAt5 - a.precisionAt5); + .sort((a, b) => b.requiredCoverageAt5 - a.requiredCoverageAt5); console.log( `\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━`, @@ -223,23 +177,23 @@ function main() { if (passing.length === 0) { console.log('\nNo configs passed all thresholds.'); - console.log('\nTop 5 by P@5:'); + console.log('\nTop 5 by Coverage@5:'); allResults - .sort((a, b) => b.precisionAt5 - a.precisionAt5) + .sort((a, b) => b.requiredCoverageAt5 - a.requiredCoverageAt5) .slice(0, 5) .forEach((r, i) => { console.log( - ` ${i + 1}. P@5=${(r.precisionAt5 * 100).toFixed(1)}% zero-hits=${r.zeroHits} isolation=${(r.isolationScore * 100).toFixed(0)}% config=${JSON.stringify(r.config)}`, + ` ${i + 1}. Coverage@5=${(r.requiredCoverageAt5 * 100).toFixed(1)}% zero-hits=${r.zeroHits} isolation=${(r.isolationScore * 100).toFixed(0)}% config=${JSON.stringify(r.config)}`, ); }); } else { console.log( - `\n${'Rank'.padEnd(5)} ${'P@5'.padEnd(8)} ${'Zero-hits'.padEnd(10)} ${'Isolation'.padEnd(10)} Config`, + `\n${'Rank'.padEnd(5)} ${'Coverage@5'.padEnd(8)} ${'Zero-hits'.padEnd(10)} ${'Isolation'.padEnd(10)} Config`, ); console.log('─'.repeat(80)); passing.forEach((r, i) => { console.log( - ` ${String(i + 1).padEnd(4)} ${(r.precisionAt5 * 100).toFixed(1).padEnd(8)}% ${String(r.zeroHits).padEnd(10)} ${(r.isolationScore * 100).toFixed(0).padEnd(10)}% cW=${r.config.cosineWeight} K=${r.config.preRerankK} iW=${r.config.importanceFinalWeight}`, + ` ${String(i + 1).padEnd(4)} ${(r.requiredCoverageAt5 * 100).toFixed(1).padEnd(8)}% ${String(r.zeroHits).padEnd(10)} ${(r.isolationScore * 100).toFixed(0).padEnd(10)}% cW=${r.config.cosineWeight} K=${r.config.preRerankK} iW=${r.config.importanceFinalWeight}`, ); }); @@ -251,7 +205,7 @@ function main() { ` importanceFinalWeight: ${best.config.importanceFinalWeight}`, ); console.log( - ` P@5: ${(best.precisionAt5 * 100).toFixed(1)}%`, + ` Coverage@5: ${(best.requiredCoverageAt5 * 100).toFixed(1)}%`, ); } diff --git a/test/benchmark/history.ts b/test/benchmark/history.ts index 2c92f8fa..86ab405e 100644 --- a/test/benchmark/history.ts +++ b/test/benchmark/history.ts @@ -75,6 +75,8 @@ export function compareReports( current: BenchmarkReport, previous: BenchmarkReport, ): string { + if (current.metricVersion !== previous.metricVersion) + return 'Reports use incompatible metric versions; rerun or rescore original results.'; const lines: string[] = []; lines.push(''); lines.push('═══ BENCHMARK COMPARISON ═══'); @@ -89,20 +91,20 @@ export function compareReports( const metrics: Array<{ name: string; - prev: number; - curr: number; + prev: number | null; + curr: number | null; higherBetter: boolean; }> = [ { - name: 'Precision@5', - prev: previous.overallPrecisionAt5, - curr: current.overallPrecisionAt5, + name: 'Required coverage@5', + prev: previous.overallRequiredCoverageAt5, + curr: current.overallRequiredCoverageAt5, higherBetter: true, }, { - name: 'Recall@20', - prev: previous.overallRecallAt20, - curr: current.overallRecallAt20, + name: 'Expected coverage@20', + prev: previous.overallExpectedCoverageAt20, + curr: current.overallExpectedCoverageAt20, higherBetter: true, }, { @@ -130,6 +132,10 @@ export function compareReports( lines.push(' ├─────────────────┼──────────┼──────────┼──────────┤'); for (const m of metrics) { + if (m.prev === null || m.curr === null) { + lines.push(`${m.name}: N/A (no labels)`); + continue; + } const delta = m.curr - m.prev; const deltaStr = delta >= 0 diff --git a/test/benchmark/recall-benchmark-dream.e2e-spec.ts b/test/benchmark/recall-benchmark-dream.e2e-spec.ts index 4d54163e..05e207f2 100644 --- a/test/benchmark/recall-benchmark-dream.e2e-spec.ts +++ b/test/benchmark/recall-benchmark-dream.e2e-spec.ts @@ -10,7 +10,7 @@ * * Thresholds: * - Isolation score = 100% (zero tolerance for cross-tenant leaks) - * - Precision@5 >= 95% + * - Required coverage@5 >= 95% * - No Dream Cycle stage returns internal errors * - No must_top5 query has 0 hits * @@ -32,7 +32,7 @@ import { EmbeddingService as EmbeddingGeneratorService } from '../../src/embeddi import { LLMService } from '../../src/llm/llm.service'; import type { LLMMessage } from '../../src/llm/llm.interface'; import { - PRECISION_AT_5_THRESHOLD, + REQUIRED_COVERAGE_AT_5_THRESHOLD, scoreQuery, buildReport, formatReport, @@ -327,36 +327,22 @@ describe('Recall Benchmark (Post-Dream-Cycle)', () => { expect(isolationFailures).toHaveLength(0); }); - it('should meet post-dream-cycle precision thresholds (Precision@5 >= 95%)', () => { + it('should meet post-dream-cycle coverage thresholds (Required coverage@5 >= 95%)', () => { if (allScores.length === 0) { console.warn('No scores to check thresholds against'); return; } - const usingRealEmbeddings = - process.env.BENCHMARK_REAL_EMBEDDINGS === 'true'; - - if (!usingRealEmbeddings) { - const avgP5 = - allScores.reduce((s, q) => s + q.precisionAt5, 0) / allScores.length; - console.log( - `Using cached embeddings — Precision@5 = ${(avgP5 * 100).toFixed(1)}% (thresholds relaxed)`, - ); - console.log( - ' Set BENCHMARK_REAL_EMBEDDINGS=true to enforce post-DC precision thresholds.', - ); - return; - } - + // This suite also always requests real embeddings; enforce its gate. const { sha, branch } = getGitInfo(); const report = buildReport(allScores, sha, branch); // Post-dream-cycle threshold: 95% (keeps production recall quality honest) - expect(report.overallPrecisionAt5).toBeGreaterThanOrEqual( - PRECISION_AT_5_THRESHOLD, + expect(report.overallRequiredCoverageAt5).toBeGreaterThanOrEqual( + REQUIRED_COVERAGE_AT_5_THRESHOLD, ); - // Log any zero-hit queries (aspirational — P@5 threshold is the hard gate). + // Log any zero-hit queries (aspirational — Coverage@5 threshold is the hard gate). const zeroHitQueries = allScores.filter( (s) => s.details.expectedTop5.length > 0 && s.details.top5Hits.length === 0, diff --git a/test/benchmark/recall-benchmark.e2e-spec.ts b/test/benchmark/recall-benchmark.e2e-spec.ts index f3c8556a..6d887fea 100644 --- a/test/benchmark/recall-benchmark.e2e-spec.ts +++ b/test/benchmark/recall-benchmark.e2e-spec.ts @@ -6,8 +6,8 @@ * * Thresholds: * - Isolation score = 100% (zero tolerance for cross-tenant leaks) - * - Precision@5 >= 95% - * - No must_top5 query has 0 hits + * - Required coverage@5 >= 95% + * - Explicit no-answer judgments must return no records */ import { INestApplication } from '@nestjs/common'; @@ -22,13 +22,12 @@ import type { GoldQuery } from '../fixtures/types'; import { PrismaService } from '../../src/prisma/prisma.service'; import { EmbeddingService as EmbeddingGeneratorService } from '../../src/embedding/embedding.service'; import { - PRECISION_AT_5_THRESHOLD, + REQUIRED_COVERAGE_AT_5_THRESHOLD, scoreQuery, buildReport, formatReport, checkThresholds, type QueryScore, - type BenchmarkReport, } from './scoring'; import { saveReport, getGitInfo } from './history'; import { generateCorpusEmbeddings } from '../helpers/generate-embeddings'; @@ -91,25 +90,21 @@ describe('Recall Benchmark', () => { .set(headers) .send({ query: query.query, limit: 20 }) .expect((r) => { - // Accept 200, 201, or 400 (for edge case queries like empty string) - if (r.status !== 200 && r.status !== 201 && r.status !== 400) { + // Empty input is handled above; other API failures must fail the run. + if (r.status !== 200 && r.status !== 201) { throw new Error( `Unexpected status ${r.status} for query "${query.query}": ${JSON.stringify(r.body)}`, ); } }); - if (res.status === 400) { - return []; - } - const body = res.body as { memories?: Array<{ id: string }> }; return (body.memories ?? []).map((m) => m.id); } catch (error) { console.error( `Query execution failed for [${query.id}]: ${(error as Error).message}`, ); - return []; + throw error; } } @@ -163,7 +158,7 @@ describe('Recall Benchmark', () => { expect(isolationFailures).toHaveLength(0); }); - it('should meet precision thresholds (with real embeddings)', () => { + it('should meet coverage thresholds (with real embeddings)', () => { if (allScores.length === 0) { console.warn('No scores to check thresholds against'); return; @@ -172,43 +167,14 @@ describe('Recall Benchmark', () => { const { sha, branch } = getGitInfo(); const report = buildReport(allScores, sha, branch); - // When using CachedEmbeddingService (hash-based stubs), precision will - // be low because vectors are not semantically meaningful. These thresholds - // only apply when using real embeddings (e.g., in CI with a real embedding provider). - // createTestApp() swaps EmbeddingService for CachedEmbeddingService by default. - // Only enforce precision thresholds when explicitly using real embeddings. - const usingRealEmbeddings = - process.env.BENCHMARK_REAL_EMBEDDINGS === 'true'; - - if (usingRealEmbeddings) { - // Precision@5 >= 95% - expect(report.overallPrecisionAt5).toBeGreaterThanOrEqual( - PRECISION_AT_5_THRESHOLD, - ); - - // Log any must_top5 queries with 0 hits (aspirational — not a hard gate). - // P@5 threshold above already captures overall quality. - // These edge cases are tracked for improvement but don't block CI. - const zeroHitQueries = allScores.filter( - (s) => - s.details.expectedTop5.length > 0 && - s.details.top5Hits.length === 0, - ); - if (zeroHitQueries.length > 0) { - const ids = zeroHitQueries.map((q) => q.queryId).join(', '); - console.warn( - `⚠️ Zero-hit queries (${zeroHitQueries.length}): ${ids}`, - ); - } - } else { - // With cached embeddings, just log the baseline - console.log( - `⚠️ Using cached embeddings — precision@5 = ${(report.overallPrecisionAt5 * 100).toFixed(1)}% (thresholds relaxed)`, - ); - console.log( - ' Set EMBEDDING_PROVIDER to a real provider to enforce precision thresholds.', - ); - } + // This suite always generates real embeddings. The gate cannot be disabled + // by omitting an unrelated environment flag. + expect(allScores).toHaveLength(GOLD_QUERIES.length); + expect(report.overallRequiredCoverageAt5).not.toBeNull(); + expect(report.overallRequiredCoverageAt5).toBeGreaterThanOrEqual( + REQUIRED_COVERAGE_AT_5_THRESHOLD, + ); + expect(checkThresholds(allScores)).toBe(true); }); }); }); diff --git a/test/benchmark/scoring.spec.ts b/test/benchmark/scoring.spec.ts new file mode 100644 index 00000000..e2574ab2 --- /dev/null +++ b/test/benchmark/scoring.spec.ts @@ -0,0 +1,71 @@ +import { scoreQuery, buildReport, checkThresholds } from './scoring'; +import type { GoldQuery } from '../fixtures/types'; +const query: GoldQuery = { + id: 'q', + query: 'fact', + user: 'a', + must_top5: ['gold'], + must_absent: ['forbidden'], + category: 'semantic', +}; +describe('honest benchmark accounting', () => { + it('names a required hit plus four unjudged results coverage, never precision', () => { + const score = scoreQuery(query, [ + 'gold', + 'noise1', + 'noise2', + 'noise3', + 'noise4', + ]); + expect(score.requiredCoverageAt5).toBe(1); + expect(score).not.toHaveProperty('precisionAt5'); + }); + it('excludes unlabelled queries and reports each label denominator', () => { + const report = buildReport( + [scoreQuery(query, []), scoreQuery({ ...query, must_top5: [] }, [])], + 'sha', + 'branch', + ); + expect(report.overallRequiredCoverageAt5).toBe(0); + expect(report.labelledTop5Count).toBe(1); + expect(report.unlabelledCount).toBe(1); + expect(report.overallExpectedCoverageAt20).toBeNull(); + expect(report.overallMrr).toBe(0); + }); + it('separates explicit negative judgments from absent labels', () => { + const negative = scoreQuery( + { ...query, must_top5: [], expect_no_answer: true }, + ['noise'], + ); + expect(negative.noAnswerPassed).toBe(false); + expect(negative.passed).toBe(false); + expect(checkThresholds([scoreQuery(query, ['gold']), negative])).toBe( + false, + ); + expect( + scoreQuery({ ...query, must_top5: [] }, ['noise']).noAnswerPassed, + ).toBeNull(); + }); + it('uses the first required hit for reciprocal rank', () => { + expect( + scoreQuery({ ...query, must_top5: ['gold', 'other'] }, [ + 'noise', + 'gold', + 'other', + ]).mrr, + ).toBe(0.5); + }); + it('does not let unlabelled queries inflate a failing 95% quality gate', () => { + const scores = Array.from({ length: 19 }, () => + scoreQuery(query, ['gold']), + ); + scores.push(scoreQuery(query, []), scoreQuery(query, [])); + scores.push( + ...Array.from({ length: 100 }, () => + scoreQuery({ ...query, must_top5: [] }, []), + ), + ); + expect(checkThresholds(scores)).toBe(false); + expect(checkThresholds([])).toBe(false); + }); +}); diff --git a/test/benchmark/scoring.ts b/test/benchmark/scoring.ts index 19bcef9e..3c0fb095 100644 --- a/test/benchmark/scoring.ts +++ b/test/benchmark/scoring.ts @@ -1,27 +1,18 @@ -/** - * Benchmark scoring utilities. - * - * Calculates precision@5, recall@20, MRR, and isolation scores - * for recall benchmark queries. - */ - +/** Label coverage is not precision: the fixture lists required items, not exhaustive relevance judgments. */ import type { GoldQuery } from '../fixtures/types'; -export const PRECISION_AT_5_THRESHOLD = 0.95; - +// Retain the original quality bar; exclude unlabelled cases instead of awarding them 1. +export const REQUIRED_COVERAGE_AT_5_THRESHOLD = 0.95; export interface QueryScore { queryId: string; category: string; passed: boolean; - /** How many must_top5 items appeared in actual top 5 */ - precisionAt5: number; - /** How many should_top20 items appeared in actual top 20 */ - recallAt20: number; - /** Mean reciprocal rank for must_top5 items */ - mrr: number; - /** True if no must_absent items appeared */ + requiredCoverageAt5: number | null; + expectedCoverageAt20: number | null; + /** Reciprocal rank of the FIRST required item, not the mean over required items. */ + mrr: number | null; + noAnswerPassed: boolean | null; isolationPassed: boolean; - /** Details for debugging failures */ details: { query: string; user: string; @@ -33,304 +24,193 @@ export interface QueryScore { top20Hits: string[]; }; } - export interface CategoryScore { category: string; queryCount: number; + labelledTop5Count: number; + labelledTop20Count: number; + noAnswerCount: number; + unlabelledCount: number; passed: number; failed: number; - avgPrecisionAt5: number; - avgRecallAt20: number; - avgMrr: number; + avgRequiredCoverageAt5: number | null; + avgExpectedCoverageAt20: number | null; + avgMrr: number | null; isolationScore: number; } - export interface BenchmarkReport { + metricVersion: 2; timestamp: string; gitSha: string; branch: string; totalQueries: number; + labelledTop5Count: number; + labelledTop20Count: number; + noAnswerCount: number; + unlabelledCount: number; passedQueries: number; failedQueries: number; - overallPrecisionAt5: number; - overallRecallAt20: number; - overallMrr: number; + overallRequiredCoverageAt5: number | null; + overallExpectedCoverageAt20: number | null; + overallMrr: number | null; overallIsolationScore: number; categoryScores: CategoryScore[]; queryScores: QueryScore[]; failedQueryDetails: QueryScore[]; thresholdsPassed: boolean; } - -/** - * Score a single query's results against its gold expectations. - */ export function scoreQuery(query: GoldQuery, resultIds: string[]): QueryScore { + if ( + query.expect_no_answer && + (query.must_top5.length || query.should_top20?.length) + ) { + throw new Error(`Contradictory no-answer labels: ${query.id}`); + } const top5 = resultIds.slice(0, 5); const top20 = resultIds.slice(0, 20); - - // Precision@5: fraction of must_top5 items found in top 5 - const top5Hits = query.must_top5.filter((id) => top5.includes(id)); - const precisionAt5 = - query.must_top5.length > 0 ? top5Hits.length / query.must_top5.length : 1.0; - - // Recall@20: fraction of should_top20 items found in top 20 - const shouldTop20 = query.should_top20 ?? []; - const top20Hits = shouldTop20.filter((id) => top20.includes(id)); - const recallAt20 = - shouldTop20.length > 0 ? top20Hits.length / shouldTop20.length : 1.0; - - // MRR: mean reciprocal rank for must_top5 items - let mrr: number; - if (query.must_top5.length > 0) { - const reciprocalRanks = query.must_top5.map((id) => { - const rank = resultIds.indexOf(id); - return rank >= 0 ? 1 / (rank + 1) : 0; - }); - mrr = - reciprocalRanks.reduce((sum, rr) => sum + rr, 0) / query.must_top5.length; - } else { - mrr = 1.0; // no must_top5 = perfect by default - } - - // Isolation: any must_absent item in results is a hard fail + const expectedTop5 = [...new Set(query.must_top5)]; + const expectedTop20 = [...new Set(query.should_top20 ?? [])]; + const top5Hits = expectedTop5.filter((id) => top5.includes(id)); + const top20Hits = expectedTop20.filter((id) => top20.includes(id)); + const firstRank = resultIds.findIndex((id) => expectedTop5.includes(id)); const mustAbsentViolations = query.must_absent.filter((id) => resultIds.includes(id), ); const isolationPassed = mustAbsentViolations.length === 0; - - // Overall pass: isolation + at least some hits if expected - const passed = - isolationPassed && (query.must_top5.length === 0 || top5Hits.length > 0); - + const noAnswerPassed = query.expect_no_answer ? resultIds.length === 0 : null; return { queryId: query.id, category: query.category, - passed, - precisionAt5, - recallAt20, - mrr, + passed: + isolationPassed && + noAnswerPassed !== false && + (!expectedTop5.length || top5Hits.length > 0), + requiredCoverageAt5: expectedTop5.length + ? top5Hits.length / expectedTop5.length + : null, + expectedCoverageAt20: expectedTop20.length + ? top20Hits.length / expectedTop20.length + : null, + mrr: expectedTop5.length ? (firstRank < 0 ? 0 : 1 / (firstRank + 1)) : null, + noAnswerPassed, isolationPassed, details: { query: query.query, user: query.user, - expectedTop5: query.must_top5, - expectedTop20: shouldTop20, - actualIds: resultIds.slice(0, 20), + expectedTop5, + expectedTop20, + actualIds: top20, mustAbsentViolations, top5Hits, top20Hits, }, }; } - -/** - * Aggregate scores by category. - */ +export function averageLabelled(values: (number | null)[]): number | null { + const labelled = values.filter((v): v is number => v !== null); + return labelled.length + ? labelled.reduce((s, v) => s + v, 0) / labelled.length + : null; +} +function counts(scores: QueryScore[]) { + return { + labelledTop5Count: scores.filter((s) => s.requiredCoverageAt5 !== null) + .length, + labelledTop20Count: scores.filter((s) => s.expectedCoverageAt20 !== null) + .length, + noAnswerCount: scores.filter((s) => s.noAnswerPassed !== null).length, + unlabelledCount: scores.filter( + (s) => + s.requiredCoverageAt5 === null && + s.expectedCoverageAt20 === null && + s.noAnswerPassed === null, + ).length, + }; +} export function aggregateByCategory(scores: QueryScore[]): CategoryScore[] { - const categories = new Map(); - for (const score of scores) { - const list = categories.get(score.category) ?? []; - list.push(score); - categories.set(score.category, list); - } - - const result: CategoryScore[] = []; - for (const [category, catScores] of categories) { - const passed = catScores.filter((s) => s.passed).length; - result.push({ + return [...new Set(scores.map((s) => s.category))].sort().map((category) => { + const group = scores.filter((s) => s.category === category); + const passed = group.filter((s) => s.passed).length; + return { category, - queryCount: catScores.length, + queryCount: group.length, + ...counts(group), passed, - failed: catScores.length - passed, - avgPrecisionAt5: avg(catScores.map((s) => s.precisionAt5)), - avgRecallAt20: avg(catScores.map((s) => s.recallAt20)), - avgMrr: avg(catScores.map((s) => s.mrr)), + failed: group.length - passed, + avgRequiredCoverageAt5: averageLabelled( + group.map((s) => s.requiredCoverageAt5), + ), + avgExpectedCoverageAt20: averageLabelled( + group.map((s) => s.expectedCoverageAt20), + ), + avgMrr: averageLabelled(group.map((s) => s.mrr)), isolationScore: - catScores.filter((s) => s.isolationPassed).length / catScores.length, - }); - } - - return result.sort((a, b) => a.category.localeCompare(b.category)); + group.filter((s) => s.isolationPassed).length / group.length, + }; + }); +} +export function checkThresholds(scores: QueryScore[]): boolean { + const coverage = averageLabelled(scores.map((s) => s.requiredCoverageAt5)); + return ( + scores.length > 0 && + coverage !== null && + coverage >= REQUIRED_COVERAGE_AT_5_THRESHOLD && + scores.every((s) => s.isolationPassed && s.noAnswerPassed !== false) + ); } - -/** - * Build the full benchmark report. - */ export function buildReport( scores: QueryScore[], gitSha: string, branch: string, ): BenchmarkReport { - const categoryScores = aggregateByCategory(scores); - const failedQueryDetails = scores.filter((s) => !s.passed); - return { + metricVersion: 2, timestamp: new Date().toISOString(), gitSha, branch, totalQueries: scores.length, + ...counts(scores), passedQueries: scores.filter((s) => s.passed).length, - failedQueries: failedQueryDetails.length, - overallPrecisionAt5: avg(scores.map((s) => s.precisionAt5)), - overallRecallAt20: avg(scores.map((s) => s.recallAt20)), - overallMrr: avg(scores.map((s) => s.mrr)), - overallIsolationScore: - scores.filter((s) => s.isolationPassed).length / scores.length, - categoryScores, + failedQueries: scores.filter((s) => !s.passed).length, + overallRequiredCoverageAt5: averageLabelled( + scores.map((s) => s.requiredCoverageAt5), + ), + overallExpectedCoverageAt20: averageLabelled( + scores.map((s) => s.expectedCoverageAt20), + ), + overallMrr: averageLabelled(scores.map((s) => s.mrr)), + overallIsolationScore: scores.length + ? scores.filter((s) => s.isolationPassed).length / scores.length + : 0, + categoryScores: aggregateByCategory(scores), queryScores: scores, - failedQueryDetails, + failedQueryDetails: scores.filter((s) => !s.passed), thresholdsPassed: checkThresholds(scores), }; } - -/** - * Check if all thresholds pass. - */ -export function checkThresholds(scores: QueryScore[]): boolean { - const isolationScore = - scores.filter((s) => s.isolationPassed).length / scores.length; - const precisionAt5 = avg(scores.map((s) => s.precisionAt5)); - - // Any must_top5 query with 0 hits? - const zeroHitQueries = scores.filter( - (s) => s.details.expectedTop5.length > 0 && s.details.top5Hits.length === 0, - ); - - // Zero-hit queries are tracked as warnings — P@5 threshold is the hard gate. - if (zeroHitQueries.length > 0) { - console.warn( - `⚠️ ${zeroHitQueries.length} zero-hit queries (tracked, not blocking): ${zeroHitQueries.map((q) => q.queryId).join(', ')}`, - ); - } - - return ( - isolationScore >= 1.0 && // Zero tolerance for isolation failures - precisionAt5 >= PRECISION_AT_5_THRESHOLD // At least 95% precision - ); -} - -/** - * Format the report for console output. - */ export function formatReport(report: BenchmarkReport): string { - const lines: string[] = []; - - lines.push(''); - lines.push( - '╔══════════════════════════════════════════════════════════════╗', - ); - lines.push('║ ENGRAM RECALL BENCHMARK REPORT ║'); - lines.push( - '╚══════════════════════════════════════════════════════════════╝', - ); - lines.push(''); - lines.push(` Git SHA: ${report.gitSha}`); - lines.push(` Branch: ${report.branch}`); - lines.push(` Timestamp: ${report.timestamp}`); - lines.push(''); - - // Overall scores - lines.push('┌─────────────────────────────────────────────────────────────┐'); - lines.push('│ OVERALL SCORES │'); - lines.push('├─────────────────────────────────────────────────────────────┤'); - lines.push( - `│ Total Queries: ${report.totalQueries} │`.slice( - 0, - 62, - ) + '│', - ); - lines.push( - `│ Passed: ${report.passedQueries} / ${report.totalQueries} (${pct(report.passedQueries / report.totalQueries)})` + - ' '.repeat(40), - ); - lines.push( - `│ Precision@5: ${pct(report.overallPrecisionAt5)} ${report.overallPrecisionAt5 >= PRECISION_AT_5_THRESHOLD ? '✅' : '❌'} (threshold: ${pct(PRECISION_AT_5_THRESHOLD)})`, - ); - lines.push(`│ Recall@20: ${pct(report.overallRecallAt20)}`); - lines.push(`│ MRR: ${report.overallMrr.toFixed(4)}`); - lines.push( - `│ Isolation: ${pct(report.overallIsolationScore)} ${report.overallIsolationScore >= 1.0 ? '✅' : '❌'} (threshold: 100%)`, - ); - lines.push('└─────────────────────────────────────────────────────────────┘'); - lines.push(''); - - // Category breakdown - lines.push( - '┌──────────────────┬───────┬────────┬──────────┬──────────┬──────────┬───────────┐', - ); - lines.push( - '│ Category │ Total │ Passed │ Prec@5 │ Rec@20 │ MRR │ Isolation │', - ); - lines.push( - '├──────────────────┼───────┼────────┼──────────┼──────────┼──────────┼───────────┤', - ); - for (const cat of report.categoryScores) { - lines.push( - `│ ${pad(cat.category, 16)} │ ${pad(String(cat.queryCount), 5)} │ ${pad(String(cat.passed), 6)} │ ${pad(pct(cat.avgPrecisionAt5), 8)} │ ${pad(pct(cat.avgRecallAt20), 8)} │ ${pad(cat.avgMrr.toFixed(4), 8)} │ ${pad(pct(cat.isolationScore), 9)} │`, - ); - } - lines.push( - '└──────────────────┴───────┴────────┴──────────┴──────────┴──────────┴───────────┘', - ); - lines.push(''); - - // Failed queries - if (report.failedQueryDetails.length > 0) { - lines.push(`❌ FAILED QUERIES (${report.failedQueryDetails.length}):`); - lines.push(''); - for (const q of report.failedQueryDetails) { - lines.push( - ` [${q.queryId}] "${q.details.query}" (user: ${q.details.user})`, - ); - if (!q.isolationPassed) { - lines.push( - ` ⛔ ISOLATION FAILURE: ${q.details.mustAbsentViolations.join(', ')}`, - ); - } - if ( - q.details.expectedTop5.length > 0 && - q.details.top5Hits.length === 0 - ) { - lines.push( - ` ⚠️ ZERO HITS: expected ${q.details.expectedTop5.join(', ')} in top 5`, - ); - } - if (q.precisionAt5 < 1.0 && q.details.expectedTop5.length > 0) { - lines.push( - ` 📊 Precision@5: ${pct(q.precisionAt5)} — hit: [${q.details.top5Hits.join(', ')}], missed: [${q.details.expectedTop5.filter((id) => !q.details.top5Hits.includes(id)).join(', ')}]`, - ); - } - lines.push( - ` 📋 Actual top 5: [${q.details.actualIds.slice(0, 5).join(', ')}]`, - ); - lines.push(''); - } - } - - // Final verdict - lines.push(''); - if (report.thresholdsPassed) { - lines.push('✅ ALL THRESHOLDS PASSED'); - } else { - lines.push('❌ THRESHOLDS NOT MET — benchmark FAILED'); - } - lines.push(''); - - return lines.join('\n'); -} - -// ── Helpers ───────────────────────────────────────────────────── - -function avg(values: number[]): number { - if (values.length === 0) return 0; - return values.reduce((s, v) => s + v, 0) / values.length; -} - -function pct(value: number): string { - return `${(value * 100).toFixed(1)}%`; -} - -function pad(str: string, len: number): string { - return str.padEnd(len); + const pct = (n: number | null) => + n === null ? 'N/A (no labels)' : `${(100 * n).toFixed(2)}%`; + return [ + `Engram recall benchmark (metric version ${report.metricVersion})`, + `${report.gitSha} (${report.branch}) at ${report.timestamp}`, + `Queries: ${report.totalQueries}; case assertions passed: ${report.passedQueries}`, + `Required-item coverage@5: ${pct(report.overallRequiredCoverageAt5)} (${report.labelledTop5Count} labelled; gate 95%)`, + `Expected-item coverage@20: ${pct(report.overallExpectedCoverageAt20)} (${report.labelledTop20Count} labelled)`, + `MRR of first required hit: ${report.overallMrr?.toFixed(4) ?? 'N/A'}`, + `Explicit no-answer cases: ${report.noAnswerCount}; unlabelled relevance cases: ${report.unlabelledCount}`, + `Forbidden-ID isolation checks: ${pct(report.overallIsolationScore)}`, + 'These partial labels do not measure precision or exhaustive recall. Unlabelled cases contribute only forbidden-ID checks.', + ...report.categoryScores.map( + (c) => + `${c.category}: coverage@5 ${pct(c.avgRequiredCoverageAt5)} (n=${c.labelledTop5Count}), coverage@20 ${pct(c.avgExpectedCoverageAt20)} (n=${c.labelledTop20Count})`, + ), + ...report.failedQueryDetails.map( + (q) => + `FAIL ${q.queryId}: expected [${q.details.expectedTop5.join(', ')}], returned [${q.details.actualIds.slice(0, 5).join(', ')}], forbidden [${q.details.mustAbsentViolations.join(', ')}], no-answer=${q.noAnswerPassed}`, + ), + report.thresholdsPassed + ? 'Quality thresholds passed' + : 'Quality thresholds FAILED', + ].join('\n'); } diff --git a/test/fixtures/types.ts b/test/fixtures/types.ts index 4bda00eb..07307a92 100644 --- a/test/fixtures/types.ts +++ b/test/fixtures/types.ts @@ -41,6 +41,8 @@ export interface FixtureUser { } export interface GoldQuery { + /** Explicit negative judgment; empty must_top5 alone is NOT a no-answer label. */ + expect_no_answer?: boolean; /** Unique query ID */ id: string; /** The search query text */ From 4d817c04caa61f48f0522d498923d812aaa88663 Mon Sep 17 00:00:00 2001 From: heybeaux Date: Sun, 6 Sep 2026 18:31:02 -0700 Subject: [PATCH 4/6] test(audit): verify preservation writes through ensemble retrieval --- scripts/audit/preservation-postgres.cjs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/scripts/audit/preservation-postgres.cjs b/scripts/audit/preservation-postgres.cjs index 7af1cae6..dd0ee041 100644 --- a/scripts/audit/preservation-postgres.cjs +++ b/scripts/audit/preservation-postgres.cjs @@ -7,6 +7,7 @@ require('ts-node').register({ transpileOnly: true, project: path.join(repo, 'tsc const { Client } = require('pg'); const { DreamCycleConsolidationStage } = require(path.join(repo, 'src/consolidation/stages/dream-cycle-consolidation.stage.ts')); const { EmbeddingWriteService } = require(path.join(repo, 'src/vector/embedding-write.service.ts')); +const { PgVectorEnsembleProvider } = require(path.join(repo, 'src/ensemble/pgvector-ensemble.provider.ts')); const db = `engram_preservation_test_${process.pid}`; const config = { host: '/tmp', database: 'postgres' }; (async () => { @@ -74,7 +75,14 @@ const config = { host: '/tmp', database: 'postgres' }; assert.equal(eligible.length, fail ? 3 : 4); assert.deepEqual(eligible.filter(r => r.id !== 'new').map(r => r.id), ['a','b','c']); assert.ok(eligible.every(r => r.dimensions === dims)); - console.log(JSON.stringify({ name, eligible: eligible.length, sourcesPreserved: 3, rollback: fail })); + // Exercise the integrated ensemble reader against the preservation writer's + // committed model rows, including rollback paths and a different owner. + const provider = new PgVectorEnsembleProvider(prisma); + const search = { modelId: model, embedding: Array(dims).fill(0.1), userId: 'u', limit: 10 }; + const retrieved = await provider.queryByModel(search); + assert.deepEqual(retrieved.map(r => r.memoryId).sort(), eligible.map(r => r.id).sort()); + assert.deepEqual(await provider.queryByModel({ ...search, userId: 'other-owner' }), []); + console.log(JSON.stringify({ name, eligible: eligible.length, sourcesPreserved: 3, rollback: fail, ensembleRetrieved: retrieved.length, otherOwnerRetrieved: 0 })); } await tx.query("UPDATE memories SET consolidated_into='missing-replacement' WHERE id IN ('a','b','c')"); const inventory = require('node:fs').readFileSync(path.join(__dirname, 'consolidation-reconciliation.sql'), 'utf8') From ab0226d412f7ad88a5c00e849de7325a85003443 Mon Sep 17 00:00:00 2001 From: heybeaux Date: Sun, 6 Sep 2026 19:10:05 -0700 Subject: [PATCH 5/6] fix(recall): reward actual use and preserve temporal query meaning --- src/memory/memory-query.service.spec.ts | 26 ++++++++ src/memory/memory-query.service.ts | 4 +- .../recall-weight-usage.service.spec.ts | 63 ++++++++++--------- src/memory/recall-weight.service.ts | 20 +++--- 4 files changed, 71 insertions(+), 42 deletions(-) diff --git a/src/memory/memory-query.service.spec.ts b/src/memory/memory-query.service.spec.ts index 1ceefd0d..103546cf 100644 --- a/src/memory/memory-query.service.spec.ts +++ b/src/memory/memory-query.service.spec.ts @@ -653,6 +653,32 @@ describe('MemoryQueryService', () => { }); describe('temporal path — reranking query selection', () => { + it('embeds the original temporal meaning without a reranker and retains date eligibility', async () => { + const start = new Date(Date.now() - 7 * 86400000); + const end = new Date(); + temporalParser.parse.mockReturnValue({ + semanticQuery: 'What did I work on?', + temporalFilter: { expression: 'last week', start, end }, + } as any); + await service.recall(userId, { + query: 'What did I work on last week?', + limit: 5, + }); + expect(embedding.generateForRecall).toHaveBeenCalledWith( + 'What did I work on last week?', + ); + expect(prisma.memory.findMany).toHaveBeenCalledWith( + expect.objectContaining({ + where: expect.objectContaining({ + OR: [ + { observedAt: { gte: start, lte: end } }, + { observedAt: null, createdAt: { gte: start, lte: end } }, + ], + }), + }), + ); + }); + it('should pass original query (with temporal expression) to reranker on temporal path', async () => { const mockRerankService = { rerank: jest.fn().mockResolvedValue([{ index: 0, score: 0.9 }]), diff --git a/src/memory/memory-query.service.ts b/src/memory/memory-query.service.ts index cb2150a9..67d22fea 100644 --- a/src/memory/memory-query.service.ts +++ b/src/memory/memory-query.service.ts @@ -116,7 +116,9 @@ export class MemoryQueryService { const now = new Date(); const parsed = this.temporalParser.parse(dto.query, now); const hasTemporalIntent = parsed.temporalFilter !== null; - const searchQuery = parsed.semanticQuery; + // Date filters constrain eligibility; the full phrase still carries meaning + // for semantic ranking, especially when no cross-encoder is available. + const searchQuery = hasTemporalIntent ? dto.query : parsed.semanticQuery; if (hasTemporalIntent) { this.logger.log('[Recall] Temporal intent detected:', { diff --git a/src/memory/recall-weight-usage.service.spec.ts b/src/memory/recall-weight-usage.service.spec.ts index c8f83e58..d2ce4d52 100644 --- a/src/memory/recall-weight-usage.service.spec.ts +++ b/src/memory/recall-weight-usage.service.spec.ts @@ -93,12 +93,12 @@ describe('RecallWeightService — Usage Weighting (ENG-27)', () => { }); describe('usageSignal', () => { - it('should return 0 for memories below minimum retrieval threshold', () => { + it('should return 0 for memories with no actual use', () => { const memory = createMemory({ retrievalCount: 2, usedCount: 0 }); expect(service.usageSignal(memory)).toBe(0); }); - it('should return > 0 for memories with enough retrievals', () => { + it('should return > 0 for memories with actual use', () => { const memory = createMemory({ retrievalCount: 10, usedCount: 5, @@ -107,34 +107,39 @@ describe('RecallWeightService — Usage Weighting (ENG-27)', () => { expect(service.usageSignal(memory)).toBeGreaterThan(0); }); - it('should weight usedCount higher than retrievalCount', () => { - const highUsed = createMemory({ - retrievalCount: 5, - usedCount: 10, - lastRetrievedAt: new Date(), - }); - const highRetrieved = createMemory({ - retrievalCount: 25, // same raw total: 10*2 + 5 = 25, vs 0*2 + 25 = 25 - usedCount: 0, - lastRetrievedAt: new Date(), - }); - // With usedCountMultiplier=2: highUsed = 10*2 + 5 = 25, highRetrieved = 0*2 + 25 = 25 - // But with usedCount=10 vs 0, let's make the numbers different - const moreUsed = createMemory({ - retrievalCount: 3, - usedCount: 10, - lastRetrievedAt: new Date(), - }); - const moreRetrieved = createMemory({ - retrievalCount: 20, - usedCount: 0, - lastRetrievedAt: new Date(), + it('does not promote repeatedly retrieved distractions over a relevant unused memory', async () => { + const relevant = { ...createMemory({ id: 'relevant' }), score: 0.51 }; + const distraction = { + ...createMemory({ + id: 'distraction', + retrievalCount: 1000, + lastRetrievedAt: new Date(), + }), + score: 0.48, + }; + const ranked = await service.applyUsageWeighting([distraction, relevant]); + expect(ranked.map((memory) => memory.id)).toEqual([ + 'relevant', + 'distraction', + ]); + expect(service.usageSignal(distraction)).toBe(0); + }); + + it('uses actual usage even without retrieval history and ignores later exposures', () => { + const memory = createMemory({ + usedCount: 5, + retrievalCount: 0, + lastUsedAt: new Date(Date.now() - DAY_MS), }); - // moreUsed raw = 10*2 + 3 = 23 - // moreRetrieved raw = 0*2 + 20 = 20 - expect(service.usageSignal(moreUsed)).toBeGreaterThan( - service.usageSignal(moreRetrieved), - ); + const signal = service.usageSignal(memory); + expect(signal).toBeGreaterThan(0); + expect( + service.usageSignal({ + ...memory, + retrievalCount: 1000, + lastRetrievedAt: new Date(), + }), + ).toBe(signal); }); it('should decay with time since last use', () => { diff --git a/src/memory/recall-weight.service.ts b/src/memory/recall-weight.service.ts index 94bf0b62..a7064b5e 100644 --- a/src/memory/recall-weight.service.ts +++ b/src/memory/recall-weight.service.ts @@ -16,13 +16,13 @@ export interface UsageWeightConfig { usageWeight: number; /** Half-life for recency decay in days (default 14) */ recencyHalfLifeDays: number; - /** Multiplier for usedCount vs retrievalCount (default 2) */ + /** Multiplier for actual use count (default 2) */ usedCountMultiplier: number; /** Boost factor for positive feedback (default 1.5) */ feedbackBoost: number; /** Penalty factor for negative feedback (default 0.5) */ feedbackPenalty: number; - /** Minimum retrievals before usage weighting kicks in (default 3) */ + /** @deprecated Retained for config compatibility; actual use needs no retrieval threshold. */ minRetrievals: number; } @@ -134,29 +134,25 @@ export class RecallWeightService { * ENG-27: Calculate a usage-based signal for a memory. * * Combines: - * - usedCount (weighted higher — actual usage > mere retrieval) - * - retrievalCount + * - usedCount (actual use, not appearances in recall results) * - recency decay (recent usage matters more) * * Returns a value in [0, 1] representing usage strength. */ usageSignal(memory: Memory): number { - // Don't boost memories without enough retrieval data (cold-start protection) - if (memory.retrievalCount < this.usageConfig.minRetrievals) return 0; + // Retrieval is exposure, not evidence of use. Rewarding it creates a + // feedback loop in which repeatedly surfaced distractions gain relevance. + if (!(memory.usedCount > 0)) return 0; const now = Date.now(); const { usedCountMultiplier, recencyHalfLifeDays } = this.usageConfig; - // Raw usage score: usedCount is weighted higher than retrievalCount - const rawUsage = - memory.usedCount * usedCountMultiplier + memory.retrievalCount; + const rawUsage = memory.usedCount * usedCountMultiplier; // Recency decay: exp(-lambda * days) where lambda = ln(2) / halfLife const lambda = Math.LN2 / recencyHalfLifeDays; const lastUsedTime = - memory.lastUsedAt?.getTime() ?? - memory.lastRetrievedAt?.getTime() ?? - memory.createdAt.getTime(); + memory.lastUsedAt?.getTime() ?? memory.createdAt.getTime(); const daysSinceUse = (now - lastUsedTime) / DAY_MS; const recencyDecay = Math.exp(-lambda * daysSinceUse); From ff615b3a3d10cc5632161aa490255219cd2e2fb9 Mon Sep 17 00:00:00 2001 From: heybeaux Date: Sun, 6 Sep 2026 19:29:55 -0700 Subject: [PATCH 6/6] fix(ci): parse current recall quality report markers --- .github/workflows/benchmark.yml | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 610fd0ca..1a8e1890 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -136,16 +136,18 @@ jobs: } // Extract the report section - const reportStart = output.indexOf('╔══'); - const reportEnd = output.lastIndexOf('benchmark PASSED') !== -1 - ? output.lastIndexOf('benchmark PASSED') + 20 - : output.lastIndexOf('benchmark FAILED') + 20; + const reportStart = output.indexOf('Engram recall benchmark (metric version '); + const successMarker = 'Quality thresholds passed'; + const failureMarker = 'Quality thresholds FAILED'; + const reportEnd = output.lastIndexOf(successMarker) !== -1 + ? output.lastIndexOf(successMarker) + successMarker.length + : output.lastIndexOf(failureMarker) + failureMarker.length; const report = reportStart >= 0 ? output.slice(reportStart, reportEnd > reportStart ? reportEnd : undefined) : 'Report not found in output'; - const passed = output.includes('ALL THRESHOLDS PASSED') && !output.includes('THRESHOLDS NOT MET') && !output.includes('Benchmark FAILED'); + const passed = output.includes('Quality thresholds passed') && !output.includes('Quality thresholds FAILED') && !output.includes('Benchmark FAILED'); const emoji = passed ? '✅' : '❌'; const body = `## ${emoji} Recall Benchmark Results @@ -199,11 +201,11 @@ jobs: echo "::error::Benchmark output file not found — benchmark step likely failed" exit 1 fi - if ! grep -q "ALL THRESHOLDS PASSED" benchmark-output.txt; then - echo "::error::Benchmark did not report ALL THRESHOLDS PASSED" + if ! grep -q "Quality thresholds passed" benchmark-output.txt; then + echo "::error::Benchmark did not report Quality thresholds passed" exit 1 fi - if grep -Eq "THRESHOLDS NOT MET|Benchmark FAILED" benchmark-output.txt; then + if grep -Eq "Quality thresholds FAILED|Benchmark FAILED" benchmark-output.txt; then echo "::error::Benchmark thresholds not met" exit 1 fi