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

src/vectorStore.js

@@ -0,0 +1,61 @@
+import { Client } from "chromadb";
+import { OpenAIEmbeddings } from "openai";
+
+export class ChromaVectorStore {
+  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.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 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),
+    });
+  }
+
+  async query(queryText, topK = 5) {
+    if (!this.collection) await this.init();
+    const embedding = await this._embedTexts([queryText]);
+    const results = await this.collection.query({
+      queryEmbeddings: embedding,
+      nResults: topK,
+    });
+    return results.documents.map((doc, idx) => ({
+      content: doc,
+      score: results.distances[idx],
+      metadata: results.metadatas[idx],
+    }));
+  }
+
+  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);
+  }
+}