import { ChromaClient } from "chromadb";
import { OpenAIEmbeddings } from "langchain/embeddings/openai";
import { OpenAI } from "langchain/llms/openai";

const client = new ChromaClient({
  path: process.env.CHROMA_DB_PATH || "./chromadb",
});

const embeddings = new OpenAIEmbeddings({
  openAIApiKey: process.env.OPENAI_API_KEY,
});

export async function addDocument(collectionName, text, metadata = {}) {
  const collection = await client.getOrCreateCollection({
    name: collectionName,
  });
  const embedding = await embeddings.embedQuery(text);
  await collection.add({
    documents: [text],
    embeddings: [embedding],
    metadatas: [metadata],
  });
}

export async function getSimilarDocuments(collectionName, query, k = 5) {
  const collection = await client.getOrCreateCollection({
    name: collectionName,
  });
  const embedding = await embeddings.embedQuery(query);
  const results = await collection.query({
    queryEmbeddings: [embedding],
    nResults: k,
  });
  return results.ids[0].map((id, idx) => ({
    id,
    score: results.scores[0][idx],
    document: results.documents[0][idx],
    metadata: results.metadatas[0][idx],
  }));
}