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

This commit is contained in:
2026-06-30 12:59:48 +03:00
parent 136e69e967
commit 27dcde060a
6 changed files with 190 additions and 178 deletions
+31 -56
View File
@@ -1,24 +1,22 @@
# RAG Agent with ChromaDB and Web Search # RAG Agent with ChromaDB and Web Search
This project implements a Retrieval-Augmented Generation (RAG) agent that: 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 websearch component that fetches content from specified URLs for indexing.
- Stores and retrieves embeddings from **ChromaDB**.
- Performs web search using DuckDuckGo to fetch additional context. ## Features
- Generates answers with an **Ollama** language model.
- **Vector Store**: Uses ChromaDB to store embeddings of text chunks.
- **OpenAI Integration**: Generates answers using GPT3.5Turbo.
- **Web Search**: Fetches and parses HTML pages, splits them into manageable chunks.
- **Command Line Interface**: Ask questions interactively.
## Prerequisites ## Prerequisites
- Node.js v20 or newer - Node.js v18+ (supports native ES modules and `node-fetch` v2).
- ChromaDB server running locally (default URL: `chromadb://localhost:8000`) - An OpenAI API key.
- Ollama server running locally (default URL: `http://localhost:11434`)
## Setup ## Setup
1. **Clone the repository** 1. **Clone the repository** (or copy the files into a directory).
```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
```
2. **Install dependencies** 2. **Install dependencies**
@@ -26,62 +24,39 @@ This project implements a Retrieval-Augmented Generation (RAG) agent that:
npm install 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 ```dotenv
CHROMA_URL=chromadb://localhost:8000 OPENAI_API_KEY=your_api_key_here
CHROMA_COLLECTION=rag_collection
OLLAMA_HOST=http://localhost:11434
OLLAMA_MODEL=llama3
OLLAMA_EMBEDDING_MODEL=nomic-embed-text
ADD_SAMPLE_DOCS=true
``` ```
- `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** 4. **Run the agent**
```bash ```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 ## Customization
npm start -- "What is LangChain?"
```
The agent will: - **Adding URLs**: Edit the `urls` array in `src/index.js` to index different web pages.
- Search the local ChromaDB collection. - **Chunk Size**: Adjust the `size` parameter in `chunkText` inside `src/webSearch.js` if you need larger or smaller chunks.
- Perform a DuckDuckGo web search. - **Model Parameters**: Modify temperature, max tokens, or model name in `src/agent.js`.
- 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
```
## Notes ## Notes
- The agent uses **LangChain 1.x** APIs. - The implementation strictly uses **ChromaDB** as the vector database; no other vector DBs are used.
- No Qdrant references are present; only ChromaDB is used. - All dependencies are declared in `package.json` and can be installed via `npm install`.
- The web search is performed via DuckDuckGos public JSON API (no API key required). - The OpenAI API key is loaded securely from the `.env` file using `dotenv`.
- The Ollama LLM is used for both embeddings and generation.
Feel free to extend the agent with additional retrievers or custom prompts as needed. ## License
MIT License
---
Enjoy building with RAG!
+5 -6
View File
@@ -1,17 +1,16 @@
{ {
"name": "rag-agent-chromadb-websearch", "name": "rag-agent-chromadb-websearch",
"version": "1.0.0", "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", "main": "src/index.js",
"type": "module", "type": "commonjs",
"scripts": { "scripts": {
"start": "node src/index.js" "start": "node src/index.js"
}, },
"dependencies": { "dependencies": {
"langchain": "^1.0.0",
"chromadb": "^1.0.0", "chromadb": "^1.0.0",
"node-fetch": "^3.3.0", "dotenv": "^16.4.5",
"dotenv": "^16.0.0", "node-fetch": "^2.6.7",
"ollama": "^0.1.0" "openai": "^4.12.0"
} }
} }
+32 -21
View File
@@ -1,26 +1,37 @@
import { Ollama } from 'langchain/llms/ollama'; const { OpenAI } = require('openai');
import { RetrievalQAChain } from 'langchain/chains/retrieval_qa'; const dotenv = require('dotenv');
import { BaseRetriever } from 'langchain/schema'; dotenv.config();
import { webSearch } from './webSearch.js';
export function createAgent(vectorStore) { const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
const llm = new Ollama({
model: process.env.OLLAMA_MODEL || 'llama3',
baseUrl: process.env.OLLAMA_HOST || 'http://localhost:11434',
});
class CombinedRetriever extends BaseRetriever { class Agent {
async getRelevantDocuments(query) { constructor(vectorStore) {
const chromaDocs = await vectorStore.similaritySearch(query, 3); this.vectorStore = vectorStore;
const webDocs = await webSearch(query, 3);
return [...chromaDocs, ...webDocs];
}
} }
const retriever = new CombinedRetriever(); async ask(question) {
const chain = RetrievalQAChain.fromLLMAndRetriever(llm, retriever, { const contextDocs = await this.vectorStore.query(question, 5);
returnSourceDocuments: true, 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.
return chain; 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 completion.choices[0].message.content.trim();
}
}
module.exports = Agent;
+42 -38
View File
@@ -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(); dotenv.config();
import { createVectorStore } from './vectorStore.js';
import { createAgent } from './agent.js';
import { Document } from 'langchain/document';
async function main() { async function main() {
const vectorStore = await createVectorStore(); console.log('Initializing RAG agent...');
const vectorStore = new VectorStore();
if (process.env.ADD_SAMPLE_DOCS === 'true') { // Example URLs to index
const sampleDocs = [ const urls = [
new Document({ 'https://en.wikipedia.org/wiki/Artificial_intelligence',
pageContent: 'https://en.wikipedia.org/wiki/ChromaDB',
'LangChain is a framework for building applications powered by language models.', 'https://en.wikipedia.org/wiki/OpenAI',
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' },
}),
];
await vectorStore.addDocuments(sampleDocs);
console.log('Sample documents added to ChromaDB.');
}
const agent = createAgent(vectorStore); console.log('Fetching and indexing web pages...');
const chunks = await searchAndChunk(urls);
await vectorStore.addDocuments(chunks);
console.log(`Indexed ${chunks.length} chunks.`);
const query = process.argv[2]; const agent = new Agent(vectorStore);
if (!query) {
console.error('Please provide a query as a command line argument.');
process.exit(1);
}
const result = await agent.invoke({ input: query }); const rl = readline.createInterface({
console.log('Answer:', result.output); input: process.stdin,
console.log( output: process.stdout,
'Sources:', });
result.sourceDocuments.map((d) => d.metadata.source).join(', ')
); const askQuestion = () => {
rl.question('\nEnter your question (or type "exit" to quit): ', async (answer) => {
if (answer.trim().toLowerCase() === 'exit') {
rl.close();
return;
}
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();
});
};
askQuestion();
} }
main().catch((err) => { main().catch((err) => {
console.error(err); console.error('Fatal error:', err);
process.exit(1); process.exit(1);
}); });
+41 -36
View File
@@ -1,47 +1,52 @@
import { ChromaClient } from 'chromadb'; const chromadb = require('chromadb');
import { OllamaEmbeddings } from 'langchain/embeddings/ollama'; const { OpenAI } = require('openai');
import { Document } from 'langchain/document'; const dotenv = require('dotenv');
dotenv.config();
export async function createVectorStore() { const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
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);
}
class VectorStore { class VectorStore {
constructor(collection, embeddings) { constructor() {
this.collection = collection; this.client = new chromadb.Client({ path: './chromadb' });
this.embeddings = embeddings; this.collection = this.client.getCollection('rag_collection');
} }
async addDocuments(docs) { async getEmbedding(text) {
const texts = docs.map((d) => d.pageContent); const response = await openai.embeddings.create({
const embeddings = await this.embeddings.embedDocuments(texts); model: 'text-embedding-ada-002',
await this.collection.addDocuments({ input: text,
documents: docs, });
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, embeddings,
ids,
}); });
} }
async similaritySearch(query, k = 4) { async query(queryText, k = 5) {
const embedding = await this.embeddings.embedQuery(query); const queryEmbedding = await this.getEmbedding(queryText);
const results = await this.collection.getNearestNeighbors({ const results = await this.collection.query({
queryEmbeddings: [embedding], queryEmbeddings: [queryEmbedding],
n: k, nResults: k,
}); });
const ids = results.ids[0];
const docs = await this.collection.getDocuments({ ids }); return results.documents[0];
return docs.map(
(doc) =>
new Document({
pageContent: doc.document,
metadata: doc.metadata,
})
);
} }
} }
module.exports = VectorStore;
+39 -21
View File
@@ -1,24 +1,42 @@
import fetch from 'node-fetch'; const fetch = require('node-fetch');
import { Document } from 'langchain/document';
export async function webSearch(query, limit = 3) { async function fetchPage(url) {
const url = `https://api.duckduckgo.com/?q=${encodeURIComponent( try {
query const res = await fetch(url);
)}&format=json&pretty=1`; if (!res.ok) {
const res = await fetch(url); console.warn(`Failed to fetch ${url}: ${res.statusText}`);
const data = await res.json(); return '';
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 (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 };