Update app.py

app.py

@@ -1,99 +1,391 @@
 import os
 import time
-from fastapi import FastAPI, Request, HTTPException, Depends
+import secrets
+
+from fastapi import (
+    FastAPI,
+    Request,
+    HTTPException,
+    Depends,
+)
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
 
-app = FastAPI(title="Qwen 2.5 Coder Public Secure Server")
+
+# ============================================================
+# APPLICATION
+# ============================================================
+
+app = FastAPI(
+    title="Varsh LLM",
+    description="Secure Qwen 2.5 Coder 7B OpenAI-compatible API",
+    version="1.0.0",
+)
+
+
+# ============================================================
+# CORS
+# ============================================================
 
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_credentials=True,
+    allow_credentials=False,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# 🔐 YOUR PREFERRED CUSTOM API KEY
-CUSTOM_API_KEY = "varsshancoder"
 
-security = HTTPBearer()
+# ============================================================
+# CONFIGURATION
+# ============================================================
+
+MODEL_ID = "qwen2.5-coder-7b-instruct-q4_k_m"
+
+MODEL_REPO = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF"
+
+MODEL_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf"
+
+# Read API key from Hugging Face Space Secret
+CUSTOM_API_KEY = os.getenv("CUSTOM_API_KEY")
+
+if not CUSTOM_API_KEY:
+    raise RuntimeError(
+        "CUSTOM_API_KEY is not configured. "
+        "Add CUSTOM_API_KEY in Hugging Face Space Settings > Secrets."
+    )
+
+
+# ============================================================
+# SECURITY
+# ============================================================
 
-def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
-    if credentials.credentials != CUSTOM_API_KEY:
+security = HTTPBearer(auto_error=False)
+
+
+def verify_api_key(
+    credentials: HTTPAuthorizationCredentials | None = Depends(security),
+):
+    """
+    Validate:
+
+    Authorization: Bearer <CUSTOM_API_KEY>
+    """
+
+    if credentials is None:
         raise HTTPException(
             status_code=401,
-            detail="Invalid or missing custom API Key."
+            detail="Missing API key.",
+            headers={"WWW-Authenticate": "Bearer"},
         )
-    return credentials.credentials
 
-MODEL_REPO = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF"
-MODEL_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf"
+    if credentials.scheme.lower() != "bearer":
+        raise HTTPException(
+            status_code=401,
+            detail="Authentication scheme must be Bearer.",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    if not secrets.compare_digest(
+        credentials.credentials,
+        CUSTOM_API_KEY,
+    ):
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid API key.",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    return True
+
+
+# ============================================================
+# MODEL LOADING
+# ============================================================
 
-print("Checking and downloading Qwen 2.5 Coder 7B GGUF model...")
-model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
+print("=" * 70)
+print("Starting Varsh LLM")
+print("=" * 70)
+
+print(f"Model repository: {MODEL_REPO}")
+print(f"Model file:       {MODEL_FILE}")
+print("Checking/downloading model...")
+
+model_path = hf_hub_download(
+    repo_id=MODEL_REPO,
+    filename=MODEL_FILE,
+)
+
+print(f"Model path: {model_path}")
+print("Loading Qwen 2.5 Coder 7B...")
 
-print(f"Loading model into memory from: {model_path}")
 llm = Llama(
     model_path=model_path,
-    n_ctx=4096,      
-    n_threads=2,     
+
+    # Context window
+    n_ctx=4096,
+
+    # Your HF Space has 2 vCPU
+    n_threads=2,
+
+    # Batch size
     n_batch=512,
-    verbose=False
+
+    # CPU inference
+    verbose=False,
 )
-print("Model loaded successfully and secured!")
 
-@app.post("/v1/chat/completions", dependencies=[Depends(verify_api_key)])
+print("Model loaded successfully.")
+print("=" * 70)
+
+
+# ============================================================
+# HEALTH / ROOT
+# ============================================================
+
+@app.get("/")
+async def root():
+    return {
+        "status": "healthy",
+        "security": "enabled",
+        "model": MODEL_ID,
+        "endpoint": "/v1/chat/completions",
+    }
+
+
+@app.get("/health")
+async def health():
+    return {
+        "status": "healthy",
+        "model": MODEL_ID,
+    }
+
+
+# ============================================================
+# CHAT COMPLETIONS
+# ============================================================
+
+@app.post(
+    "/v1/chat/completions",
+    dependencies=[Depends(verify_api_key)],
+)
 async def chat_completions(request: Request):
-    body = await request.json()
-    messages = body.get("messages", [])
+
+    # --------------------------------------------------------
+    # Parse JSON
+    # --------------------------------------------------------
+
+    try:
+        body = await request.json()
+    except Exception:
+        raise HTTPException(
+            status_code=400,
+            detail="Invalid JSON request body.",
+        )
+
+    # --------------------------------------------------------
+    # Extract parameters
+    # --------------------------------------------------------
+
+    messages = body.get("messages")
+
+    if not messages:
+        raise HTTPException(
+            status_code=400,
+            detail="The 'messages' field is required.",
+        )
+
+    requested_model = body.get("model", MODEL_ID)
+
+    # Accept our model ID
+    if requested_model != MODEL_ID:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Unsupported model. Use '{MODEL_ID}'.",
+        )
+
     temperature = body.get("temperature", 0.2)
+
     max_tokens = body.get("max_tokens", 1024)
 
+    # Safety limits
+    try:
+        temperature = float(temperature)
+        max_tokens = int(max_tokens)
+    except (TypeError, ValueError):
+        raise HTTPException(
+            status_code=400,
+            detail="temperature must be a number and max_tokens must be an integer.",
+        )
+
+    if temperature < 0:
+        temperature = 0.0
+
+    if temperature > 2:
+        temperature = 2.0
+
+    if max_tokens < 1:
+        max_tokens = 1
+
+    if max_tokens > 4096:
+        max_tokens = 4096
+
+    # --------------------------------------------------------
+    # Build Qwen ChatML prompt
+    # --------------------------------------------------------
+
     formatted_prompt = ""
-    for msg in messages:
-        role = msg.get("role", "user")
-        content = msg.get("content", "")
-        formatted_prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
+
+    for message in messages:
+
+        role = message.get("role", "user")
+
+        content = message.get("content", "")
+
+        # Handle simple OpenAI-style text messages
+        if isinstance(content, list):
+
+            text_parts = []
+
+            for item in content:
+                if isinstance(item, dict):
+                    if item.get("type") == "text":
+                        text_parts.append(
+                            item.get("text", "")
+                        )
+
+            content = "\n".join(text_parts)
+
+        if not isinstance(content, str):
+            content = str(content)
+
+        formatted_prompt += (
+            f"<|im_start|>{role}\n"
+            f"{content}"
+            f"<|im_end|>\n"
+        )
+
     formatted_prompt += "<|im_start|>assistant\n"
 
-    response = llm(
-        formatted_prompt,
-        max_tokens=max_tokens,
-        stop=["<|im_end|>", "<|im_start|>"],
-        temperature=temperature
+    # --------------------------------------------------------
+    # Generate
+    # --------------------------------------------------------
+
+    start_time = time.time()
+
+    try:
+
+        response = llm(
+            formatted_prompt,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            stop=[
+                "<|im_end|>",
+                "<|im_start|>",
+            ],
+        )
+
+    except Exception as e:
+
+        print(f"Generation error: {e}")
+
+        raise HTTPException(
+            status_code=500,
+            detail="Model generation failed.",
+        )
+
+    generation_time = time.time() - start_time
+
+    # --------------------------------------------------------
+    # IMPORTANT:
+    # llama.cpp returns choices as a LIST
+    # --------------------------------------------------------
+
+    try:
+        text_output = response["choices"][0]["text"]
+    except (KeyError, IndexError, TypeError):
+
+        print(f"Unexpected model response: {response}")
+
+        raise HTTPException(
+            status_code=500,
+            detail="Invalid response received from model.",
+        )
+
+    # --------------------------------------------------------
+    # Token information
+    # --------------------------------------------------------
+
+    usage = response.get("usage", {})
+
+    prompt_tokens = usage.get(
+        "prompt_tokens",
+        0,
+    )
+
+    completion_tokens = usage.get(
+        "completion_tokens",
+        0,
+    )
+
+    total_tokens = usage.get(
+        "total_tokens",
+        prompt_tokens + completion_tokens,
     )
 
-    text_output = response["choices"]["text"]
-    
-    return JSONResponse(content={
-        "id": f"chatcmpl-{int(time.time())}",
-        "object": "chat.completion",
-        "created": int(time.time()),
-        "model": "qwen2.5-coder-7b",
-        "choices": [{
-            "index": 0,
-            "message": {
-                "role": "assistant",
-                "content": text_output
+    # --------------------------------------------------------
+    # OpenAI-compatible response
+    # --------------------------------------------------------
+
+    created = int(time.time())
+
+    return JSONResponse(
+        content={
+            "id": f"chatcmpl-{created}",
+
+            "object": "chat.completion",
+
+            "created": created,
+
+            "model": MODEL_ID,
+
+            "choices": [
+                {
+                    "index": 0,
+
+                    "message": {
+                        "role": "assistant",
+                        "content": text_output,
+                    },
+
+                    "finish_reason": "stop",
+                }
+            ],
+
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
             },
-            "finish_reason": "stop"
-        }],
-        "usage": {
-            "prompt_tokens": len(formatted_prompt) // 4, 
-            "completion_tokens": len(text_output) // 4,
-            "total_tokens": (len(formatted_prompt) + len(text_output)) // 4
+
+            "x_generation_time_seconds": round(
+                generation_time,
+                3,
+            ),
         }
-    })
+    )
 
-@app.get("/")
-def index():
-    return {
-        "status": "healthy", 
-        "security": "enabled",
-        "model": "qwen2.5-coder-7b-instruct-q4_k_m", 
-        "endpoint": "/v1/chat/completions"
-    }
+
+# ============================================================
+# STARTUP MESSAGE
+# ============================================================
+
+print("=" * 70)
+print("Varsh LLM API READY")
+print("Listening on: http://0.0.0.0:7860")
+print("Endpoint:     /v1/chat/completions")
+print("Authentication: Bearer API key")
+print("=" * 70)