fix(memory): scope semantic search before limiting

This commit is contained in:
Dragan Spiridonov
2026-08-03 18:15:22 +02:00
parent c547283f1a
commit 0a98f6a1ba
6 changed files with 55 additions and 18 deletions
+6 -2
View File
@@ -237,10 +237,14 @@ export class HybridMemoryBackend implements MemoryBackend {
* Vector similarity search
* Now uses unified memory's persistent vector storage
*/
async vectorSearch(embedding: number[], k: number): Promise<VectorSearchResult[]> {
async vectorSearch(
embedding: number[],
k: number,
keyPrefix?: string
): Promise<VectorSearchResult[]> {
this.ensureInitialized();
const results = await this.unifiedMemory!.vectorSearch(embedding, k);
const results = await this.unifiedMemory!.vectorSearch(embedding, k, undefined, keyPrefix);
// Convert to VectorSearchResult format
return results.map(r => ({
+2 -2
View File
@@ -298,8 +298,8 @@ export interface MemoryBackend extends Initializable, Disposable {
/** Search by pattern (`options.namespace` mirrors `get`). */
search(pattern: string, limit?: number, options?: RetrieveOptions): Promise<string[]>;
/** Vector similarity search (HNSW) */
vectorSearch(embedding: number[], k: number): Promise<VectorSearchResult[]>;
/** Vector similarity search (HNSW), optionally scoped by logical key prefix. */
vectorSearch(embedding: number[], k: number, keyPrefix?: string): Promise<VectorSearchResult[]>;
/** Store vector embedding */
storeVector(key: string, embedding: number[], metadata?: unknown): Promise<void>;
+6 -1
View File
@@ -109,10 +109,15 @@ export class InMemoryBackend implements MemoryBackend {
return results;
}
async vectorSearch(embedding: number[], k: number): Promise<VectorSearchResult[]> {
async vectorSearch(
embedding: number[],
k: number,
keyPrefix?: string
): Promise<VectorSearchResult[]> {
const results: VectorSearchResult[] = [];
for (const [key, entry] of this.vectors.entries()) {
if (keyPrefix && !key.startsWith(keyPrefix)) continue;
const score = cosineSimilarity(embedding, entry.embedding);
results.push({ key, score, metadata: entry.metadata });
}
+24 -6
View File
@@ -851,16 +851,30 @@ export class UnifiedMemoryManager {
return result.changes > 0;
}
async vectorSearch(
query: number[], k: number = 10, namespace?: string
): Promise<Array<{ id: string; score: number; metadata?: unknown }>> {
async vectorSearch(
query: number[], k: number = 10, namespace?: string, keyPrefix?: string
): Promise<Array<{ id: string; score: number; metadata?: unknown }>> {
this.ensureInitialized();
if (!this.vectorsLoaded) {
await this.loadVectorIndex();
}
const results = this.vectorIndex.search(query, k * 2);
// Logical namespaces are encoded in vector ids (`namespace:key`), while
// the vectors table namespace identifies the physical backend. Widen the
// ANN query until enough logical matches are found so global top-K results
// cannot hide a valid namespaced hit.
const indexSize = this.vectorIndex.size();
let candidateCount = Math.min(indexSize, Math.max(k * 2, 1));
let results = this.vectorIndex.search(query, candidateCount);
while (
keyPrefix &&
results.filter(result => result.id.startsWith(keyPrefix)).length < k &&
candidateCount < indexSize
) {
candidateCount = Math.min(indexSize, candidateCount * 2);
results = this.vectorIndex.search(query, candidateCount);
}
if (results.length === 0) return [];
const ids = results.map(r => r.id);
@@ -875,7 +889,7 @@ export class UnifiedMemoryManager {
const filteredResults: Array<{ id: string; score: number; metadata?: unknown }> = [];
for (const result of results) {
const row = metadataMap.get(result.id);
if (row && row.namespace === namespace) {
if (row && row.namespace === namespace && (!keyPrefix || result.id.startsWith(keyPrefix))) {
filteredResults.push({
id: result.id, score: result.score,
metadata: row.metadata ? safeJsonParse(row.metadata) : undefined,
@@ -886,7 +900,11 @@ export class UnifiedMemoryManager {
return filteredResults;
}
return results.slice(0, k).map(result => {
const scopedResults = keyPrefix
? results.filter(result => result.id.startsWith(keyPrefix)).slice(0, k)
: results.slice(0, k);
return scopedResults.map(result => {
const row = metadataMap.get(result.id);
return {
id: result.id, score: result.score,
+6 -6
View File
@@ -245,12 +245,12 @@ export async function handleMemoryQuery(
try {
// Generate real 384-dim transformer embedding for accurate cosine similarity search
const embedding = await computeRealEmbedding(params.pattern);
const vectorResults = await kernel!.memory.vectorSearch(embedding, limit + offset);
// Filter by namespace if specified
const filtered = namespace !== 'default'
? vectorResults.filter(r => r.key.startsWith(`${namespace}:`))
: vectorResults;
const keyPrefix = namespace !== 'default' ? `${namespace}:` : undefined;
const filtered = await kernel!.memory.vectorSearch(
embedding,
limit + offset,
keyPrefix
);
const paginatedResults = filtered.slice(offset, offset + limit);
+11 -1
View File
@@ -393,7 +393,17 @@ describe('UnifiedMemoryManager', () => {
});
});
describe('vectorSearch', () => {
describe('vectorSearch', () => {
it('scopes candidates by logical key prefix before applying the limit', async () => {
await manager.vectorStore('other:closer', [1, 0, 0], 'qe-kernel');
await manager.vectorStore('target:only', [0.9, 0.1, 0], 'qe-kernel');
const results = await manager.vectorSearch([1, 0, 0], 1, undefined, 'target:');
expect(results).toHaveLength(1);
expect(results[0].id).toBe('target:only');
});
beforeEach(async () => {
// Store orthogonal vectors
await manager.vectorStore('v1', [1, 0, 0], 'default', { axis: 'x' });