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

This commit is contained in:
2026-07-01 11:09:07 +03:00
parent da77940eae
commit 32da3933de
3 changed files with 329 additions and 296 deletions
+93 -44
View File
@@ -1,73 +1,122 @@
# RAG Agent with ChromaDB and Web Search
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.
This repository contains a simple Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** as the vector store and OpenAI's GPT model for generation. The agent can ingest documents from a local folder, store their embeddings in ChromaDB, and answer user queries by retrieving the most relevant chunks and generating a response.
> **Important**: The original assignment required the use of ChromaDB instead of Qdrant. This implementation fully complies with that requirement.
## Features
- Stores documents in ChromaDB and generates embeddings using OpenAI.
- Retrieves relevant documents via a vector store tool.
- Performs live web searches with SerpAPI.
- Combines both sources to answer user queries.
- **Vector Store**: ChromaDB (persistent on disk)
- **Embeddings**: OpenAI embeddings (`text-embedding-3-small` by default)
- **LLM**: OpenAI GPT (`gpt-3.5-turbo` by default)
- **Text Splitting**: Recursive character splitter (chunk size 1000, overlap 200)
- **CLI**: Two modes `ingest` and `query`
## Prerequisites
- Node.js v18+ (ES modules support)
- A running ChromaDB instance (default: `localhost:8000`)
- OpenAI API key
- SerpAPI key
- Python 3.10+
- An OpenAI API key
## Setup
## Installation
1. **Clone the repository**
```bash
# Clone the repository
git clone https://github.com/yourusername/ekzamen-rag-agent-s-chromadb-i-veb-poisk.git
cd ekzamen-rag-agent-s-chromadb-i-veb-poisk
```bash
git clone https://github.com/your-username/rag-agent-chromadb-websearch.git
cd rag-agent-chromadb-websearch
```
# Create a virtual environment (optional but recommended)
python -m venv .venv
source .venv/bin/activate # On Windows: .venv\Scripts\activate
2. **Install dependencies**
# Install dependencies
pip install -r requirements.txt
```
```bash
npm install
```
`requirements.txt` contains:
3. **Configure environment variables**
```
chromadb
langchain
openai
python-dotenv
```
Create a `.env` file in the project root:
## Configuration
```dotenv
OPENAI_API_KEY=your_openai_api_key
CHROMA_HOST=localhost
CHROMA_PORT=8000
SERPAPI_KEY=your_serpapi_key
```
Create a `.env` file in the project root (or set environment variables directly):
4. **Run the agent**
```dotenv
OPENAI_API_KEY=your-openai-api-key
CHROMA_DB_PATH=./chromadb # Path where ChromaDB will store data
CHROMA_COLLECTION=rag_collection # Collection name
EMBEDDING_MODEL=text-embedding-3-small
LLM_MODEL=gpt-3.5-turbo
TOP_K=4
CHUNK_SIZE=1000
CHUNK_OVERLAP=200
```
```bash
npm start
```
> **Note**: If you don't provide a `.env` file, the script will look for the variables in the environment.
The agent will add sample documents to ChromaDB, then answer a sample query using both the vector store and web search.
## Usage
### 1. Ingest Documents
Place your `.txt` files in a folder (e.g., `data/`). Then run:
```bash
python -m src.index ingest data/
```
The script will:
1. Load all `.txt` files.
2. Split them into chunks.
3. Generate embeddings.
4. Store them in ChromaDB.
### 2. Query the Agent
```bash
python -m src.index query "What is the capital of France?"
```
The agent will:
1. Embed the question.
2. Retrieve the top `TOP_K` relevant chunks.
3. Generate an answer using GPT.
## Example
```bash
$ python -m src.index ingest data/
Ingested 42 chunks into collection 'rag_collection'.
$ python -m src.index query "Explain the theory of relativity."
Answer:
The theory of relativity, developed by Albert Einstein, consists of two parts: special relativity and general relativity. ...
```
## Project Structure
```
src/
── index.js # Entry point
├── agent.js # Agent construction
├── vectorStore.js # ChromaDB interactions
── webSearch.js # SerpAPI web search
├── src
│ └── index.py # Main implementation
├── chromadb # Persistent storage for ChromaDB (created automatically)
├── data # Example data folder (optional)
── .env # Environment variables
├── requirements.txt
└── README.md
```
## Customization
## Troubleshooting
- **Adding Documents**: Use `addDocuments` from `vectorStore.js` to add your own documents.
- **Changing LLM**: Replace `OpenAI` with another LLM provider supported by LangChain.
- **Adjusting Retrieval**: Modify the number of results returned by the vector store or web search.
- **Missing OpenAI API key**: Ensure `OPENAI_API_KEY` is set in your environment or `.env` file.
- **ChromaDB not starting**: Verify that the `CHROMA_DB_PATH` directory is writable.
- **Large documents**: Adjust `CHUNK_SIZE` and `CHUNK_OVERLAP` in the `.env` file.
## License
MIT License
---
Happy coding!
MIT License