Update app.py

app.py

@@ -33,18 +33,34 @@ 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."
     )
 
 
 # ============================================================
-# APP
+# CONTEXT SETTINGS
+# ============================================================
+
+# 16 GB RAM + 7B Q4 CPU deployment.
+# 8192 is a safer target than 16384.
+CONTEXT_SIZE = 8192
+
+# Reserve room for the model's answer.
+RESERVED_OUTPUT_TOKENS = 1024
+
+# Maximum amount of input context we allow.
+MAX_INPUT_TOKENS = (
+    CONTEXT_SIZE - RESERVED_OUTPUT_TOKENS
+)
+
+
+# ============================================================
+# FASTAPI
 # ============================================================
 
 app = FastAPI(
     title=APP_NAME,
-    version="3.0.0",
+    version="4.0.0",
 )
 
 app.add_middleware(
@@ -91,12 +107,16 @@ def verify_api_key(
 # MODEL
 # ============================================================
 
-print("=" * 65)
+print("=" * 70)
 print("Starting Varsh LLM")
-print("=" * 65)
+print("=" * 70)
 
 print(f"Model repository: {MODEL_REPO}")
 print(f"Model file:       {MODEL_FILE}")
+print(f"Context size:     {CONTEXT_SIZE}")
+print(f"Max input:        {MAX_INPUT_TOKENS}")
+print(f"Reserved output:  {RESERVED_OUTPUT_TOKENS}")
+
 print("Checking/downloading model...")
 
 model_path = hf_hub_download(
@@ -110,10 +130,13 @@ print("Loading Qwen 2.5 Coder 7B...")
 llm = Llama(
     model_path=model_path,
 
-    # 2 vCPU / 16 GB RAM Space
-    n_ctx=4096,
+    # CPU deployment
+    n_ctx=CONTEXT_SIZE,
+
     n_threads=2,
     n_threads_batch=2,
+
+    # Moderate batch for 2 vCPU.
     n_batch=256,
 
     use_mmap=True,
@@ -124,19 +147,21 @@ llm = Llama(
 
 print("Model loaded successfully.")
 
-print("=" * 65)
+print("=" * 70)
 print("VARSH LLM API READY")
-print("=" * 65)
+print("=" * 70)
 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)
+print(f"Context:            {CONTEXT_SIZE}")
+print("=" * 70)
 
 
-# llama.cpp should not receive simultaneous generations
+# llama.cpp generation is serialized because this is
+# a 2-vCPU CPU deployment.
 generation_lock = asyncio.Lock()
 
 
@@ -145,42 +170,13 @@ generation_lock = asyncio.Lock()
 # ============================================================
 
 def content_to_text(content: Any) -> str:
-    """
-    Convert OpenAI/Claude-style content into a plain string.
-
-    Examples:
-
-        "hello"
-
-    becomes:
-
-        "hello"
-
-    And:
-
-        [
-            {
-                "type": "text",
-                "text": "hello"
-            }
-        ]
-
-    becomes:
-
-        "hello"
-
-    Image/audio blocks are ignored because this model is
-    text-only.
-    """
 
     if content is None:
         return ""
 
-    # Normal string
     if isinstance(content, str):
         return content
 
-    # OpenAI content array
     if isinstance(content, list):
 
         parts = []
@@ -196,40 +192,39 @@ def content_to_text(content: Any) -> str:
 
             item_type = item.get("type")
 
-            # Standard text block
             if item_type == "text":
-                text = item.get("text")
+
+                text = item.get("text", "")
 
                 if isinstance(text, str):
                     parts.append(text)
 
                 continue
 
-            # Some clients may omit type
-            if isinstance(item.get("text"), str):
-                parts.append(item["text"])
-
-                continue
-
-            # Ignore image/audio content.
-            if item_type in (
-                "image",
-                "image_url",
-                "input_image",
-                "audio",
-                "input_audio",
+            # Some clients use:
+            # {"text": "..."}
+            if isinstance(
+                item.get("text"),
+                str,
             ):
-                continue
+                parts.append(
+                    item["text"]
+                )
 
         return "\n".join(parts)
 
-    # Dict content
     if isinstance(content, dict):
 
-        if isinstance(content.get("text"), str):
+        if isinstance(
+            content.get("text"),
+            str,
+        ):
             return content["text"]
 
-        if isinstance(content.get("content"), str):
+        if isinstance(
+            content.get("content"),
+            str,
+        ):
             return content["content"]
 
         return json.dumps(
@@ -245,10 +240,6 @@ def content_to_text(content: Any) -> str:
 # ============================================================
 
 def normalize_messages(messages: Any) -> list:
-    """
-    Convert all incoming messages into llama.cpp-safe
-    text-only messages.
-    """
 
     if not isinstance(messages, list):
         raise HTTPException(
@@ -278,34 +269,28 @@ def normalize_messages(messages: Any) -> list:
             )
         )
 
-        # IMPORTANT:
-        # llama.cpp must receive string content.
-        normalized_message = {
+        normalized = {
             "role": role,
             "content": content,
         }
 
-        # Preserve name when present.
         if isinstance(
             message.get("name"),
             str,
         ):
-            normalized_message["name"] = (
+            normalized["name"] = (
                 message["name"]
             )
 
-        # Preserve tool call ID if present.
         if isinstance(
             message.get("tool_call_id"),
             str,
         ):
-            normalized_message["tool_call_id"] = (
+            normalized["tool_call_id"] = (
                 message["tool_call_id"]
             )
 
-        result.append(
-            normalized_message
-        )
+        result.append(normalized)
 
     if not result:
         raise HTTPException(
@@ -317,39 +302,304 @@ def normalize_messages(messages: Any) -> list:
 
 
 # ============================================================
-# PARAMETERS
+# TOKEN COUNTING
 # ============================================================
 
-def safe_float(
-    value: Any,
-    default: float,
-) -> float:
+def count_message_tokens(messages: list) -> int:
+    """
+    Use llama.cpp's actual tokenizer instead of estimating
+    tokens using character length.
+    """
 
     try:
-        return float(value)
-    except Exception:
-        return default
 
+        # llama.cpp's chat template will ultimately create
+        # a prompt. We use a conservative approximation here.
+        total = 0
 
-def safe_int(
-    value: Any,
-    default: int,
-) -> int:
+        for message in messages:
+
+            role = message.get(
+                "role",
+                "",
+            )
+
+            content = message.get(
+                "content",
+                "",
+            )
+
+            text = (
+                f"<|im_start|>{role}\n"
+                f"{content}"
+                f"<|im_end|>\n"
+            )
+
+            tokens = llm.tokenize(
+                text.encode(
+                    "utf-8"
+                ),
+                add_bos=False,
+            )
+
+            total += len(tokens)
+
+        # Assistant generation marker.
+        assistant_tokens = llm.tokenize(
+            b"<|im_start|>assistant\n",
+            add_bos=False,
+        )
+
+        total += len(
+            assistant_tokens
+        )
+
+        return total
 
-    try:
-        return int(value)
     except Exception:
-        return default
+
+        # Safe fallback.
+        characters = sum(
+            len(
+                str(
+                    m.get(
+                        "content",
+                        "",
+                    )
+                )
+            )
+            for m in messages
+        )
+
+        return max(
+            1,
+            characters // 4,
+        )
 
 
 # ============================================================
-# ID
+# CONTEXT MANAGEMENT
+# ============================================================
+
+def fit_messages_to_context(
+    messages: list,
+    max_tokens: int,
+) -> tuple[list, int]:
+
+    if not messages:
+        return messages, 0
+
+    current_tokens = count_message_tokens(
+        messages
+    )
+
+    if current_tokens <= max_tokens:
+        return messages, current_tokens
+
+    print(
+        f"Context too large: "
+        f"{current_tokens} tokens > "
+        f"{max_tokens}. "
+        f"Trimming old messages."
+    )
+
+    # --------------------------------------------------------
+    # Preserve system messages.
+    # --------------------------------------------------------
+
+    system_messages = [
+        m
+        for m in messages
+        if m.get("role") == "system"
+    ]
+
+    non_system_messages = [
+        m
+        for m in messages
+        if m.get("role") != "system"
+    ]
+
+    # --------------------------------------------------------
+    # Keep newest messages first.
+    # --------------------------------------------------------
+
+    kept = list(system_messages)
+
+    # Add newest conversation messages backwards.
+    selected = []
+
+    for message in reversed(
+        non_system_messages
+    ):
+
+        candidate = (
+            system_messages
+            + list(reversed(
+                selected + [message]
+            ))
+        )
+
+        tokens = count_message_tokens(
+            candidate
+        )
+
+        if tokens <= max_tokens:
+            selected.append(message)
+        else:
+            break
+
+    selected.reverse()
+
+    final_messages = (
+        system_messages
+        + selected
+    )
+
+    final_tokens = count_message_tokens(
+        final_messages
+    )
+
+    # --------------------------------------------------------
+    # If system prompt itself is too large,
+    # truncate the system content.
+    # --------------------------------------------------------
+
+    if final_tokens > max_tokens:
+
+        print(
+            "System context is too large. "
+            "Truncating system message."
+        )
+
+        new_messages = []
+
+        remaining = max_tokens
+
+        for message in final_messages:
+
+            if (
+                message.get("role")
+                != "system"
+            ):
+                continue
+
+            content = message.get(
+                "content",
+                "",
+            )
+
+            # Tokenize system content.
+            tokens = llm.tokenize(
+                content.encode(
+                    "utf-8"
+                ),
+                add_bos=False,
+            )
+
+            if len(tokens) <= remaining:
+
+                new_messages.append(
+                    message
+                )
+
+                remaining -= len(tokens)
+
+            else:
+
+                # Approximate character cutoff.
+                ratio = (
+                    remaining
+                    / max(
+                        len(tokens),
+                        1,
+                    )
+                )
+
+                cutoff = max(
+                    1,
+                    int(
+                        len(content)
+                        * ratio
+                    ),
+                )
+
+                truncated = content[
+                    :cutoff
+                ]
+
+                new_messages.append(
+                    {
+                        "role": "system",
+                        "content": truncated,
+                    }
+                )
+
+                remaining = 0
+                break
+
+        # Add as many newest messages as fit.
+        for message in selected:
+
+            candidate = (
+                new_messages
+                + [message]
+            )
+
+            if (
+                count_message_tokens(
+                    candidate
+                )
+                <= max_tokens
+            ):
+                new_messages.append(
+                    message
+                )
+
+        final_messages = new_messages
+
+        final_tokens = count_message_tokens(
+            final_messages
+        )
+
+    print(
+        f"Context after trimming: "
+        f"{final_tokens} tokens"
+    )
+
+    return (
+        final_messages,
+        final_tokens,
+    )
+
+
+# ============================================================
+# HELPERS
 # ============================================================
 
 def make_id():
     return f"chatcmpl-{uuid.uuid4().hex}"
 
 
+def safe_int(
+    value,
+    default,
+):
+    try:
+        return int(value)
+    except Exception:
+        return default
+
+
+def safe_float(
+    value,
+    default,
+):
+    try:
+        return float(value)
+    except Exception:
+        return default
+
+
 # ============================================================
 # ROOT
 # ============================================================
@@ -359,12 +609,13 @@ async def root():
 
     return {
         "status": "healthy",
-        "service": "Varsh LLM",
+        "service": APP_NAME,
         "security": "enabled",
         "model": MODEL_ID,
         "endpoint": "/v1/chat/completions",
         "models_endpoint": "/v1/models",
         "streaming": True,
+        "context_size": CONTEXT_SIZE,
     }
 
 
@@ -378,6 +629,7 @@ async def health():
     return {
         "status": "ok",
         "model": MODEL_ID,
+        "context_size": CONTEXT_SIZE,
     }
 
 
@@ -405,7 +657,7 @@ async def models():
 
 
 # ============================================================
-# CHAT COMPLETIONS
+# CHAT
 # ============================================================
 
 @app.post(
@@ -417,7 +669,7 @@ async def chat_completions(
 ):
 
     # --------------------------------------------------------
-    # Parse JSON
+    # JSON
     # --------------------------------------------------------
 
     try:
@@ -455,10 +707,7 @@ async def chat_completions(
         MODEL_ID,
     )
 
-    if requested_model in (
-        MODEL_ID,
-        f"varsh/{MODEL_ID}",
-    ):
+    if requested_model == f"varsh/{MODEL_ID}":
         requested_model = MODEL_ID
 
     # --------------------------------------------------------
@@ -496,7 +745,6 @@ async def chat_completions(
         )
     )
 
-    # Reasonable limits for your CPU Space.
     temperature = max(
         0.0,
         min(
@@ -513,18 +761,48 @@ async def chat_completions(
         ),
     )
 
-    max_tokens = max(
-        1,
-        min(
-            max_tokens,
-            4096,
-        ),
-    )
-
     # --------------------------------------------------------
-    # STREAMING
+    # Context management
     # --------------------------------------------------------
 
+    messages, prompt_tokens = (
+        fit_messages_to_context(
+            messages,
+            MAX_INPUT_TOKENS,
+        )
+    )
+
+    # Never request more than the remaining context.
+    available_output = (
+        CONTEXT_SIZE
+        - prompt_tokens
+        - 16
+    )
+
+    if available_output < 1:
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                "Prompt is too large even after "
+                "context trimming."
+            ),
+        )
+
+    max_tokens = min(
+        max_tokens,
+        available_output,
+    )
+
+    print(
+        f"Request: prompt={prompt_tokens}, "
+        f"max_tokens={max_tokens}, "
+        f"stream={stream}"
+    )
+
+    # ========================================================
+    # STREAM
+    # ========================================================
+
     if stream:
 
         chat_id = make_id()
@@ -533,13 +811,7 @@ async def chat_completions(
         async def stream_generator():
 
             # ------------------------------------------------
-            # 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.
+            # Immediate first SSE event.
             # ------------------------------------------------
 
             first_chunk = {
@@ -570,18 +842,16 @@ async def chat_completions(
 
             try:
 
-                # --------------------------------------------
-                # Generate
-                # --------------------------------------------
-
                 async with generation_lock:
 
-                    result = llm.create_chat_completion(
-                        messages=messages,
-                        temperature=temperature,
-                        top_p=top_p,
-                        max_tokens=max_tokens,
-                        stream=True,
+                    result = (
+                        llm.create_chat_completion(
+                            messages=messages,
+                            temperature=temperature,
+                            top_p=top_p,
+                            max_tokens=max_tokens,
+                            stream=True,
+                        )
                     )
 
                     for chunk in result:
@@ -609,11 +879,6 @@ async def chat_completions(
                         if content is None:
                             content = ""
 
-                        # ------------------------------------
-                        # Every streamed event must be valid
-                        # SSE JSON.
-                        # ------------------------------------
-
                         output_chunk = {
                             "id": chat_id,
                             "object": "chat.completion.chunk",
@@ -641,9 +906,9 @@ async def chat_completions(
                             + "\n\n"
                         )
 
-                # --------------------------------------------
-                # Final chunk
-                # --------------------------------------------
+                # ------------------------------------------------
+                # Final event.
+                # ------------------------------------------------
 
                 final_chunk = {
                     "id": chat_id,
@@ -673,7 +938,7 @@ async def chat_completions(
             except Exception as exc:
 
                 print(
-                    "STREAM GENERATION ERROR:",
+                    "STREAM ERROR:",
                     repr(exc),
                 )
 
@@ -687,9 +952,10 @@ async def chat_completions(
                             "index": 0,
                             "delta": {
                                 "content": (
-                                    f"\n[Generation error: "
-                                    f"{str(exc)}]"
-                                ),
+                                    "\n[Generation error: "
+                                    + str(exc)
+                                    + "]"
+                                )
                             },
                             "finish_reason": "stop",
                         }
@@ -718,19 +984,21 @@ async def chat_completions(
         )
 
     # ========================================================
-    # NON-STREAMING
+    # NON-STREAM
     # ========================================================
 
     try:
 
         async with generation_lock:
 
-            result = llm.create_chat_completion(
-                messages=messages,
-                temperature=temperature,
-                top_p=top_p,
-                max_tokens=max_tokens,
-                stream=False,
+            result = (
+                llm.create_chat_completion(
+                    messages=messages,
+                    temperature=temperature,
+                    top_p=top_p,
+                    max_tokens=max_tokens,
+                    stream=False,
+                )
             )
 
         choices = result.get(
@@ -786,7 +1054,7 @@ async def chat_completions(
                 "usage": {
                     "prompt_tokens": usage.get(
                         "prompt_tokens",
-                        0,
+                        prompt_tokens,
                     ),
                     "completion_tokens": usage.get(
                         "completion_tokens",
@@ -794,7 +1062,7 @@ async def chat_completions(
                     ),
                     "total_tokens": usage.get(
                         "total_tokens",
-                        0,
+                        prompt_tokens,
                     ),
                 },
             }