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

src/vectorStore.js

@@ -1,40 +1,69 @@
-const { ChromaClient } = require('chromadb');
+import { Client } from "@chromadb/chromadb";
 
-class VectorStore {
+/**
+ * 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 class VectorStore {
   constructor() {
-    this.client = new ChromaClient(); // uses local storage by default
+    this.client = new Client();
     this.collection = null;
   }
 
-  async init(collectionName = 'default') {
+  async init() {
     this.collection = await this.client.getOrCreateCollection({
-      name: collectionName,
-      metadata: { hnsw: { efConstruction: 200, M: 16 } }
+      name: "rag_collection",
     });
   }
 
-  async add(embeddings, metadatas, ids) {
+  /**
+   * Adds an array of documents to the collection.
+   * @param {Array<{id: string, text: string}>} docs
+   */
+  async addDocuments(docs) {
     if (!this.collection) {
-      throw new Error('Collection not initialized. Call init() first.');
+      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,
       embeddings,
-      metadatas,
-      ids
+      documents,
     });
   }
 
-  async query(queryEmbedding, nResults = 5) {
+  /**
+   * 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('Collection not initialized. Call init() first.');
+      throw new Error("VectorStore not initialized. Call init() first.");
     }
+    const queryEmbedding = embed(queryText);
     const results = await this.collection.query({
       queryEmbeddings: [queryEmbedding],
       nResults,
-      include: ['metadatas', 'documents']
     });
-    return results;
+    // 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],
+    }));
   }
-}
-
-module.exports = { VectorStore };
+}