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

Uploaded Python 3

File details

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

File metadata

  • Download URL: retrievalmind-0.1.0.tar.gz
  • Upload date:
  • Size: 10.6 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.0.tar.gz
Algorithm Hash digest
SHA256 f5a1d6160f50d0611bec45de9ff81698188c8f49e9de80f7dcb0178a61fa03ec
MD5 fac70c4f0a921be419210d82a6cf6323
BLAKE2b-256 a9904f4bcc035a6e96ae766ee0bad84a924bfe918856921028787e86d5c22604

See more details on using hashes here.

File details

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

File metadata

  • Download URL: RetrievalMind-0.1.0-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.0-py3-none-any.whl
Algorithm Hash digest
SHA256 f93027ceebc9cd5444cd728f667bc925df194c2cec2b1687f5f803ea50e79028
MD5 b23780eb61c8591535736d55c1fafefc
BLAKE2b-256 16952c1faaf154a29383b0eae65950fbeb30e0f61228796f4e255b9063067a2a

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