Add: Memory with multi user support with MongoDB

Author: ToastCoder

dependencies.txt

@@ -7,4 +7,5 @@ langchain-community
 langchain-text-splitters
 langchain-ollama
 pymupdf
-protobuf==3.20.3
+protobuf==3.20.3
+pymongo

main.py

@@ -3,8 +3,37 @@
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain_ollama import OllamaEmbeddings
+from pymongo import MongoClient
+import datetime
 import ollama
 
+# MongoDB Setup
+print("[INFO] Connecting to MongoDB...")
+try:
+    # Connecting to default local MongoDB Port
+    mongo_client = MongoClient("mongodb://localhost:27017/", serverSelectionTimeoutMS=2000)
+
+    # Forcing a connection test
+    mongo_client.server_info()
+    db = mongo_client["hr_assistant"]
+    chat_collection = db["chat_history"]
+    print("[OK] Connected to MongoDB!")
+except Exception as e:
+    print("[ERROR] Could not connect to MongoDB. Is the server running?")
+    exit(1)
+
+# Dynamic Session ID for handling users and their past questions
+print("\nWelcome to HR Buddy")
+SESSION_ID = input("Please enter your name to login: ").strip()
+
+# Fallback support as a guest session
+if not SESSION_ID:
+    SESSION_ID = "guest_user"
+    print("[INFO] No name provided. Logging in as 'guest_user'.")
+else:
+    print(f"[INFO] Hello, {SESSION_ID}! Fetching your records...")
+
+
 # Load and Chunk the Data
 print("[INFO] Reading HR Policy PDF...")
 loader = PyMuPDFLoader("rag_source/nasscom-hr-manual-2016.pdf") 
@@ -15,7 +44,8 @@
 chunks = text_splitter.split_documents(pages)
 
 # Embed and Store in Vector DB
-print("[INFO] Embedding text and building Chroma database...")
+print(f"[INFO] Chunking complete. Found {len(chunks)} text chunks.")
+print("[INFO] Sending chunks to Ollama to create embeddings...")
 # fast, local nomic model to turn text into math
 embeddings = OllamaEmbeddings(model="nomic-embed-text")
 vectorstore = Chroma.from_documents(documents=chunks, embedding=embeddings)
@@ -30,7 +60,14 @@
     while True:
         user_input = input("\n:> ")
 
-        # Semantic Search: Find the relevant policy paragraphs
+        # Get last 10 messages from the user
+        past_messages_cursor = chat_collection.find({"session_id": SESSION_ID}).sort("timestamp", -1).limit(10)
+
+        # Convert cursor to list and reverse for reading it from start to end
+        past_messages = list(past_messages_cursor)[::-1]
+        history_text = "\n".join([f"{msg['role']}: {msg['content']}" for msg in past_messages])
+
+        # Semantic Search: Find the relevant policy paragraphs from the RAG source
         docs = retriever.invoke(user_input)
         context_block = "\n\n".join([doc.page_content for doc in docs])
 
@@ -41,9 +78,13 @@
         
         If the Context does not contain the exact answer, simply reply: "I cannot find the answer to that in the current policy."
         
-        Context: 
+        --- PAST CONVERSATION MEMORY ---
+        {history_text}
+        
+        --- CURRENT HR CONTEXT --- 
         {context_block}
         
+        --- CURRENT QUESTION ---
         Question: {user_input}
         Answer:"""
 
@@ -52,7 +93,13 @@
             {'role': 'user', 'content': rag_prompt}
         ])
 
-        print(f"\nAI: {response['message']['content']}")
+        ai_answer = response['message']['content']
+        print(f"\nAI: {ai_answer}")
+
+        chat_collection.insert_many([
+            {"session_id": SESSION_ID, "role": "User", "content": user_input, "timestamp": datetime.datetime.now(datetime.timezone.utc)},
+            {"session_id": SESSION_ID, "role": "AI", "content": ai_answer, "timestamp": datetime.datetime.now(datetime.timezone.utc)}
+        ])
 
 except KeyboardInterrupt:
     print("\n\n[INFO] Closing Chat.")

run.sh

@@ -85,6 +85,10 @@ else
     echo "[OK] Python recognizes 'pymupdf'."
 fi
 
+# Check and configure MongoDB
+chmod +x scripts/mongodb_cfg.sh
+./scripts/mongodb_cfg.sh
+
 # Required Models for RAG
 echo "[INFO] Verifying AI models (Llama3.2 & Nomic)..."
 ollama pull llama3.2

scripts/mongodb_cfg.sh

@@ -0,0 +1,49 @@
+# System MongoDB Engine Check & Install
+echo "[INFO] Checking for MongoDB..."
+if ! command -v mongod &> /dev/null && ! command -v mongo &> /dev/null; then
+    echo "[WARNING] MongoDB not found. Detecting OS to install..."
+    
+    # Detect the operating system
+    if [ -f /etc/os-release ]; then
+        . /etc/os-release
+        case "$ID_LIKE" in
+            *fedora*|*rhel*|*centos*)
+                echo "[INFO] Detected RedHat/Fedora base. Installing via dnf..."
+                sudo dnf install -y mongodb-server
+                SVC_NAME="mongod"
+                ;;
+            *debian*|*ubuntu*)
+                echo "[INFO] Detected Debian/Ubuntu base. Installing via apt..."
+                sudo apt-get update
+                sudo apt-get install -y mongodb
+                SVC_NAME="mongodb"
+                ;;
+            *archlinux*)
+                echo "[INFO] Detected Arch base. Installing via pacman..."
+                sudo pacman -S --noconfirm mongodb
+                SVC_NAME="mongodb"
+                ;;
+            *)
+                echo "[ERROR] Auto-install for this OS is not supported. Please install MongoDB manually."
+                exit 1
+                ;;
+        esac
+        
+        echo "[OK] MongoDB installed successfully."
+        
+        # Automatically enable and start the background service
+        echo "[INFO] Starting MongoDB service..."
+        sudo systemctl enable --now $SVC_NAME
+    else
+        echo "[ERROR] Could not detect OS. Please install MongoDB manually."
+        exit 1
+    fi
+else
+    echo "[OK] MongoDB is already installed."
+    
+    # Ensure the service is actually awake and running
+    if ! systemctl is-active --quiet mongod && ! systemctl is-active --quiet mongodb; then
+         echo "[INFO] Waking up stopped MongoDB service..."
+         sudo systemctl start mongod || sudo systemctl start mongodb
+    fi
+fi