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";