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

src/vectorStore.js

@@ -1,69 +1,47 @@
-import { Client } from "@chromadb/chromadb";
+import { ChromaClient } from 'chromadb';
+import { OllamaEmbeddings } from 'langchain/embeddings/ollama';
+import { Document } from 'langchain/document';
 
-/**
- * Simple embedding function that converts text into a fixed-length numeric vector.
- * This is a placeholder and should be replaced with a real embedding model for production use.
- */
-function embed(text) {
-  const vector = Array.from(text)
-    .map((c) => c.charCodeAt(0))
-    .slice(0, 10);
-  while (vector.length < 10) {
-    vector.push(0);
-  }
-  return vector;
+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);
 }
 
-export class VectorStore {
-  constructor() {
-    this.client = new Client();
-    this.collection = null;
+class VectorStore {
+  constructor(collection, embeddings) {
+    this.collection = collection;
+    this.embeddings = embeddings;
   }
 
-  async init() {
-    this.collection = await this.client.getOrCreateCollection({
-      name: "rag_collection",
-    });
-  }
-
-  /**
-   * Adds an array of documents to the collection.
-   * @param {Array<{id: string, text: string}>} docs
-   */
   async addDocuments(docs) {
-    if (!this.collection) {
-      throw new Error("VectorStore not initialized. Call init() first.");
-    }
-    const ids = docs.map((d) => d.id);
-    const embeddings = docs.map((d) => embed(d.text));
-    const documents = docs.map((d) => d.text);
-    await this.collection.add({
-      ids,
+    const texts = docs.map((d) => d.pageContent);
+    const embeddings = await this.embeddings.embedDocuments(texts);
+    await this.collection.addDocuments({
+      documents: docs,
       embeddings,
-      documents,
     });
   }
 
-  /**
-   * Queries the collection for the most relevant documents.
-   * @param {string} queryText
-   * @param {number} nResults
-   * @returns {Promise<Array<{id: string, text: string, score: number}>>}
-   */
-  async query(queryText, nResults = 3) {
-    if (!this.collection) {
-      throw new Error("VectorStore not initialized. Call init() first.");
-    }
-    const queryEmbedding = embed(queryText);
-    const results = await this.collection.query({
-      queryEmbeddings: [queryEmbedding],
-      nResults,
+  async similaritySearch(query, k = 4) {
+    const embedding = await this.embeddings.embedQuery(query);
+    const results = await this.collection.getNearestNeighbors({
+      queryEmbeddings: [embedding],
+      n: k,
     });
-    // results is an array of objects with ids, documents, and scores
-    return results[0].ids.map((id, idx) => ({
-      id,
-      text: results[0].documents[idx],
-      score: results[0].distances[idx],
-    }));
+    const ids = results.ids[0];
+    const docs = await this.collection.getDocuments({ ids });
+    return docs.map(
+      (doc) =>
+        new Document({
+          pageContent: doc.document,
+          metadata: doc.metadata,
+        })
+    );
   }
 }