Upload 5 files

Dockerfile

@@ -0,0 +1,17 @@
+# Use a stable Python base
+FROM python:3.10-slim
+
+# set working dir
+WORKDIR /app
+
+# copy files
+COPY . /app
+
+# install dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# expose default Space port (7860 is HF default)
+EXPOSE 7860
+
+# start the app (use PORT env var set by HF - default 7860)
+CMD ["sh", "-c", "uvicorn server:app --host 0.0.0.0 --port ${PORT:-7860} --proxy-headers"]

README.md

@@ -1,10 +1,6 @@
----
-title: Server
-emoji: 🏆
-colorFrom: green
-colorTo: blue
-sdk: docker
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+sdk: docker
+app_port: 7860
+---
+# Akatsuki (FastAPI WebSocket chat)
+... normal README content ...

requirements.txt

@@ -0,0 +1,5 @@
+fastapi
+uvicorn[standard]
+httpx
+python-dotenv
+websockets

server.py

@@ -0,0 +1,536 @@
+#!/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()})
+
+@app.get("/stats")
+async def get_stats(): return JSONResponse({"active_users": len(connected_users)})
+
+@app.websocket("/ws")
+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
+
+
+@app.on_event("startup")
+async def startup_event():
+    load_state()
+    asyncio.create_task(state_cleanup_task())
+    asyncio.create_task(heartbeat_task())
+
+@app.on_event("shutdown")
+def shutdown_event(): save_state()
+@app.get("/")
+async def root(): return HTMLResponse("<h1>Akatsuki Server</h1>")

start.sh

@@ -0,0 +1,2 @@
+#!/bin/sh
+exec uvicorn server:app --host 0.0.0.0 --port ${PORT:-7860} --proxy-headers