serving/server.py

"""
AEC AI Reader - FastAPI Server (Production-Ready for Hugging Face Spaces)
Menangani 3 masalah deployment utama:
1. Model tidak bisa ada di repo Space (>1GB) → download dari HF Hub saat startup
2. SQLite cache hilang saat Space restart → preseed ulang dari Dataset repo saat startup
3. Tidak ada auth → OpenAI-compatible Bearer Token enforcement
4. Cold start timeout → /health endpoint agar client tahu kapan siap
"""

import os
import secrets
import time
import asyncio
import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Depends, Request
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
from typing import Optional, List, Any, Dict

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("aec-server")

# --- Konfigurasi dari Environment Variables (set di HF Space Secrets) ---
API_KEY = os.getenv("AEC_API_KEY", "aec-local-dev-key")
MODEL_REPO_ID   = os.getenv("MODEL_REPO_ID", "")        # e.g., "yourusername/qwen3-4b-aec-gguf" di HF Hub
MODEL_FILENAME  = os.getenv("MODEL_FILENAME", "qwen3-4b-instruct-q4_k_m.gguf")
GRAMMAR_PATH    = os.getenv("GRAMMAR_PATH", "serving/grammar.gbnf")
CACHE_DB_PATH   = os.getenv("CACHE_DB_PATH", "/tmp/chain_cache.sqlite")

# GitLab sebagai sumber Dataset (free tier, no CI needed, hanya raw file API)
GITLAB_TOKEN      = os.getenv("GITLAB_TOKEN", "")
GITLAB_PROJECT_ID = os.getenv("GITLAB_PROJECT_ID", "")
GITLAB_BRANCH     = os.getenv("GITLAB_BRANCH", "main")
GITLAB_DATASET_PATH = os.getenv("GITLAB_DATASET_PATH", "dataset/output/training_data_v2.jsonl")

engine = None
startup_ready = False
startup_error = None


# --- Security: Constant-time API Key Comparison (anti-timing-attack) ---
security = HTTPBearer()

def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
    is_valid = secrets.compare_digest(credentials.credentials, API_KEY)
    if not is_valid:
        raise HTTPException(status_code=403, detail="Invalid API key")
    return credentials.credentials


# --- Startup: Download model + preseed cache ---
async def load_engine_background():
    global engine, startup_ready, startup_error
    try:
        model_path = MODEL_FILENAME  # default: sudah ada di repo (dev mode)

        if MODEL_REPO_ID:
            log.info(f"[Startup] Downloading model from HF Hub: {MODEL_REPO_ID}/{MODEL_FILENAME}")
            from huggingface_hub import hf_hub_download
            model_path = hf_hub_download(
                repo_id=MODEL_REPO_ID,
                filename=MODEL_FILENAME,
                local_dir="/tmp/models"
            )
            log.info(f"[Startup] Model downloaded to: {model_path}")

        # Import engine setelah model tersedia
        from serving.inference import AECInferenceEngine
        engine = AECInferenceEngine(
            model_path=model_path,
            grammar_path=GRAMMAR_PATH,
            cache_db_path=CACHE_DB_PATH
        )

        # Preseed cache dari GitLab Raw API (free tier, tidak perlu CI aktif)
        if GITLAB_PROJECT_ID and GITLAB_TOKEN and engine.cache.stats()["total_entries"] < 100:
            import urllib.request
            import urllib.parse

            encoded_path = urllib.parse.quote(GITLAB_DATASET_PATH, safe="")
            gitlab_url = (
                f"https://gitlab.com/api/v4/projects/{GITLAB_PROJECT_ID}"
                f"/repository/files/{encoded_path}/raw?ref={GITLAB_BRANCH}"
            )
            log.info(f"[Startup] Downloading dataset from GitLab: project {GITLAB_PROJECT_ID}")

            req = urllib.request.Request(gitlab_url)
            req.add_header("PRIVATE-TOKEN", GITLAB_TOKEN)

            dataset_local = "/tmp/training_data_v2.jsonl"
            with urllib.request.urlopen(req, timeout=120) as resp, \
                 open(dataset_local, "wb") as f:
                f.write(resp.read())

            log.info(f"[Startup] Dataset downloaded. Preseeding cache...")
            engine.cache.preseed_from_dataset(dataset_local)

        startup_ready = True
        log.info("[Startup] AEC AI Engine ready.")

    except Exception as e:
        startup_error = str(e)
        log.error(f"[Startup] FAILED: {e}")


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Jalankan loading di background agar HF tidak kill proses karena startup timeout
    asyncio.create_task(load_engine_background())
    yield


# --- App ---
app = FastAPI(
    title="AEC AI Reader API",
    version="2.0.0",
    lifespan=lifespan,
    docs_url=None,   # Sembunyikan Swagger di produksi
    redoc_url=None
)


# --- OpenAI-Compatible Request/Response Models ---
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: str = "aec-reader"
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.3
    max_tokens: Optional[int] = 1024
    stream: Optional[bool] = False

class ChatCompletionChoice(BaseModel):
    index: int
    message: ChatMessage
    finish_reason: str

class ChatCompletionResponse(BaseModel):
    id: str
    object: str = "chat.completion"
    model: str = "aec-reader"
    choices: List[ChatCompletionChoice]
    usage: Dict[str, Any] = {}


# --- Endpoints ---
@app.get("/", response_class=HTMLResponse)
async def root():
    status = "LOADING..." if not startup_ready else "ONLINE"
    color = "#f59e0b" if not startup_ready else "#10b981"
    return f"""
    <html><head><title>AEC AI Engine</title></head>
    <body style='background:#111;color:#fff;font-family:monospace;display:flex;align-items:center;justify-content:center;height:100vh;margin:0'>
    <div style='text-align:center'>
    <div style='width:14px;height:14px;border-radius:50%;background:{color};display:inline-block;margin-right:8px'></div>
    AEC AI Engine &nbsp;<strong style='color:{color}'>{status}</strong>
    <br><small style='color:#666;margin-top:8px;display:block'>Endpoint: POST /v1/chat/completions</small>
    </div></body></html>
    """

@app.get("/health")
async def health():
    return {
        "ready": startup_ready,
        "error": startup_error,
        "cache_entries": engine.cache.stats()["total_entries"] if engine else 0
    }

@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
async def chat_completions(
    req: ChatCompletionRequest,
    _key: str = Depends(verify_api_key)
):
    # Jika engine masih loading, kembalikan 503 bukan hang
    if not startup_ready:
        raise HTTPException(
            status_code=503,
            detail=f"Engine loading. Startup error: {startup_error}" if startup_error else "Engine still loading, retry in 30s"
        )

    # Ambil konten pesan terakhir dari user sebagai instruksi
    user_msg = next(
        (m.content for m in reversed(req.messages) if m.role == "user"),
        None
    )
    if not user_msg:
        raise HTTPException(status_code=400, detail="No user message found")

    try:
        result = engine.process_instruction(user_msg)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    import json
    import time
    response_text = json.dumps(result["output"], ensure_ascii=False, indent=2)

    return ChatCompletionResponse(
        id=f"aec-{int(time.time())}",
        model="aec-reader",
        choices=[
            ChatCompletionChoice(
                index=0,
                message=ChatMessage(role="assistant", content=response_text),
                finish_reason="stop"
            )
        ],
        usage={
            "source": result["source"],
            "output_type": result["output_type"],
            "cache_similarity": result.get("similarity", 0)
        }
    )


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("serving.server:app", host="0.0.0.0", port=7860, reload=False)