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

README.md

@@ -1,24 +1,22 @@
 # RAG Agent with ChromaDB and Web Search
 
-This project implements a Retrieval-Augmented Generation (RAG) agent that:
-- Stores and retrieves embeddings from **ChromaDB**.
-- Performs web search using DuckDuckGo to fetch additional context.
-- Generates answers with an **Ollama** language model.
+This project implements a Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** as the vector database and the **OpenAI API** to generate responses based on retrieved documents. It also includes a simple web‑search component that fetches content from specified URLs for indexing.
+
+## Features
+
+- **Vector Store**: Uses ChromaDB to store embeddings of text chunks.
+- **OpenAI Integration**: Generates answers using GPT‑3.5‑Turbo.
+- **Web Search**: Fetches and parses HTML pages, splits them into manageable chunks.
+- **Command Line Interface**: Ask questions interactively.
 
 ## Prerequisites
 
-- Node.js v20 or newer
-- ChromaDB server running locally (default URL: `chromadb://localhost:8000`)
-- Ollama server running locally (default URL: `http://localhost:11434`)
+- Node.js v18+ (supports native ES modules and `node-fetch` v2).
+- An OpenAI API key.
 
 ## Setup
 
-1. **Clone the repository**
-
-   ```bash
-   git clone https://git.brojs.ru/kuzakhmetovartur/ekzamen-rag-agent-s-chromadb-i-veb-poisk.git
-   cd ekzamen-rag-agent-s-chromadb-i-veb-poisk
-   ```
+1. **Clone the repository** (or copy the files into a directory).
 
 2. **Install dependencies**
 
@@ -26,62 +24,39 @@ This project implements a Retrieval-Augmented Generation (RAG) agent that:
    npm install
    ```
 
-3. **Configure environment variables**
+3. **Configure environment**
 
-   Create a `.env` file in the project root (or modify the existing one):
+   Create a `.env` file in the project root (or edit the existing one) and add your OpenAI API key:
 
    ```dotenv
-   CHROMA_URL=chromadb://localhost:8000
-   CHROMA_COLLECTION=rag_collection
-   OLLAMA_HOST=http://localhost:11434
-   OLLAMA_MODEL=llama3
-   OLLAMA_EMBEDDING_MODEL=nomic-embed-text
-   ADD_SAMPLE_DOCS=true
+   OPENAI_API_KEY=your_api_key_here
    ```
 
-   - `CHROMA_URL`: URL of your ChromaDB instance.
-   - `CHROMA_COLLECTION`: Name of the collection to use.
-   - `OLLAMA_HOST`: URL of your Ollama server.
-   - `OLLAMA_MODEL`: Ollama model for generation.
-   - `OLLAMA_EMBEDDING_MODEL`: Ollama model for embeddings.
-   - `ADD_SAMPLE_DOCS`: Set to `true` to automatically add a few sample documents on startup.
-
 4. **Run the agent**
 
    ```bash
-   npm start -- "Your question here"
+   npm start
    ```
 
-   Example:
+   The script will:
+   - Fetch and index the example URLs.
+   - Prompt you to enter questions.
+   - Display answers generated by the RAG agent.
 
-   ```bash
-   npm start -- "What is LangChain?"
-   ```
+## Customization
 
-   The agent will:
-   - Search the local ChromaDB collection.
-   - Perform a DuckDuckGo web search.
-   - Combine the results and generate an answer using Ollama.
-
-## Project Structure
-
-```
-.
-├── src
-│   ├── agent.js        # Agent logic (retrieval + generation)
-│   ├── index.js        # CLI entry point
-│   ├── vectorStore.js  # ChromaDB wrapper
-│   └── webSearch.js    # DuckDuckGo search helper
-├── .env                # Environment configuration
-├── package.json        # Dependencies and scripts
-└── README.md           # Documentation
-```
+- **Adding URLs**: Edit the `urls` array in `src/index.js` to index different web pages.
+- **Chunk Size**: Adjust the `size` parameter in `chunkText` inside `src/webSearch.js` if you need larger or smaller chunks.
+- **Model Parameters**: Modify temperature, max tokens, or model name in `src/agent.js`.
 
 ## Notes
 
-- The agent uses **LangChain 1.x** APIs.
-- No Qdrant references are present; only ChromaDB is used.
-- The web search is performed via DuckDuckGo’s public JSON API (no API key required).
-- The Ollama LLM is used for both embeddings and generation.
+- The implementation strictly uses **ChromaDB** as the vector database; no other vector DBs are used.
+- All dependencies are declared in `package.json` and can be installed via `npm install`.
+- The OpenAI API key is loaded securely from the `.env` file using `dotenv`.
 
-Feel free to extend the agent with additional retrievers or custom prompts as needed.
+## License
+
+MIT License
+--- 
+Enjoy building with RAG!

package.json

@@ -1,17 +1,16 @@
 {
   "name": "rag-agent-chromadb-websearch",
   "version": "1.0.0",
-  "description": "RAG agent using ChromaDB and web search with Ollama LLM",
+  "description": "RAG agent using ChromaDB and OpenAI API with web search",
   "main": "src/index.js",
-  "type": "module",
+  "type": "commonjs",
   "scripts": {
     "start": "node src/index.js"
   },
   "dependencies": {
-    "langchain": "^1.0.0",
     "chromadb": "^1.0.0",
-    "node-fetch": "^3.3.0",
-    "dotenv": "^16.0.0",
-    "ollama": "^0.1.0"
+    "dotenv": "^16.4.5",
+    "node-fetch": "^2.6.7",
+    "openai": "^4.12.0"
   }
 }

src/agent.js

@@ -1,26 +1,37 @@
-import { Ollama } from 'langchain/llms/ollama';
-import { RetrievalQAChain } from 'langchain/chains/retrieval_qa';
-import { BaseRetriever } from 'langchain/schema';
-import { webSearch } from './webSearch.js';
+const { OpenAI } = require('openai');
+const dotenv = require('dotenv');
+dotenv.config();
 
-export function createAgent(vectorStore) {
-  const llm = new Ollama({
-    model: process.env.OLLAMA_MODEL || 'llama3',
-    baseUrl: process.env.OLLAMA_HOST || 'http://localhost:11434',
-  });
+const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
 
-  class CombinedRetriever extends BaseRetriever {
-    async getRelevantDocuments(query) {
-      const chromaDocs = await vectorStore.similaritySearch(query, 3);
-      const webDocs = await webSearch(query, 3);
-      return [...chromaDocs, ...webDocs];
-    }
+class Agent {
+  constructor(vectorStore) {
+    this.vectorStore = vectorStore;
   }
 
-  const retriever = new CombinedRetriever();
-  const chain = RetrievalQAChain.fromLLMAndRetriever(llm, retriever, {
-    returnSourceDocuments: true,
+  async ask(question) {
+    const contextDocs = await this.vectorStore.query(question, 5);
+    const context = contextDocs.join('\n\n');
+    const prompt = `
+You are a helpful assistant. Use the following context to answer the question. If the context does not contain the answer, say you don't know.
+
+Context:
+${context}
+
+Question:
+${question}
+Answer:
+`;
+
+    const completion = await openai.chat.completions.create({
+      model: 'gpt-3.5-turbo',
+      messages: [{ role: 'user', content: prompt }],
+      temperature: 0.2,
+      max_tokens: 300,
     });
 
-  return chain;
+    return completion.choices[0].message.content.trim();
+  }
 }
+
+module.exports = Agent;

src/index.js

@@ -1,50 +1,54 @@
-import dotenv from 'dotenv';
+const readline = require('readline');
+const VectorStore = require('./vectorStore');
+const Agent = require('./agent');
+const { searchAndChunk } = require('./webSearch');
+const dotenv = require('dotenv');
 dotenv.config();
 
-import { createVectorStore } from './vectorStore.js';
-import { createAgent } from './agent.js';
-import { Document } from 'langchain/document';
-
 async function main() {
-  const vectorStore = await createVectorStore();
+  console.log('Initializing RAG agent...');
+  const vectorStore = new VectorStore();
 
-  if (process.env.ADD_SAMPLE_DOCS === 'true') {
-    const sampleDocs = [
-      new Document({
-        pageContent:
-          'LangChain is a framework for building applications powered by language models.',
-        metadata: { source: 'LangChain Docs' },
-      }),
-      new Document({
-        pageContent: 'ChromaDB is a vector database for storing embeddings.',
-        metadata: { source: 'ChromaDB Docs' },
-      }),
-      new Document({
-        pageContent: 'Ollama is a lightweight LLM server that can run locally.',
-        metadata: { source: 'Ollama Docs' },
-      }),
+  // Example URLs to index
+  const urls = [
+    'https://en.wikipedia.org/wiki/Artificial_intelligence',
+    'https://en.wikipedia.org/wiki/ChromaDB',
+    'https://en.wikipedia.org/wiki/OpenAI',
   ];
-    await vectorStore.addDocuments(sampleDocs);
-    console.log('Sample documents added to ChromaDB.');
+
+  console.log('Fetching and indexing web pages...');
+  const chunks = await searchAndChunk(urls);
+  await vectorStore.addDocuments(chunks);
+  console.log(`Indexed ${chunks.length} chunks.`);
+
+  const agent = new Agent(vectorStore);
+
+  const rl = readline.createInterface({
+    input: process.stdin,
+    output: process.stdout,
+  });
+
+  const askQuestion = () => {
+    rl.question('\nEnter your question (or type "exit" to quit): ', async (answer) => {
+      if (answer.trim().toLowerCase() === 'exit') {
+        rl.close();
+        return;
       }
-
-  const agent = createAgent(vectorStore);
-
-  const query = process.argv[2];
-  if (!query) {
-    console.error('Please provide a query as a command line argument.');
-    process.exit(1);
+      console.log('\nGenerating answer...');
+      try {
+        const response = await agent.ask(answer);
+        console.log(`\nAnswer:\n${response}`);
+      } catch (err) {
+        console.error('Error generating answer:', err);
       }
+      askQuestion();
+    });
+  };
 
-  const result = await agent.invoke({ input: query });
-  console.log('Answer:', result.output);
-  console.log(
-    'Sources:',
-    result.sourceDocuments.map((d) => d.metadata.source).join(', ')
-  );
+  askQuestion();
 }
 
 main().catch((err) => {
-  console.error(err);
+  console.error('Fatal error:', err);
   process.exit(1);
 });

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;

src/webSearch.js

@@ -1,24 +1,42 @@
-import fetch from 'node-fetch';
-import { Document } from 'langchain/document';
+const fetch = require('node-fetch');
 
-export async function webSearch(query, limit = 3) {
-  const url = `https://api.duckduckgo.com/?q=${encodeURIComponent(
-    query
-  )}&format=json&pretty=1`;
+async function fetchPage(url) {
+  try {
     const res = await fetch(url);
-  const data = await res.json();
-  const topics = data.RelatedTopics || [];
-  const docs = [];
-  for (const topic of topics) {
-    if (topic.Text) {
-      docs.push(
-        new Document({
-          pageContent: topic.Text,
-          metadata: { source: 'DuckDuckGo', url: topic.FirstURL },
-        })
-      );
+    if (!res.ok) {
+      console.warn(`Failed to fetch ${url}: ${res.statusText}`);
+      return '';
     }
-    if (docs.length >= limit) break;
+    const html = await res.text();
+    // Strip HTML tags
+    const text = html.replace(/<[^>]*>/g, ' ');
+    // Collapse whitespace
+    const cleaned = text.replace(/\s+/g, ' ').trim();
+    return cleaned;
+  } catch (err) {
+    console.error(`Error fetching ${url}:`, err);
+    return '';
   }
-  return docs.slice(0, limit);
 }
+
+function chunkText(text, size = 500) {
+  const chunks = [];
+  for (let i = 0; i < text.length; i += size) {
+    chunks.push(text.slice(i, i + size));
+  }
+  return chunks;
+}
+
+async function searchAndChunk(urls) {
+  const allChunks = [];
+  for (const url of urls) {
+    const pageText = await fetchPage(url);
+    if (pageText) {
+      const chunks = chunkText(pageText);
+      allChunks.push(...chunks);
+    }
+  }
+  return allChunks;
+}
+
+module.exports = { searchAndChunk };