paif/store.js

'use strict';

const path = require('path');
const fs = require('fs');
const os = require('os');
const { randomUUID } = require('crypto');
const { LocalIndex } = require('vectra');

const BASE_DIR = path.join(os.homedir(), '.memory-bridge');
const INDEXES_DIR = path.join(BASE_DIR, 'indexes');

// Map of agent_id -> LocalIndex instance
const agentIndexes = new Map();

// Legacy single-tenant index (for backward compat / master access)
let legacyIndex = null;
const LEGACY_STORE_DIR = path.join(BASE_DIR, 'index');

async function init() {
  // Ensure indexes directory exists
  fs.mkdirSync(INDEXES_DIR, { recursive: true });

  // Initialize legacy index for backward compatibility
  fs.mkdirSync(LEGACY_STORE_DIR, { recursive: true });
  legacyIndex = new LocalIndex(LEGACY_STORE_DIR);
  if (!(await legacyIndex.isIndexCreated())) {
    await legacyIndex.createIndex();
    console.log(`Legacy vector index : created at ${LEGACY_STORE_DIR}`);
  } else {
    console.log(`Legacy vector index : loaded from ${LEGACY_STORE_DIR}`);
  }
}

// Get or create an agent's index
async function getAgentIndex(agentId) {
  if (!agentId) {
    // No agent specified - use legacy index
    return legacyIndex;
  }

  if (agentIndexes.has(agentId)) {
    return agentIndexes.get(agentId);
  }

  const agentDir = path.join(INDEXES_DIR, agentId);
  fs.mkdirSync(agentDir, { recursive: true });

  const index = new LocalIndex(agentDir);
  if (!(await index.isIndexCreated())) {
    await index.createIndex();
    console.log(`Agent index created : ${agentId} at ${agentDir}`);
  }

  agentIndexes.set(agentId, index);
  return index;
}

async function add(agentId, { text, tags = [], source = '', vector }) {
  const index = await getAgentIndex(agentId);
  const id = randomUUID();
  const createdAt = new Date().toISOString();
  await index.insertItem({ id, vector, metadata: { id, text, tags, source, createdAt, agentId } });
  return { id, createdAt };
}

async function query(agentId, { vector, limit = 10, tags = [] }) {
  const index = await getAgentIndex(agentId);
  const topK = tags.length ? Math.min(limit * 3, 300) : limit;
  const results = await index.queryItems(vector, topK);
  let items = results;
  if (tags.length) {
    items = results.filter((r) =>
      tags.some((t) => (r.item.metadata.tags || []).includes(t))
    );
  }
  return items.slice(0, limit).map((r) => ({
    ...r.item.metadata,
    score: r.score,
  }));
}

async function remove(agentId, id) {
  const index = await getAgentIndex(agentId);
  await index.deleteItem(id);
}

async function list(agentId, { page = 1, pageSize = 50 } = {}) {
  const index = await getAgentIndex(agentId);
  const all = await index.listItems();
  const total = all.length;
  const start = (page - 1) * pageSize;
  const data = all.slice(start, start + pageSize).map((item) => item.metadata);
  return { data, total, page, pageSize, pages: Math.ceil(total / pageSize) };
}

async function count(agentId) {
  if (!agentId) {
    // Return total across all indexes for health check
    let total = 0;
    try {
      const agents = fs.readdirSync(INDEXES_DIR).filter(d => {
        const stat = fs.statSync(path.join(INDEXES_DIR, d));
        return stat.isDirectory();
      });
      for (const agent of agents) {
        const idx = await getAgentIndex(agent);
        const all = await idx.listItems();
        total += all.length;
      }
    } catch {
      // Ignore errors
    }
    return total;
  }

  const index = await getAgentIndex(agentId);
  const all = await index.listItems();
  return all.length;
}

// Get stats for all agents (master only)
async function stats() {
  const result = {
    totalAgents: 0,
    totalMemories: 0,
    agents: {}
  };

  try {
    const agents = fs.readdirSync(INDEXES_DIR).filter(d => {
      const stat = fs.statSync(path.join(INDEXES_DIR, d));
      return stat.isDirectory();
    });

    result.totalAgents = agents.length;

    for (const agent of agents) {
      const idx = await getAgentIndex(agent);
      const all = await idx.listItems();
      result.agents[agent] = all.length;
      result.totalMemories += all.length;
    }
  } catch {
    // Ignore errors
  }

  return result;
}

module.exports = { init, add, query, remove, list, count, stats };