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

src/agent.js

@@ -0,0 +1,35 @@
+import { ChromaVectorStore } from "./vectorStore.js";
+import { fetchWebContent } from "./webSearch.js";
+import { OpenAI } from "openai";
+
+export class RAGAgent {
+  constructor() {
+    this.vectorStore = new ChromaVectorStore();
+    this.openai = new OpenAI({
+      apiKey: process.env.OPENAI_API_KEY,
+    });
+  }
+
+  async ingestFromUrl(url) {
+    const content = await fetchWebContent(url);
+    if (!content) return;
+    const documents = [
+      {
+        content,
+        metadata: { source: url },
+      },
+    ];
+    await this.vectorStore.addDocuments(documents);
+  }
+
+  async ask(question) {
+    const relevant = await this.vectorStore.query(question, 3);
+    const context = relevant.map((r) => r.content).join("\n---\n");
+    const prompt = `You are an assistant. Use the following context to answer the question.\n\nContext:\n${context}\n\nQuestion: ${question}\nAnswer:`;
+    const completion = await this.openai.chat.completions.create({
+      model: "gpt-4o-mini",
+      messages: [{ role: "user", content: prompt }],
+    });
+    return completion.choices[0].message.content.trim();
+  }
+}

src/index.js

@@ -0,0 +1,25 @@
+import dotenv from "dotenv";
+import { RAGAgent } from "./agent.js";
+
+dotenv.config();
+
+async function main() {
+  const agent = new RAGAgent();
+
+  // Example ingestion
+  const url = "https://raw.githubusercontent.com/openai/openai-node/main/README.md";
+  console.log(`Ingesting content from ${url}...`);
+  await agent.ingestFromUrl(url);
+  console.log("Ingestion complete.");
+
+  // Example question
+  const question = "What is the purpose of the OpenAI Node.js library?";
+  console.log(`\nAsking: ${question}`);
+  const answer = await agent.ask(question);
+  console.log(`\nAnswer:\n${answer}`);
+}
+
+main().catch((err) => {
+  console.error(err);
+  process.exit(1);
+});

src/vectorStore.js

@@ -0,0 +1,61 @@
+import { Client } from "chromadb";
+import { OpenAIEmbeddings } from "openai";
+
+export class ChromaVectorStore {
+  constructor() {
+    const host = process.env.CHROMA_HOST || "localhost";
+    const port = process.env.CHROMA_PORT || "8000";
+    this.client = new Client({ path: `http://${host}:${port}` });
+    this.collectionName = "rag_collection";
+    this.collection = null;
+  }
+
+  async init() {
+    const collections = await this.client.getCollections();
+    const exists = collections.some((c) => c.name === this.collectionName);
+    if (!exists) {
+      this.collection = await this.client.createCollection({
+        name: this.collectionName,
+        metadata: { hnsw: { ef_construction: 128, M: 64 } },
+      });
+    } else {
+      this.collection = await this.client.getCollection({
+        name: this.collectionName,
+      });
+    }
+  }
+
+  async addDocuments(documents) {
+    if (!this.collection) await this.init();
+    const embeddings = await this._embedTexts(documents.map((d) => d.content));
+    const ids = documents.map((_, idx) => `doc_${Date.now()}_${idx}`);
+    await this.collection.add({
+      ids,
+      embeddings,
+      documents: documents.map((d) => d.content),
+      metadatas: documents.map((d) => d.metadata),
+    });
+  }
+
+  async query(queryText, topK = 5) {
+    if (!this.collection) await this.init();
+    const embedding = await this._embedTexts([queryText]);
+    const results = await this.collection.query({
+      queryEmbeddings: embedding,
+      nResults: topK,
+    });
+    return results.documents.map((doc, idx) => ({
+      content: doc,
+      score: results.distances[idx],
+      metadata: results.metadatas[idx],
+    }));
+  }
+
+  async _embedTexts(texts) {
+    const openai = new OpenAIEmbeddings({
+      apiKey: process.env.OPENAI_API_KEY,
+    });
+    const embeddings = await openai.embedTexts(texts);
+    return embeddings.data.map((d) => d.embedding);
+  }
+}

src/webSearch.js

@@ -0,0 +1,15 @@
+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 "";
+  }
+}