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

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:

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:

   export OPENAI_API_KEY="sk-..."
   

2. Prepare Documents
   Place all .txt files you want to index in a directory, e.g., data/.

Usage

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

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.