import asyncio import json import os import threading import time import uuid from concurrent.futures import ThreadPoolExecutor from fastapi import Depends, FastAPI, Header, HTTPException, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, StreamingResponse from server.auth import verify_key from server.db import ( add_message, create_conversation, get_messages, init_db, init_users_table, init_moderation_table, list_moderation_events, list_conversations, get_conversation, create_user, get_user_by_email, get_user_by_id, ) from server.inference import generate_tokens, load_model from server.schemas import ( ChatRequest, AgentRequest, ModelCard, ModelList, RegisterRequest, LoginRequest, TokenResponse, UserResponse, ModerationRequest, ModerationResponse, ) from server.user_auth import ( hash_password, verify_password, create_access_token, decode_access_token, ) from server.limits import ( MAX_BODY_BYTES, MAX_CONCURRENT_GENERATIONS, GENERATION_QUE_TIMEOUT, RATE_LIMIT_CHAT, RATE_LIMIT_AGENT, RATE_LIMIT_AUTH, RATE_LIMIT_READ, ) from server.ratelimit import init_redis, close_redis, rate_limit, identity_from_request from server.safety import moderate_input, moderate_output, injection, audit import server.safety.config as safety_config def _default_checkpoint(): """Prefer sft_ckpt.pt (the blended chat+agent model). agent_ckpt.pt is the retired, over-fit two-stage model, kept only as a fallback if no sft exists.""" base = os.path.join(os.path.dirname(__file__), "..", "ml", "checkpoints") sft = os.path.join(base, "sft_ckpt.pt") agent = os.path.join(base, "agent_ckpt.pt") return sft if os.path.exists(sft) else agent CHECKPOINT = os.environ.get("ABYSS_CHECKPOINT") or _default_checkpoint() def _resolve_checkpoint() -> str: """Locate the model checkpoint: explicit env > local file > Hugging Face Hub. On a deploy like Hugging Face Spaces the weights are NOT baked into the image (kept small), so when ABYSS_HF_REPO is set and nothing is on disk we pull the checkpoint from the Hub once at boot (cached afterwards). """ env_path = os.environ.get("ABYSS_CHECKPOINT") if env_path and os.path.exists(env_path): return env_path if os.path.exists(CHECKPOINT): return CHECKPOINT repo = os.environ.get("ABYSS_HF_REPO") if repo: from huggingface_hub import hf_hub_download filename = os.environ.get("ABYSS_HF_CHECKPOINT", "sft_ckpt.pt") print(f"[startup] downloading {filename} from HF Hub repo '{repo}' ...") path = hf_hub_download(repo_id=repo, filename=filename) print(f"[startup] checkpoint ready: {path}") return path return CHECKPOINT # nothing found -> load_model raises a clear error app = FastAPI(title="Abyss Inference Server", version="0.1.0") # CORS — defaults to "*" (safe here because auth is a Bearer token, not cookies). # Set ABYSS_CORS_ORIGINS to a comma-separated allowlist (e.g. your Cloudflare # Pages URL) to tighten it in production. _cors_env = os.environ.get("ABYSS_CORS_ORIGINS", "*").strip() _allow_origins = ["*"] if _cors_env in ("", "*") else [ o.strip() for o in _cors_env.split(",") if o.strip() ] app.add_middleware( CORSMiddleware, allow_origins=_allow_origins, allow_methods=["*"], allow_headers=["*"], ) # ----- BODY-SIZE GUARD + RATE-LIMIT RESPONSE HEADERS (PHASE S) ----- @app.middleware("http") async def guard_middleware(request: Request, call_next): cl = request.headers.get("content-length") if cl and cl.isdigit() and int(cl) > MAX_BODY_BYTES: return JSONResponse( {"detail": f"Request body too large (max {MAX_BODY_BYTES} bytes)"}, status_code=413, ) response = await call_next(request) headers = getattr(request.state, "rate_headers", None) if headers: for k, v in headers.items(): response.headers[k] = v return response # ----- GLOBAL GENERATION CONCURRENCY CAP (PHASE S) ----- _gen_semaphore = asyncio.Semaphore(MAX_CONCURRENT_GENERATIONS) async def _acquire_slot() -> None: try: await asyncio.wait_for( _gen_semaphore.acquire(), timeout=GENERATION_QUE_TIMEOUT ) except asyncio.TimeoutError: raise HTTPException( status_code=503, detail="Server is busy. Please try again shortly" ) _executor = ThreadPoolExecutor(max_workers=1) # ----- PHASE W: SAFETY HELPERS ----- ADMIN_KEY = os.environ.get("ABYSS_ADMIN_KEY", "") def require_admin(x_admin_key: str = Header(default="")): """Gate the moderation review endpoint. Locked unless ABYSS_ADMIN_KEY is set.""" if not ADMIN_KEY or x_admin_key != ADMIN_KEY: raise HTTPException(status_code=403, detail="Admin access required") return True def _static_stream(message, conv_id, completion_id, created, obj="chat.completion.chunk"): """Emit a fixed message (safety refusal) in the normal SSE chunk shape.""" async def gen(): yield "data: " + json.dumps({ "id": completion_id, "object": obj, "created": created, "model": "abyss-1", "conversation_id": conv_id, "choices": [{"index": 0, "delta": {"content": message}, "finish_reason": None}], }) + "\n\n" yield "data: " + json.dumps({ "id": completion_id, "object": obj, "created": created, "model": "abyss-1", "conversation_id": conv_id, "choices": [{"index": 0, "delta": {}, "finish_reason": "content_filter"}], }) + "\n\n" yield "data: [DONE]\n\n" return StreamingResponse(gen(), media_type="text/event-stream") # ----- STARTUP / SHUTDOWN ----- @app.on_event("startup") async def startup() -> None: init_db() init_users_table() init_moderation_table() await init_redis() load_model(_resolve_checkpoint()) @app.on_event("shutdown") async def shutdown() -> None: await close_redis() # ----- HEALTH ----- @app.get("/health") def health(): return {"status": "ok", "model": "abyss-1"} # ----- MODELS ----- @app.get("/v1/models", dependencies=[Depends(verify_key), Depends(rate_limit(*RATE_LIMIT_READ, scope="read"))]) def list_models(): return ModelList(data=[ModelCard(id="abyss-1")]) # ----- AUTH ----- @app.post("/auth/register", response_model=TokenResponse, dependencies=[Depends(rate_limit(*RATE_LIMIT_AUTH, scope="auth"))]) def register(req: RegisterRequest): if get_user_by_email(req.email): raise HTTPException(status_code=400, detail="Email already registered") hashed = hash_password(req.password) user_id = create_user(req.username, req.email, hashed) token = create_access_token(user_id) return TokenResponse(access_token=token, username=req.username) @app.post("/auth/login", response_model=TokenResponse, dependencies=[Depends(rate_limit(*RATE_LIMIT_AUTH, scope="auth"))]) def login(req: LoginRequest): user = get_user_by_email(req.email) if not user or not verify_password(req.password, user["hashed_password"]): raise HTTPException(status_code=401, detail="Invalid email or password") token = create_access_token(user["id"]) return TokenResponse(access_token=token, username=user["username"]) @app.get("/auth/me", response_model=UserResponse, dependencies=[Depends(rate_limit(*RATE_LIMIT_READ, scope="read"))]) def me(token: str = Depends(verify_key)): user_id = decode_access_token(token) if not user_id: raise HTTPException(status_code=401, detail="Invalid token") user = get_user_by_id(user_id) if not user: raise HTTPException(status_code=404, detail="User not found") return user # ----- CONVERSATIONS ----- @app.get("/v1/conversations", dependencies=[Depends(verify_key), Depends(rate_limit(*RATE_LIMIT_READ, scope="read"))]) def get_conversations(): return list_conversations() @app.get("/v1/conversations/{conversation_id}", dependencies=[Depends(verify_key), Depends(rate_limit(*RATE_LIMIT_READ, scope="read"))]) def get_conv(conversation_id: str): conv = get_conversation(conversation_id) if conv is None: raise HTTPException(status_code=404, detail="Conversation not found") return conv # ----- CHAT COMPLETIONS (SSE) ----- @app.post( "/v1/chat/completions", dependencies=[ Depends(verify_key), Depends(rate_limit(*RATE_LIMIT_CHAT, scope="chat")), ], ) async def chat_completions(req: ChatRequest, request: Request): messages = [m.model_dump() for m in req.messages] identity = identity_from_request(request) user_text = "\n".join(m["content"] for m in messages if m["role"] == "user") # ----- PHASE W: INPUT MODERATION (before consuming a generation slot) ----- in_mod = moderate_input(user_text) if in_mod.flagged: audit.record_result("input", "chat", in_mod, identity=identity, snippet=user_text) if in_mod.action in (safety_config.ACTION_BLOCK, safety_config.ACTION_SOFT_REFUSE): conv_id = req.conversation_id or create_conversation( title=(user_text[:60] or "New Chat") ) for m in messages: add_message(conv_id, m["role"], m["content"]) refusal = safety_config.message_for(in_mod.action, in_mod.top_category) add_message(conv_id, "assistant", refusal) return _static_stream( refusal, conv_id, f"chatcmpl-{uuid.uuid4().hex[:12]}", int(time.time()) ) await _acquire_slot() # Phase S: concurrency cap try: # RESOLVE / CREATE CONVERSATION conv_id = req.conversation_id if conv_id is None: first_user = next( (m["content"] for m in messages if m["role"] == "user"), "New Chat" ) conv_id = create_conversation(title=first_user[:60]) # PERSIST INCOMING MESSAGES for m in messages: add_message(conv_id, m["role"], m["content"]) completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) except Exception: _gen_semaphore.release() # release if we fail before streaming raise async def event_stream(): try: loop = asyncio.get_event_loop() queue: asyncio.Queue = asyncio.Queue() stop_event = threading.Event() def _run(): try: for token in generate_tokens( messages, max_tokens=req.max_tokens, temperature=req.temperature, top_k=req.top_k, top_p=req.top_p, ): if stop_event.is_set(): # Phase W: halt on unsafe output break asyncio.run_coroutine_threadsafe( queue.put(("token", token)), loop ) except Exception as exc: asyncio.run_coroutine_threadsafe( queue.put(("error", str(exc))), loop ) finally: asyncio.run_coroutine_threadsafe( queue.put(("done", None)), loop ) _executor.submit(_run) streamed: list[str] = [] since_check = 0 while True: kind, value = await queue.get() if kind == "error": yield f'data: {{"error": "{value}"}}\n\n' break if kind == "done": text = "".join(streamed) out_mod = moderate_output(text) if out_mod.flagged: audit.record_result( "output", "chat", out_mod, identity=identity, conversation_id=conv_id, snippet=text, ) add_message(conv_id, "assistant", text) final_chunk = { "id": completion_id, "object": "chat.completion.chunk", "created": created, "model": "abyss-1", "conversation_id": conv_id, "choices": [ {"index": 0, "delta": {}, "finish_reason": "stop"} ], } yield f"data: {json.dumps(final_chunk)}\n\n" yield "data: [DONE]\n\n" break # ----- PHASE W: ROLLING OUTPUT MODERATION ----- streamed.append(value) since_check += len(value) if since_check >= 48: since_check = 0 out_mod = moderate_output("".join(streamed)) if out_mod.action == safety_config.ACTION_BLOCK: stop_event.set() audit.record_result( "output", "chat", out_mod, identity=identity, conversation_id=conv_id, snippet="".join(streamed), ) add_message(conv_id, "assistant", safety_config.OUTPUT_BLOCK_MESSAGE) yield "data: " + json.dumps({ "id": completion_id, "object": "chat.completion.chunk", "created": created, "model": "abyss-1", "conversation_id": conv_id, "choices": [{"index": 0, "delta": {"content": " " + safety_config.OUTPUT_BLOCK_MESSAGE}, "finish_reason": "content_filter"}], }) + "\n\n" yield "data: [DONE]\n\n" break chunk = { "id": completion_id, "object": "chat.completion.chunk", "created": created, "model": "abyss-1", "conversation_id": conv_id, "choices": [ {"index": 0, "delta": {"content": value}, "finish_reason": None} ], } yield f"data: {json.dumps(chunk)}\n\n" finally: _gen_semaphore.release() # Phase S: free the slot return StreamingResponse(event_stream(), media_type="text/event-stream") # ----- AGENT CHAT (SSE) ----- @app.post( "/v1/agent/chat", dependencies=[ Depends(verify_key), Depends(rate_limit(*RATE_LIMIT_AGENT, scope="agent")), ], ) async def agent_chat(req: AgentRequest, request: Request): from agent.core.loop import run_agent messages = [m.model_dump() for m in req.messages] identity = identity_from_request(request) user_text = "\n".join(m["content"] for m in messages if m["role"] == "user") # ----- PHASE W: INPUT MODERATION + PROMPT-INJECTION SCREEN ----- in_mod = moderate_input(user_text) if in_mod.flagged: audit.record_result("input", "agent", in_mod, identity=identity, snippet=user_text) inj = injection.detect(user_text) if inj.flagged: audit.record_injection("agent", inj, identity=identity, snippet=user_text) refuse = None if in_mod.action in (safety_config.ACTION_BLOCK, safety_config.ACTION_SOFT_REFUSE): refuse = safety_config.message_for(in_mod.action, in_mod.top_category) elif inj.score >= 0.8: refuse = ("I can't follow instructions that try to override my safety " "guidelines or reveal my system configuration.") if refuse is not None: conv_id = req.conversation_id or create_conversation( title=f"[Agent] {user_text[:55] or 'Agent Chat'}" ) for m in messages: add_message(conv_id, m["role"], m["content"]) add_message(conv_id, "assistant", refuse) return _static_stream( refuse, conv_id, f"agentcmpl-{uuid.uuid4().hex[:12]}", int(time.time()), obj="agent.completion.chunk", ) await _acquire_slot() # Phase S: concurrency cap try: conv_id = req.conversation_id if conv_id is None: first_user = next( (m["content"] for m in messages if m["role"] == "user"), "Agent Chat" ) conv_id = create_conversation(title=f"[Agent] {first_user[:55]}") for m in messages: add_message(conv_id, m["role"], m["content"]) completion_id = f"agentcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) except Exception: _gen_semaphore.release() raise async def event_stream(): try: loop = asyncio.get_event_loop() queue: asyncio.Queue = asyncio.Queue() stop_event = threading.Event() def _run(): try: for kind, value in run_agent( messages, generate_tokens, max_tokens=req.max_tokens, temperature=req.temperature, top_k=req.top_k, top_p=req.top_p, ): if stop_event.is_set(): # Phase W: halt on unsafe output break asyncio.run_coroutine_threadsafe( queue.put((kind, value)), loop ) except Exception as exc: asyncio.run_coroutine_threadsafe( queue.put(("error", str(exc))), loop ) finally: asyncio.run_coroutine_threadsafe( queue.put(("end", None)), loop ) _executor.submit(_run) streamed: list[str] = [] since_check = 0 while True: kind, value = await queue.get() if kind == "error": yield f'data: {{"error": "{value}"}}\n\n' break if kind == "end": final_chunk = { "id": completion_id, "object": "agent.completion.chunk", "created": created, "model": "abyss-1", "conversation_id": conv_id, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], } yield f"data: {json.dumps(final_chunk)}\n\n" yield "data: [DONE]\n\n" break if kind == "done": out_mod = moderate_output(value) if out_mod.flagged: audit.record_result( "output", "agent", out_mod, identity=identity, conversation_id=conv_id, snippet=value, ) stored = safety_config.OUTPUT_BLOCK_MESSAGE if out_mod.blocked else value add_message(conv_id, "assistant", stored) continue # ----- PHASE W: ROLLING OUTPUT MODERATION (model tokens only) ----- if kind == "token": streamed.append(value) since_check += len(value) if since_check >= 48: since_check = 0 out_mod = moderate_output("".join(streamed)) if out_mod.action == safety_config.ACTION_BLOCK: stop_event.set() audit.record_result( "output", "agent", out_mod, identity=identity, conversation_id=conv_id, snippet="".join(streamed), ) add_message(conv_id, "assistant", safety_config.OUTPUT_BLOCK_MESSAGE) yield "data: " + json.dumps({ "id": completion_id, "object": "agent.completion.chunk", "created": created, "model": "abyss-1", "conversation_id": conv_id, "choices": [{"index": 0, "delta": {"content": " " + safety_config.OUTPUT_BLOCK_MESSAGE, "type": "token"}, "finish_reason": "content_filter"}], }) + "\n\n" yield "data: [DONE]\n\n" break chunk = { "id": completion_id, "object": "agent.completion.chunk", "created": created, "model": "abyss-1", "conversation_id": conv_id, "choices": [{ "index": 0, "delta": {"content": value, "type": kind}, "finish_reason": None, }], } yield f"data: {json.dumps(chunk)}\n\n" finally: _gen_semaphore.release() return StreamingResponse(event_stream(), media_type="text/event-stream") # ----- MODERATIONS (PHASE W) ----- @app.post( "/v1/moderations", response_model=ModerationResponse, dependencies=[ Depends(verify_key), Depends(rate_limit(*RATE_LIMIT_READ, scope="read")), ], ) def moderations(req: ModerationRequest, request: Request): """Standalone moderation endpoint - classify a piece of text directly.""" res = moderate_input(req.input) if res.flagged: audit.record_result( "api", "moderations", res, identity=identity_from_request(request), snippet=req.input, ) return ModerationResponse(**res.to_dict()) @app.get("/v1/admin/moderation", dependencies=[Depends(require_admin)]) def admin_moderation(limit: int = 100, min_severity: str | None = None): """Review the moderation audit trail. Requires the X-Admin-Key header.""" return list_moderation_events(limit=limit, min_severity=min_severity)