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
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 websearch 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 GPT3.5Turbo.
- **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 DuckDuckGos 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!
+5 -6
View File
@@ -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"
}
}
+32 -21
View File
@@ -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.
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();
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' },
}),
];
await vectorStore.addDocuments(sampleDocs);
console.log('Sample documents added to ChromaDB.');
}
// 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',
];
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];
if (!query) {
console.error('Please provide a query as a command line argument.');
process.exit(1);
}
const agent = new Agent(vectorStore);
const result = await agent.invoke({ input: query });
console.log('Answer:', result.output);
console.log(
'Sources:',
result.sourceDocuments.map((d) => d.metadata.source).join(', ')
);
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;
}
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) => {
console.error(err);
console.error('Fatal error:', err);
process.exit(1);
});
+41 -36
View File
@@ -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;
+39 -21
View File
@@ -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`;
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 },
})
);
async function fetchPage(url) {
try {
const res = await fetch(url);
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 };