Update main.py

main.py

@@ -2,18 +2,23 @@ import uuid
 import json
 import asyncio
 import httpx
+import sqlite3
+import time
 from pathlib import Path
+from datetime import datetime
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, UploadFile, File, HTTPException
-from fastapi.responses import HTMLResponse
+from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 from fastapi.middleware.cors import CORSMiddleware
+from typing import List, Dict, Set
 
-app = FastAPI(title="Antaram Chat AI")
+app = FastAPI(title="Antaram Chat AI Pro")
 
 # --- Configuration ---
 AI_URL = "https://sarveshpatel-unsloth0-6bqwen.hf.space/chat/raw"
 MAX_TOKENS = 30000
+SYSTEM_PROMPT = "You are Antaram AI. Answer concisely, accurately, and professionally. Use emojis very sparingly (max 1 per response). Do not hallucinate."
 
 # CORS
 app.add_middleware(
@@ -32,163 +37,212 @@ UPLOAD_DIR.mkdir(exist_ok=True)
 templates = Jinja2Templates(directory="templates")
 app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads")
 
-# State
-active_rooms = {}  # {room_id: [websocket_objects]}
-room_files = {}
+# --- Database Setup (SQLite) ---
+DB_PATH = "chat.db"
+
+def init_db():
+    with sqlite3.connect(DB_PATH) as conn:
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS messages (
+                id TEXT PRIMARY KEY,
+                room_id TEXT,
+                username TEXT,
+                content TEXT,
+                file_url TEXT,
+                file_type TEXT,
+                reply_to_id TEXT,
+                reply_content TEXT,
+                timestamp REAL
+            )
+        """)
+        conn.commit()
+
+init_db()
+
+# --- State ---
+active_rooms: Dict[str, List[WebSocket]] = {}
+active_users: Dict[str, Set[str]] = {} # room_id -> set of usernames
+
+# --- Helpers ---
+
+def save_message(room_id, username, content, file_data=None, reply_to_id=None, reply_content=None):
+    msg_id = str(uuid.uuid4())
+    ts = time.time()
+    file_url = file_data['file_url'] if file_data else None
+    file_type = file_data['file_type'] if file_data else None
+    
+    with sqlite3.connect(DB_PATH) as conn:
+        conn.execute(
+            "INSERT INTO messages VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
+            (msg_id, room_id, username, content, file_url, file_type, reply_to_id, reply_content, ts)
+        )
+    return {
+        "id": msg_id, "username": username, "text": content, 
+        "file": file_data, "reply_to": reply_to_id, "reply_content": reply_content,
+        "timestamp": ts, "type": "message"
+    }
+
+def get_history(room_id, limit=50):
+    with sqlite3.connect(DB_PATH) as conn:
+        conn.row_factory = sqlite3.Row
+        cursor = conn.execute(
+            "SELECT * FROM messages WHERE room_id = ? ORDER BY timestamp ASC LIMIT ?", 
+            (room_id, limit)
+        )
+        rows = cursor.fetchall()
+        
+    history = []
+    for row in rows:
+        file_data = None
+        if row['file_url']:
+            file_data = {"file_url": row['file_url'], "file_type": row['file_type'], "original_name": "File"}
+            
+        history.append({
+            "type": "message",
+            "id": row['id'],
+            "username": row['username'],
+            "text": row['content'],
+            "file": file_data,
+            "reply_to": row['reply_to_id'],
+            "reply_content": row['reply_content'],
+            "timestamp": row['timestamp']
+        })
+    return history
+
+# --- AI Logic ---
+
+async def stream_antaram_ai(room_id: str, prompt: str, context_msg: str = None):
+    """Streams AI response. If context_msg is provided (reply), it's added to prompt."""
+    
+    # 1. Notify Start
+    await broadcast_to_room(room_id, json.dumps({"type": "ai_start", "username": "Antaram AI"}))
+
+    final_prompt = f"{SYSTEM_PROMPT}\n\n"
+    if context_msg:
+        final_prompt += f"Context (User replied to this): {context_msg}\n"
+    
+    clean_user_prompt = prompt.replace("@antaram.ai", "").strip()
+    final_prompt += f"User Query: {clean_user_prompt}"
+
+    full_response = ""
+    
+    try:
+        async with httpx.AsyncClient(timeout=60) as client:
+            async with client.stream("POST", AI_URL, json={"prompt": final_prompt, "max_tokens": MAX_TOKENS}) as response:
+                async for chunk in response.aiter_text():
+                    clean_chunk = chunk.replace("<think>", "").replace("</think>", "")
+                    if clean_chunk:
+                        full_response += clean_chunk
+                        await broadcast_to_room(room_id, json.dumps({
+                            "type": "ai_chunk", "chunk": clean_chunk
+                        }))
+    except Exception as e:
+        await broadcast_to_room(room_id, json.dumps({"type": "error", "message": "AI Unreachable"}))
+
+    # 2. Notify End & Save AI response to DB
+    save_message(room_id, "Antaram AI", full_response)
+    await broadcast_to_room(room_id, json.dumps({"type": "ai_end"}))
 
 # --- Routes ---
 
 @app.get("/", response_class=HTMLResponse)
 async def home(request: Request):
-    """Home page"""
     return templates.TemplateResponse("index.html", {"request": request, "room_id": None})
 
 @app.get("/room/{room_id}", response_class=HTMLResponse)
 async def dynamic_room(request: Request, room_id: str):
-    """Dynamic route for sharing links directly"""
     return templates.TemplateResponse("index.html", {"request": request, "room_id": room_id.upper()})
 
 @app.post("/create-room")
 async def create_room():
     room_id = str(uuid.uuid4())[:8].upper()
-    active_rooms[room_id] = []
-    room_files[room_id] = []
     return {"room_id": room_id, "success": True}
 
 @app.post("/upload-file/{room_id}")
 async def upload_file(room_id: str, file: UploadFile = File(...)):
-    # Allow upload if room is active OR just created
-    if room_id not in room_files:
-        if room_id in active_rooms:
-            room_files[room_id] = []
-        else:
-            # Auto-create entry if connecting via file upload first (rare edge case)
-            room_files[room_id] = []
-    
     file_ext = Path(file.filename).suffix
     unique_name = f"{uuid.uuid4().hex}{file_ext}"
     file_path = UPLOAD_DIR / unique_name
-    
     content = await file.read()
     with open(file_path, "wb") as f:
         f.write(content)
     
-    file_info = {
-        "filename": unique_name,
-        "original_name": file.filename,
-        "file_type": file.content_type or "application/octet-stream",
-        "size": len(content),
-        "file_url": f"/uploads/{unique_name}"
-    }
-    room_files[room_id].append(file_info)
-    return {"success": True, "file_info": file_info}
-
-# --- AI Logic ---
-
-async def stream_antaram_ai(room_id: str, prompt: str):
-    """Streams AI response to a specific room"""
-    
-    # 1. Notify room that AI is thinking
-    await broadcast_to_room(room_id, json.dumps({
-        "type": "ai_start",
-        "username": "Antaram AI"
-    }))
-
-    full_response = ""
-    
-    try:
-        async with httpx.AsyncClient(timeout=120) as client:
-            # Removing the trigger word for cleaner prompt
-            clean_prompt = prompt.replace("@antaram.ai", "").strip()
-            
-            async with client.stream("POST", AI_URL, json={"prompt": clean_prompt, "max_tokens": MAX_TOKENS}) as response:
-                async for chunk in response.aiter_text():
-                    # Strip <think> tags as requested
-                    clean_chunk = chunk.replace("<think>", "").replace("</think>", "")
-                    if clean_chunk:
-                        full_response += clean_chunk
-                        # Stream chunk to WebSocket
-                        await broadcast_to_room(room_id, json.dumps({
-                            "type": "ai_chunk",
-                            "username": "Antaram AI",
-                            "chunk": clean_chunk
-                        }))
-    except Exception as e:
-        await broadcast_to_room(room_id, json.dumps({
-            "type": "error",
-            "message": f"AI Error: {str(e)}"
-        }))
-
-    # 2. Notify finished
-    await broadcast_to_room(room_id, json.dumps({
-        "type": "ai_end",
-        "username": "Antaram AI"
-    }))
+    return {"success": True, "file_info": {
+        "file_url": f"/uploads/{unique_name}",
+        "file_type": file.content_type,
+        "original_name": file.filename
+    }}
 
 # --- WebSocket ---
 
 @app.websocket("/ws/{room_id}")
 async def websocket_endpoint(websocket: WebSocket, room_id: str):
     room_id = room_id.upper()
+    await websocket.accept()
     
-    # Init room if accessed via direct link and not created yet
     if room_id not in active_rooms:
         active_rooms[room_id] = []
-        room_files[room_id] = []
-    
-    await websocket.accept()
+        active_users[room_id] = set()
+
     active_rooms[room_id].append(websocket)
     
-    # Notify Join
-    await broadcast_to_room(room_id, json.dumps({
-        "type": "system",
-        "message": "User joined.",
-        "userCount": len(active_rooms[room_id])
-    }))
+    # 1. Send History First
+    history = get_history(room_id)
+    await websocket.send_text(json.dumps({"type": "history", "data": history}))
     
     try:
         while True:
             data = await websocket.receive_text()
-            message_data = json.loads(data)
+            msg_data = json.loads(data)
             
-            user_text = message_data.get("text", "")
-            username = message_data.get("username", "Anonymous")
+            msg_type = msg_data.get("type")
+            username = msg_data.get("username", "Guest")
             
-            # 1. Broadcast User Message Immediately
-            await broadcast_to_room(room_id, json.dumps({
-                "type": message_data.get("type", "message"),
-                "username": username,
-                "text": user_text,
-                "file": message_data.get("file", None)
-            }))
-
-            # 2. Check for AI Trigger
-            if "@antaram.ai" in user_text.lower():
-                # Start AI processing in background (non-blocking)
-                asyncio.create_task(stream_antaram_ai(room_id, user_text))
-                    
-    except WebSocketDisconnect:
-        if room_id in active_rooms and websocket in active_rooms[room_id]:
-            active_rooms[room_id].remove(websocket)
-            # Cleanup if empty
-            if not active_rooms[room_id]:
-                # Optional: Keep room alive for a bit or delete
-                pass 
-            else:
+            # Handle Join (to update user list)
+            if msg_type == "join":
+                active_users[room_id].add(username)
                 await broadcast_to_room(room_id, json.dumps({
-                    "type": "system",
-                    "message": "User left.",
-                    "userCount": len(active_rooms[room_id])
+                    "type": "system", 
+                    "message": f"{username} joined.",
+                    "users": list(active_users[room_id])
                 }))
+                continue
+
+            # Handle Message
+            if msg_type == "message":
+                text = msg_data.get("text", "")
+                reply_to = msg_data.get("reply_to")
+                reply_content = msg_data.get("reply_content")
+                
+                # Save to DB
+                saved_msg = save_message(room_id, username, text, msg_data.get("file"), reply_to, reply_content)
+                
+                # Broadcast
+                await broadcast_to_room(room_id, json.dumps(saved_msg))
+
+                # Check AI
+                if "@antaram.ai" in text.lower():
+                    # If user is replying to a message AND mentioning AI, pass that context
+                    asyncio.create_task(stream_antaram_ai(room_id, text, reply_content))
+
+    except WebSocketDisconnect:
+        if room_id in active_rooms:
+            if websocket in active_rooms[room_id]:
+                active_rooms[room_id].remove(websocket)
+            # We don't remove user immediately from set to keep history context, 
+            # or you can remove if you want strict online status.
+            await broadcast_to_room(room_id, json.dumps({
+                "type": "system", 
+                "message": "User left.",
+                "users": list(active_users[room_id])
+            }))
 
 async def broadcast_to_room(room_id: str, message: str):
     if room_id in active_rooms:
-        # Create a copy of the list to iterate safely
         for connection in list(active_rooms[room_id]):
             try:
                 await connection.send_text(message)
-            except RuntimeError:
-                # Connection likely closed
+            except:
                 pass
 
 if __name__ == "__main__":