feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'
This commit is contained in:
@@ -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
|
||||||
|
|
||||||
|
1. **Clone the repository**
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
# Clone the repository
|
git clone https://github.com/your-username/rag-agent-chromadb-websearch.git
|
||||||
git clone https://git.brojs.ru/kuzakhmetovartur/ekzamen-rag-agent-s-chromadb-i-veb-poisk.git
|
cd rag-agent-chromadb-websearch
|
||||||
cd ekzamen-rag-agent-s-chromadb-i-veb-poisk
|
|
||||||
|
|
||||||
# Create a virtual environment (recommended)
|
|
||||||
python -m venv .venv
|
|
||||||
source .venv/bin/activate # On Windows use `.venv\Scripts\activate`
|
|
||||||
|
|
||||||
# Install dependencies
|
|
||||||
pip install -r requirements.txt
|
|
||||||
```
|
```
|
||||||
|
|
||||||
`requirements.txt` contains:
|
2. **Install dependencies**
|
||||||
|
|
||||||
```
|
|
||||||
openai
|
|
||||||
chromadb
|
|
||||||
duckduckgo-search
|
|
||||||
beautifulsoup4
|
|
||||||
requests
|
|
||||||
```
|
|
||||||
|
|
||||||
## Environment Variables
|
|
||||||
|
|
||||||
| Variable | Description | Example |
|
|
||||||
|----------|-------------|---------|
|
|
||||||
| `OPENAI_API_KEY` | Your OpenAI API key | `sk-...` |
|
|
||||||
| `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
|
```bash
|
||||||
python src/index.py ingest /path/to/text/files
|
npm install
|
||||||
```
|
```
|
||||||
|
|
||||||
The script will read all `.txt` files, split them into chunks, embed them, and store them in ChromaDB.
|
3. **Configure environment variables**
|
||||||
|
|
||||||
### Ask a Question
|
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
|
```bash
|
||||||
python src/index.py ask "What is the capital of France?"
|
npm start
|
||||||
```
|
```
|
||||||
|
|
||||||
The agent will:
|
The agent will add sample documents to ChromaDB, then answer a sample query using both the vector store and web search.
|
||||||
|
|
||||||
1. Query the local vector store for relevant passages.
|
## Project Structure
|
||||||
2. If none are found above the similarity threshold, perform a DuckDuckGo search.
|
|
||||||
3. Combine the retrieved context into a prompt.
|
|
||||||
4. Call OpenAI’s `gpt-3.5-turbo` to generate an answer.
|
|
||||||
|
|
||||||
## Example
|
```
|
||||||
|
src/
|
||||||
```bash
|
├── index.js # Entry point
|
||||||
$ python src/index.py ingest ./data
|
├── agent.js # Agent construction
|
||||||
INFO:root:Added 12 documents to collection 'rag_collection'.
|
├── vectorStore.js # ChromaDB interactions
|
||||||
|
└── webSearch.js # SerpAPI web search
|
||||||
$ 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
|
## Customization
|
||||||
|
|
||||||
Unit tests are provided in the `tests/` directory. To run them:
|
- **Adding Documents**: Use `addDocuments` from `vectorStore.js` to add your own documents.
|
||||||
|
- **Changing LLM**: Replace `OpenAI` with another LLM provider supported by LangChain.
|
||||||
```bash
|
- **Adjusting Retrieval**: Modify the number of results returned by the vector store or web search.
|
||||||
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
@@ -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"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
+49
-32
@@ -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,
|
||||||
class Agent {
|
openAIApiKey: process.env.OPENAI_API_KEY,
|
||||||
constructor(vectorStore) {
|
|
||||||
this.vectorStore = vectorStore;
|
|
||||||
}
|
|
||||||
|
|
||||||
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 completion.choices[0].message.content.trim();
|
/**
|
||||||
}
|
* Tool that retrieves relevant documents from ChromaDB.
|
||||||
|
*/
|
||||||
|
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');
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
module.exports = Agent;
|
/**
|
||||||
|
* Tool that performs a web search.
|
||||||
|
*/
|
||||||
|
const webSearchTool = new Tool({
|
||||||
|
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');
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
const tools = [vectorStoreTool, webSearchTool];
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Create and return a LangChain 1.x AgentExecutor.
|
||||||
|
*/
|
||||||
|
export async function createAgent() {
|
||||||
|
const agent = await AgentExecutor.fromLLMAndTools(llm, tools, {
|
||||||
|
verbose: true,
|
||||||
|
});
|
||||||
|
return agent;
|
||||||
|
}
|
||||||
+16
-42
@@ -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);
|
||||||
});
|
});
|
||||||
+48
-42
@@ -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',
|
||||||
class VectorStore {
|
port: parseInt(process.env.CHROMA_PORT, 10) || 8000,
|
||||||
constructor() {
|
|
||||||
this.client = new chromadb.Client({ path: './chromadb' });
|
|
||||||
this.collection = this.client.getCollection('rag_collection');
|
|
||||||
}
|
|
||||||
|
|
||||||
async getEmbedding(text) {
|
|
||||||
const response = await openai.embeddings.create({
|
|
||||||
model: 'text-embedding-ada-002',
|
|
||||||
input: text,
|
|
||||||
});
|
});
|
||||||
return response.data[0].embedding;
|
|
||||||
|
const embeddings = new OpenAIEmbeddings({
|
||||||
|
openAIApiKey: process.env.OPENAI_API_KEY,
|
||||||
|
});
|
||||||
|
|
||||||
|
const COLLECTION_NAME = 'rag_collection';
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Ensure the collection exists in ChromaDB.
|
||||||
|
*/
|
||||||
|
async function ensureCollection() {
|
||||||
|
const collections = await chroma.listCollections();
|
||||||
|
if (!collections.includes(COLLECTION_NAME)) {
|
||||||
|
await chroma.createCollection({ name: COLLECTION_NAME });
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async addDocuments(chunks) {
|
/**
|
||||||
const documents = [];
|
* Add an array of documents to the vector store.
|
||||||
const embeddings = [];
|
* @param {string[]} docs
|
||||||
const ids = [];
|
*/
|
||||||
|
export async function addDocuments(docs) {
|
||||||
for (const chunk of chunks) {
|
await ensureCollection();
|
||||||
const embedding = await this.getEmbedding(chunk);
|
const ids = docs.map((_, idx) => `doc-${Date.now()}-${idx}`);
|
||||||
documents.push(chunk);
|
const embeddingsResult = await embeddings.embedDocuments(docs);
|
||||||
embeddings.push(embedding);
|
await chroma.add({
|
||||||
ids.push(`${Date.now()}-${Math.random()}`);
|
collection_name: COLLECTION_NAME,
|
||||||
}
|
|
||||||
|
|
||||||
await this.collection.add({
|
|
||||||
documents,
|
|
||||||
embeddings,
|
|
||||||
ids,
|
ids,
|
||||||
|
documents: docs,
|
||||||
|
embeddings: embeddingsResult,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
async query(queryText, k = 5) {
|
/**
|
||||||
const queryEmbedding = await this.getEmbedding(queryText);
|
* Query the vector store for the most relevant documents.
|
||||||
const results = await this.collection.query({
|
* @param {string} queryText
|
||||||
queryEmbeddings: [queryEmbedding],
|
* @param {number} nResults
|
||||||
nResults: k,
|
* @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
|
||||||
return results.documents[0];
|
const docs = results.documents || [];
|
||||||
|
return { documents: docs };
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|
||||||
module.exports = VectorStore;
|
|
||||||
+17
-40
@@ -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({
|
||||||
|
api_key: process.env.SERPAPI_KEY,
|
||||||
|
engine: 'google',
|
||||||
|
});
|
||||||
|
const results = await search.get({ q: query, num: 5 });
|
||||||
|
// Return only the organic results with title and link
|
||||||
|
return (results.organic_results || []).map((r) => ({
|
||||||
|
title: r.title,
|
||||||
|
link: r.link,
|
||||||
|
}));
|
||||||
}
|
}
|
||||||
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 '';
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
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 };
|
|
||||||
Reference in New Issue
Block a user