import { ChromaClient } from 'chromadb';
import { OllamaEmbeddings } from 'langchain/embeddings/ollama';
import { Document } from 'langchain/document';

export async function createVectorStore() {
  const chroma = new ChromaClient({ path: process.env.CHROMA_URL });
  const collection = await chroma.getOrCreateCollection({
    name: process.env.CHROMA_COLLECTION,
  });
  const embeddings = new OllamaEmbeddings({
    model: process.env.OLLAMA_EMBEDDING_MODEL || 'nomic-embed-text',
  });
  return new VectorStore(collection, embeddings);
}

class VectorStore {
  constructor(collection, embeddings) {
    this.collection = collection;
    this.embeddings = embeddings;
  }

  async addDocuments(docs) {
    const texts = docs.map((d) => d.pageContent);
    const embeddings = await this.embeddings.embedDocuments(texts);
    await this.collection.addDocuments({
      documents: docs,
      embeddings,
    });
  }

  async similaritySearch(query, k = 4) {
    const embedding = await this.embeddings.embedQuery(query);
    const results = await this.collection.getNearestNeighbors({
      queryEmbeddings: [embedding],
      n: k,
    });
    const ids = results.ids[0];
    const docs = await this.collection.getDocuments({ ids });
    return docs.map(
      (doc) =>
        new Document({
          pageContent: doc.document,
          metadata: doc.metadata,
        })
    );
  }
}