feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'

src/vectorStore.js

@@ -1,61 +1,40 @@
-import { Client } from "chromadb";
-import { OpenAIEmbeddings } from "openai";
+const { ChromaClient } = require('chromadb');
 
-export class ChromaVectorStore {
+class VectorStore {
   constructor() {
-    const host = process.env.CHROMA_HOST || "localhost";
-    const port = process.env.CHROMA_PORT || "8000";
-    this.client = new Client({ path: `http://${host}:${port}` });
-    this.collectionName = "rag_collection";
+    this.client = new ChromaClient(); // uses local storage by default
     this.collection = null;
   }
 
-  async init() {
-    const collections = await this.client.getCollections();
-    const exists = collections.some((c) => c.name === this.collectionName);
-    if (!exists) {
-      this.collection = await this.client.createCollection({
-        name: this.collectionName,
-        metadata: { hnsw: { ef_construction: 128, M: 64 } },
-      });
-    } else {
-      this.collection = await this.client.getCollection({
-        name: this.collectionName,
-      });
+  async init(collectionName = 'default') {
+    this.collection = await this.client.getOrCreateCollection({
+      name: collectionName,
+      metadata: { hnsw: { efConstruction: 200, M: 16 } }
+    });
+  }
+
+  async add(embeddings, metadatas, ids) {
+    if (!this.collection) {
+      throw new Error('Collection not initialized. Call init() first.');
     }
-  }
-
-  async addDocuments(documents) {
-    if (!this.collection) await this.init();
-    const embeddings = await this._embedTexts(documents.map((d) => d.content));
-    const ids = documents.map((_, idx) => `doc_${Date.now()}_${idx}`);
     await this.collection.add({
-      ids,
       embeddings,
-      documents: documents.map((d) => d.content),
-      metadatas: documents.map((d) => d.metadata),
+      metadatas,
+      ids
     });
   }
 
-  async query(queryText, topK = 5) {
-    if (!this.collection) await this.init();
-    const embedding = await this._embedTexts([queryText]);
+  async query(queryEmbedding, nResults = 5) {
+    if (!this.collection) {
+      throw new Error('Collection not initialized. Call init() first.');
+    }
     const results = await this.collection.query({
-      queryEmbeddings: embedding,
-      nResults: topK,
+      queryEmbeddings: [queryEmbedding],
+      nResults,
+      include: ['metadatas', 'documents']
     });
-    return results.documents.map((doc, idx) => ({
-      content: doc,
-      score: results.distances[idx],
-      metadata: results.metadatas[idx],
-    }));
+    return results;
   }
+}
 
-  async _embedTexts(texts) {
-    const openai = new OpenAIEmbeddings({
-      apiKey: process.env.OPENAI_API_KEY,
-    });
-    const embeddings = await openai.embedTexts(texts);
-    return embeddings.data.map((d) => d.embedding);
-  }
-}
+module.exports = { VectorStore };