Update app.py

app.py

@@ -1,11 +1,14 @@
 import os
 import time
 import uuid
+import json
+import threading
 
 from fastapi import FastAPI, Request, HTTPException, Depends
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-from fastapi.responses import JSONResponse
+
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
 
@@ -15,43 +18,51 @@ from llama_cpp import Llama
 # ============================================================
 
 APP_NAME = "Varsh LLM"
+
 MODEL_REPO = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF"
 MODEL_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf"
 
+# This MUST exactly match the model ID expected by OmniRoute.
 MODEL_ID = "qwen2.5-coder-7b-instruct-q4_k_m"
 
-# IMPORTANT:
-# Set CUSTOM_API_KEY in:
-# Hugging Face Space → Settings → Secrets
-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."
-    )
-
-
-# ============================================================
-# SERVER SETTINGS
-# ============================================================
-
 HOST = "0.0.0.0"
 PORT = int(os.getenv("PORT", "7860"))
 
-# Your current HF Space has 2 vCPU.
+# Your current HF Space = 2 vCPU.
 # Keep this at 2 unless you upgrade the Space.
 N_THREADS = int(os.getenv("N_THREADS", "2"))
 
-# Prompt-processing batch size.
+# llama.cpp prompt processing batch size.
 N_BATCH = int(os.getenv("N_BATCH", "512"))
 
-# Context size.
+# Context window.
 N_CTX = int(os.getenv("N_CTX", "4096"))
 
 
 # ============================================================
-# FASTAPI
+# API KEY
+# ============================================================
+
+# Set this in:
+#
+# Hugging Face Space
+# -> Settings
+# -> Secrets
+# -> CUSTOM_API_KEY
+#
+# DO NOT hard-code the secret in this file.
+
+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."
+    )
+
+
+# ============================================================
+# FASTAPI APPLICATION
 # ============================================================
 
 app = FastAPI(
@@ -61,6 +72,10 @@ app = FastAPI(
 )
 
 
+# ============================================================
+# CORS
+# ============================================================
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -84,7 +99,9 @@ def verify_api_key(
         raise HTTPException(
             status_code=401,
             detail="Missing API key.",
-            headers={"WWW-Authenticate": "Bearer"},
+            headers={
+                "WWW-Authenticate": "Bearer"
+            },
         )
 
     if credentials.scheme.lower() != "bearer":
@@ -112,6 +129,7 @@ print("=" * 70)
 
 print(f"Model repository : {MODEL_REPO}")
 print(f"Model file       : {MODEL_FILE}")
+print(f"Model ID         : {MODEL_ID}")
 print(f"Threads          : {N_THREADS}")
 print(f"Batch size       : {N_BATCH}")
 print(f"Context size     : {N_CTX}")
@@ -124,31 +142,62 @@ model_path = hf_hub_download(
 )
 
 print(f"Model path: {model_path}")
+
 print("Loading Qwen 2.5 Coder 7B...")
 
 llm = Llama(
     model_path=model_path,
+
+    # Context window
     n_ctx=N_CTX,
+
+    # CPU threads
     n_threads=N_THREADS,
+
+    # Prompt processing batch
     n_batch=N_BATCH,
+
+    # Let llama.cpp use its Qwen chat template
+    chat_format=None,
+
     verbose=False,
 )
 
 print("Model loaded successfully.")
 
+
+# ============================================================
+# GENERATION LOCK
+# ============================================================
+
+# The HF Space has one model instance on CPU.
+#
+# This prevents multiple requests from trying to generate
+# through the same llama.cpp instance simultaneously.
+
+generation_lock = threading.Lock()
+
+
+# ============================================================
+# STARTUP INFORMATION
+# ============================================================
+
 print("=" * 70)
 print("Varsh LLM API READY")
 print("=" * 70)
+
 print(f"Listening on: http://{HOST}:{PORT}")
 print("Root endpoint: /")
+print("Health endpoint: /health")
 print("Models endpoint: /v1/models")
 print("Chat endpoint: /v1/chat/completions")
 print("Authentication: Bearer API key")
+print("Streaming: ENABLED")
 print("=" * 70)
 
 
 # ============================================================
-# ROOT / HEALTH
+# ROOT
 # ============================================================
 
 @app.get("/")
@@ -160,19 +209,25 @@ def root():
         "model": MODEL_ID,
         "endpoint": "/v1/chat/completions",
         "models_endpoint": "/v1/models",
+        "streaming": True,
     }
 
 
+# ============================================================
+# HEALTH
+# ============================================================
+
 @app.get("/health")
 def health():
     return {
         "status": "healthy",
+        "service": APP_NAME,
         "model": MODEL_ID,
     }
 
 
 # ============================================================
-# OPENAI-COMPATIBLE MODELS ENDPOINT
+# OPENAI /v1/models
 # ============================================================
 
 @app.get(
@@ -193,11 +248,16 @@ def list_models():
     }
 
 
+# ============================================================
+# OPENAI /v1/models/{model_id}
+# ============================================================
+
 @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,
@@ -212,6 +272,59 @@ def get_model(model_id: str):
     }
 
 
+# ============================================================
+# MESSAGE VALIDATION
+# ============================================================
+
+def clean_messages(messages):
+
+    if not isinstance(messages, list):
+        raise HTTPException(
+            status_code=400,
+            detail="'messages' must be a list.",
+        )
+
+    cleaned = []
+
+    for message in messages:
+
+        if not isinstance(message, dict):
+            raise HTTPException(
+                status_code=400,
+                detail="Each message must be an object.",
+            )
+
+        role = message.get("role", "user")
+        content = message.get("content", "")
+
+        if content is None:
+            content = ""
+
+        if not isinstance(content, str):
+            raise HTTPException(
+                status_code=400,
+                detail=(
+                    "Message content must be a string. "
+                    "Multimodal content is not supported."
+                ),
+            )
+
+        cleaned.append(
+            {
+                "role": role,
+                "content": content,
+            }
+        )
+
+    if not cleaned:
+        raise HTTPException(
+            status_code=400,
+            detail="No messages were provided.",
+        )
+
+    return cleaned
+
+
 # ============================================================
 # CHAT COMPLETIONS
 # ============================================================
@@ -228,6 +341,7 @@ async def chat_completions(request: Request):
 
     try:
         body = await request.json()
+
     except Exception:
         raise HTTPException(
             status_code=400,
@@ -235,118 +349,356 @@ async def chat_completions(request: Request):
         )
 
     # --------------------------------------------------------
-    # Extract parameters
+    # Model
     # --------------------------------------------------------
 
-    requested_model = body.get("model", MODEL_ID)
+    requested_model = body.get(
+        "model",
+        MODEL_ID,
+    )
 
     if requested_model != MODEL_ID:
+
         raise HTTPException(
             status_code=404,
             detail=f"Model '{requested_model}' not found.",
         )
 
-    messages = body.get("messages")
+    # --------------------------------------------------------
+    # Messages
+    # --------------------------------------------------------
 
-    if not messages or not isinstance(messages, list):
-        raise HTTPException(
-            status_code=400,
-            detail="The 'messages' field is required and must be a list.",
-        )
+    messages = clean_messages(
+        body.get("messages")
+    )
 
-    temperature = body.get("temperature", 0.2)
-    max_tokens = body.get("max_tokens", 1024)
+    # --------------------------------------------------------
+    # Parameters
+    # --------------------------------------------------------
 
-    # Optional OpenAI-compatible parameters
-    top_p = body.get("top_p", 0.95)
-    stop = body.get("stop", None)
+    temperature = body.get(
+        "temperature",
+        0.2,
+    )
 
-    stream = body.get("stream", False)
+    max_tokens = body.get(
+        "max_tokens",
+        1024,
+    )
 
-    # Current implementation returns normal JSON.
-    # Streaming can be added separately once basic routing is stable.
-    if stream:
-        raise HTTPException(
-            status_code=400,
-            detail="Streaming is not enabled on this deployment.",
-        )
+    top_p = body.get(
+        "top_p",
+        0.95,
+    )
+
+    stop = body.get(
+        "stop",
+        None,
+    )
+
+    stream = body.get(
+        "stream",
+        False,
+    )
 
     # --------------------------------------------------------
-    # Sanitize messages
+    # Optional seed
     # --------------------------------------------------------
 
-    clean_messages = []
+    seed = body.get(
+        "seed",
+        None,
+    )
 
-    for message in messages:
-        if not isinstance(message, dict):
-            continue
+    # --------------------------------------------------------
+    # Request metadata
+    # --------------------------------------------------------
 
-        role = message.get("role", "user")
-        content = message.get("content", "")
+    request_id = (
+        "chatcmpl-"
+        + uuid.uuid4().hex
+    )
 
-        if content is None:
-            content = ""
+    created = int(
+        time.time()
+    )
 
-        # Handle normal string content.
-        # For now, multimodal content arrays are not supported.
-        if not isinstance(content, str):
-            raise HTTPException(
-                status_code=400,
-                detail="Message content must be a string.",
-            )
+    # ========================================================
+    # STREAMING
+    # ========================================================
 
-        clean_messages.append(
-            {
-                "role": role,
-                "content": content,
-            }
-        )
+    if stream:
 
-    if not clean_messages:
-        raise HTTPException(
-            status_code=400,
-            detail="No valid messages were provided.",
-        )
+        def generate_stream():
 
-    # --------------------------------------------------------
-    # Generate response
-    # --------------------------------------------------------
+            start_time = time.perf_counter()
+
+            try:
+
+                # ------------------------------------------------
+                # Lock model generation
+                # ------------------------------------------------
+
+                with generation_lock:
+
+                    stream_kwargs = {
+                        "messages": messages,
+                        "temperature": float(temperature),
+                        "max_tokens": int(max_tokens),
+                        "top_p": float(top_p),
+                        "stream": True,
+                    }
+
+                    if stop is not None:
+                        stream_kwargs["stop"] = stop
+
+                    if seed is not None:
+                        stream_kwargs["seed"] = int(seed)
+
+                    response_stream = (
+                        llm.create_chat_completion(
+                            **stream_kwargs
+                        )
+                    )
+
+                    first_chunk = True
+
+                    for chunk in response_stream:
+
+                        choices = chunk.get(
+                            "choices",
+                            [],
+                        )
+
+                        if not choices:
+                            continue
+
+                        choice = choices[0]
+
+                        delta = choice.get(
+                            "delta",
+                            {},
+                        )
+
+                        finish_reason = choice.get(
+                            "finish_reason",
+                            None,
+                        )
+
+                        # ----------------------------------------
+                        # First chunk
+                        # ----------------------------------------
+
+                        if first_chunk:
+
+                            first_chunk = False
+
+                            first_data = {
+                                "id": request_id,
+                                "object": "chat.completion.chunk",
+                                "created": created,
+                                "model": MODEL_ID,
+                                "choices": [
+                                    {
+                                        "index": 0,
+                                        "delta": {
+                                            "role": "assistant"
+                                        },
+                                        "finish_reason": None,
+                                    }
+                                ],
+                            }
+
+                            yield (
+                                "data: "
+                                + json.dumps(
+                                    first_data,
+                                    ensure_ascii=False,
+                                )
+                                + "\n\n"
+                            )
+
+                        # ----------------------------------------
+                        # Content chunk
+                        # ----------------------------------------
+
+                        content = delta.get(
+                            "content",
+                            "",
+                        )
+
+                        if content:
+
+                            data = {
+                                "id": request_id,
+                                "object": "chat.completion.chunk",
+                                "created": created,
+                                "model": MODEL_ID,
+                                "choices": [
+                                    {
+                                        "index": 0,
+                                        "delta": {
+                                            "content": content
+                                        },
+                                        "finish_reason": None,
+                                    }
+                                ],
+                            }
+
+                            yield (
+                                "data: "
+                                + json.dumps(
+                                    data,
+                                    ensure_ascii=False,
+                                )
+                                + "\n\n"
+                            )
+
+                        # ----------------------------------------
+                        # Finish reason
+                        # ----------------------------------------
+
+                        if finish_reason is not None:
+
+                            final_data = {
+                                "id": request_id,
+                                "object": "chat.completion.chunk",
+                                "created": created,
+                                "model": MODEL_ID,
+                                "choices": [
+                                    {
+                                        "index": 0,
+                                        "delta": {},
+                                        "finish_reason": finish_reason,
+                                    }
+                                ],
+                            }
+
+                            yield (
+                                "data: "
+                                + json.dumps(
+                                    final_data,
+                                    ensure_ascii=False,
+                                )
+                                + "\n\n"
+                            )
+
+                # ------------------------------------------------
+                # DONE
+                # ------------------------------------------------
+
+                yield "data: [DONE]\n\n"
+
+                elapsed = (
+                    time.perf_counter()
+                    - start_time
+                )
+
+                print(
+                    f"Stream {request_id} "
+                    f"completed in "
+                    f"{elapsed:.2f}s"
+                )
+
+            except Exception as e:
+
+                print(
+                    f"Streaming error "
+                    f"{request_id}: {e}"
+                )
+
+                error_data = {
+                    "error": {
+                        "message": str(e),
+                        "type": "server_error",
+                    }
+                }
 
-    request_id = f"chatcmpl-{uuid.uuid4().hex}"
+                yield (
+                    "data: "
+                    + json.dumps(
+                        error_data,
+                        ensure_ascii=False,
+                    )
+                    + "\n\n"
+                )
+
+                yield "data: [DONE]\n\n"
+
+        return StreamingResponse(
+            generate_stream(),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",
+            },
+        )
 
-    created = int(time.time())
+    # ========================================================
+    # NON-STREAMING
+    # ========================================================
 
     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,
-        )
+
+        with generation_lock:
+
+            completion_kwargs = {
+                "messages": messages,
+                "temperature": float(temperature),
+                "max_tokens": int(max_tokens),
+                "top_p": float(top_p),
+                "stream": False,
+            }
+
+            if stop is not None:
+                completion_kwargs["stop"] = stop
+
+            if seed is not None:
+                completion_kwargs["seed"] = int(seed)
+
+            response = (
+                llm.create_chat_completion(
+                    **completion_kwargs
+                )
+            )
 
     except Exception as e:
-        print(f"Inference error: {e}")
+
+        print(
+            f"Inference error "
+            f"{request_id}: {e}"
+        )
 
         raise HTTPException(
             status_code=500,
             detail=f"Model inference failed: {str(e)}",
         )
 
-    generation_time = time.perf_counter() - start_time
+    elapsed = (
+        time.perf_counter()
+        - start_time
+    )
 
     # --------------------------------------------------------
-    # Extract generated text
+    # Extract response
     # --------------------------------------------------------
 
     try:
+
         choice = response["choices"][0]
 
-        message = choice["message"]
+        message = choice.get(
+            "message",
+            {},
+        )
 
-        text_output = message.get("content", "")
+        text_output = message.get(
+            "content",
+            "",
+        )
 
         finish_reason = choice.get(
             "finish_reason",
@@ -354,18 +706,27 @@ async def chat_completions(request: Request):
         )
 
     except Exception as e:
-        print(f"Response parsing error: {e}")
+
+        print(
+            f"Response parsing error: {e}"
+        )
 
         raise HTTPException(
             status_code=500,
-            detail="Invalid response returned by the model.",
+            detail=(
+                "Invalid response returned "
+                "by the model."
+            ),
         )
 
     # --------------------------------------------------------
     # Usage
     # --------------------------------------------------------
 
-    usage = response.get("usage", {})
+    usage = response.get(
+        "usage",
+        {},
+    )
 
     prompt_tokens = usage.get(
         "prompt_tokens",
@@ -386,7 +747,7 @@ async def chat_completions(request: Request):
         f"Request {request_id} | "
         f"prompt={prompt_tokens} | "
         f"completion={completion_tokens} | "
-        f"time={generation_time:.2f}s"
+        f"time={elapsed:.2f}s"
     )
 
     # --------------------------------------------------------
@@ -415,7 +776,7 @@ async def chat_completions(request: Request):
                 "total_tokens": total_tokens,
             },
             "x_generation_time_seconds": round(
-                generation_time,
+                elapsed,
                 3,
             ),
         }