Blame for server/main.py

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)