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

src/vectorStore.js

@@ -1,47 +1,52 @@
-import { ChromaClient } from 'chromadb';
-import { OllamaEmbeddings } from 'langchain/embeddings/ollama';
-import { Document } from 'langchain/document';
+const chromadb = require('chromadb');
+const { OpenAI } = require('openai');
+const dotenv = require('dotenv');
+dotenv.config();
 
-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);
-}
+const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
 
 class VectorStore {
-  constructor(collection, embeddings) {
-    this.collection = collection;
-    this.embeddings = embeddings;
+  constructor() {
+    this.client = new chromadb.Client({ path: './chromadb' });
+    this.collection = this.client.getCollection('rag_collection');
   }
 
-  async addDocuments(docs) {
-    const texts = docs.map((d) => d.pageContent);
-    const embeddings = await this.embeddings.embedDocuments(texts);
-    await this.collection.addDocuments({
-      documents: docs,
+  async getEmbedding(text) {
+    const response = await openai.embeddings.create({
+      model: 'text-embedding-ada-002',
+      input: text,
+    });
+    return response.data[0].embedding;
+  }
+
+  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 similaritySearch(query, k = 4) {
-    const embedding = await this.embeddings.embedQuery(query);
-    const results = await this.collection.getNearestNeighbors({
-      queryEmbeddings: [embedding],
-      n: k,
+  async query(queryText, k = 5) {
+    const queryEmbedding = await this.getEmbedding(queryText);
+    const results = await this.collection.query({
+      queryEmbeddings: [queryEmbedding],
+      nResults: 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,
-        })
-    );
+
+    return results.documents[0];
   }
-}
+}
+
+module.exports = VectorStore;