# 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.