Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """ | |
| Akatsuki - simple ephemeral single-room WebSocket chat server. | |
| """ | |
| import os | |
| import asyncio | |
| import json | |
| import uuid | |
| import secrets | |
| import re | |
| import logging | |
| from datetime import datetime, timedelta, timezone | |
| from typing import Dict, Optional, Any, Set, List | |
| from fastapi import FastAPI, WebSocket, WebSocketDisconnect, status | |
| from fastapi.responses import HTMLResponse, JSONResponse | |
| from fastapi.middleware.cors import CORSMiddleware | |
| import httpx | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| # --- Basic Logging Setup --- | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | |
| APP_NAME = "Akatsuki" | |
| ADMIN_USERNAME = os.environ.get("ADMIN_USERNAME", "Aadish") | |
| ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "Aadish20m") | |
| GROQ_API_KEY = os.environ.get("GROQ_API_KEY") | |
| GROQ_ENDPOINT = "https://api.groq.com/openai/v1/chat/completions" | |
| # --- AI Model Configuration --- | |
| MODEL_ALIASES = { | |
| "gpt": "openai/gpt-oss-120b", | |
| "llama": "meta-llama/llama-4-maverick-17b-128e-instruct", | |
| "deepseek": "deepseek-r1-distill-llama-70b", | |
| "qwen": "qwen/qwen3-32b", | |
| "compound": "compound-beta", | |
| } | |
| DEFAULT_MODEL = "compound-beta" | |
| # --- Whitelist Configuration --- | |
| WHITELISTED_USERS = {"Prakhar", "Priyanshu", "Summit", "Aditya", "Yuvraj", "Yash"} | |
| TEMPORARY_WHITELIST = set() | |
| app = FastAPI(title=APP_NAME) | |
| app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]) | |
| # --- State Management --- | |
| STATE_FILE = "server_state.json" | |
| TEMP_STATE_FILE = "server_state.json.tmp" | |
| LATEST_BROADCAST: Optional[Dict[str, Any]] = None | |
| class User: | |
| def __init__(self, websocket: WebSocket, username: str, role: str): | |
| self.ws = websocket | |
| self.username = username | |
| self.role = role | |
| self.session_id = uuid.uuid4().hex | |
| self.last_pong_at = datetime.now(timezone.utc) | |
| connected_users: Dict[str, User] = {} | |
| bans: Dict[str, Optional[str]] = {} | |
| mutes: Dict[str, str] = {} | |
| user_tags: Dict[str, str] = {} # username.lower() -> tag | |
| connected_lock = asyncio.Lock() | |
| state_lock = asyncio.Lock() | |
| def ts_iso(dt: Optional[datetime] = None) -> str: | |
| dt = dt or datetime.now(timezone.utc) | |
| return dt.isoformat() | |
| def ts_from_iso(iso_str: str) -> Optional[datetime]: | |
| try: | |
| return datetime.fromisoformat(iso_str.replace('Z', '+00:00')) | |
| except (ValueError, TypeError): | |
| return None | |
| # --- Persistence --- | |
| def save_state(): | |
| logging.info("Attempting to save state to disk...") | |
| try: | |
| with open(TEMP_STATE_FILE, "w") as f: | |
| json.dump({"bans": bans, "mutes": mutes, "tags": user_tags}, f, indent=2) | |
| os.replace(TEMP_STATE_FILE, STATE_FILE) | |
| logging.info("State saved successfully.") | |
| except (IOError, os.error) as e: | |
| logging.error(f"Failed to save state: {e}") | |
| async def async_save_state(): | |
| async with state_lock: | |
| loop = asyncio.get_running_loop() | |
| await loop.run_in_executor(None, save_state) | |
| def cleanup_expired_state(): | |
| now = datetime.now(timezone.utc) | |
| changed = False | |
| for u, exp_str in list(bans.items()): | |
| if exp_str and (expiry := ts_from_iso(exp_str)) and expiry <= now: | |
| del bans[u] | |
| changed = True | |
| for u, exp_str in list(mutes.items()): | |
| if (expiry := ts_from_iso(exp_str)) and expiry <= now: | |
| del mutes[u] | |
| changed = True | |
| if changed: save_state() | |
| def load_state(): | |
| global bans, mutes, user_tags | |
| if not os.path.exists(STATE_FILE): return | |
| try: | |
| with open(STATE_FILE, "r") as f: | |
| data = json.load(f) | |
| bans = data.get("bans", {}) | |
| mutes = data.get("mutes", {}) | |
| user_tags = data.get("tags", {}) | |
| cleanup_expired_state() | |
| logging.info("Server state loaded successfully.") | |
| except (IOError, json.JSONDecodeError) as e: | |
| logging.error(f"Failed to load state: {e}") | |
| async def state_cleanup_task(): | |
| while True: | |
| await asyncio.sleep(60 * 5) | |
| logging.info("Running periodic state cleanup task...") | |
| try: | |
| loop = asyncio.get_running_loop() | |
| await loop.run_in_executor(None, cleanup_expired_state) | |
| except Exception as e: | |
| logging.error(f"Error in state_cleanup_task: {e}", exc_info=True) | |
| # --- Business Logic Helpers --- | |
| async def is_banned(username: str) -> Optional[str]: | |
| username_lower = username.lower() | |
| for ban_user, ban_expiry_val in bans.items(): | |
| if ban_user.lower() == username_lower: | |
| if ban_expiry_val is None: return "You are permanently banned." | |
| expiry_dt = ts_from_iso(ban_expiry_val) | |
| if expiry_dt and expiry_dt > datetime.now(timezone.utc): | |
| remaining = expiry_dt - datetime.now(timezone.utc) | |
| return f"You are banned for another {int(remaining.total_seconds() / 60)} minutes." | |
| elif expiry_dt: | |
| del bans[ban_user] | |
| await async_save_state() | |
| return None | |
| return None | |
| async def is_muted(username: str) -> Optional[str]: | |
| username_lower = username.lower() | |
| for mute_user, mute_expiry_val in mutes.items(): | |
| if mute_user.lower() == username_lower: | |
| expiry_dt = ts_from_iso(mute_expiry_val) | |
| if expiry_dt and expiry_dt > datetime.now(timezone.utc): | |
| remaining = expiry_dt - datetime.now(timezone.utc) | |
| return f"You are muted for another {int(remaining.total_seconds())}s." | |
| elif expiry_dt: | |
| del mutes[mute_user] | |
| await async_save_state() | |
| return None | |
| return None | |
| # --- User & Session Helpers --- | |
| async def find_user_by_name(username: str) -> Optional[User]: | |
| username_lower = username.lower() | |
| async with connected_lock: | |
| for user in connected_users.values(): | |
| if user.username.lower() == username_lower: | |
| return user | |
| return None | |
| async def is_username_in_use(username: str) -> bool: | |
| return await find_user_by_name(username) is not None | |
| # --- Messaging Helpers --- | |
| async def safe_send(ws: WebSocket, obj: dict): | |
| try: await ws.send_text(json.dumps(obj)) | |
| except (WebSocketDisconnect, ConnectionError, RuntimeError): pass | |
| async def broadcast(obj: dict, exclude_ids: Optional[Set[str]] = None): | |
| exclude = exclude_ids or set() | |
| connections_to_send = [] | |
| async with connected_lock: | |
| for sid, u in connected_users.items(): | |
| if sid not in exclude: | |
| connections_to_send.append(u.ws) | |
| tasks = [safe_send(ws, obj) for ws in connections_to_send] | |
| if tasks: await asyncio.gather(*tasks) | |
| async def send_to_user(username: str, obj: dict) -> bool: | |
| user = await find_user_by_name(username) | |
| if user: | |
| await safe_send(user.ws, obj) | |
| return True | |
| return False | |
| def get_users_list() -> List[Dict[str, Any]]: | |
| return [{"name": u.username, "role": u.role, "tag": user_tags.get(u.username.lower())} for u in connected_users.values()] | |
| # --- Groq AI Helper --- | |
| async def call_groq_api(user_prompt: str, system_prompt: Optional[str] = None, model: str = DEFAULT_MODEL) -> str: | |
| if not GROQ_API_KEY: | |
| await asyncio.sleep(0.5) | |
| return f"[Simulated AI Response for {model}] You asked: '{user_prompt[:100]}...'" | |
| logging.info(f"Making Groq API call. Model: {model}, Endpoint: {GROQ_ENDPOINT}") | |
| headers = {"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"} | |
| payload = { | |
| "messages": [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}], | |
| "model": model | |
| } | |
| if model == "openai/gpt-oss-120b": | |
| payload["tools"] = [{"type": "browser_search"}, {"type": "code_interpreter"}] | |
| logging.info(f"Request Payload: {json.dumps(payload)}") | |
| try: | |
| async with httpx.AsyncClient(timeout=60.0) as client: | |
| resp = await client.post(GROQ_ENDPOINT, headers=headers, json=payload) | |
| logging.info(f"Groq API Response Status: {resp.status_code}") | |
| resp.raise_for_status() | |
| content = resp.json()["choices"][0]["message"]["content"] | |
| # Strip out the <think>...</think> block before returning | |
| cleaned_content = re.sub(r'<think>.*?</think>\s*', '', content, flags=re.DOTALL) | |
| return cleaned_content.strip() | |
| except httpx.HTTPStatusError as e: | |
| logging.error(f"Groq API HTTP Status Error: {e.response.status_code} Body: {e.response.text}") | |
| return f"[AI API Error] Status {e.response.status_code}" | |
| except Exception as e: | |
| logging.error(f"Groq API call failed: {e}", exc_info=True) | |
| return f"[AI API Error] Could not connect: {e}" | |
| # --- Command Handlers --- | |
| def parse_command_args(args: List[str]) -> (Set[str], List[str]): | |
| users, remaining_args = set(), [] | |
| for arg in args: | |
| if arg.startswith('@'): users.add(arg[1:]) | |
| else: remaining_args.append(arg) | |
| return users, remaining_args | |
| async def handle_admin_command(admin_user: User, raw_cmd: str): | |
| global LATEST_BROADCAST | |
| parts = raw_cmd.strip().split() | |
| cmd, args = parts[0].lower(), parts[1:] | |
| changed_state = False | |
| if cmd == "/login": | |
| target_users, _ = parse_command_args(args) | |
| if not target_users: | |
| return await safe_send(admin_user.ws, {"type": "system", "text": "Usage: /login @username"}) | |
| for username in target_users: | |
| TEMPORARY_WHITELIST.add(username.lower()) | |
| await safe_send(admin_user.ws, {"type": "system", "text": f"User '{username}' has been temporarily whitelisted to join."}) | |
| return | |
| if cmd == "/clearall": | |
| LATEST_BROADCAST = None | |
| await broadcast({"type": "clear_chat"}) | |
| await broadcast({"type": "system", "text": f"Chat history cleared by {admin_user.username}."}) | |
| return | |
| if cmd in ("/broadcast", "/b"): | |
| message = " ".join(args) | |
| if not message: | |
| await safe_send(admin_user.ws, {"type": "system", "text": "Usage: /broadcast <message>"}) | |
| else: | |
| broadcast_payload = {"type": "broadcast", "from": admin_user.username, "text": message, "ts": ts_iso()} | |
| LATEST_BROADCAST = broadcast_payload | |
| await broadcast(broadcast_payload) | |
| return | |
| if cmd == "/clearbroadcast": | |
| LATEST_BROADCAST = None | |
| await broadcast({"type": "clear_broadcast"}) | |
| return | |
| if cmd == "/tag": | |
| target_users, tag_parts = parse_command_args(args) | |
| tag_arg = next((arg for arg in tag_parts if arg.startswith('--')), None) | |
| if not target_users or not tag_arg: | |
| return await safe_send(admin_user.ws, {"type": "system", "text": "Usage: /tag @user --tagname"}) | |
| tag = tag_arg[2:] | |
| if not tag: | |
| return await safe_send(admin_user.ws, {"type": "system", "text": "Tag name cannot be empty."}) | |
| for username in target_users: | |
| user_tags[username.lower()] = tag | |
| target_user_obj = await find_user_by_name(username) | |
| role = target_user_obj.role if target_user_obj else "user" | |
| await broadcast({"type": "system", "text": f"{admin_user.username} gave the tag '{tag}' to {username}."}) | |
| await broadcast({"type": "user_update", "user": {"name": username, "role": role, "tag": tag}}) | |
| if target_users: | |
| await async_save_state() | |
| return | |
| if cmd == "/removetag": | |
| target_users, _ = parse_command_args(args) | |
| if not target_users: | |
| return await safe_send(admin_user.ws, {"type": "system", "text": "Usage: /removetag @user"}) | |
| a_tag_was_removed = False | |
| for username in target_users: | |
| if username.lower() in user_tags: | |
| del user_tags[username.lower()] | |
| a_tag_was_removed = True | |
| target_user_obj = await find_user_by_name(username) | |
| role = target_user_obj.role if target_user_obj else "user" | |
| await broadcast({"type": "system", "text": f"{admin_user.username} removed the tag from {username}."}) | |
| await broadcast({"type": "user_update", "user": {"name": username, "role": role, "tag": None}}) | |
| if a_tag_was_removed: | |
| await async_save_state() | |
| return | |
| target_users, remaining_args = parse_command_args(args) | |
| if not target_users: | |
| if cmd in ("/kick", "/ban", "/unban", "/mute", "/unmute"): | |
| return await safe_send(admin_user.ws, {"type": "system", "text": "You must specify at least one user with @mention."}) | |
| else: | |
| return # Other commands might not need a user | |
| duration_min = int(remaining_args.pop(0)) if remaining_args and remaining_args[0].isdigit() else None | |
| reason = " ".join(remaining_args) or "No reason specified." | |
| for username in target_users: | |
| target_user = await find_user_by_name(username) | |
| if cmd == "/kick" and target_user: | |
| await safe_send(target_user.ws, {"type": "system", "text": f"Kicked by admin. Reason: {reason}"}) | |
| await target_user.ws.close(code=status.WS_1008_POLICY_VIOLATION) | |
| await broadcast({"type": "system", "text": f"{username} was kicked by {admin_user.username}."}) | |
| elif cmd == "/ban": | |
| changed_state = True | |
| bans[username] = ts_iso(datetime.now(timezone.utc) + timedelta(minutes=duration_min)) if duration_min else None | |
| d_str = f"for {duration_min} minutes" if duration_min else "permanently" | |
| await broadcast({"type": "system", "text": f"{username} was banned {d_str} by {admin_user.username}."}) | |
| if target_user: | |
| await safe_send(target_user.ws, {"type": "system", "text": f"You have been banned {d_str}."}) | |
| await target_user.ws.close(code=status.WS_1008_POLICY_VIOLATION) | |
| elif cmd == "/unban": | |
| user_to_unban = next((u for u in bans if u.lower() == username.lower()), None) | |
| if user_to_unban: | |
| del bans[user_to_unban] | |
| changed_state = True | |
| await broadcast({"type": "system", "text": f"{username} was unbanned by {admin_user.username}."}) | |
| elif cmd == "/mute": | |
| changed_state = True | |
| d_min = duration_min or 5 | |
| mutes[username] = ts_iso(datetime.now(timezone.utc) + timedelta(minutes=d_min)) | |
| await broadcast({"type": "system", "text": f"{username} was muted for {d_min} min by {admin_user.username}."}) | |
| if target_user: await send_to_user(username, {"type": "system", "text": f"You are muted for {d_min} min."}) | |
| elif cmd == "/unmute": | |
| user_to_unmute = next((u for u in mutes if u.lower() == username.lower()), None) | |
| if user_to_unmute: | |
| del mutes[user_to_unmute] | |
| changed_state = True | |
| await broadcast({"type": "system", "text": f"{username} was unmuted by {admin_user.username}."}) | |
| if changed_state: await async_save_state() | |
| async def handle_message(user: User, data: dict): | |
| typ = data.get("type") | |
| if typ == "pong": | |
| user.last_pong_at = datetime.now(timezone.utc) | |
| return | |
| if typ in ("message", "pm"): | |
| if mute_reason := await is_muted(user.username): | |
| await safe_send(user.ws, {"type": "system", "text": mute_reason}) | |
| return | |
| if typ == "message": | |
| await broadcast({"type": "message", "id": data.get("id", str(uuid.uuid4())), "from": user.username, "text": data.get("text", ""), "ts": ts_iso()}) | |
| elif typ == "command": | |
| if user.role == "admin": await handle_admin_command(user, data.get("raw", "")) | |
| else: await safe_send(user.ws, {"type": "system", "text": "You do not have permission to use admin commands."}) | |
| elif typ == "nick": | |
| new_nick = (data.get("toNick") or "").strip() | |
| if new_nick and 1 <= len(new_nick) <= 32 and not await is_username_in_use(new_nick) and new_nick.lower() != ADMIN_USERNAME.lower(): | |
| old_nick = user.username | |
| user.username = new_nick | |
| users_list = [] | |
| async with connected_lock: | |
| users_list = get_users_list() | |
| await broadcast({"type": "system", "text": f"{old_nick} is now known as {new_nick}."}) | |
| await broadcast({"type": "users", "users": users_list}) | |
| else: | |
| await safe_send(user.ws, {"type": "system", "text": "Invalid or taken nickname."}) | |
| elif typ == "pm": | |
| recipients, text = data.get("to", []), data.get("text", "") | |
| if recipients and text: | |
| pm_payload = {**data, "from": user.username} | |
| for r_name in recipients: | |
| if not await send_to_user(r_name, pm_payload): | |
| await safe_send(user.ws, {"type": "system", "text": f"User '{r_name}' not found."}) | |
| await safe_send(user.ws, pm_payload) | |
| elif typ == "ai": | |
| prompt_parts = data.get("text", "").split() | |
| model_id = DEFAULT_MODEL | |
| model_alias = "compound" | |
| final_prompt_list = prompt_parts | |
| if prompt_parts and prompt_parts[0].startswith('--'): | |
| flag = prompt_parts[0][2:].lower() | |
| if flag in MODEL_ALIASES: | |
| model_alias = flag | |
| model_id = MODEL_ALIASES[flag] | |
| final_prompt_list = prompt_parts[1:] | |
| prompt_text = " ".join(final_prompt_list) | |
| if not prompt_text: | |
| return await safe_send(user.ws, {"type": "system", "text": "Usage: /ai [--model_alias] <prompt>"}) | |
| await broadcast({"type": "system", "text": f"{user.username} is asking the AI ({model_alias})..."}) | |
| system_prompt = "You are a helpful and concise assistant integrated into a terminal-based chat application. Your answers should be brief and to the point. Use markdown for formatting if necessary." | |
| response = await call_groq_api(prompt_text, system_prompt=system_prompt, model=model_id) | |
| await broadcast({"type": "ai_resp", "id": str(uuid.uuid4()), "from": "AI", "text": response, "ts": ts_iso()}) | |
| async def get_stats(): return JSONResponse({"active_users": len(connected_users)}) | |
| async def websocket_endpoint(ws: WebSocket): | |
| user: Optional[User] = None | |
| await ws.accept() | |
| try: | |
| auth_msg = await asyncio.wait_for(ws.receive_json(), timeout=15.0) | |
| if auth_msg.get("type") != "auth": | |
| return await safe_send(ws, {"type": "auth_failed", "reason": "First message must be auth"}) | |
| req_name = (auth_msg.get("username") or "").strip() | |
| req_name_lower = req_name.lower() | |
| if not req_name or len(req_name) > 32: | |
| return await safe_send(ws, {"type": "auth_failed", "reason": "Username must be 1-32 characters"}) | |
| if ban_reason := await is_banned(req_name): | |
| return await safe_send(ws, {"type": "auth_failed", "reason": ban_reason}) | |
| if await is_username_in_use(req_name): | |
| return await safe_send(ws, {"type": "auth_failed", "reason": "Username is already in use"}) | |
| is_whitelisted = any(u.lower() == req_name_lower for u in WHITELISTED_USERS) | |
| is_temp_whitelisted = req_name_lower in TEMPORARY_WHITELIST | |
| is_test_user = re.match(r'^test\d+$', req_name_lower) is not None | |
| is_admin_user = req_name_lower == ADMIN_USERNAME.lower() | |
| if not (is_whitelisted or is_temp_whitelisted or is_test_user or is_admin_user): | |
| return await safe_send(ws, {"type": "auth_failed", "reason": "You are not authorized to join this server."}) | |
| role = "user" | |
| if is_admin_user: | |
| password = auth_msg.get("password") | |
| if secrets.compare_digest(password or "", ADMIN_PASSWORD): | |
| role = "admin" | |
| else: return await safe_send(ws, {"type": "auth_failed", "reason": "Invalid admin credentials"}) | |
| user = User(websocket=ws, username=req_name, role=role) | |
| if req_name_lower in TEMPORARY_WHITELIST: | |
| TEMPORARY_WHITELIST.remove(req_name_lower) | |
| users_list = [] | |
| async with connected_lock: | |
| connected_users[user.session_id] = user | |
| users_list = get_users_list() | |
| await safe_send(ws, {"type": "auth_ok", "username": user.username, "role": user.role}) | |
| await broadcast({"type": "user_join", "user": {"name": user.username, "role": user.role, "tag": user_tags.get(user.username.lower())}}, exclude_ids={user.session_id}) | |
| await safe_send(ws, {"type": "users", "users": users_list}) | |
| if LATEST_BROADCAST: | |
| await safe_send(ws, LATEST_BROADCAST) | |
| welcome_system_prompt = f"You are a witty and welcoming chatbot for the '{APP_NAME}' chat server. Your task is to generate a very short, cool, and slightly mysterious welcome message for a new user named '{user.username}'. Keep it under 15 words. Be creative." | |
| welcome_message = await call_groq_api(f"A user named {user.username} has joined.", system_prompt=welcome_system_prompt, model="llama-3.1-8b-instant") | |
| await safe_send(ws, {"type": "system", "text": welcome_message}) | |
| while True: | |
| data = await ws.receive_json() | |
| await handle_message(user, data) | |
| except (WebSocketDisconnect, asyncio.TimeoutError, json.JSONDecodeError): pass | |
| finally: | |
| if user: | |
| async with connected_lock: | |
| if user.session_id in connected_users: del connected_users[user.session_id] | |
| user_list_payload = {"type": "users", "users": get_users_list()} | |
| leave_payload = {"type": "system", "text": f"{user.username} has left the chat."} | |
| await broadcast(user_list_payload) | |
| await broadcast(leave_payload) | |
| # --- Heartbeat Task --- | |
| PING_INTERVAL = 30 | |
| PONG_TIMEOUT = 65 | |
| async def heartbeat_task(): | |
| while True: | |
| await asyncio.sleep(PING_INTERVAL) | |
| now = datetime.now(timezone.utc) | |
| users_to_check = [] | |
| async with connected_lock: | |
| users_to_check = list(connected_users.values()) | |
| for user in users_to_check: | |
| if (now - user.last_pong_at).total_seconds() > PONG_TIMEOUT: | |
| logging.warning(f"User {user.username} timed out due to no pong. Disconnecting.") | |
| try: | |
| await user.ws.close(code=status.WS_1008_POLICY_VIOLATION, reason="Connection timed out") | |
| except RuntimeError: | |
| pass # Already closed | |
| else: | |
| try: | |
| await user.ws.send_text(json.dumps({"type": "ping"})) | |
| except (WebSocketDisconnect, ConnectionError, RuntimeError): | |
| pass | |
| async def startup_event(): | |
| load_state() | |
| asyncio.create_task(state_cleanup_task()) | |
| asyncio.create_task(heartbeat_task()) | |
| def shutdown_event(): save_state() | |
| async def root(): return HTMLResponse("<h1>Akatsuki Server</h1>") | |