server/main.py

import os
from typing import List
from fastapi import FastAPI, HTTPException, Depends, Security
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security.api_key import APIKeyHeader
from pydantic import BaseModel
import torch
import yaml

from model.transformer import Transformer, ModelArgs
from tokenizer.tokenizer import BPETokenizer
from inference.engine import InferenceEngine

# --------------------------------------------------
# FastAPI App
# --------------------------------------------------

app = FastAPI(title="Production LLM API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --------------------------------------------------
# API Key
# --------------------------------------------------

API_KEY_NAME = "access_token"
api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)

async def get_api_key(api_key: str = Security(api_key_header)):
    if api_key == "default_secret_key" or not os.getenv("REQUIRE_AUTH"):
        return api_key
    raise HTTPException(status_code=403, detail="Invalid API key")

# --------------------------------------------------
# Global engine
# --------------------------------------------------

engine = None

# --------------------------------------------------
# Request models
# --------------------------------------------------

class ChatMessage(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: str = "custom-llm-1b"
    messages: List[ChatMessage]
    temperature: float = 0.7
    top_p: float = 0.9
    max_tokens: int = 128   # ✅ SAFE for CPU
    stream: bool = False

# --------------------------------------------------
# Startup: Load model from YAML
# --------------------------------------------------

@app.on_event("startup")
def startup_event():
    global engine

    print("🔄 Starting model load...")

    device = torch.device("cpu")

    # ✅ Load config path from env
    config_path = os.getenv("MODEL_CONFIG", "config/custom.yaml")
    print(f"📄 Using config: {config_path}")

    if not os.path.exists(config_path):
        raise RuntimeError(f"Config file not found: {config_path}")

    # ✅ Load YAML
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    if "model" not in config:
        raise RuntimeError("Config missing 'model' section")

    # ✅ Build model from config
    params = ModelArgs(**config["model"])

    model = Transformer(params).to(device)
    model.eval()

    # ✅ Load weights (if available)
    
    base_dir = os.getcwd()
    model_path = os.path.join(base_dir, "model", "model.pt")

    print("📂 Current directory:", base_dir)
    print("📦 Checking model path:", model_path)
    print("📦 Exists:", os.path.exists(model_path))

    
    if os.path.exists(model_path):
        print(f"📦 Loading weights from {model_path}")
    
        state_dict = torch.load(model_path, map_location=device)

    # ✅ Fix key mismatch: remove "model." prefix
        new_state_dict = {}
        for key, value in state_dict.items():
            if key.startswith("model."):
                new_key = key[len("model."):]  # remove prefix
            else:
                new_key = key
            new_state_dict[new_key] = value

    # ✅ Load safely
        result = model.load_state_dict(new_state_dict, strict=False)

        print("✅ Model weights loaded successfully")
        print("Missing keys:", result.missing_keys)
        print("Unexpected keys:", result.unexpected_keys)
    else:
        print("⚠️ No model weights found — using random weights")

    # ✅ Tokenizer
    tokenizer_path = config.get("tokenizer_path", "tokenizer/")
    tokenizer = BPETokenizer(tokenizer_path)

    engine = InferenceEngine(model, tokenizer)

    print("✅ Model loaded successfully!")

# --------------------------------------------------
# Health
# --------------------------------------------------

@app.get("/health")
def health():
    return {"status": "ok"}

# --------------------------------------------------
# Models endpoint
# --------------------------------------------------

@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [
            {"id": "custom-llm-1b", "object": "model"}
        ]
    }

# --------------------------------------------------
# Chat API (FIXED ✅)
# --------------------------------------------------

@app.post("/v1/chat/completions")
async def chat_completions(
    request: ChatCompletionRequest,
    api_key: str = Depends(get_api_key)
):

    if engine is None:
        raise HTTPException(status_code=500, detail="Model not loaded")

    # ✅ SIMPLE prompt (safe for tokenizer)
    prompt = ""
    for msg in request.messages:
        prompt += msg.content + "\n"

    # ✅ SAFE generation (prevents 500 crash)
    try:
        response_text = engine.generate(
            prompt,
            max_new_tokens=request.max_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            stream=False
        )
    except Exception as e:
        import traceback
        error_msg = traceback.format_exc()
        print("❌ ERROR:", error_msg)

        return {
            "id": "error",
            "object": "chat.completion",
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": f"Error:\n{error_msg}"
                },
                "finish_reason": "error"
            }]
        }

    return {
        "id": "chatcmpl-123",
        "object": "chat.completion",
        "model": request.model,
        "choices": [{
            "index": 0,
            "message": {
                "role": "assistant",
                "content": response_text,
            },
            "finish_reason": "stop"
        }],
        "usage": {
            "prompt_tokens": len(prompt),
            "completion_tokens": len(response_text),
            "total_tokens": len(prompt) + len(response_text)
        }
    }