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

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.1.tar.gz (10.7 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.1-py3-none-any.whl (11.2 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: retrievalmind-0.1.1.tar.gz
  • Upload date:
  • Size: 10.7 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.1.tar.gz
Algorithm Hash digest
SHA256 dbffc31c85a1d77fff1ffaa1760c1424fae355394d94c749d108f39279674e8d
MD5 891342f0e08678acd09a3362ea19547e
BLAKE2b-256 5abfdc2f1cd71150869285c5808b575ffcbbd3cec0e0723a16c35ed862dc540c

See more details on using hashes here.

File details

Details for the file RetrievalMind-0.1.1-py3-none-any.whl.

File metadata

  • Download URL: RetrievalMind-0.1.1-py3-none-any.whl
  • Upload date:
  • Size: 11.2 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.1-py3-none-any.whl
Algorithm Hash digest
SHA256 b643c962a25d1c5d9d0e8610dc2b0b61f6395deb97a0f20d4ba96533fd426d1a
MD5 77ae67d4159c4155a13840b6bcf8e934
BLAKE2b-256 e374873e6cc34731fea5f27d7aa7f78e8c24b767dfc66d9963321ea1fc85d6a3

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