From ae5b4de3d9c7b9d4ebed098c0697e03710360d94 Mon Sep 17 00:00:00 2001 From: Lum1104 Date: Sat, 14 Mar 2026 23:16:10 +0800 Subject: [PATCH] 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 --- .../src/__tests__/embedding-search.test.ts | 92 +++++++++++++++++++ packages/core/src/embedding-search.ts | 83 +++++++++++++++++ packages/core/src/index.ts | 5 + 3 files changed, 180 insertions(+) create mode 100644 packages/core/src/__tests__/embedding-search.test.ts create mode 100644 packages/core/src/embedding-search.ts diff --git a/packages/core/src/__tests__/embedding-search.test.ts b/packages/core/src/__tests__/embedding-search.test.ts new file mode 100644 index 0000000..4952c79 --- /dev/null +++ b/packages/core/src/__tests__/embedding-search.test.ts @@ -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 = { + 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); + }); + }); +}); diff --git a/packages/core/src/embedding-search.ts b/packages/core/src/embedding-search.ts new file mode 100644 index 0000000..71192ca --- /dev/null +++ b/packages/core/src/embedding-search.ts @@ -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; + + constructor(nodes: GraphNode[], embeddings: Record) { + 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; + } +} diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 695e63c..7615d6d 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -43,3 +43,8 @@ export { type PluginConfig, type PluginEntry, } from "./plugins/discovery.js"; +export { + SemanticSearchEngine, + cosineSimilarity, + type SemanticSearchOptions, +} from "./embedding-search.js";