Skip to main content

A custom Retrieval-Augmented Generation (RAG) framework for AI Agent applications.

Project description

RetrievalMind

RetrievalMind is an end-to-end Retrieval-Augmented Generation (RAG) pipeline framework designed to efficiently ingest, embed, store, and retrieve information from documents. This project serves as a personal learning and educational framework for implementing RAG functionalities in AI agents.


Table of Contents


Project Overview

RetrievalMind allows you to:

  1. Ingest Documents – Load PDFs or text files into a structured format for processing.
  2. Generate Embeddings – Convert textual content into vector representations using pre-trained models.
  3. Vector Store Management – Store and manage embeddings in ChromaDB for efficient retrieval.
  4. Information Retrieval – Retrieve top relevant content for user queries using a similarity-based search.

This framework can be extended and reused for AI projects requiring knowledge retrieval or RAG-based reasoning, making it ideal for personal learning and experimentation.


Features

  • Supports PDF and text file ingestion with customizable chunking.
  • Generates semantic embeddings using all-miniLM-L6-v2 model.
  • Persistent vector storage with ChromaDB.
  • Flexible retrieval pipeline for natural language queries.
  • Modular structure for easy integration into AI agents and projects.

Project Structure

RetrievalMind/
├── data
│   ├── pdf
│   │   └── Company_Policy_Document.pdf
│   ├── text_files
│   │   ├── large_language_model.txt
│   │   ├── machine_learning.txt
│   │   └── python.txt
│   └── vector_store
│       ├── <UUID>/
│       └── chroma.sqlite3
├── main.py
├── src
│   ├── data_ingestion
│   │   ├── pdf_ingestor.py
│   │   └── text_ingestor.py
│   ├── embeddings_manager
│   │   └── embedding_manager.py
│   ├── rag_retriver
│   │   └── retriver.py
│   └── vector_store_manager
│       └── vector_store.py
  • data/pdf/ – Stores PDF documents for ingestion.
  • data/text_files/ – Stores plain text documents.
  • data/vector_store/ – ChromaDB storage for embeddings.
  • src/ – Core modular code for ingestion, embeddings, retrieval, and vector store management.
  • main.py – Demonstrates the full RAG pipeline workflow.

Installation

  1. Clone the repository:
git clone https://github.com/Himanshu7921/RetrievalMind
  1. Install dependencies:
pip install -r requirements.txt

Dependencies include: langchain, chromadb, sentence-transformers, PyMuPDF, etc.


Usage

from src.data_ingestion import PDFDocumentIngestor
from src.embeddings_manager import EmbeddingManager
from src.vector_store_manager import VectorStore
from src.rag_retriver import Retrieval

# Step 1: Load and chunk PDF
pdf_ingestor = PDFDocumentIngestor(file_path="data/pdf/Company_Policy_Document.pdf", loader_type='mu')
pdf_loader = pdf_ingestor.load_document()
document_chunks = pdf_loader.load()
chunk_texts = [chunk.page_content for chunk in document_chunks]

# Step 2: Generate embeddings
embedding_manager = EmbeddingManager()
chunk_embeddings = embedding_manager.generate_embeddings(chunk_texts)

# Step 3: Store chunks in vector store
vector_store = VectorStore(collection_name="company_policy_collection", persist_directory="data/vector_store", document_type="PDF")
vector_store.add_document(documents=document_chunks, embeddings=chunk_embeddings)

# Step 4: Initialize retrieval pipeline
retrieval_pipeline = Retrieval(vector_store=vector_store, embedding_manager=embedding_manager)
query_text = "What is the termination policy of the company?"
retrieved_results = retrieval_pipeline.retrieve(query=query_text, top_k=1, score_threshold=0)

if retrieved_results:
    print("Retrieved Content:\n", retrieved_results[0]['content'])
else:
    print("No relevant documents found.")

Example Output

[INFO] Vector store ready: 'company_policy_collection'
[INFO] Current document count: 4
[INFO] Successfully added 4 documents to 'company_policy_collection'.
[INFO] Total documents in collection: 8
Retrieved Document Content:
7. Customer Support Policy
7.1 Customer queries will be acknowledged within 24 hours of receipt.
...
8. Termination Policy
8.1 Employment may be terminated due to misconduct, non-performance, or policy violation.
...

Installation

From PyPI

Install the latest release directly from PyPI:

pip install RetrievalMind==0.1.3

Dependencies include: langchain, chromadb, sentence-transformers, PyMuPDF, etc.


Usage

Once installed, you can use RetrievalMind as a package:

from RetrievalMind.data_ingestion import PDFDocumentIngestor
from RetrievalMind.embeddings_manager import EmbeddingManager
from RetrievalMind.vector_store_manager import VectorStore
from RetrievalMind.rag_retriver import Retrieval

# Step 1: Load and chunk PDF
pdf_ingestor = PDFDocumentIngestor(file_path="data/pdf/Company_Policy_Document.pdf", loader_type='mu')
pdf_loader = pdf_ingestor.load_document()
document_chunks = pdf_loader.load()
chunk_texts = [chunk.page_content for chunk in document_chunks]

# Step 2: Generate embeddings
embedding_manager = EmbeddingManager()
chunk_embeddings = embedding_manager.generate_embeddings(chunk_texts)

# Step 3: Store chunks in vector store
vector_store = VectorStore(collection_name="company_policy_collection", persist_directory="data/vector_store", document_type="PDF")
vector_store.add_document(documents=document_chunks, embeddings=chunk_embeddings)

# Step 4: Initialize retrieval pipeline
retrieval_pipeline = Retrieval(vector_store=vector_store, embedding_manager=embedding_manager)
query_text = "What is the termination policy of the company?"
retrieved_results = retrieval_pipeline.retrieve(query=query_text, top_k=1, score_threshold=0)

if retrieved_results:
    print("Retrieved Content:\n", retrieved_results[0]['content'])
else:
    print("No relevant documents found.")

Future Enhancements

  • Support for multi-format document ingestion (Word, CSV).
  • Advanced query ranking with hybrid embeddings + keyword search.
  • Integration with LLMs for generative responses using retrieved knowledge.
  • Dashboard UI for visualizing vector store and retrieval results.

License

This project is for personal learning and educational purposes. No commercial use permitted.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

retrievalmind-0.1.3.tar.gz (11.9 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

retrievalmind-0.1.3-py3-none-any.whl (12.3 kB view details)

Uploaded Python 3

File details

Details for the file retrievalmind-0.1.3.tar.gz.

File metadata

  • Download URL: retrievalmind-0.1.3.tar.gz
  • Upload date:
  • Size: 11.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.7

File hashes

Hashes for retrievalmind-0.1.3.tar.gz
Algorithm Hash digest
SHA256 ac31c29ec80bdc199d617596063180dd801fdfcd01d4365cafd13102ce4055ad
MD5 0d8ca516c1e5830eddb769366972ef48
BLAKE2b-256 c73f11882319755e7982eb9a5b680ef2a87f3f25fce5fb178072d0974ec520fd

See more details on using hashes here.

File details

Details for the file retrievalmind-0.1.3-py3-none-any.whl.

File metadata

  • Download URL: retrievalmind-0.1.3-py3-none-any.whl
  • Upload date:
  • Size: 12.3 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.7

File hashes

Hashes for retrievalmind-0.1.3-py3-none-any.whl
Algorithm Hash digest
SHA256 96fab8e7ca5c1ccf653ac7a9d07a4c584bdf4b0a6017c9e1a131f05fd0459389
MD5 a7536370a52cfef4ecc060480de7f5ca
BLAKE2b-256 c84ead7cf4703e07b94d07eb4912fa876927e6ac68f6b5c57dc2c21c2c19316e

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page