94 lines
2.5 KiB
Markdown
94 lines
2.5 KiB
Markdown
# RAG Agent with ChromaDB
|
||
|
||
This project implements a simple Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** as the vector store.
|
||
The agent loads text documents, indexes them with embeddings, and answers user questions by retrieving relevant passages and generating a response with an OpenAI LLM.
|
||
|
||
## Features
|
||
|
||
- **ChromaDB** persistence for fast similarity search.
|
||
- OpenAI embeddings (`text-embedding-3-small`) for vector representation.
|
||
- OpenAI LLM (`gpt-4o-mini` by default) for answer generation.
|
||
- Simple command‑line interface to index documents and ask questions.
|
||
- Backward‑compatible API: `RAGAgent` exposes `add_documents`, `ask`, `get_document_count`, and `clear_store`.
|
||
|
||
## Requirements
|
||
|
||
```text
|
||
chromadb==0.4.24
|
||
langchain==0.1.13
|
||
openai==1.12.0
|
||
tqdm==4.66.1
|
||
pydantic==2.6.3
|
||
python-dotenv==1.0.1
|
||
```
|
||
|
||
Install them with:
|
||
|
||
```bash
|
||
pip install -r requirements.txt
|
||
```
|
||
|
||
## Setup
|
||
|
||
1. **OpenAI API Key**
|
||
The agent uses OpenAI services for embeddings and LLM.
|
||
Set your key in an environment variable:
|
||
|
||
```bash
|
||
export OPENAI_API_KEY="sk-..."
|
||
```
|
||
|
||
2. **Prepare Documents**
|
||
Place all `.txt` files you want to index in a directory, e.g., `data/`.
|
||
|
||
## Usage
|
||
|
||
```bash
|
||
python -m src.main --docs data/ --question "What is the capital of France?"
|
||
```
|
||
|
||
### Arguments
|
||
|
||
| Argument | Description | Default |
|
||
|----------|-------------|---------|
|
||
| `--docs` | Path to directory with `.txt` files. | **Required** |
|
||
| `--question` | The question to ask the agent. | **Required** |
|
||
| `--persist` | Directory where ChromaDB stores its data. | `./chromadb` |
|
||
| `--model` | OpenAI LLM model to use. | `gpt-4o-mini` |
|
||
| `--k` | Number of documents to retrieve for RAG. | `4` |
|
||
|
||
The first run will index all documents. Subsequent runs reuse the persisted index.
|
||
|
||
## API
|
||
|
||
```python
|
||
from src.vector_store import ChromaDBVectorStore
|
||
from src.agent import RAGAgent
|
||
from langchain.schema import Document
|
||
|
||
# Create vector store
|
||
store = ChromaDBVectorStore(persist_directory="./chromadb")
|
||
|
||
# Add documents
|
||
docs = [Document(page_content="Hello world", metadata={"source": "greeting.txt"})]
|
||
store.add_documents(docs)
|
||
|
||
# Create agent
|
||
agent = RAGAgent(vector_store=store)
|
||
|
||
# Ask a question
|
||
answer = agent.ask("What is this?")
|
||
print(answer)
|
||
```
|
||
|
||
## Testing
|
||
|
||
The project includes no automated tests, but you can manually verify:
|
||
|
||
1. Run the CLI with a small set of documents.
|
||
2. Ask a question that should be answered using the indexed content.
|
||
3. Verify that the answer references the correct context.
|
||
|
||
## License
|
||
|
||
MIT License. |