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

This commit is contained in:
2026-06-30 13:11:18 +03:00
parent 27dcde060a
commit 99e229be28
4 changed files with 458 additions and 132 deletions
+85 -37
View File
@@ -1,62 +1,110 @@
# RAG Agent with ChromaDB and Web Search
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.
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.
## 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.
- **Local Retrieval** Store and query embeddings in a persistent ChromaDB collection.
- **Web Search Fallback** If local retrieval fails to find relevant context, the agent performs a DuckDuckGo search and uses the snippets.
- **OpenAI Integration** Uses OpenAI embeddings (`text-embedding-ada-002`) and the `gpt-3.5-turbo` model for generation.
- **CLI** Simple command line interface for ingesting documents and asking questions.
## Prerequisites
- Node.js v18+ (supports native ES modules and `node-fetch` v2).
- An OpenAI API key.
- Python 3.9+
- An OpenAI API key
- (Optional) Internet access for web search
## Setup
## Installation
1. **Clone the repository** (or copy the files into a directory).
```bash
# 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
2. **Install dependencies**
# Create a virtual environment (recommended)
python -m venv .venv
source .venv/bin/activate # On Windows use `.venv\Scripts\activate`
```bash
npm install
```
# Install dependencies
pip install -r requirements.txt
```
3. **Configure environment**
`requirements.txt` contains:
Create a `.env` file in the project root (or edit the existing one) and add your OpenAI API key:
```
openai
chromadb
duckduckgo-search
beautifulsoup4
requests
```
```dotenv
OPENAI_API_KEY=your_api_key_here
```
## Environment Variables
4. **Run the agent**
| 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` |
```bash
npm start
```
Set them in your shell or create a `.env` file and load with `dotenv` (optional).
The script will:
- Fetch and index the example URLs.
- Prompt you to enter questions.
- Display answers generated by the RAG agent.
## Usage
## Customization
### Ingest Documents
- **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`.
Place your plain text files (`.txt`) in a folder, then run:
## Notes
```bash
python src/index.py ingest /path/to/text/files
```
- 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`.
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
MIT License
---
Enjoy building with RAG!
MIT License