rag

import { LocalIndex } from "vectra";
import { join } from "path";
import { mkdir } from "fs/promises";

export interface ChunkMetadata {
  text: string;
  source: string;
  fileName: string;
  chunkIndex: number;
  charStart: number;
  charEnd: number;
  collection: string;
  ingestedAt: string;
  [key: string]: string | number | boolean;
}

const _indexes = new Map<string, LocalIndex<ChunkMetadata>>();

export function getIndex(dataPath: string, collection: string): LocalIndex<ChunkMetadata> {
  const key = `${dataPath}::${collection}`;
  let idx = _indexes.get(key);
  if (!idx) {
    idx = new LocalIndex<ChunkMetadata>(join(dataPath, collection));
    _indexes.set(key, idx);
  }
  return idx;
}

export async function ensureIndex(idx: LocalIndex<ChunkMetadata>): Promise<void> {
  await mkdir(idx.folderPath, { recursive: true });
  if (!(await idx.isIndexCreated())) {
    await idx.createIndex({ version: 1 });
  }
}

export async function insertChunks(
  idx: LocalIndex<ChunkMetadata>,
  chunks: Array<{ vector: number[]; metadata: ChunkMetadata }>,
): Promise<void> {
  await ensureIndex(idx);
  await idx.batchInsertItems(chunks.map(c => ({ vector: c.vector, metadata: c.metadata })));
}

export async function queryChunks(
  idx: LocalIndex<ChunkMetadata>,
  vector: number[],
  topK: number,
): Promise<Array<{ score: number; metadata: ChunkMetadata }>> {
  await ensureIndex(idx);
  const results = await idx.queryItems(vector, "", topK);
  return results.map(r => ({ score: r.score, metadata: r.item.metadata }));
}

export async function deleteBySource(
  idx: LocalIndex<ChunkMetadata>,
  sourcePath: string,
): Promise<number> {
  await ensureIndex(idx);
  const items = await idx.listItemsByMetadata({ source: { $eq: sourcePath } });
  if (items.length === 0) return 0;
  await idx.deleteItems(items.map(i => i.id));
  return items.length;
}

export async function listSources(
  idx: LocalIndex<ChunkMetadata>,
): Promise<Array<{ source: string; fileName: string; chunks: number; ingestedAt: string }>> {
  await ensureIndex(idx);
  const items = await idx.listItems();
  const map = new Map<string, { fileName: string; chunks: number; ingestedAt: string }>();
  for (const item of items) {
    const { source, fileName, ingestedAt } = item.metadata;
    const existing = map.get(source);
    if (!existing) {
      map.set(source, { fileName, chunks: 1, ingestedAt });
    } else {
      existing.chunks++;
    }
  }
  return Array.from(map.entries())
    .map(([source, v]) => ({ source, ...v }))
    .sort((a, b) => a.source.localeCompare(b.source));
}

export async function indexStats(
  idx: LocalIndex<ChunkMetadata>,
): Promise<{ documents: number; chunks: number }> {
  await ensureIndex(idx);
  const items = await idx.listItems();
  const sources = new Set(items.map(i => i.metadata.source));
  return { documents: sources.size, chunks: items.length };
}