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, 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 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( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Setup BASE_DIR = Path(__file__).resolve().parent UPLOAD_DIR = BASE_DIR / "uploads" UPLOAD_DIR.mkdir(exist_ok=True) templates = Jinja2Templates(directory="templates") app.mount("/uploads", StaticFiles(directory="uploads"), name="uploads") # --- 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("", "").replace("", "") 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): 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): 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() return {"room_id": room_id, "success": True} @app.post("/upload-file/{room_id}") async def upload_file(room_id: str, file: UploadFile = File(...)): 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) 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() if room_id not in active_rooms: active_rooms[room_id] = [] active_users[room_id] = set() active_rooms[room_id].append(websocket) # 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() msg_data = json.loads(data) msg_type = msg_data.get("type") username = msg_data.get("username", "Guest") # 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": 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: for connection in list(active_rooms[room_id]): try: await connection.send_text(message) except: pass if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)