mirror of
https://github.com/proffesor-for-testing/agentic-qe.git
synced 2026-09-19 08:45:47 +08:00
fix(memory): scope semantic search before limiting
This commit is contained in:
@@ -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 => ({
|
||||
|
||||
@@ -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>;
|
||||
|
||||
@@ -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 });
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
|
||||
|
||||
@@ -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' });
|
||||
|
||||
Reference in New Issue
Block a user