mirror of
https://github.com/Egonex-AI/Understand-Anything.git
synced 2026-06-22 10:58:03 +08:00
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:
co-authored by
Claude Opus 4.6
parent
6224761359
commit
ae5b4de3d9
@@ -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);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -43,3 +43,8 @@ export {
|
||||
type PluginConfig,
|
||||
type PluginEntry,
|
||||
} from "./plugins/discovery.js";
|
||||
export {
|
||||
SemanticSearchEngine,
|
||||
cosineSimilarity,
|
||||
type SemanticSearchOptions,
|
||||
} from "./embedding-search.js";
|
||||
|
||||
Reference in New Issue
Block a user