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

src/agent.js

@@ -1,13 +1,28 @@
-const { embed } = require('./utils');
+import { VectorStore } from "./vectorStore.js";
+import { webSearch } from "./webSearch.js";
 
-async function answerQuestion(question, vectorStore) {
-  const questionEmbedding = embed(question);
-  const results = await vectorStore.query(questionEmbedding, 3);
-  const contexts = results[0].metadatas.map(m => m.text).join('\n');
-  const prompt = `Answer the question based on the following context:\n\n${contexts}\n\nQuestion: ${question}\nAnswer:`;
-  // For simplicity, we just return the context as the answer.
-  // In a real scenario, you would pass the prompt to a language model.
-  return contexts;
-}
+/**
+ * A simple RAG agent that retrieves relevant documents from the vector store
+ * and optionally performs a web search if no relevant documents are found.
+ */
+export class Agent {
+  constructor(vectorStore) {
+    this.vectorStore = vectorStore;
+  }
 
-module.exports = { answerQuestion };
+  /**
+   * Processes a user query and returns the best answer.
+   * @param {string} query
+   * @returns {Promise<string>}
+   */
+  async answer(query) {
+    const results = await this.vectorStore.query(query, 3);
+    if (results.length > 0 && results[0].score < 0.5) {
+      // Return the most relevant document text
+      return results[0].text;
+    }
+    // Fallback to web search
+    const html = await webSearch(query);
+    return `No relevant local documents found. Here is the raw web search result:\n${html}`;
+  }
+}

src/index.js

@@ -1,27 +1,29 @@
-const { VectorStore } = require('./vectorStore');
-const { answerQuestion } = require('./agent');
-const { webSearch } = require('./search');
-require('dotenv').config();
+import { VectorStore } from "./vectorStore.js";
+import { Agent } from "./agent.js";
 
-(async () => {
-  const vectorStore = new VectorStore();
-  await vectorStore.init('rag_collection');
+async function main() {
+  const store = new VectorStore();
+  await store.init();
 
-  // Example usage: add some documents
+  // Sample documents to index
   const docs = [
-    { text: 'ChromaDB is a vector database.', id: 'doc1' },
-    { text: 'It supports similarity search.', id: 'doc2' }
+    { id: "1", text: "ChromaDB is a fast, lightweight vector database." },
+    { id: "2", text: "It supports in-memory and persistent storage." },
+    { id: "3", text: "You can use it with various embedding models." },
   ];
-  const embeddings = docs.map(d => require('./utils').embed(d.text));
-  const metadatas = docs.map(d => ({ id: d.id, text: d.text }));
-  await vectorStore.add(embeddings, metadatas, docs.map(d => d.id));
 
-  // Example question
-  const question = 'What is ChromaDB?';
-  const answer = await answerQuestion(question, vectorStore);
-  console.log('Answer:', answer);
+  await store.addDocuments(docs);
 
-  // Example web search
-  const results = await webSearch('ChromaDB documentation');
-  console.log('Web search results:', results);
-})();
+  const agent = new Agent(store);
+
+  const query = process.argv[2] || "What is ChromaDB?";
+  console.log(`Query: ${query}`);
+  const answer = await agent.answer(query);
+  console.log("\nAnswer:");
+  console.log(answer);
+}
+
+main().catch((err) => {
+  console.error(err);
+  process.exit(1);
+});

src/test.js

@@ -0,0 +1,15 @@
+import { VectorStore } from "./vectorStore.js";
+
+async function testVectorStore() {
+  const store = new VectorStore();
+  await store.init();
+  const docs = [
+    { id: "a", text: "Hello world" },
+    { id: "b", text: "Goodbye world" },
+  ];
+  await store.addDocuments(docs);
+  const results = await store.query("Hello", 2);
+  console.log("Test results:", results);
+}
+
+testVectorStore().catch((err) => console.error(err));

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 };
+}

src/webSearch.js

@@ -1,15 +1,17 @@
 import fetch from "node-fetch";
 
-export async function fetchWebContent(url) {
-  try {
-    const response = await fetch(url);
-    if (!response.ok) {
-      throw new Error(`HTTP error ${response.status}`);
-    }
-    const text = await response.text();
-    return text;
-  } catch (err) {
-    console.error(`Failed to fetch ${url}: ${err.message}`);
-    return "";
+/**
+ * Performs a simple web search using DuckDuckGo's HTML interface.
+ * This is a lightweight example and does not use an official API.
+ * @param {string} query
+ * @returns {Promise<string>} The raw HTML of the search results page.
+ */
+export async function webSearch(query) {
+  const url = `https://duckduckgo.com/html/?q=${encodeURIComponent(query)}`;
+  const response = await fetch(url);
+  if (!response.ok) {
+    throw new Error(`Web search failed with status ${response.status}`);
   }
+  const html = await response.text();
+  return html;
 }