app.py

"""
Qwen3.6-27B MTP TQ3_4S — Custom HTML chat server for Hugging Face Spaces.
Free CPU tier (2 vCPUs, 16 GB RAM).

Architecture:
  - FastAPI server on port 7860
  - GET /         → HTML chat interface
  - POST /api/chat → SSE stream of reasoning + response tokens + live metrics
  - No Gradio — pure HTML/CSS/JS frontend

Reasoning (Qwen3 thinking mode):
  The model's embedded chat template defaults to thinking ENABLED. We don't
  pass chat_template_kwargs (not supported in llama-cpp-python 0.3.32).
  The model emits reasoning tokens in delta.reasoning_content automatically.
"""

from __future__ import annotations

import asyncio
import json
import os
import re
import time
import traceback
import uuid
from threading import Lock
from typing import Iterator, List

import uvicorn
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
from sse_starlette.sse import EventSourceResponse

# =============================================================================
# Configuration
# =============================================================================
REPO_ID = "openbmb/MiniCPM5-1B-GGUF"
FILENAME = "MiniCPM5-1B-Q8_0.gguf"

# --- CPU / memory tuning (free tier = 2 vCPUs, 16 GB RAM) ------------------
# MiniCPM5-1B is only 1.07 GB (Q8_0) — tons of headroom on 16 GB RAM.
#
# 100k context window — MiniCPM5 supports up to 262k (n_ctx_train).
# KV cache at q8_0 for 100k context ≈ 2.3 GB (fits easily in 16 GB).
#
#   N_CTX = 102400     : 100k token context (model supports up to 262k)
#   N_BATCH = 512      : prompt-eval batch size
#   KV_CACHE = q8_0    : optimal for CPU (memory bandwidth bound)
N_THREADS = 2
N_CTX = 102400
N_BATCH = 512
KV_CACHE_DTYPE = "q8_0"
N_GPU_LAYERS = 0

# --- Generation defaults ---------------------------------------------------
# Max output = 30k tokens. At ~12 tok/s, 30k tokens ≈ 42 minutes.
# Default is 4096 (reasonable for chat); users can bump to 30k via slider.
MAX_TOKENS_DEFAULT = 4096
TEMPERATURE_DEFAULT = 0.7
TOP_P_DEFAULT = 0.9
REPEAT_PENALTY_DEFAULT = 1.05

PORT = int(os.environ.get("PORT", 7860))

SYSTEM_PROMPT_DEFAULT = (
    "You are MiniCPM5, a helpful and concise assistant. "
    "Answer in clear, well-structured prose."
)

# ChatML template — the model's embedded template already handles thinking,
# but we set this explicitly as a fallback for models that omit it.
CHAT_TEMPLATE = (
    "{% for message in messages %}"
    "{% if message['role'] == 'system' %}"
    "<|im_start|>system\n{{ message['content'] }}<|im_end|>\n"
    "{% elif message['role'] == 'user' %}"
    "<|im_start|>user\n{{ message['content'] }}<|im_end|>\n"
    "{% elif message['role'] == 'assistant' %}"
    "<|im_start|>assistant\n{{ message['content'] }}<|im_end|>\n"
    "{% endif %}"
    "{% endfor %}"
    "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
)

# =============================================================================
# Quant label parser
# =============================================================================
def _parse_quant(filename: str) -> str:
    m = re.search(r"(Q[0-9]+_[A-Z]+|IQ[0-9]+_[A-Z]+|TQ[0-9]+_[A-Z0-9]+|F16|F32)", filename, re.IGNORECASE)
    return m.group(1).upper() if m else "unknown"


# =============================================================================
# Model download + load
# =============================================================================
print(f"[boot] Downloading {FILENAME} from {REPO_ID} ...", flush=True)
model_path = hf_hub_download(
    repo_id=REPO_ID,
    filename=FILENAME,
    repo_type="model",
)
print(f"[boot] Model cached at: {model_path}", flush=True)

model_size_mb = os.path.getsize(model_path) / (1024 * 1024)
quant_label = _parse_quant(FILENAME)
print(f"[boot] Model file size: {model_size_mb:.1f} MB ({quant_label})", flush=True)
if model_size_mb < 100:
    raise RuntimeError(
        f"Model file is suspiciously small ({model_size_mb:.1f} MB). "
        f"Try clearing the HF cache and restarting."
    )

import llama_cpp as _llama_cpp_mod
print(f"[boot] llama-cpp-python version: {_llama_cpp_mod.__version__}", flush=True)

print("[boot] Loading Llama model (this can take 1-2 minutes) ...", flush=True)

LOAD_ATTEMPTS = [
    # MiniCPM5 has its own embedded chat template — do NOT override it.
    # q8_0 KV cache is optimal for CPU (tested: fp16 was 2x slower due to
    # memory bandwidth bottleneck).
    ("q8_0 KV cache, model's embedded template", dict(
        kv_cache_dtype=KV_CACHE_DTYPE,
    )),
    ("fp16 KV cache, model's embedded template", dict(
        kv_cache_dtype="fp16",
    )),
    ("bare minimum (all defaults)", dict()),
]

llm = None
last_error = None
for label, extra_kwargs in LOAD_ATTEMPTS:
    print(f"[boot] Attempt: {label}", flush=True)
    try:
        llm = Llama(
            model_path=model_path,
            n_threads=N_THREADS,
            n_ctx=N_CTX,
            n_batch=N_BATCH,
            n_gpu_layers=N_GPU_LAYERS,
            use_mlock=False,
            use_mmap=False,     # eager-load 1GB model into RAM (avoids 73s
                                # mmap cold-start on first request over HF's
                                # network-attached storage)
            verbose=False,      # suppress per-request debug output (was True
                                # for debugging load issues — no longer needed)
            **extra_kwargs,
        )
        print(f"[boot] SUCCESS with: {label}", flush=True)
        break
    except Exception as exc:
        last_error = exc
        print(f"[boot] FAILED with: {label}", flush=True)
        print(f"[boot]   error: {exc}", flush=True)
        print("-" * 70, flush=True)
        llm = None
        continue

if llm is None:
    raise RuntimeError(f"All model load attempts failed. Last error: {last_error}")

print("[boot] Model ready. Starting FastAPI server ...", flush=True)

LLM_LOCK = Lock()


# =============================================================================
# FastAPI app
# =============================================================================
app = FastAPI(title="Qwen3.6-27B TQ3_4S Chat")

# CORS — allow browser-based apps and tools to call the API directly.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

MODEL_ID = FILENAME.replace(".gguf", "")


def _build_messages(history: List[dict], system_prompt: str) -> List[dict]:
    """
    Build the messages list for create_chat_completion.

    The HTML frontend sends the full conversation history (including the
    latest user message) via the `history` field. We prepend a system
    message if one isn't already present in the history.

    The Qwen3.6 chat template requires at least one user message —
    otherwise it raises "No user query found in messages."
    """
    messages: List[dict] = []

    # Check if the history already starts with a system message.
    has_system = (
        len(history) > 0
        and isinstance(history[0], dict)
        and history[0].get("role") == "system"
    )

    if not has_system:
        messages.append({
            "role": "system",
            "content": system_prompt or SYSTEM_PROMPT_DEFAULT,
        })

    for msg in history:
        if isinstance(msg, dict) and msg.get("role") and msg.get("content") is not None:
            messages.append({
                "role": msg["role"],
                "content": msg["content"],
            })

    # Safety check: the Qwen3.6 template raises "No user query found"
    # if there's no user message. Log a warning if we're about to send
    # a messages list with no user role.
    has_user = any(m.get("role") == "user" for m in messages)
    if not has_user:
        print(f"[warn] _build_messages: no user message in history "
              f"(roles: {[m['role'] for m in messages]})", flush=True)

    return messages


def _run_llm_stream(
    messages: List[dict],
    temperature: float,
    top_p: float,
    max_tokens: int,
    repeat_penalty: float,
    queue: asyncio.Queue,
    loop: asyncio.AbstractEventLoop,
) -> None:
    """
    Synchronous worker that runs llama-cpp-python's blocking create_chat_completion
    in a background thread. Pushes events to the asyncio queue so the async SSE
    generator can yield them immediately without blocking the event loop.

    This is CRITICAL for streaming — if we ran create_chat_completion directly
    in the async path, each token would block the event loop and the SSE
    response wouldn't flush until the entire generation finished.
    """
    def _put(event):
        # Thread-safe put into the asyncio queue.
        loop.call_soon_threadsafe(queue.put_nowait, event)

    try:
        _put({"type": "status", "stage": "processing_prompt"})

        state = {
            "prompt_start": time.perf_counter(),
            "first_token": None,
            "first_reasoning_token": None,
            "first_response_token": None,
            "last_reasoning_token": None,
            "last_response_token": None,
            "last_token": None,
            "reasoning_tokens": 0,
            "response_tokens": 0,
        }

        def _tps(first, last, count):
            if first is None or last is None or count == 0:
                return 0.0
            elapsed = last - first
            return count / elapsed if elapsed > 0 else 0.0

        # State machine for parsing <think> tags from content.
        # MiniCPM5 (and some other models) emit reasoning as <think>...</think>
        # inside the regular content field, NOT as a separate reasoning_content
        # field. We parse these tags in real-time and split into reasoning vs
        # response events.
        #
        # States:
        #   "thinking"  — inside <think> block, emit as reasoning
        #   "responding" — after </think>, emit as response
        #   "buffering"  — haven't seen <think> yet, buffering to check
        think_state = "buffering"
        content_buffer = ""

        with LLM_LOCK:
            stream = llm.create_chat_completion(
                messages=messages,
                max_tokens=int(max_tokens),
                temperature=float(temperature),
                top_p=float(top_p),
                repeat_penalty=float(repeat_penalty),
                stream=True,
            )
            for chunk in stream:
                delta = (chunk.get("choices") or [{}])[0].get("delta", {}) or {}
                now = time.perf_counter()

                # Check for reasoning_content field first (Qwen3 / GPT-oss style)
                reasoning_token = delta.get("reasoning_content")
                if reasoning_token:
                    if state["first_token"] is None:
                        state["first_token"] = now
                    if state["first_reasoning_token"] is None:
                        state["first_reasoning_token"] = now
                    state["last_reasoning_token"] = now
                    state["last_token"] = now
                    state["reasoning_tokens"] += 1
                    _put({
                        "type": "reasoning",
                        "token": reasoning_token,
                        "count": state["reasoning_tokens"],
                        "tps": _tps(
                            state["first_reasoning_token"],
                            state["last_reasoning_token"],
                            state["reasoning_tokens"],
                        ),
                    })
                    continue

                # Content tokens — may contain <think> tags (MiniCPM5 style)
                token = delta.get("content")
                if not token:
                    continue

                # Parse <think> tags from the token stream
                content_buffer += token

                while content_buffer:
                    if think_state == "buffering":
                        # Check if we have <think> in the buffer
                        if "<think>" in content_buffer:
                            # Emit anything before <think> as response (rare)
                            idx = content_buffer.index("<think>")
                            if idx > 0:
                                pre = content_buffer[:idx]
                                if state["first_token"] is None:
                                    state["first_token"] = now
                                if state["first_response_token"] is None:
                                    state["first_response_token"] = now
                                state["last_response_token"] = now
                                state["last_token"] = now
                                state["response_tokens"] += 1
                                _put({
                                    "type": "response",
                                    "token": pre,
                                    "count": state["response_tokens"],
                                    "tps": _tps(
                                        state["first_response_token"],
                                        state["last_response_token"],
                                        state["response_tokens"],
                                    ),
                                })
                            content_buffer = content_buffer[idx + 7:]  # skip <think>
                            think_state = "thinking"
                            if state["first_token"] is None:
                                state["first_token"] = now
                            if state["first_reasoning_token"] is None:
                                state["first_reasoning_token"] = now
                            state["last_reasoning_token"] = now
                            state["last_token"] = now
                            state["reasoning_tokens"] += 1
                            _put({
                                "type": "reasoning",
                                "token": "",
                                "count": state["reasoning_tokens"],
                                "tps": _tps(
                                    state["first_reasoning_token"],
                                    state["last_reasoning_token"],
                                    state["reasoning_tokens"],
                                ),
                            })
                        elif len(content_buffer) > 7:
                            # Not enough to contain <think> — emit as response
                            # but keep last 7 chars in case "<think>" spans chunks
                            safe = content_buffer[:-7]
                            content_buffer = content_buffer[-7:]
                            if safe and state["first_token"] is None:
                                state["first_token"] = now
                            if safe and state["first_response_token"] is None:
                                state["first_response_token"] = now
                            if safe:
                                state["last_response_token"] = now
                                state["last_token"] = now
                                state["response_tokens"] += 1
                                _put({
                                    "type": "response",
                                    "token": safe,
                                    "count": state["response_tokens"],
                                    "tps": _tps(
                                        state["first_response_token"],
                                        state["last_response_token"],
                                        state["response_tokens"],
                                    ),
                                })
                            break
                        else:
                            break  # buffer too short, wait for more

                    elif think_state == "thinking":
                        # Check for </think> in the buffer
                        if "</think>" in content_buffer:
                            idx = content_buffer.index("</think>")
                            reasoning_text = content_buffer[:idx]
                            if reasoning_text:
                                if state["first_token"] is None:
                                    state["first_token"] = now
                                state["last_reasoning_token"] = now
                                state["last_token"] = now
                                state["reasoning_tokens"] += 1
                                _put({
                                    "type": "reasoning",
                                    "token": reasoning_text,
                                    "count": state["reasoning_tokens"],
                                    "tps": _tps(
                                        state["first_reasoning_token"],
                                        state["last_reasoning_token"],
                                        state["reasoning_tokens"],
                                    ),
                                })
                            content_buffer = content_buffer[idx + 8:]  # skip </think>
                            think_state = "responding"
                        elif len(content_buffer) > 8:
                            # Emit as reasoning but keep last 8 chars
                            safe = content_buffer[:-8]
                            content_buffer = content_buffer[-8:]
                            if safe:
                                state["last_reasoning_token"] = now
                                state["last_token"] = now
                                state["reasoning_tokens"] += 1
                                _put({
                                    "type": "reasoning",
                                    "token": safe,
                                    "count": state["reasoning_tokens"],
                                    "tps": _tps(
                                        state["first_reasoning_token"],
                                        state["last_reasoning_token"],
                                        state["reasoning_tokens"],
                                    ),
                                })
                            break
                        else:
                            break  # buffer too short

                    elif think_state == "responding":
                        # Everything goes to response
                        if state["first_response_token"] is None:
                            state["first_response_token"] = now
                        state["last_response_token"] = now
                        state["last_token"] = now
                        state["response_tokens"] += 1
                        _put({
                            "type": "response",
                            "token": content_buffer,
                            "count": state["response_tokens"],
                            "tps": _tps(
                                state["first_response_token"],
                                state["last_response_token"],
                                state["response_tokens"],
                            ),
                        })
                        content_buffer = ""
                        break

        # Final metrics
        total_tokens = state["reasoning_tokens"] + state["response_tokens"]
        ttft_ms = (state["first_token"] - state["prompt_start"]) * 1000 if state["first_token"] else 0
        gen_time = (state["last_token"] - state["first_token"]) if state["first_token"] and state["last_token"] else 0
        prompt_time = (state["first_token"] - state["prompt_start"]) if state["first_token"] else 0
        total_time = prompt_time + gen_time
        total_tps = total_tokens / gen_time if gen_time > 0 else 0
        reasoning_tps = _tps(state["first_reasoning_token"], state["last_reasoning_token"], state["reasoning_tokens"])
        response_tps = _tps(state["first_response_token"], state["last_response_token"], state["response_tokens"])

        _put({
            "type": "metrics",
            "ttft_ms": round(ttft_ms, 0),
            "prompt_time_s": round(prompt_time, 2),
            "gen_time_s": round(gen_time, 2),
            "total_time_s": round(total_time, 2),
            "reasoning_tokens": state["reasoning_tokens"],
            "response_tokens": state["response_tokens"],
            "total_tokens": total_tokens,
            "reasoning_tps": round(reasoning_tps, 2),
            "response_tps": round(response_tps, 2),
            "total_tps": round(total_tps, 2),
        })
        _put({"type": "done"})

    except Exception as exc:
        traceback.print_exc()
        _put({"type": "error", "message": str(exc)})


async def _stream_chat_async(
    messages: List[dict],
    temperature: float,
    top_p: float,
    max_tokens: int,
    repeat_penalty: float,
) -> asyncio.Queue:
    """
    Launch the blocking LLM stream in a background thread and return an
    asyncio.Queue that yields events as they arrive. Each token is available
    immediately — the event loop stays responsive for SSE flushing.
    """
    queue: asyncio.Queue = asyncio.Queue()
    loop = asyncio.get_event_loop()
    # Run the sync worker in a thread pool. The worker pushes events to the
    # queue via call_soon_threadsafe, so the async side can await them.
    asyncio.get_event_loop().run_in_executor(
        None,
        _run_llm_stream,
        messages, temperature, top_p, max_tokens, repeat_penalty,
        queue, loop,
    )
    return queue


@app.get("/")
async def index() -> HTMLResponse:
    return HTMLResponse(content=HTML_PAGE)


@app.post("/api/chat")
async def chat(request: Request):
    body = await request.json()

    messages = _build_messages(
        body.get("history", []),
        body.get("system_prompt", SYSTEM_PROMPT_DEFAULT),
    )

    queue = await _stream_chat_async(
        messages=messages,
        temperature=body.get("temperature", TEMPERATURE_DEFAULT),
        top_p=body.get("top_p", TOP_P_DEFAULT),
        max_tokens=body.get("max_tokens", MAX_TOKENS_DEFAULT),
        repeat_penalty=body.get("repeat_penalty", REPEAT_PENALTY_DEFAULT),
    )

    async def _event_generator():
        while True:
            event = await queue.get()
            yield {"data": json.dumps(event)}
            if event.get("type") in ("done", "error"):
                break

    return EventSourceResponse(_event_generator())


@app.get("/api/health")
async def health():
    return JSONResponse({
        "status": "ok",
        "model": FILENAME,
        "repo": REPO_ID,
        "quant": quant_label,
        "size_mb": round(model_size_mb, 1),
        "llama_cpp_version": _llama_cpp_mod.__version__,
        "n_threads": N_THREADS,
        "n_ctx": N_CTX,
        "kv_cache_dtype": KV_CACHE_DTYPE,
    })


# =============================================================================
# OpenAI-compatible API (/v1/*)
#
# These endpoints implement the OpenAI Chat Completions API spec so any
# client that speaks OpenAI (Python `openai` library, LangChain, curl, etc.)
# can connect directly:
#
#   from openai import OpenAI
#   client = OpenAI(base_url="https://YOUR_SPACE.hf.space/v1", api_key="x")
#   resp = client.chat.completions.create(
#       model="Qwen3.6-27B-MTP-TQ3_4S",
#       messages=[{"role": "user", "content": "Hello"}],
#       stream=True,
#   )
#
# Streaming format follows OpenAI SSE conventions:
#   data: {"choices":[{"delta":{"content":"..."}}]}
#   data: {"choices":[{"delta":{"reasoning_content":"..."}}]}  ← Qwen3 thinking
#   data: [DONE]
#
# Non-streaming returns a standard chat.completion object.
# =============================================================================
@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [{
            "id": MODEL_ID,
            "object": "model",
            "created": int(time.time()),
            "owned_by": REPO_ID.split("/")[0],
        }],
    }


@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    body = await request.json()

    messages = body.get("messages", [])
    stream = body.get("stream", False)
    temperature = body.get("temperature", TEMPERATURE_DEFAULT)
    top_p = body.get("top_p", TOP_P_DEFAULT)
    # OpenAI uses max_tokens (legacy) or max_completion_tokens (newer)
    max_tokens = body.get("max_completion_tokens") or body.get("max_tokens") or MAX_TOKENS_DEFAULT
    repeat_penalty = body.get("repeat_penalty", REPEAT_PENALTY_DEFAULT)
    stop = body.get("stop")
    # Normalize stop to a list (OpenAI accepts string or list)
    if isinstance(stop, str):
        stop = [stop]

    completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    created = int(time.time())
    model_name = body.get("model", MODEL_ID)

    if stream:
        queue = await _stream_chat_async(messages, temperature, top_p, max_tokens, repeat_penalty)

        async def _openai_sse():
            # Initial role delta (OpenAI convention)
            yield {"data": json.dumps({
                "id": completion_id,
                "object": "chat.completion.chunk",
                "created": created,
                "model": model_name,
                "choices": [{
                    "index": 0,
                    "delta": {"role": "assistant"},
                    "finish_reason": None,
                }],
            })}

            while True:
                event = await queue.get()
                etype = event.get("type")

                if etype == "reasoning":
                    yield {"data": json.dumps({
                        "id": completion_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": model_name,
                        "choices": [{
                            "index": 0,
                            "delta": {"reasoning_content": event["token"]},
                            "finish_reason": None,
                        }],
                    })}
                elif etype == "response":
                    yield {"data": json.dumps({
                        "id": completion_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": model_name,
                        "choices": [{
                            "index": 0,
                            "delta": {"content": event["token"]},
                            "finish_reason": None,
                        }],
                    })}
                elif etype == "done":
                    yield {"data": json.dumps({
                        "id": completion_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": model_name,
                        "choices": [{
                            "index": 0,
                            "delta": {},
                            "finish_reason": "stop",
                        }],
                    })}
                    yield {"data": "[DONE]"}
                    break
                elif etype == "error":
                    yield {"data": json.dumps({
                        "error": {"message": event.get("message", "unknown error")},
                    })}
                    break

        return EventSourceResponse(_openai_sse())

    else:
        # Non-streaming: collect all tokens from the queue, return a single JSON.
        queue = await _stream_chat_async(messages, temperature, top_p, max_tokens, repeat_penalty)
        reasoning_content = ""
        content = ""
        final_metrics = None

        while True:
            event = await queue.get()
            etype = event.get("type")
            if etype == "reasoning":
                reasoning_content += event["token"]
            elif etype == "response":
                content += event["token"]
            elif etype == "metrics":
                final_metrics = event
            elif etype == "error":
                return JSONResponse(
                    status_code=500,
                    content={"error": {"message": event.get("message", "unknown error")}},
                )
            elif etype == "done":
                break

        prompt_tokens = 0
        try:
            # Count prompt tokens using the model's tokenizer
            full_prompt = " ".join(m.get("content", "") for m in messages)
            prompt_tokens = len(llm.tokenize(full_prompt.encode("utf-8")))
        except Exception:
            pass

        completion_tokens = (final_metrics or {}).get("total_tokens", 0)

        message_obj = {
            "role": "assistant",
            "content": content,
        }
        if reasoning_content:
            message_obj["reasoning_content"] = reasoning_content

        return {
            "id": completion_id,
            "object": "chat.completion",
            "created": created,
            "model": model_name,
            "choices": [{
                "index": 0,
                "message": message_obj,
                "finish_reason": "stop",
            }],
            "usage": {
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens,
            },
        }


# =============================================================================
# HTML / CSS / JS — custom chat interface
# =============================================================================
HTML_PAGE = r"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>MiniCPM5-1B · CPU Chat</title>
<style>
  :root {
    --bg: #0f0f0f;
    --surface: #1a1a1a;
    --surface-hover: #242424;
    --border: #2a2a2a;
    --text: #e0e0e0;
    --text-dim: #888;
    --accent: #6b8afd;
    --user-bubble: #2563eb;
    --asst-bubble: #1e1e1e;
    --reasoning-bg: #161616;
    --reasoning-border: #333;
    --success: #22c55e;
    --error: #ef4444;
  }
  * { box-sizing: border-box; margin: 0; padding: 0; }
  body {
    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
    background: var(--bg); color: var(--text);
    display: flex; flex-direction: column; height: 100vh; overflow: hidden;
  }
  header {
    padding: 10px 20px; background: var(--surface);
    border-bottom: 1px solid var(--border);
    display: flex; align-items: center; justify-content: space-between;
    flex-shrink: 0;
  }
  header h1 { font-size: 16px; font-weight: 600; }
  header .meta { font-size: 12px; color: var(--text-dim); }
  #main { display: flex; flex: 1; overflow: hidden; }
  #chat-area { flex: 1; display: flex; flex-direction: column; overflow: hidden; }
  #messages {
    flex: 1; overflow-y: auto; padding: 20px;
    display: flex; flex-direction: column; gap: 16px;
  }
  .msg { max-width: 75%; word-wrap: break-word; }
  .msg.user {
    align-self: flex-end;
    background: var(--user-bubble); color: white;
    padding: 10px 14px; border-radius: 16px 16px 4px 16px;
  }
  .msg.asst {
    align-self: flex-start;
    background: var(--asst-bubble); border: 1px solid var(--border);
    padding: 12px 16px; border-radius: 16px 16px 16px 4px;
  }
  .msg.asst .content { line-height: 1.6; }
  .msg.asst .content p { margin-bottom: 8px; }
  .msg.asst .content code {
    background: #0a0a0a; padding: 2px 5px; border-radius: 3px;
    font-family: "SF Mono", Monaco, monospace; font-size: 13px;
  }
  .msg.asst .content pre {
    background: #0a0a0a; padding: 10px; border-radius: 6px;
    overflow-x: auto; margin: 8px 0;
  }
  .msg.asst .content pre code { background: none; padding: 0; }
  .reasoning {
    background: var(--reasoning-bg); border: 1px solid var(--reasoning-border);
    border-radius: 8px; margin-bottom: 10px; overflow: hidden;
  }
  .reasoning summary {
    padding: 8px 12px; cursor: pointer; font-size: 13px;
    color: var(--text-dim); user-select: none;
  }
  .reasoning summary b { color: var(--accent); }
  .reasoning[open] summary { border-bottom: 1px solid var(--reasoning-border); }
  .reasoning .body {
    padding: 10px 12px; font-size: 13px; color: var(--text-dim);
    line-height: 1.5; white-space: pre-wrap; max-height: 400px; overflow-y: auto;
  }
  .metrics-bar {
    margin-top: 10px; padding-top: 8px; border-top: 1px solid var(--border);
    font-size: 11px; color: var(--text-dim); font-family: monospace;
  }
  #input-area {
    padding: 12px 20px; background: var(--surface);
    border-top: 1px solid var(--border); display: flex; gap: 10px;
    flex-shrink: 0;
  }
  #input {
    flex: 1; background: var(--bg); border: 1px solid var(--border);
    color: var(--text); padding: 10px 14px; border-radius: 8px;
    font-size: 14px; font-family: inherit; resize: none;
    min-height: 44px; max-height: 120px;
  }
  #input:focus { outline: none; border-color: var(--accent); }
  #send {
    background: var(--accent); color: white; border: none;
    padding: 0 20px; border-radius: 8px; font-size: 14px; font-weight: 600;
    cursor: pointer; white-space: nowrap;
  }
  #send:disabled { opacity: 0.5; cursor: not-allowed; }
  #send.stop { background: var(--error); }
  #sidebar {
    width: 280px; background: var(--surface); border-left: 1px solid var(--border);
    padding: 16px; overflow-y: auto; flex-shrink: 0;
  }
  #sidebar h3 { font-size: 13px; text-transform: uppercase; color: var(--text-dim); margin-bottom: 10px; }
  #sidebar .setting { margin-bottom: 14px; }
  #sidebar label { display: block; font-size: 12px; margin-bottom: 4px; color: var(--text-dim); }
  #sidebar input[type=text], #sidebar input[type=number] {
    width: 100%; background: var(--bg); border: 1px solid var(--border);
    color: var(--text); padding: 6px 8px; border-radius: 4px; font-size: 13px;
  }
  #sidebar input[type=range] { width: 100%; }
  #sidebar .val { float: right; font-family: monospace; font-size: 12px; color: var(--accent); }
  #live-metrics {
    background: var(--bg); border: 1px solid var(--border);
    border-radius: 6px; padding: 10px; margin-top: 12px;
    font-family: monospace; font-size: 12px;
  }
  #live-metrics .row { display: flex; justify-content: space-between; padding: 2px 0; }
  #live-metrics .label { color: var(--text-dim); }
  #live-metrics .value { color: var(--accent); }
  .examples { display: flex; flex-wrap: wrap; gap: 6px; margin-bottom: 8px; }
  .example {
    background: var(--surface); border: 1px solid var(--border);
    padding: 4px 10px; border-radius: 12px; font-size: 12px;
    cursor: pointer; color: var(--text-dim);
  }
  .example:hover { background: var(--surface-hover); color: var(--text); }
  @media (max-width: 768px) {
    #sidebar { display: none; }
    .msg { max-width: 90%; }
  }
  /* Markdown content styling */
  .msg.asst .content h1, .msg.asst .content h2, .msg.asst .content h3 {
    margin: 12px 0 6px; color: #fff; font-weight: 600;
  }
  .msg.asst .content h1 { font-size: 1.4em; }
  .msg.asst .content h2 { font-size: 1.2em; }
  .msg.asst .content h3 { font-size: 1.1em; }
  .msg.asst .content ul, .msg.asst .content ol {
    margin: 6px 0; padding-left: 24px;
  }
  .msg.asst .content li { margin: 3px 0; }
  .msg.asst .content blockquote {
    border-left: 3px solid var(--accent); padding-left: 12px;
    margin: 8px 0; color: var(--text-dim); font-style: italic;
  }
  .msg.asst .content table {
    border-collapse: collapse; margin: 8px 0; width: 100%;
  }
  .msg.asst .content th, .msg.asst .content td {
    border: 1px solid var(--border); padding: 6px 10px; text-align: left;
  }
  .msg.asst .content th { background: var(--surface); }
  .msg.asst .content a { color: var(--accent); }
  .msg.asst .content hr { border: none; border-top: 1px solid var(--border); margin: 12px 0; }
  .msg.asst .content p:last-child { margin-bottom: 0; }
</style>
<script src="https://cdn.jsdelivr.net/npm/marked@12.0.2/marked.min.js"></script>
</head>
<body>

<header>
  <h1>🤖 MiniCPM5-1B (Q8_0) · CPU</h1>
  <div class="meta">100k context · Max output 30k · Hybrid Reasoning ON · OpenAI API at <code>/v1/chat/completions</code></div>
</header>

<div id="main">
  <div id="chat-area">
    <div id="messages"></div>
    <div class="examples" id="examples">
      <div class="example" onclick="useExample(this)">Explain transformer attention in 3 sentences</div>
      <div class="example" onclick="useExample(this)">Write a Python dedup function</div>
      <div class="example" onclick="useExample(this)">Three habits for senior DevOps</div>
      <div class="example" onclick="useExample(this)">Draft a bug-fix release note</div>
    </div>
    <div id="input-area">
      <textarea id="input" placeholder="Message MiniCPM5-1B... (Enter to send, Shift+Enter for newline)" rows="1"></textarea>
      <button id="send" onclick="sendOrStop()">Send</button>
    </div>
  </div>

  <div id="sidebar">
    <h3>⚙️ Settings</h3>
    <div class="setting">
      <label>System prompt</label>
      <input type="text" id="system-prompt" value="You are Qwen3.6, a helpful and concise assistant. Answer in clear, well-structured prose.">
    </div>
    <div class="setting">
      <label>Temperature <span class="val" id="temp-val">0.70</span></label>
      <input type="range" id="temperature" min="0" max="2" step="0.05" value="0.7" oninput="document.getElementById('temp-val').textContent=parseFloat(this.value).toFixed(2)">
    </div>
    <div class="setting">
      <label>Top-p <span class="val" id="topp-val">0.90</span></label>
      <input type="range" id="top-p" min="0" max="1" step="0.01" value="0.9" oninput="document.getElementById('topp-val').textContent=parseFloat(this.value).toFixed(2)">
    </div>
    <div class="setting">
      <label>Max tokens <span class="val" id="maxt-val">4096</span></label>
      <input type="range" id="max-tokens" min="256" max="30000" step="256" value="4096" oninput="document.getElementById('maxt-val').textContent=this.value">
    </div>
    <div class="setting">
      <label>Repeat penalty <span class="val" id="rp-val">1.05</span></label>
      <input type="range" id="repeat-penalty" min="0.8" max="2" step="0.01" value="1.05" oninput="document.getElementById('rp-val').textContent=parseFloat(this.value).toFixed(2)">
    </div>

    <h3>📊 Live Metrics</h3>
    <div id="live-metrics">
      <div class="row"><span class="label">Status</span><span class="value" id="m-status">idle</span></div>
      <div class="row"><span class="label">TTFT</span><span class="value" id="m-ttft">—</span></div>
      <div class="row"><span class="label">Reasoning tok/s</span><span class="value" id="m-rtps">—</span></div>
      <div class="row"><span class="label">Response tok/s</span><span class="value" id="m-vtps">—</span></div>
      <div class="row"><span class="label">Total tok/s</span><span class="value" id="m-ttps">—</span></div>
      <div class="row"><span class="label">Reasoning count</span><span class="value" id="m-rcount">0</span></div>
      <div class="row"><span class="label">Response count</span><span class="value" id="m-vcount">0</span></div>
      <div class="row"><span class="label">Total tokens</span><span class="value" id="m-total">0</span></div>
      <div class="row"><span class="label">Elapsed</span><span class="value" id="m-elapsed">—</span></div>
      <div class="row"><span class="label">Gen time</span><span class="value" id="m-gen">—</span></div>
      <div class="row"><span class="label">Total time</span><span class="value" id="m-total-time">—</span></div>
    </div>
  </div>
</div>

<script>
let history = [];
let isGenerating = false;
let currentController = null;
let reasoningText = '';
let responseText = '';
let promptStart = 0;
let elapsedTimer = null;

const messagesEl = document.getElementById('messages');
const inputEl = document.getElementById('input');
const sendBtn = document.getElementById('send');

inputEl.addEventListener('keydown', e => {
  if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); sendOrStop(); }
});

function useExample(el) {
  inputEl.value = el.textContent;
  inputEl.focus();
}

function setMetrics(id, val) {
  const el = document.getElementById(id);
  if (el) el.textContent = val;
}

function startElapsedTimer() {
  stopElapsedTimer();
  promptStart = Date.now();
  elapsedTimer = setInterval(() => {
    const elapsed = ((Date.now() - promptStart) / 1000).toFixed(1);
    setMetrics('m-elapsed', elapsed + 's');
  }, 100);
}

function stopElapsedTimer() {
  if (elapsedTimer) {
    clearInterval(elapsedTimer);
    elapsedTimer = null;
  }
}

function resetLiveMetrics() {
  setMetrics('m-status', 'starting...');
  setMetrics('m-ttft', '—');
  setMetrics('m-rtps', '—');
  setMetrics('m-vtps', '—');
  setMetrics('m-ttps', '—');
  setMetrics('m-rcount', '0');
  setMetrics('m-vcount', '0');
  setMetrics('m-total', '0');
  setMetrics('m-elapsed', '0.0s');
  setMetrics('m-gen', '—');
  setMetrics('m-total-time', '—');
  startElapsedTimer();
}

function escapeHtml(s) {
  return s.replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;');
}

// Full markdown rendering via marked.js (loaded from CDN in <head>).
// Falls back to escaped text if marked isn't available.
function renderMarkdown(text) {
  if (typeof marked !== 'undefined') {
    try {
      return marked.parse(text, { breaks: true, gfm: true });
    } catch (e) {
      console.error('marked error:', e);
    }
  }
  return escapeHtml(text).replace(/\n/g, '<br>');
}

function addMessage(role, content) {
  const div = document.createElement('div');
  div.className = 'msg ' + role;
  if (role === 'user') {
    div.innerHTML = escapeHtml(content);
  } else {
    div.innerHTML = `<div class="content">${renderMarkdown(content)}</div>`;
  }
  messagesEl.appendChild(div);
  messagesEl.scrollTop = messagesEl.scrollHeight;
  return div;
}

function updateAsstBubble(reasoning, response, isThinking) {
  let bubble = document.getElementById('current-asst');
  if (!bubble) {
    bubble = document.createElement('div');
    bubble.className = 'msg asst';
    bubble.id = 'current-asst';
    messagesEl.appendChild(bubble);
  }

  // Check if we need to (re)create the structure. We avoid replacing
  // innerHTML on every token because that resets scroll positions.
  // Instead, we create the structure once and only update text content.
  let reasoningEl = bubble.querySelector('details.reasoning');
  let reasoningBody = reasoningEl ? reasoningEl.querySelector('.body') : null;
  let contentEl = bubble.querySelector('.content');

  // (Re)create reasoning block if presence changed
  const needReasoning = !!reasoning;
  const hasReasoning = !!reasoningEl;
  if (needReasoning && !hasReasoning) {
    reasoningEl = document.createElement('details');
    reasoningEl.className = 'reasoning';
    reasoningEl.open = true;
    const summary = document.createElement('summary');
    summary.innerHTML = '<b>💭 Thinking...</b>';
    reasoningBody = document.createElement('div');
    reasoningBody.className = 'body';
    reasoningEl.appendChild(summary);
    reasoningEl.appendChild(reasoningBody);
    bubble.insertBefore(reasoningEl, bubble.firstChild);
  }
  if (!needReasoning && hasReasoning) {
    reasoningEl.remove();
    reasoningEl = null;
    reasoningBody = null;
  }

  // Update reasoning summary label (Thinking... vs Reasoning)
  if (reasoningEl) {
    const summary = reasoningEl.querySelector('summary');
    if (summary) {
      const label = isThinking ? '💭 Thinking...' : '💭 Reasoning';
      summary.innerHTML = `<b>${label}</b>`;
    }
  }

  // Update reasoning body text WITHOUT resetting scroll
  if (reasoningBody) {
    reasoningBody.textContent = reasoning;
    // Auto-scroll reasoning to bottom while thinking (follow new tokens)
    if (isThinking) {
      reasoningBody.scrollTop = reasoningBody.scrollHeight;
    }
  }

  // (Re)create content block if presence changed
  const needContent = !!response || !!reasoning;  // show placeholder during thinking
  const hasContent = !!contentEl;
  if (needContent && !hasContent) {
    contentEl = document.createElement('div');
    contentEl.className = 'content';
    bubble.appendChild(contentEl);
  }

  // Update content
  if (contentEl) {
    if (response) {
      contentEl.innerHTML = renderMarkdown(response);
      contentEl.style.color = '';
      contentEl.style.fontStyle = '';
    } else if (reasoning) {
      contentEl.innerHTML = 'generating response...';
      contentEl.style.color = '#888';
      contentEl.style.fontStyle = 'italic';
    }
  }

  // Auto-scroll the main messages container to follow the response,
  // but only when we're past the reasoning phase (response is streaming).
  // During reasoning, the reasoning box handles its own scroll.
  if (response) {
    messagesEl.scrollTop = messagesEl.scrollHeight;
  }
}

function finalizeAsstBubble(metrics) {
  let bubble = document.getElementById('current-asst');
  if (!bubble) return;
  bubble.removeAttribute('id');

  // Collapse the reasoning block now that we're done
  const details = bubble.querySelector('details.reasoning');
  if (details) details.removeAttribute('open');

  // Append metrics bar with TPS + total time prominently
  if (metrics) {
    const m = document.createElement('div');
    m.className = 'metrics-bar';
    m.innerHTML = `
      <div style="display:flex;gap:16px;flex-wrap:wrap;">
        <span>📊 <b>${metrics.total_tokens}</b> tokens</span>
        <span>🧠 ${metrics.reasoning_tokens} thinking</span>
        <span>💬 ${metrics.response_tokens} response</span>
      </div>
      <div style="display:flex;gap:16px;flex-wrap:wrap;margin-top:4px;">
        <span>⚡ TTFT <b>${metrics.ttft_ms}ms</b></span>
        <span>🧠 <b>${metrics.reasoning_tps}</b> tok/s (reasoning)</span>
        <span>💬 <b>${metrics.response_tps}</b> tok/s (response)</span>
        <span>⏱️ <b>${metrics.total_time_s}s</b> total</span>
      </div>
    `;
    bubble.appendChild(m);
  }
}

async function sendOrStop() {
  if (isGenerating) {
    if (currentController) currentController.abort();
    return;
  }
  const msg = inputEl.value.trim();
  if (!msg) return;

  inputEl.value = '';
  inputEl.style.height = 'auto';
  addMessage('user', msg);
  history.push({ role: 'user', content: msg });

  reasoningText = '';
  responseText = '';
  promptStart = Date.now();
  resetLiveMetrics();
  updateAsstBubble('', '', true);

  isGenerating = true;
  sendBtn.textContent = 'Stop';
  sendBtn.classList.add('stop');

  currentController = new AbortController();

  try {
    const resp = await fetch('/api/chat', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        history: history,
        system_prompt: document.getElementById('system-prompt').value,
        temperature: parseFloat(document.getElementById('temperature').value),
        top_p: parseFloat(document.getElementById('top-p').value),
        max_tokens: parseInt(document.getElementById('max-tokens').value),
        repeat_penalty: parseFloat(document.getElementById('repeat-penalty').value),
      }),
      signal: currentController.signal,
    });

    const reader = resp.body.getReader();
    const decoder = new TextDecoder();
    let buffer = '';
    let finalMetrics = null;

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      buffer += decoder.decode(value, { stream: true });

      const lines = buffer.split('\n');
      buffer = lines.pop();

      for (const line of lines) {
        if (!line.startsWith('data: ')) continue;
        try {
          const ev = JSON.parse(line.slice(6));
          handleEvent(ev);
          if (ev.type === 'metrics') finalMetrics = ev;
          if (ev.type === 'done' || ev.type === 'error') {
            finalizeAsstBubble(finalMetrics);
            if (ev.type === 'error') {
              setMetrics('m-status', 'error');
            }
          }
        } catch (e) {}
      }
    }

    // Add the full response to history
    history.push({ role: 'assistant', content: responseText });

  } catch (err) {
    stopElapsedTimer();
    if (err.name === 'AbortError') {
      setMetrics('m-status', 'stopped');
      history.push({ role: 'assistant', content: responseText + ' [stopped]' });
    } else {
      setMetrics('m-status', 'error: ' + err.message);
    }
    finalizeAsstBubble(null);
  } finally {
    stopElapsedTimer();
    isGenerating = false;
    sendBtn.textContent = 'Send';
    sendBtn.classList.remove('stop');
    currentController = null;
  }
}

function handleEvent(ev) {
  switch (ev.type) {
    case 'status':
      setMetrics('m-status', ev.stage === 'processing_prompt' ? 'processing prompt...' : ev.stage);
      break;
    case 'reasoning':
      reasoningText += ev.token;
      setMetrics('m-rcount', ev.count);
      setMetrics('m-rtps', ev.tps.toFixed(2));
      const rTotal = ev.count + parseInt(document.getElementById('m-vcount').textContent);
      setMetrics('m-total', rTotal);
      updateAsstBubble(reasoningText, responseText, true);
      break;
    case 'response':
      responseText += ev.token;
      setMetrics('m-vcount', ev.count);
      setMetrics('m-vtps', ev.tps.toFixed(2));
      const vTotal = parseInt(document.getElementById('m-rcount').textContent) + ev.count;
      setMetrics('m-total', vTotal);
      if (reasoningText) {
        updateAsstBubble(reasoningText, responseText, false);
      } else {
        updateAsstBubble('', responseText, false);
      }
      setMetrics('m-status', 'streaming response');
      break;
    case 'metrics':
      stopElapsedTimer();
      setMetrics('m-status', 'done');
      setMetrics('m-ttft', ev.ttft_ms + 'ms');
      setMetrics('m-gen', ev.gen_time_s + 's');
      setMetrics('m-total-time', ev.total_time_s + 's');
      setMetrics('m-rtps', ev.reasoning_tps.toFixed(2));
      setMetrics('m-vtps', ev.response_tps.toFixed(2));
      setMetrics('m-ttps', ev.total_tps.toFixed(2));
      setMetrics('m-elapsed', ev.total_time_s + 's');
      break;
    case 'error':
      stopElapsedTimer();
      setMetrics('m-status', 'error');
      responseText += '\n\n[error] ' + ev.message;
      updateAsstBubble(reasoningText, responseText, false);
      break;
    case 'done':
      stopElapsedTimer();
      break;
  }
}

// Auto-resize textarea
inputEl.addEventListener('input', () => {
  inputEl.style.height = 'auto';
  inputEl.style.height = Math.min(inputEl.scrollHeight, 120) + 'px';
});
</script>
</body>
</html>
"""


# =============================================================================
# Launch
# =============================================================================
if __name__ == "__main__":
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=PORT,
        log_level="info",
    )