claude-flow
Version:
Ruflo - Enterprise AI agent orchestration for Claude Code. Deploy 60+ specialized agents in coordinated swarms with self-learning, fault-tolerant consensus, vector memory, and MCP integration
191 lines (159 loc) • 5.24 kB
text/typescript
export interface HnswSearchResult {
id: string;
score: number;
}
export class HnswLite {
private vectors = new Map<string, Float32Array>();
private neighbors = new Map<string, Set<string>>();
private readonly dimensions: number;
private readonly maxNeighbors: number;
private readonly efConstruction: number;
private readonly metric: string;
constructor(dimensions: number, m: number, efConstruction: number, metric: string) {
this.dimensions = dimensions;
this.maxNeighbors = m;
this.efConstruction = efConstruction;
this.metric = metric;
}
get size(): number {
return this.vectors.size;
}
add(id: string, vector: Float32Array): void {
this.vectors.set(id, vector);
if (this.vectors.size === 1) {
this.neighbors.set(id, new Set());
return;
}
const nearest = this.findNearest(vector, this.maxNeighbors);
const neighborSet = new Set<string>();
for (const n of nearest) {
neighborSet.add(n.id);
const nNeighbors = this.neighbors.get(n.id);
if (nNeighbors) {
nNeighbors.add(id);
if (nNeighbors.size > this.maxNeighbors * 2) {
this.pruneNeighbors(n.id);
}
}
}
this.neighbors.set(id, neighborSet);
}
remove(id: string): void {
this.vectors.delete(id);
const myNeighbors = this.neighbors.get(id);
if (myNeighbors) {
for (const nId of myNeighbors) {
this.neighbors.get(nId)?.delete(id);
}
}
this.neighbors.delete(id);
}
search(query: Float32Array, k: number, threshold?: number): HnswSearchResult[] {
if (this.vectors.size === 0) return [];
if (this.vectors.size <= k * 2) {
return this.bruteForce(query, k, threshold);
}
const visited = new Set<string>();
const candidates: HnswSearchResult[] = [];
let entryId: string | undefined;
let bestScore = -1;
for (const [id] of this.vectors) {
const score = this.similarity(query, this.vectors.get(id)!);
if (score > bestScore) {
bestScore = score;
entryId = id;
}
if (visited.size >= Math.min(this.efConstruction, this.vectors.size)) break;
visited.add(id);
candidates.push({ id, score });
}
if (entryId) {
const queue = [entryId];
let idx = 0;
while (idx < queue.length && visited.size < this.efConstruction * 2) {
const currentId = queue[idx++];
const currentNeighbors = this.neighbors.get(currentId);
if (!currentNeighbors) continue;
for (const nId of currentNeighbors) {
if (visited.has(nId)) continue;
visited.add(nId);
const vec = this.vectors.get(nId);
if (!vec) continue;
const score = this.similarity(query, vec);
candidates.push({ id: nId, score });
queue.push(nId);
}
}
}
candidates.sort((a, b) => b.score - a.score);
let filtered = candidates;
if (threshold !== undefined) {
filtered = filtered.filter(c => c.score >= threshold);
}
return filtered.slice(0, k);
}
private bruteForce(query: Float32Array, k: number, threshold?: number): HnswSearchResult[] {
const results: HnswSearchResult[] = [];
for (const [id, vec] of this.vectors) {
const score = this.similarity(query, vec);
if (threshold !== undefined && score < threshold) continue;
results.push({ id, score });
}
results.sort((a, b) => b.score - a.score);
return results.slice(0, k);
}
private findNearest(query: Float32Array, k: number): HnswSearchResult[] {
return this.bruteForce(query, k);
}
private pruneNeighbors(id: string): void {
const myNeighbors = this.neighbors.get(id);
if (!myNeighbors) return;
const vec = this.vectors.get(id);
if (!vec) return;
const scored: HnswSearchResult[] = [];
for (const nId of myNeighbors) {
const nVec = this.vectors.get(nId);
if (!nVec) continue;
scored.push({ id: nId, score: this.similarity(vec, nVec) });
}
scored.sort((a, b) => b.score - a.score);
const keep = new Set(scored.slice(0, this.maxNeighbors).map(s => s.id));
for (const nId of myNeighbors) {
if (!keep.has(nId)) {
myNeighbors.delete(nId);
}
}
}
private similarity(a: Float32Array, b: Float32Array): number {
if (this.metric === 'dot') return dotProduct(a, b);
if (this.metric === 'euclidean') return 1 / (1 + euclideanDistance(a, b));
return cosineSimilarity(a, b);
}
}
export function cosineSimilarity(a: Float32Array, b: Float32Array): number {
let dot = 0;
let normA = 0;
let normB = 0;
for (let i = 0; i < a.length; i++) {
dot += a[i] * b[i];
normA += a[i] * a[i];
normB += b[i] * b[i];
}
if (normA === 0 || normB === 0) return 0;
return dot / (Math.sqrt(normA) * Math.sqrt(normB));
}
function dotProduct(a: Float32Array, b: Float32Array): number {
let sum = 0;
for (let i = 0; i < a.length; i++) {
sum += a[i] * b[i];
}
return sum;
}
function euclideanDistance(a: Float32Array, b: Float32Array): number {
let sum = 0;
for (let i = 0; i < a.length; i++) {
const d = a[i] - b[i];
sum += d * d;
}
return Math.sqrt(sum);
}