multi-role

import { readFile, writeFile, mkdir } from "fs/promises";
import { join } from "path";

interface Document {
  id: string;
  text: string;
  source: string;
  embedding: number[];
  addedAt: string;
}

function cosineSimilarity(a: number[], b: number[]): number {
  if (a.length !== b.length || a.length === 0) return 0;
  let dot = 0, normA = 0, 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];
  }
  const denom = Math.sqrt(normA) * Math.sqrt(normB);
  return denom === 0 ? 0 : dot / denom;
}

export class VectorStore {
  private docs: Document[] = [];
  private storePath: string;

  constructor(workingDir: string) {
    this.storePath = join(workingDir, ".multi-role-vectors.json");
  }

  async load(): Promise<void> {
    try {
      const raw = await readFile(this.storePath, "utf-8");
      this.docs = JSON.parse(raw);
    } catch {
      this.docs = [];
    }
  }

  async save(): Promise<void> {
    await mkdir(join(this.storePath, ".."), { recursive: true });
    await writeFile(this.storePath, JSON.stringify(this.docs), "utf-8");
  }

  async add(text: string, source: string, embedding: number[]): Promise<string> {
    const id = `doc_${Date.now()}_${Math.random().toString(36).slice(2, 7)}`;
    this.docs.push({ id, text, source, embedding, addedAt: new Date().toISOString() });
    await this.save();
    return id;
  }

  search(queryEmbedding: number[], topK: number = 3): Array<{ id: string; source: string; text: string; score: number }> {
    if (this.docs.length === 0) return [];

    return this.docs
      .map(doc => ({ id: doc.id, source: doc.source, text: doc.text, score: cosineSimilarity(queryEmbedding, doc.embedding) }))
      .sort((a, b) => b.score - a.score)
      .slice(0, topK);
  }

  remove(id: string): boolean {
    const before = this.docs.length;
    this.docs = this.docs.filter(d => d.id !== id);
    return this.docs.length < before;
  }

  list(): Array<{ id: string; source: string; chars: number; addedAt: string }> {
    return this.docs.map(d => ({ id: d.id, source: d.source, chars: d.text.length, addedAt: d.addedAt }));
  }

  count(): number {
    return this.docs.length;
  }
}