Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
-import time
 import json
+import time
 import uuid
 import asyncio
 from typing import Any
@@ -15,10 +15,11 @@ from llama_cpp import Llama
 
 
 # ============================================================
-# CONFIGURATION
+# CONFIG
 # ============================================================
 
 APP_NAME = "Varsh LLM"
+
 MODEL_ID = "qwen2.5-coder-7b-instruct-q4_k_m"
 
 MODEL_REPO = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF"
@@ -27,27 +28,23 @@ MODEL_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf"
 HOST = "0.0.0.0"
 PORT = int(os.getenv("PORT", "7860"))
 
-# IMPORTANT:
-# Add CUSTOM_API_KEY in:
-# Hugging Face Space -> Settings -> Secrets
-#
-# Do NOT hard-code the key 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."
+        "Add CUSTOM_API_KEY in Hugging Face Space "
+        "Settings > Secrets."
     )
 
 
 # ============================================================
-# FASTAPI
+# APP
 # ============================================================
 
 app = FastAPI(
     title=APP_NAME,
-    version="2.0.0",
+    version="3.0.0",
 )
 
 app.add_middleware(
@@ -58,12 +55,16 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-
 security = HTTPBearer(auto_error=False)
 
 
+# ============================================================
+# AUTH
+# ============================================================
+
 def verify_api_key(
-    credentials: HTTPAuthorizationCredentials = Depends(security),
+    credentials: HTTPAuthorizationCredentials =
+    Depends(security),
 ):
     if credentials is None:
         raise HTTPException(
@@ -90,12 +91,13 @@ def verify_api_key(
 # MODEL
 # ============================================================
 
-print("=" * 60)
+print("=" * 65)
 print("Starting Varsh LLM")
-print("=" * 60)
+print("=" * 65)
+
 print(f"Model repository: {MODEL_REPO}")
 print(f"Model file:       {MODEL_FILE}")
-print("Downloading/checking model...")
+print("Checking/downloading model...")
 
 model_path = hf_hub_download(
     repo_id=MODEL_REPO,
@@ -108,15 +110,12 @@ print("Loading Qwen 2.5 Coder 7B...")
 llm = Llama(
     model_path=model_path,
 
-    # Your Space has 2 vCPUs / 16 GB RAM.
+    # 2 vCPU / 16 GB RAM Space
     n_ctx=4096,
     n_threads=2,
     n_threads_batch=2,
-
-    # Moderate batch size for CPU deployment.
     n_batch=256,
 
-    # Keep memory usage reasonable.
     use_mmap=True,
     use_mlock=False,
 
@@ -124,97 +123,131 @@ llm = Llama(
 )
 
 print("Model loaded successfully.")
-print("=" * 60)
-print("Varsh LLM API READY")
-print("=" * 60)
-print(f"Listening on: {HOST}:{PORT}")
-print("Root endpoint: /")
-print("Models endpoint: /v1/models")
-print("Chat endpoint: /v1/chat/completions")
-print("Authentication: Bearer API key")
-print("=" * 60)
-
-
-# llama.cpp generation should be serialized on a small CPU instance.
+
+print("=" * 65)
+print("VARSH LLM API READY")
+print("=" * 65)
+print(f"Listening on:       {HOST}:{PORT}")
+print("Root endpoint:      /")
+print("Models endpoint:    /v1/models")
+print("Chat endpoint:      /v1/chat/completions")
+print("Authentication:     Bearer API key")
+print("Streaming:          enabled")
+print("=" * 65)
+
+
+# llama.cpp should not receive simultaneous generations
 generation_lock = asyncio.Lock()
 
 
 # ============================================================
-# HELPERS
+# CONTENT NORMALIZATION
 # ============================================================
 
-def normalize_content(content: Any) -> str:
+def content_to_text(content: Any) -> str:
     """
-    Converts OpenAI-style message content into plain text.
+    Convert OpenAI/Claude-style content into a plain string.
+
+    Examples:
+
+        "hello"
 
-    Supports:
+    becomes:
 
-        "Hello"
+        "hello"
 
-    and:
+    And:
 
         [
             {
                 "type": "text",
-                "text": "Hello"
+                "text": "hello"
             }
         ]
 
-    Qwen2.5-Coder here is text-only, so image/audio parts
-    cannot be processed.
+    becomes:
+
+        "hello"
+
+    Image/audio blocks are ignored because this model is
+    text-only.
     """
 
     if content is None:
         return ""
 
-    # Normal OpenAI text content
+    # Normal string
     if isinstance(content, str):
         return content
 
-    # OpenAI multimodal/content-part format
+    # OpenAI content array
     if isinstance(content, list):
-        text_parts = []
 
-        for part in content:
-            if isinstance(part, str):
-                text_parts.append(part)
+        parts = []
+
+        for item in content:
+
+            if isinstance(item, str):
+                parts.append(item)
                 continue
 
-            if not isinstance(part, dict):
+            if not isinstance(item, dict):
                 continue
 
-            part_type = part.get("type")
+            item_type = item.get("type")
+
+            # Standard text block
+            if item_type == "text":
+                text = item.get("text")
 
-            # Standard text part
-            if part_type == "text":
-                text = part.get("text", "")
                 if isinstance(text, str):
-                    text_parts.append(text)
+                    parts.append(text)
+
+                continue
 
-            # Some clients may use slightly different structures.
-            elif "text" in part and isinstance(part["text"], str):
-                text_parts.append(part["text"])
+            # Some clients may omit type
+            if isinstance(item.get("text"), str):
+                parts.append(item["text"])
 
-            # Images/audio are intentionally not processed.
-            elif part_type in {
+                continue
+
+            # Ignore image/audio content.
+            if item_type in (
                 "image",
                 "image_url",
                 "input_image",
                 "audio",
                 "input_audio",
-            }:
+            ):
                 continue
 
-        return "\n".join(text_parts)
+        return "\n".join(parts)
+
+    # Dict content
+    if isinstance(content, dict):
+
+        if isinstance(content.get("text"), str):
+            return content["text"]
+
+        if isinstance(content.get("content"), str):
+            return content["content"]
+
+        return json.dumps(
+            content,
+            ensure_ascii=False,
+        )
 
-    # Fallback for unusual clients
     return str(content)
 
 
+# ============================================================
+# MESSAGE NORMALIZATION
+# ============================================================
+
 def normalize_messages(messages: Any) -> list:
     """
-    Normalize incoming OpenAI-compatible messages so that
-    llama.cpp receives string content.
+    Convert all incoming messages into llama.cpp-safe
+    text-only messages.
     """
 
     if not isinstance(messages, list):
@@ -223,114 +256,98 @@ def normalize_messages(messages: Any) -> list:
             detail="'messages' must be an array.",
         )
 
-    normalized = []
+    result = []
 
     for message in messages:
+
         if not isinstance(message, dict):
             continue
 
-        role = message.get("role", "user")
+        role = message.get(
+            "role",
+            "user",
+        )
 
         if not isinstance(role, str):
             role = "user"
 
-        content = normalize_content(
-            message.get("content", "")
+        content = content_to_text(
+            message.get(
+                "content",
+                "",
+            )
         )
 
-        new_message = {
+        # IMPORTANT:
+        # llama.cpp must receive string content.
+        normalized_message = {
             "role": role,
             "content": content,
         }
 
-        # Preserve name if supplied.
-        if "name" in message:
-            new_message["name"] = message["name"]
-
-        # Preserve tool_call_id if supplied.
-        if "tool_call_id" in message:
-            new_message["tool_call_id"] = message["tool_call_id"]
+        # Preserve name when present.
+        if isinstance(
+            message.get("name"),
+            str,
+        ):
+            normalized_message["name"] = (
+                message["name"]
+            )
 
-        # Preserve tool calls when clients send them.
-        if "tool_calls" in message:
-            new_message["tool_calls"] = message["tool_calls"]
+        # Preserve tool call ID if present.
+        if isinstance(
+            message.get("tool_call_id"),
+            str,
+        ):
+            normalized_message["tool_call_id"] = (
+                message["tool_call_id"]
+            )
 
-        normalized.append(new_message)
+        result.append(
+            normalized_message
+        )
 
-    if not normalized:
+    if not result:
         raise HTTPException(
             status_code=400,
             detail="'messages' cannot be empty.",
         )
 
-    return normalized
-
-
-def create_chat_id() -> str:
-    return f"chatcmpl-{uuid.uuid4().hex}"
-
+    return result
 
-def create_timestamp() -> int:
-    return int(time.time())
 
+# ============================================================
+# PARAMETERS
+# ============================================================
 
-def build_non_stream_response(
-    result: dict,
-    model_id: str,
-) -> dict:
-
-    choices = result.get("choices", [])
+def safe_float(
+    value: Any,
+    default: float,
+) -> float:
 
-    if not choices:
-        text = ""
-        finish_reason = "stop"
-    else:
-        choice = choices[0]
+    try:
+        return float(value)
+    except Exception:
+        return default
 
-        message = choice.get("message", {})
 
-        text = message.get("content", "")
+def safe_int(
+    value: Any,
+    default: int,
+) -> int:
 
-        if text is None:
-            text = ""
+    try:
+        return int(value)
+    except Exception:
+        return default
 
-        finish_reason = choice.get(
-            "finish_reason",
-            "stop",
-        )
 
-    usage = result.get("usage", {})
+# ============================================================
+# ID
+# ============================================================
 
-    return {
-        "id": create_chat_id(),
-        "object": "chat.completion",
-        "created": create_timestamp(),
-        "model": model_id,
-        "choices": [
-            {
-                "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": text,
-                },
-                "finish_reason": finish_reason,
-            }
-        ],
-        "usage": {
-            "prompt_tokens": usage.get(
-                "prompt_tokens",
-                0,
-            ),
-            "completion_tokens": usage.get(
-                "completion_tokens",
-                0,
-            ),
-            "total_tokens": usage.get(
-                "total_tokens",
-                0,
-            ),
-        },
-    }
+def make_id():
+    return f"chatcmpl-{uuid.uuid4().hex}"
 
 
 # ============================================================
@@ -339,6 +356,7 @@ def build_non_stream_response(
 
 @app.get("/")
 async def root():
+
     return {
         "status": "healthy",
         "service": "Varsh LLM",
@@ -350,6 +368,19 @@ async def root():
     }
 
 
+# ============================================================
+# HEALTH
+# ============================================================
+
+@app.get("/health")
+async def health():
+
+    return {
+        "status": "ok",
+        "model": MODEL_ID,
+    }
+
+
 # ============================================================
 # MODELS
 # ============================================================
@@ -360,15 +391,13 @@ async def root():
 )
 async def models():
 
-    created = create_timestamp()
-
     return {
         "object": "list",
         "data": [
             {
                 "id": MODEL_ID,
                 "object": "model",
-                "created": created,
+                "created": int(time.time()),
                 "owned_by": "varsshancoder",
             }
         ],
@@ -383,10 +412,17 @@ async def models():
     "/v1/chat/completions",
     dependencies=[Depends(verify_api_key)],
 )
-async def chat_completions(request: Request):
+async def chat_completions(
+    request: Request,
+):
+
+    # --------------------------------------------------------
+    # Parse JSON
+    # --------------------------------------------------------
 
     try:
         body = await request.json()
+
     except Exception:
         raise HTTPException(
             status_code=400,
@@ -399,30 +435,57 @@ async def chat_completions(request: Request):
             detail="Request body must be a JSON object.",
         )
 
+    # --------------------------------------------------------
+    # Messages
+    # --------------------------------------------------------
+
     messages = normalize_messages(
-        body.get("messages", [])
+        body.get(
+            "messages",
+            [],
+        )
     )
 
+    # --------------------------------------------------------
+    # Model
+    # --------------------------------------------------------
+
     requested_model = body.get(
         "model",
         MODEL_ID,
     )
 
-    # Accept the model exposed by /v1/models.
-    # Also tolerate a provider-prefixed model ID.
-    if requested_model not in {
+    if requested_model in (
         MODEL_ID,
         f"varsh/{MODEL_ID}",
-    }:
+    ):
         requested_model = MODEL_ID
 
-    temperature = body.get(
-        "temperature",
+    # --------------------------------------------------------
+    # Parameters
+    # --------------------------------------------------------
+
+    temperature = safe_float(
+        body.get(
+            "temperature",
+            0.2,
+        ),
         0.2,
     )
 
-    max_tokens = body.get(
-        "max_tokens",
+    top_p = safe_float(
+        body.get(
+            "top_p",
+            0.95,
+        ),
+        0.95,
+    )
+
+    max_tokens = safe_int(
+        body.get(
+            "max_tokens",
+            512,
+        ),
         512,
     )
 
@@ -433,58 +496,87 @@ async def chat_completions(request: Request):
         )
     )
 
-    # Keep values safe for the 2-vCPU deployment.
-    try:
-        temperature = float(temperature)
-    except Exception:
-        temperature = 0.2
-
-    try:
-        max_tokens = int(max_tokens)
-    except Exception:
-        max_tokens = 512
-
-    if max_tokens < 1:
-        max_tokens = 1
+    # Reasonable limits for your CPU Space.
+    temperature = max(
+        0.0,
+        min(
+            temperature,
+            2.0,
+        ),
+    )
 
-    if max_tokens > 4096:
-        max_tokens = 4096
+    top_p = max(
+        0.01,
+        min(
+            top_p,
+            1.0,
+        ),
+    )
 
-    if temperature < 0:
-        temperature = 0
+    max_tokens = max(
+        1,
+        min(
+            max_tokens,
+            4096,
+        ),
+    )
 
-    if temperature > 2:
-        temperature = 2
+    # --------------------------------------------------------
+    # STREAMING
+    # --------------------------------------------------------
 
-    # Optional parameters supported by many OpenAI clients.
-    top_p = body.get("top_p", 0.95)
+    if stream:
 
-    try:
-        top_p = float(top_p)
-    except Exception:
-        top_p = 0.95
+        chat_id = make_id()
+        created = int(time.time())
 
-    if top_p <= 0:
-        top_p = 0.01
+        async def stream_generator():
 
-    if top_p > 1:
-        top_p = 1.0
+            # ------------------------------------------------
+            # FIRST EVENT
+            #
+            # Send this BEFORE model generation.
+            #
+            # This is important for Claude Code because the
+            # proxy/client must receive a real SSE event before
+            # waiting for the CPU model.
+            # ------------------------------------------------
 
-    # --------------------------------------------------------
-    # Streaming
-    # --------------------------------------------------------
+            first_chunk = {
+                "id": chat_id,
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": requested_model,
+                "choices": [
+                    {
+                        "index": 0,
+                        "delta": {
+                            "role": "assistant",
+                            "content": "",
+                        },
+                        "finish_reason": None,
+                    }
+                ],
+            }
 
-    if stream:
+            yield (
+                "data: "
+                + json.dumps(
+                    first_chunk,
+                    ensure_ascii=False,
+                )
+                + "\n\n"
+            )
 
-        chat_id = create_chat_id()
-        created = create_timestamp()
+            try:
 
-        async def event_generator():
+                # --------------------------------------------
+                # Generate
+                # --------------------------------------------
 
-            try:
                 async with generation_lock:
 
-                    result_stream = llm.create_chat_completion(
+                    result = llm.create_chat_completion(
                         messages=messages,
                         temperature=temperature,
                         top_p=top_p,
@@ -492,7 +584,7 @@ async def chat_completions(request: Request):
                         stream=True,
                     )
 
-                    for chunk in result_stream:
+                    for chunk in result:
 
                         choices = chunk.get(
                             "choices",
@@ -517,11 +609,12 @@ async def chat_completions(request: Request):
                         if content is None:
                             content = ""
 
-                        finish_reason = choice.get(
-                            "finish_reason"
-                        )
+                        # ------------------------------------
+                        # Every streamed event must be valid
+                        # SSE JSON.
+                        # ------------------------------------
 
-                        payload = {
+                        output_chunk = {
                             "id": chat_id,
                             "object": "chat.completion.chunk",
                             "created": created,
@@ -530,10 +623,11 @@ async def chat_completions(request: Request):
                                 {
                                     "index": 0,
                                     "delta": {
-                                        "role": "assistant",
                                         "content": content,
                                     },
-                                    "finish_reason": finish_reason,
+                                    "finish_reason": choice.get(
+                                        "finish_reason"
+                                    ),
                                 }
                             ],
                         }
@@ -541,45 +635,91 @@ async def chat_completions(request: Request):
                         yield (
                             "data: "
                             + json.dumps(
-                                payload,
+                                output_chunk,
                                 ensure_ascii=False,
                             )
                             + "\n\n"
                         )
 
+                # --------------------------------------------
+                # Final chunk
+                # --------------------------------------------
+
+                final_chunk = {
+                    "id": chat_id,
+                    "object": "chat.completion.chunk",
+                    "created": created,
+                    "model": requested_model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {},
+                            "finish_reason": "stop",
+                        }
+                    ],
+                }
+
+                yield (
+                    "data: "
+                    + json.dumps(
+                        final_chunk,
+                        ensure_ascii=False,
+                    )
+                    + "\n\n"
+                )
+
                 yield "data: [DONE]\n\n"
 
             except Exception as exc:
 
-                error_payload = {
-                    "error": {
-                        "message": str(exc),
-                        "type": "server_error",
-                    }
+                print(
+                    "STREAM GENERATION ERROR:",
+                    repr(exc),
+                )
+
+                error_chunk = {
+                    "id": chat_id,
+                    "object": "chat.completion.chunk",
+                    "created": created,
+                    "model": requested_model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "content": (
+                                    f"\n[Generation error: "
+                                    f"{str(exc)}]"
+                                ),
+                            },
+                            "finish_reason": "stop",
+                        }
+                    ],
                 }
 
                 yield (
                     "data: "
                     + json.dumps(
-                        error_payload,
+                        error_chunk,
                         ensure_ascii=False,
                     )
                     + "\n\n"
                 )
 
+                yield "data: [DONE]\n\n"
+
         return StreamingResponse(
-            event_generator(),
+            stream_generator(),
             media_type="text/event-stream",
             headers={
-                "Cache-Control": "no-cache",
+                "Cache-Control": "no-cache, no-transform",
                 "Connection": "keep-alive",
                 "X-Accel-Buffering": "no",
             },
         )
 
-    # --------------------------------------------------------
-    # Non-streaming
-    # --------------------------------------------------------
+    # ========================================================
+    # NON-STREAMING
+    # ========================================================
 
     try:
 
@@ -593,19 +733,78 @@ async def chat_completions(request: Request):
                 stream=False,
             )
 
-        response = build_non_stream_response(
-            result,
-            requested_model,
+        choices = result.get(
+            "choices",
+            [],
+        )
+
+        if choices:
+
+            choice = choices[0]
+
+            message = choice.get(
+                "message",
+                {},
+            )
+
+            output = message.get(
+                "content",
+                "",
+            )
+
+            finish_reason = choice.get(
+                "finish_reason",
+                "stop",
+            )
+
+        else:
+
+            output = ""
+            finish_reason = "stop"
+
+        usage = result.get(
+            "usage",
+            {},
         )
 
         return JSONResponse(
-            content=response
+            content={
+                "id": make_id(),
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": requested_model,
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": output,
+                        },
+                        "finish_reason": finish_reason,
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": usage.get(
+                        "prompt_tokens",
+                        0,
+                    ),
+                    "completion_tokens": usage.get(
+                        "completion_tokens",
+                        0,
+                    ),
+                    "total_tokens": usage.get(
+                        "total_tokens",
+                        0,
+                    ),
+                },
+            }
         )
 
     except Exception as exc:
 
         print(
-            f"Generation error: {type(exc).__name__}: {exc}"
+            "GENERATION ERROR:",
+            repr(exc),
         )
 
         raise HTTPException(
@@ -615,22 +814,11 @@ async def chat_completions(request: Request):
 
 
 # ============================================================
-# HEALTH CHECK
-# ============================================================
-
-@app.get("/health")
-async def health():
-    return {
-        "status": "ok",
-        "model": MODEL_ID,
-    }
-
-
-# ============================================================
-# STARTUP
+# MAIN
 # ============================================================
 
 if __name__ == "__main__":
+
     import uvicorn
 
     uvicorn.run(