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

This commit is contained in:
2026-06-30 16:52:49 +03:00
parent 99e229be28
commit da77940eae
6 changed files with 192 additions and 254 deletions
+51 -88
View File
@@ -1,110 +1,73 @@
# RAG Agent with ChromaDB and Web Search # RAG Agent with ChromaDB and Web Search
This project implements a Retrieval-Augmented Generation (RAG) agent that uses a local ChromaDB vector store for document retrieval and falls back to DuckDuckGo web search when the local store does not provide sufficient context. This project demonstrates a Retrieval-Augmented Generation (RAG) agent built with **LangChain 1.x**, **ChromaDB** as the vector store, and **SerpAPI** for web search integration.
## Features ## Features
- **Local Retrieval** Store and query embeddings in a persistent ChromaDB collection. - Stores documents in ChromaDB and generates embeddings using OpenAI.
- **Web Search Fallback** If local retrieval fails to find relevant context, the agent performs a DuckDuckGo search and uses the snippets. - Retrieves relevant documents via a vector store tool.
- **OpenAI Integration** Uses OpenAI embeddings (`text-embedding-ada-002`) and the `gpt-3.5-turbo` model for generation. - Performs live web searches with SerpAPI.
- **CLI** Simple command line interface for ingesting documents and asking questions. - Combines both sources to answer user queries.
## Prerequisites ## Prerequisites
- Python 3.9+ - Node.js v18+ (ES modules support)
- An OpenAI API key - A running ChromaDB instance (default: `localhost:8000`)
- (Optional) Internet access for web search - OpenAI API key
- SerpAPI key
## Installation ## Setup
```bash 1. **Clone the repository**
# Clone the repository
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
# Create a virtual environment (recommended) ```bash
python -m venv .venv git clone https://github.com/your-username/rag-agent-chromadb-websearch.git
source .venv/bin/activate # On Windows use `.venv\Scripts\activate` cd rag-agent-chromadb-websearch
```
# Install dependencies 2. **Install dependencies**
pip install -r requirements.txt
```
`requirements.txt` contains: ```bash
npm install
```
3. **Configure environment variables**
Create a `.env` file in the project root:
```dotenv
OPENAI_API_KEY=your_openai_api_key
CHROMA_HOST=localhost
CHROMA_PORT=8000
SERPAPI_KEY=your_serpapi_key
```
4. **Run the agent**
```bash
npm start
```
The agent will add sample documents to ChromaDB, then answer a sample query using both the vector store and web search.
## Project Structure
``` ```
openai src/
chromadb ├── index.js # Entry point
duckduckgo-search ├── agent.js # Agent construction
beautifulsoup4 ├── vectorStore.js # ChromaDB interactions
requests └── webSearch.js # SerpAPI web search
``` ```
## Environment Variables ## Customization
| Variable | Description | Example | - **Adding Documents**: Use `addDocuments` from `vectorStore.js` to add your own documents.
|----------|-------------|---------| - **Changing LLM**: Replace `OpenAI` with another LLM provider supported by LangChain.
| `OPENAI_API_KEY` | Your OpenAI API key | `sk-...` | - **Adjusting Retrieval**: Modify the number of results returned by the vector store or web search.
| `CHROMA_DB_PATH` | Directory where ChromaDB stores data | `./chromadb` |
| `CHROMA_COLLECTION_NAME` | Name of the collection | `rag_collection` |
| `TOP_K` | Number of top documents to retrieve | `5` |
| `SIMILARITY_THRESHOLD` | Minimum similarity to consider a document relevant | `0.5` |
| `WEB_SEARCH_MAX_RESULTS` | Max number of web snippets to fetch | `3` |
Set them in your shell or create a `.env` file and load with `dotenv` (optional).
## Usage
### Ingest Documents
Place your plain text files (`.txt`) in a folder, then run:
```bash
python src/index.py ingest /path/to/text/files
```
The script will read all `.txt` files, split them into chunks, embed them, and store them in ChromaDB.
### Ask a Question
```bash
python src/index.py ask "What is the capital of France?"
```
The agent will:
1. Query the local vector store for relevant passages.
2. If none are found above the similarity threshold, perform a DuckDuckGo search.
3. Combine the retrieved context into a prompt.
4. Call OpenAIs `gpt-3.5-turbo` to generate an answer.
## Example
```bash
$ python src/index.py ingest ./data
INFO:root:Added 12 documents to collection 'rag_collection'.
$ python src/index.py ask "Explain the theory of relativity."
Answer:
The theory of relativity, developed by Albert Einstein, consists of two parts: special relativity and general relativity. ...
```
## Testing
Unit tests are provided in the `tests/` directory. To run them:
```bash
pytest tests/
```
(If you don't have `pytest` installed, run `pip install pytest`.)
## Troubleshooting
- **No documents ingested** Ensure the folder path is correct and contains `.txt` files.
- **OpenAI errors** Verify that `OPENAI_API_KEY` is set and that you have sufficient quota.
- **Web search fails** Check your internet connection and that DuckDuckGo is reachable.
## License ## License
MIT License MIT License
---
Happy coding!
+6 -5
View File
@@ -1,16 +1,17 @@
{ {
"name": "rag-agent-chromadb-websearch", "name": "rag-agent-chromadb-websearch",
"version": "1.0.0", "version": "1.0.0",
"description": "RAG agent using ChromaDB and OpenAI API with web search", "description": "RAG agent using ChromaDB and web search with LangChain 1.x",
"main": "src/index.js", "main": "src/index.js",
"type": "commonjs", "type": "module",
"scripts": { "scripts": {
"start": "node src/index.js" "start": "node src/index.js"
}, },
"dependencies": { "dependencies": {
"chromadb": "^1.0.0", "chromadb": "^1.0.0",
"dotenv": "^16.4.5", "dotenv": "^16.0.3",
"node-fetch": "^2.6.7", "langchain": "^1.0.0",
"openai": "^4.12.0" "openai": "^3.3.0",
"serpapi": "^2.0.0"
} }
} }
+48 -31
View File
@@ -1,37 +1,54 @@
const { OpenAI } = require('openai'); import { OpenAI } from 'langchain/llms/openai';
const dotenv = require('dotenv'); import { Tool } from 'langchain/tools';
dotenv.config(); import { AgentExecutor } from 'langchain/agents';
import { query } from './vectorStore.js';
import { webSearch } from './webSearch.js';
const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY }); const llm = new OpenAI({
temperature: 0,
openAIApiKey: process.env.OPENAI_API_KEY,
});
class Agent { /**
constructor(vectorStore) { * Tool that retrieves relevant documents from ChromaDB.
this.vectorStore = vectorStore; */
} const vectorStoreTool = new Tool({
name: 'VectorStore',
description: 'Retrieve relevant documents from the vector store.',
func: async (input) => {
const result = await query(input);
if (!result.documents || result.documents.length === 0) {
return 'No relevant documents found.';
}
return result.documents.map((doc, idx) => `(${idx + 1}) ${doc}`).join('\n');
},
});
async ask(question) { /**
const contextDocs = await this.vectorStore.query(question, 5); * Tool that performs a web search.
const context = contextDocs.join('\n\n'); */
const prompt = ` const webSearchTool = new Tool({
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. name: 'WebSearch',
description: 'Search the web for up-to-date information.',
func: async (input) => {
const results = await webSearch(input);
if (!results || results.length === 0) {
return 'No web results found.';
}
return results
.map((r, idx) => `(${idx + 1}) ${r.title}: ${r.link}`)
.join('\n');
},
});
Context: const tools = [vectorStoreTool, webSearchTool];
${context}
Question: /**
${question} * Create and return a LangChain 1.x AgentExecutor.
Answer: */
`; export async function createAgent() {
const agent = await AgentExecutor.fromLLMAndTools(llm, tools, {
const completion = await openai.chat.completions.create({ verbose: true,
model: 'gpt-3.5-turbo', });
messages: [{ role: 'user', content: prompt }], return agent;
temperature: 0.2,
max_tokens: 300,
});
return completion.choices[0].message.content.trim();
}
} }
module.exports = Agent;
+16 -42
View File
@@ -1,54 +1,28 @@
const readline = require('readline'); import dotenv from 'dotenv';
const VectorStore = require('./vectorStore');
const Agent = require('./agent');
const { searchAndChunk } = require('./webSearch');
const dotenv = require('dotenv');
dotenv.config(); dotenv.config();
import { createAgent } from './agent.js';
import { addDocuments } from './vectorStore.js';
async function main() { async function main() {
console.log('Initializing RAG agent...'); // Add sample documents to the vector store
const vectorStore = new VectorStore(); const sampleDocs = [
'LangChain is a framework for building applications powered by language models.',
// Example URLs to index 'ChromaDB is an open-source vector database that stores embeddings.',
const urls = [ 'Web search can provide up-to-date information that may not be in the vector store.',
'https://en.wikipedia.org/wiki/Artificial_intelligence',
'https://en.wikipedia.org/wiki/ChromaDB',
'https://en.wikipedia.org/wiki/OpenAI',
]; ];
await addDocuments(sampleDocs);
console.log('Fetching and indexing web pages...'); const agent = await createAgent();
const chunks = await searchAndChunk(urls);
await vectorStore.addDocuments(chunks);
console.log(`Indexed ${chunks.length} chunks.`);
const agent = new Agent(vectorStore); const query = 'Explain how LangChain can use ChromaDB and web search together.';
const result = await agent.invoke({ input: query });
const rl = readline.createInterface({ console.log('\n=== Agent Response ===\n');
input: process.stdin, console.log(result.output);
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) => { main().catch((err) => {
console.error('Fatal error:', err); console.error('Error running the agent:', err);
process.exit(1); process.exit(1);
}); });
+51 -45
View File
@@ -1,52 +1,58 @@
const chromadb = require('chromadb'); import { ChromaClient } from 'chromadb';
const { OpenAI } = require('openai'); import { OpenAIEmbeddings } from 'langchain/embeddings/openai';
const dotenv = require('dotenv');
dotenv.config();
const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY }); const chroma = new ChromaClient({
host: process.env.CHROMA_HOST || 'localhost',
port: parseInt(process.env.CHROMA_PORT, 10) || 8000,
});
class VectorStore { const embeddings = new OpenAIEmbeddings({
constructor() { openAIApiKey: process.env.OPENAI_API_KEY,
this.client = new chromadb.Client({ path: './chromadb' }); });
this.collection = this.client.getCollection('rag_collection');
}
async getEmbedding(text) { const COLLECTION_NAME = 'rag_collection';
const response = await openai.embeddings.create({
model: 'text-embedding-ada-002',
input: text,
});
return response.data[0].embedding;
}
async addDocuments(chunks) { /**
const documents = []; * Ensure the collection exists in ChromaDB.
const embeddings = []; */
const ids = []; async function ensureCollection() {
const collections = await chroma.listCollections();
for (const chunk of chunks) { if (!collections.includes(COLLECTION_NAME)) {
const embedding = await this.getEmbedding(chunk); await chroma.createCollection({ name: COLLECTION_NAME });
documents.push(chunk);
embeddings.push(embedding);
ids.push(`${Date.now()}-${Math.random()}`);
}
await this.collection.add({
documents,
embeddings,
ids,
});
}
async query(queryText, k = 5) {
const queryEmbedding = await this.getEmbedding(queryText);
const results = await this.collection.query({
queryEmbeddings: [queryEmbedding],
nResults: k,
});
return results.documents[0];
} }
} }
module.exports = VectorStore; /**
* Add an array of documents to the vector store.
* @param {string[]} docs
*/
export async function addDocuments(docs) {
await ensureCollection();
const ids = docs.map((_, idx) => `doc-${Date.now()}-${idx}`);
const embeddingsResult = await embeddings.embedDocuments(docs);
await chroma.add({
collection_name: COLLECTION_NAME,
ids,
documents: docs,
embeddings: embeddingsResult,
});
}
/**
* Query the vector store for the most relevant documents.
* @param {string} queryText
* @param {number} nResults
* @returns {Promise<{documents: string[]}>}
*/
export async function query(queryText, nResults = 3) {
await ensureCollection();
const embedding = await embeddings.embedQuery(queryText);
const results = await chroma.query({
collection_name: COLLECTION_NAME,
query_embeddings: [embedding],
n_results: nResults,
});
// ChromaDB returns an array of objects; extract documents
const docs = results.documents || [];
return { documents: docs };
}
+17 -40
View File
@@ -1,42 +1,19 @@
const fetch = require('node-fetch'); import { GoogleSearchResults } from 'serpapi';
async function fetchPage(url) { /**
try { * Perform a web search using SerpAPI.
const res = await fetch(url); * @param {string} query
if (!res.ok) { * @returns {Promise<Array<{title: string, link: string}>>}
console.warn(`Failed to fetch ${url}: ${res.statusText}`); */
return ''; export async function webSearch(query) {
} const search = new GoogleSearchResults({
const html = await res.text(); api_key: process.env.SERPAPI_KEY,
// Strip HTML tags engine: 'google',
const text = html.replace(/<[^>]*>/g, ' '); });
// Collapse whitespace const results = await search.get({ q: query, num: 5 });
const cleaned = text.replace(/\s+/g, ' ').trim(); // Return only the organic results with title and link
return cleaned; return (results.organic_results || []).map((r) => ({
} catch (err) { title: r.title,
console.error(`Error fetching ${url}:`, err); link: r.link,
return ''; }));
}
} }
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 };