feat(core): add embedding-based semantic search engine

Add SemanticSearchEngine class that stores pre-computed vector embeddings
for graph nodes and performs cosine similarity search. Enables true semantic
queries (e.g., "find code that handles authentication") even when keywords
don't appear in node data. Returns the same SearchResult type as the
existing Fuse.js-based SearchEngine for compatibility.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Lum1104
2026-03-14 23:16:10 +08:00
co-authored by Claude Opus 4.6
parent 6224761359
commit ae5b4de3d9
3 changed files with 180 additions and 0 deletions
@@ -0,0 +1,92 @@
import { describe, it, expect } from "vitest";
import { SemanticSearchEngine, cosineSimilarity } from "../embedding-search.js";
import type { GraphNode } from "../types.js";
const nodes: GraphNode[] = [
{ id: "n1", type: "file", name: "auth.ts", summary: "Authentication module", tags: ["auth"], complexity: "moderate" },
{ id: "n2", type: "file", name: "db.ts", summary: "Database connection", tags: ["db"], complexity: "simple" },
{ id: "n3", type: "function", name: "login", summary: "User login handler", tags: ["auth", "login"], complexity: "moderate" },
];
// Simple unit vectors for testing
const embeddings: Record<string, number[]> = {
n1: [1, 0, 0, 0],
n2: [0, 1, 0, 0],
n3: [0.9, 0, 0.1, 0],
};
describe("embedding-search", () => {
describe("cosineSimilarity", () => {
it("returns 1 for identical vectors", () => {
expect(cosineSimilarity([1, 0, 0], [1, 0, 0])).toBeCloseTo(1);
});
it("returns 0 for orthogonal vectors", () => {
expect(cosineSimilarity([1, 0, 0], [0, 1, 0])).toBeCloseTo(0);
});
it("returns high similarity for similar vectors", () => {
const sim = cosineSimilarity([1, 0, 0], [0.9, 0.1, 0]);
expect(sim).toBeGreaterThan(0.9);
});
it("handles zero vectors", () => {
expect(cosineSimilarity([0, 0, 0], [1, 0, 0])).toBe(0);
});
});
describe("SemanticSearchEngine", () => {
it("returns results sorted by similarity", () => {
const engine = new SemanticSearchEngine(nodes, embeddings);
const queryEmbedding = [1, 0, 0, 0]; // most similar to n1 and n3
const results = engine.search(queryEmbedding);
expect(results[0].nodeId).toBe("n1");
});
it("respects limit parameter", () => {
const engine = new SemanticSearchEngine(nodes, embeddings);
const results = engine.search([1, 0, 0, 0], { limit: 2 });
expect(results).toHaveLength(2);
});
it("respects threshold parameter", () => {
const engine = new SemanticSearchEngine(nodes, embeddings);
const results = engine.search([1, 0, 0, 0], { threshold: 0.5 });
// n2 has 0 similarity, should be filtered out
const ids = results.map((r) => r.nodeId);
expect(ids).not.toContain("n2");
});
it("filters by node type", () => {
const engine = new SemanticSearchEngine(nodes, embeddings);
const results = engine.search([1, 0, 0, 0], { types: ["function"] });
expect(results.every((r) => {
const node = nodes.find((n) => n.id === r.nodeId);
return node?.type === "function";
})).toBe(true);
});
it("returns empty for nodes without embeddings", () => {
const engine = new SemanticSearchEngine(nodes, {});
const results = engine.search([1, 0, 0, 0]);
expect(results).toHaveLength(0);
});
it("hasEmbeddings returns true when embeddings exist", () => {
const engine = new SemanticSearchEngine(nodes, embeddings);
expect(engine.hasEmbeddings()).toBe(true);
});
it("hasEmbeddings returns false when empty", () => {
const engine = new SemanticSearchEngine(nodes, {});
expect(engine.hasEmbeddings()).toBe(false);
});
it("addEmbedding updates the search index", () => {
const engine = new SemanticSearchEngine(nodes, {});
expect(engine.hasEmbeddings()).toBe(false);
engine.addEmbedding("n1", [1, 0, 0, 0]);
expect(engine.hasEmbeddings()).toBe(true);
});
});
});
+83
View File
@@ -0,0 +1,83 @@
import type { GraphNode } from "./types.js";
import type { SearchResult } from "./search.js";
export interface SemanticSearchOptions {
limit?: number;
threshold?: number;
types?: string[];
}
/**
* Compute cosine similarity between two vectors.
* Returns 0 if either vector has zero magnitude.
*/
export function cosineSimilarity(a: number[], b: number[]): number {
let dot = 0;
let magA = 0;
let magB = 0;
for (let i = 0; i < a.length; i++) {
dot += a[i] * b[i];
magA += a[i] * a[i];
magB += b[i] * b[i];
}
magA = Math.sqrt(magA);
magB = Math.sqrt(magB);
if (magA === 0 || magB === 0) return 0;
return dot / (magA * magB);
}
/**
* Semantic search engine using vector embeddings.
* Stores pre-computed embeddings for graph nodes and performs
* cosine similarity search against query embeddings.
*/
export class SemanticSearchEngine {
private nodes: GraphNode[];
private embeddings: Map<string, number[]>;
constructor(nodes: GraphNode[], embeddings: Record<string, number[]>) {
this.nodes = nodes;
this.embeddings = new Map(Object.entries(embeddings));
}
hasEmbeddings(): boolean {
return this.embeddings.size > 0;
}
addEmbedding(nodeId: string, embedding: number[]): void {
this.embeddings.set(nodeId, embedding);
}
search(
queryEmbedding: number[],
options?: SemanticSearchOptions,
): SearchResult[] {
const limit = options?.limit ?? 10;
const threshold = options?.threshold ?? 0;
const typeFilter = options?.types;
const scored: Array<{ nodeId: string; score: number }> = [];
for (const node of this.nodes) {
if (typeFilter && !typeFilter.includes(node.type)) continue;
const embedding = this.embeddings.get(node.id);
if (!embedding) continue;
const similarity = cosineSimilarity(queryEmbedding, embedding);
if (similarity >= threshold) {
scored.push({ nodeId: node.id, score: 1 - similarity });
}
}
scored.sort((a, b) => a.score - b.score);
return scored.slice(0, limit);
}
updateNodes(nodes: GraphNode[]): void {
this.nodes = nodes;
}
}
+5
View File
@@ -43,3 +43,8 @@ export {
type PluginConfig,
type PluginEntry,
} from "./plugins/discovery.js";
export {
SemanticSearchEngine,
cosineSimilarity,
type SemanticSearchOptions,
} from "./embedding-search.js";