Feat: Add PostgreSQL-backed chat memory using Neon DB

bot.py

@@ -5,6 +5,8 @@ import logging
 import socket
 import time
 from io import BytesIO
+from datetime import datetime
+import psycopg2
 from dotenv import load_dotenv
 from telegram import Update
 from telegram.ext import Application, MessageHandler, filters, ContextTypes
@@ -29,6 +31,79 @@ OLLAMA_HOST = os.getenv("OLLAMA_HOST", "http://localhost:11434")
 OLLAMA_API_KEY = os.getenv("OLLAMA_API_KEY", "")  # Optional: for cloud/authenticated services
 VISION_MODEL = os.getenv("VISION_MODEL", "llava")  # Model for image analysis
 CHAT_MODEL = os.getenv("CHAT_MODEL", "mistral")    # Model for quiz generation
+DATABASE_URL = os.getenv("DATABASE_URL")  # Neon connection string
+
+
+def init_db():
+    """Initialize the database table."""
+    if not DATABASE_URL:
+        logger.warning("DATABASE_URL not set. Chat memory will be disabled.")
+        return
+
+    try:
+        conn = psycopg2.connect(DATABASE_URL)
+        cur = conn.cursor()
+        cur.execute("""
+            CREATE TABLE IF NOT EXISTS chat_history (
+                id SERIAL PRIMARY KEY,
+                chat_id BIGINT NOT NULL,
+                role VARCHAR(10) NOT NULL,
+                content TEXT NOT NULL,
+                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            );
+        """)
+        conn.commit()
+        cur.close()
+        conn.close()
+        logger.info("Database initialized successfully.")
+    except Exception as e:
+        logger.error(f"Error initializing database: {e}")
+
+
+def save_chat_message(chat_id: int, role: str, content: str):
+    """Save a chat message to the database."""
+    if not DATABASE_URL:
+        return
+
+    try:
+        conn = psycopg2.connect(DATABASE_URL)
+        cur = conn.cursor()
+        cur.execute(
+            "INSERT INTO chat_history (chat_id, role, content) VALUES (%s, %s, %s)",
+            (chat_id, role, content)
+        )
+        conn.commit()
+        cur.close()
+        conn.close()
+    except Exception as e:
+        logger.error(f"Error saving chat message: {e}")
+
+
+def get_chat_history(chat_id: int, limit: int = 10):
+    """Get recent chat history for a chat_id."""
+    if not DATABASE_URL:
+        return []
+
+    try:
+        conn = psycopg2.connect(DATABASE_URL)
+        cur = conn.cursor()
+        cur.execute(
+            """
+            SELECT role, content FROM chat_history 
+            WHERE chat_id = %s 
+            ORDER BY timestamp DESC 
+            LIMIT %s
+            """,
+            (chat_id, limit)
+        )
+        rows = cur.fetchall()
+        cur.close()
+        conn.close()
+        # Return reversed list (oldest to newest)
+        return [{"role": row[0], "content": row[1]} for row in rows][::-1]
+    except Exception as e:
+        logger.error(f"Error getting chat history: {e}")
+        return []
 
 
 async def extract_text_from_image(image_bytes: bytes) -> str:
@@ -325,8 +400,8 @@ async def handle_file(update: Update, context: ContextTypes.DEFAULT_TYPE):
         await processing_msg.edit_text(f"❌ حصل مشكلة: {str(e)[:100]}")
 
 
-async def generate_chat_response(text: str) -> str:
-    """Generate a conversational response using Ollama API."""
+async def generate_chat_response(chat_id: int, text: str) -> str:
+    """Generate a conversational response using Ollama API with history."""
     system_prompt = """You are the AI Quiz Bot Assistant. 🧠
 Your goal is to be friendly, helpful, and encourage users to study.
 - If the user greets you, reply warmly in the same language (Arabic or English).
@@ -334,13 +409,22 @@ Your goal is to be friendly, helpful, and encourage users to study.
 - Be concise and fun.
 - You are NOT generating a quiz right now, just chatting."""
 
+    # Get recent history (limit 6 messages for context)
+    history = get_chat_history(chat_id, limit=6)
+    
+    messages = [{"role": "system", "content": system_prompt}]
+    
+    # Add history
+    for msg in history:
+        messages.append(msg)
+        
+    # Add current user message
+    messages.append({"role": "user", "content": text})
+
     try:
         payload = {
             "model": CHAT_MODEL,
-            "messages": [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": text}
-            ],
+            "messages": messages,
             "stream": False
         }
         
@@ -374,9 +458,18 @@ async def handle_text(update: Update, context: ContextTypes.DEFAULT_TYPE):
         
     # If text is short (< 50 chars), treat it as conversation
     if len(text.strip()) < 50:
+        # Save user message to memory
+        save_chat_message(chat_id, "user", text)
+        
         # Show typing indicator
         await context.bot.send_chat_action(chat_id=chat_id, action="typing")
-        response = await generate_chat_response(text)
+        
+        # Generate response with history
+        response = await generate_chat_response(chat_id, text)
+        
+        # Save bot response to memory
+        save_chat_message(chat_id, "assistant", response)
+        
         await message.reply_text(response)
         return
     
@@ -513,6 +606,9 @@ def wait_for_network():
 
 def main():
     """Start the bot."""
+    # Initialize database
+    init_db()
+
     # Fix for running in background thread: Ensure an event loop exists
     try:
         loop = asyncio.get_event_loop()

requirements.txt

@@ -4,3 +4,4 @@ python-dotenv>=1.0.0
 httpx>=0.24.0
 fastapi
 uvicorn
+psycopg2-binary