Update app.py

app.py

@@ -1,80 +1,85 @@
 import os
 import time
-import secrets
+import uuid
 
-from fastapi import (
-    FastAPI,
-    Request,
-    HTTPException,
-    Depends,
-)
+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 fastapi.responses import JSONResponse
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
 
 
 # ============================================================
-# APPLICATION
+# CONFIGURATION
 # ============================================================
 
-app = FastAPI(
-    title="Varsh LLM",
-    description="Secure Qwen 2.5 Coder 7B OpenAI-compatible API",
-    version="1.0.0",
-)
+APP_NAME = "Varsh LLM"
+MODEL_REPO = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF"
+MODEL_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf"
 
+MODEL_ID = "qwen2.5-coder-7b-instruct-q4_k_m"
 
-# ============================================================
-# CORS
-# ============================================================
+# IMPORTANT:
+# Set CUSTOM_API_KEY in:
+# Hugging Face Space → Settings → Secrets
+CUSTOM_API_KEY = os.getenv("CUSTOM_API_KEY")
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=False,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+if not CUSTOM_API_KEY:
+    raise RuntimeError(
+        "CUSTOM_API_KEY is not configured. "
+        "Add CUSTOM_API_KEY in Hugging Face Space Settings > Secrets."
+    )
 
 
 # ============================================================
-# CONFIGURATION
+# SERVER SETTINGS
 # ============================================================
 
-MODEL_ID = "qwen2.5-coder-7b-instruct-q4_k_m"
+HOST = "0.0.0.0"
+PORT = int(os.getenv("PORT", "7860"))
 
-MODEL_REPO = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF"
+# Your current HF Space has 2 vCPU.
+# Keep this at 2 unless you upgrade the Space.
+N_THREADS = int(os.getenv("N_THREADS", "2"))
 
-MODEL_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf"
+# Prompt-processing batch size.
+N_BATCH = int(os.getenv("N_BATCH", "512"))
 
-# Read API key from Hugging Face Space Secret
-CUSTOM_API_KEY = os.getenv("CUSTOM_API_KEY")
+# Context size.
+N_CTX = int(os.getenv("N_CTX", "4096"))
 
-if not CUSTOM_API_KEY:
-    raise RuntimeError(
-        "CUSTOM_API_KEY is not configured. "
-        "Add CUSTOM_API_KEY in Hugging Face Space Settings > Secrets."
-    )
+
+# ============================================================
+# FASTAPI
+# ============================================================
+
+app = FastAPI(
+    title="Varsh LLM API",
+    version="1.0.0",
+    description="OpenAI-compatible Qwen 2.5 Coder API",
+)
+
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=False,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 
 # ============================================================
-# SECURITY
+# AUTHENTICATION
 # ============================================================
 
 security = HTTPBearer(auto_error=False)
 
 
 def verify_api_key(
-    credentials: HTTPAuthorizationCredentials | None = Depends(security),
+    credentials: HTTPAuthorizationCredentials = Depends(security),
 ):
-    """
-    Validate:
-
-    Authorization: Bearer <CUSTOM_API_KEY>
-    """
-
     if credentials is None:
         raise HTTPException(
             status_code=401,
@@ -85,21 +90,16 @@ def verify_api_key(
     if credentials.scheme.lower() != "bearer":
         raise HTTPException(
             status_code=401,
-            detail="Authentication scheme must be Bearer.",
-            headers={"WWW-Authenticate": "Bearer"},
+            detail="Authentication must use Bearer token.",
         )
 
-    if not secrets.compare_digest(
-        credentials.credentials,
-        CUSTOM_API_KEY,
-    ):
+    if credentials.credentials != CUSTOM_API_KEY:
         raise HTTPException(
             status_code=401,
             detail="Invalid API key.",
-            headers={"WWW-Authenticate": "Bearer"},
         )
 
-    return True
+    return credentials.credentials
 
 
 # ============================================================
@@ -110,8 +110,12 @@ print("=" * 70)
 print("Starting Varsh LLM")
 print("=" * 70)
 
-print(f"Model repository: {MODEL_REPO}")
-print(f"Model file:       {MODEL_FILE}")
+print(f"Model repository : {MODEL_REPO}")
+print(f"Model file       : {MODEL_FILE}")
+print(f"Threads          : {N_THREADS}")
+print(f"Batch size       : {N_BATCH}")
+print(f"Context size     : {N_CTX}")
+
 print("Checking/downloading model...")
 
 model_path = hf_hub_download(
@@ -124,46 +128,90 @@ print("Loading Qwen 2.5 Coder 7B...")
 
 llm = Llama(
     model_path=model_path,
-
-    # Context window
-    n_ctx=4096,
-
-    # Your HF Space has 2 vCPU
-    n_threads=2,
-
-    # Batch size
-    n_batch=512,
-
-    # CPU inference
+    n_ctx=N_CTX,
+    n_threads=N_THREADS,
+    n_batch=N_BATCH,
     verbose=False,
 )
 
 print("Model loaded successfully.")
+
+print("=" * 70)
+print("Varsh LLM API READY")
+print("=" * 70)
+print(f"Listening on: http://{HOST}:{PORT}")
+print("Root endpoint: /")
+print("Models endpoint: /v1/models")
+print("Chat endpoint: /v1/chat/completions")
+print("Authentication: Bearer API key")
 print("=" * 70)
 
 
 # ============================================================
-# HEALTH / ROOT
+# ROOT / HEALTH
 # ============================================================
 
 @app.get("/")
-async def root():
+def root():
     return {
         "status": "healthy",
+        "service": APP_NAME,
         "security": "enabled",
         "model": MODEL_ID,
         "endpoint": "/v1/chat/completions",
+        "models_endpoint": "/v1/models",
     }
 
 
 @app.get("/health")
-async def health():
+def health():
     return {
         "status": "healthy",
         "model": MODEL_ID,
     }
 
 
+# ============================================================
+# OPENAI-COMPATIBLE MODELS ENDPOINT
+# ============================================================
+
+@app.get(
+    "/v1/models",
+    dependencies=[Depends(verify_api_key)],
+)
+def list_models():
+    return {
+        "object": "list",
+        "data": [
+            {
+                "id": MODEL_ID,
+                "object": "model",
+                "created": int(time.time()),
+                "owned_by": "varsshancoder",
+            }
+        ],
+    }
+
+
+@app.get(
+    "/v1/models/{model_id}",
+    dependencies=[Depends(verify_api_key)],
+)
+def get_model(model_id: str):
+    if model_id != MODEL_ID:
+        raise HTTPException(
+            status_code=404,
+            detail=f"Model '{model_id}' not found.",
+        )
+
+    return {
+        "id": MODEL_ID,
+        "object": "model",
+        "created": int(time.time()),
+        "owned_by": "varsshancoder",
+    }
+
+
 # ============================================================
 # CHAT COMPLETIONS
 # ============================================================
@@ -190,133 +238,131 @@ async def chat_completions(request: Request):
     # Extract parameters
     # --------------------------------------------------------
 
-    messages = body.get("messages")
+    requested_model = body.get("model", MODEL_ID)
 
-    if not messages:
+    if requested_model != MODEL_ID:
         raise HTTPException(
-            status_code=400,
-            detail="The 'messages' field is required.",
+            status_code=404,
+            detail=f"Model '{requested_model}' not found.",
         )
 
-    requested_model = body.get("model", MODEL_ID)
+    messages = body.get("messages")
 
-    # Accept our model ID
-    if requested_model != MODEL_ID:
+    if not messages or not isinstance(messages, list):
         raise HTTPException(
             status_code=400,
-            detail=f"Unsupported model. Use '{MODEL_ID}'.",
+            detail="The 'messages' field is required and must be a list.",
         )
 
     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):
+    # Optional OpenAI-compatible parameters
+    top_p = body.get("top_p", 0.95)
+    stop = body.get("stop", None)
+
+    stream = body.get("stream", False)
+
+    # Current implementation returns normal JSON.
+    # Streaming can be added separately once basic routing is stable.
+    if stream:
         raise HTTPException(
             status_code=400,
-            detail="temperature must be a number and max_tokens must be an integer.",
+            detail="Streaming is not enabled on this deployment.",
         )
 
-    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
+    # Sanitize messages
     # --------------------------------------------------------
 
-    formatted_prompt = ""
+    clean_messages = []
 
     for message in messages:
+        if not isinstance(message, dict):
+            continue
 
         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 content is None:
+            content = ""
 
+        # Handle normal string content.
+        # For now, multimodal content arrays are not supported.
         if not isinstance(content, str):
-            content = str(content)
-
-        formatted_prompt += (
-            f"<|im_start|>{role}\n"
-            f"{content}"
-            f"<|im_end|>\n"
+            raise HTTPException(
+                status_code=400,
+                detail="Message content must be a string.",
+            )
+
+        clean_messages.append(
+            {
+                "role": role,
+                "content": content,
+            }
         )
 
-    formatted_prompt += "<|im_start|>assistant\n"
+    if not clean_messages:
+        raise HTTPException(
+            status_code=400,
+            detail="No valid messages were provided.",
+        )
 
     # --------------------------------------------------------
-    # Generate
+    # Generate response
     # --------------------------------------------------------
 
-    start_time = time.time()
+    request_id = f"chatcmpl-{uuid.uuid4().hex}"
 
-    try:
+    created = int(time.time())
 
-        response = llm(
-            formatted_prompt,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            stop=[
-                "<|im_end|>",
-                "<|im_start|>",
-            ],
+    start_time = time.perf_counter()
+
+    try:
+        response = llm.create_chat_completion(
+            messages=clean_messages,
+            temperature=float(temperature),
+            max_tokens=int(max_tokens),
+            top_p=float(top_p),
+            stop=stop,
         )
 
     except Exception as e:
-
-        print(f"Generation error: {e}")
+        print(f"Inference error: {e}")
 
         raise HTTPException(
             status_code=500,
-            detail="Model generation failed.",
+            detail=f"Model inference failed: {str(e)}",
         )
 
-    generation_time = time.time() - start_time
+    generation_time = time.perf_counter() - start_time
 
     # --------------------------------------------------------
-    # IMPORTANT:
-    # llama.cpp returns choices as a LIST
+    # Extract generated text
     # --------------------------------------------------------
 
     try:
-        text_output = response["choices"][0]["text"]
-    except (KeyError, IndexError, TypeError):
+        choice = response["choices"][0]
+
+        message = choice["message"]
 
-        print(f"Unexpected model response: {response}")
+        text_output = message.get("content", "")
+
+        finish_reason = choice.get(
+            "finish_reason",
+            "stop",
+        )
+
+    except Exception as e:
+        print(f"Response parsing error: {e}")
 
         raise HTTPException(
             status_code=500,
-            detail="Invalid response received from model.",
+            detail="Invalid response returned by the model.",
         )
 
     # --------------------------------------------------------
-    # Token information
+    # Usage
     # --------------------------------------------------------
 
     usage = response.get("usage", {})
@@ -336,56 +382,41 @@ async def chat_completions(request: Request):
         prompt_tokens + completion_tokens,
     )
 
+    print(
+        f"Request {request_id} | "
+        f"prompt={prompt_tokens} | "
+        f"completion={completion_tokens} | "
+        f"time={generation_time:.2f}s"
+    )
+
     # --------------------------------------------------------
     # OpenAI-compatible response
     # --------------------------------------------------------
 
-    created = int(time.time())
-
     return JSONResponse(
         content={
-            "id": f"chatcmpl-{created}",
-
+            "id": request_id,
             "object": "chat.completion",
-
             "created": created,
-
             "model": MODEL_ID,
-
             "choices": [
                 {
                     "index": 0,
-
                     "message": {
                         "role": "assistant",
                         "content": text_output,
                     },
-
-                    "finish_reason": "stop",
+                    "finish_reason": finish_reason,
                 }
             ],
-
             "usage": {
                 "prompt_tokens": prompt_tokens,
                 "completion_tokens": completion_tokens,
                 "total_tokens": total_tokens,
             },
-
             "x_generation_time_seconds": round(
                 generation_time,
                 3,
             ),
         }
-    )
-
-
-# ============================================================
-# 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)
+    )