RAG System

A simple Retrieval-Augmented Generation (RAG) demo built in Python. This project loads a local text or PDF document, splits it into chunks, embeds those chunks with sentence-transformers, stores them in ChromaDB, and answers questions with a lightweight text-generation model.

Features

• Load .txt and .pdf files through dedicated loaders
• Split source text into overlapping chunks
• Generate embeddings with all-MiniLM-L6-v2
• Store and search vectors locally with ChromaDB
• Route simple greetings to normal chat and knowledge questions to RAG mode
• Fall back to normal generation when retrieval confidence is weak

Project Structure

app.py                  # Main pipeline and interactive CLI
loaders/                # TXT and PDF document loaders
chunkers/               # Text chunking logic
embeddings/             # Embedding model wrapper
vectordb/               # Chroma persistence and search
router/                 # Query routing logic
llm/                    # Text generation wrapper
data_sample/            # Example source document
chroma_db/              # Local vector store created at runtime

How It Works

1. Load a local document.
2. Split the document into overlapping chunks.
3. Convert chunks into embeddings.
4. Save embeddings into a persistent Chroma collection.
5. Accept user questions in a loop.
6. Decide whether to use retrieval or normal chat.
7. Build a context-aware prompt when relevant content is found.

Quickstart

1. Create a virtual environment

python -m venv .venv
source .venv/bin/activate

2. Install dependencies

pip install -r requirements.txt

3. Point the app at your document

Edit file_path in app.py before running. The current default is:

file_path = "/content/data.txt"

For this repository, a working local example would be:

file_path = "data_sample/data.txt"

4. Run the app

python app.py

Type questions into the prompt and use exit to quit.

Dependencies

• chromadb
• pypdf
• sentence-transformers
• transformers and a compatible PyTorch install are also required by the generator code

Notes and Limitations

• The current app re-adds chunks to the same Chroma collection on each run. Clear chroma_db/ if you want a fresh index.
• llm/generator.py uses TinyLlama/TinyLlama-1.1B-Chat-v1.0, which may be slow on CPU-only machines.
• There is no automated test suite yet.
• Configuration is currently hardcoded in app.py, including chunk size, overlap, retrieval threshold, and file path.

Suggested Next Improvements

• Move runtime settings into environment variables or a config file
• Add pytest coverage for loaders, chunking, and routing
• Prevent duplicate vector inserts across runs
• Add support for multiple documents and top-k retrieval

License

Add a license file before publishing publicly.