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

src/vectorStore.js

@@ -1,52 +1,58 @@
-const chromadb = require('chromadb');
-const { OpenAI } = require('openai');
-const dotenv = require('dotenv');
-dotenv.config();
+import { ChromaClient } from 'chromadb';
+import { OpenAIEmbeddings } from 'langchain/embeddings/openai';
 
-const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
+const chroma = new ChromaClient({
+  host: process.env.CHROMA_HOST || 'localhost',
+  port: parseInt(process.env.CHROMA_PORT, 10) || 8000,
+});
 
-class VectorStore {
-  constructor() {
-    this.client = new chromadb.Client({ path: './chromadb' });
-    this.collection = this.client.getCollection('rag_collection');
-  }
+const embeddings = new OpenAIEmbeddings({
+  openAIApiKey: process.env.OPENAI_API_KEY,
+});
 
-  async getEmbedding(text) {
-    const response = await openai.embeddings.create({
-      model: 'text-embedding-ada-002',
-      input: text,
-    });
-    return response.data[0].embedding;
-  }
+const COLLECTION_NAME = 'rag_collection';
 
-  async addDocuments(chunks) {
-    const documents = [];
-    const embeddings = [];
-    const ids = [];
-
-    for (const chunk of chunks) {
-      const embedding = await this.getEmbedding(chunk);
-      documents.push(chunk);
-      embeddings.push(embedding);
-      ids.push(`${Date.now()}-${Math.random()}`);
-    }
-
-    await this.collection.add({
-      documents,
-      embeddings,
-      ids,
-    });
-  }
-
-  async query(queryText, k = 5) {
-    const queryEmbedding = await this.getEmbedding(queryText);
-    const results = await this.collection.query({
-      queryEmbeddings: [queryEmbedding],
-      nResults: k,
-    });
-
-    return results.documents[0];
+/**
+ * Ensure the collection exists in ChromaDB.
+ */
+async function ensureCollection() {
+  const collections = await chroma.listCollections();
+  if (!collections.includes(COLLECTION_NAME)) {
+    await chroma.createCollection({ name: COLLECTION_NAME });
   }
 }
 
-module.exports = VectorStore;
+/**
+ * Add an array of documents to the vector store.
+ * @param {string[]} docs
+ */
+export async function addDocuments(docs) {
+  await ensureCollection();
+  const ids = docs.map((_, idx) => `doc-${Date.now()}-${idx}`);
+  const embeddingsResult = await embeddings.embedDocuments(docs);
+  await chroma.add({
+    collection_name: COLLECTION_NAME,
+    ids,
+    documents: docs,
+    embeddings: embeddingsResult,
+  });
+}
+
+/**
+ * Query the vector store for the most relevant documents.
+ * @param {string} queryText
+ * @param {number} nResults
+ * @returns {Promise<{documents: string[]}>}
+ */
+export async function query(queryText, nResults = 3) {
+  await ensureCollection();
+  const embedding = await embeddings.embedQuery(queryText);
+  const results = await chroma.query({
+    collection_name: COLLECTION_NAME,
+    query_embeddings: [embedding],
+    n_results: nResults,
+  });
+  // ChromaDB returns an array of objects; extract documents
+  const docs = results.documents || [];
+  return { documents: docs };
+}