app.py

"""
LFM2.5 FastAPI Backend - ONNX Runtime Edition
==============================================
Lightweight, CPU-friendly FastAPI backend for LiquidAI LFM2.5-1.2B-Instruct.
Uses official ONNX model for fast inference without heavy PyTorch dependencies.

Features:
- ONNX Runtime for fast CPU inference (no GPU required)
- Q8 quantization for 95%+ accuracy retention
- Streaming SSE responses
- OpenAI-compatible API
- Optimized for HuggingFace Spaces (2 vCPU, 16GB RAM)
"""

import asyncio
import json
import logging
import time
import uuid
import threading
import queue  # Thread-safe queue for true streaming
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Dict, List, Optional, Union
from pathlib import Path

import numpy as np
import onnxruntime as ort
from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from huggingface_hub import hf_hub_download, list_repo_files
from pydantic import BaseModel, Field
from sse_starlette.sse import EventSourceResponse
from transformers import AutoTokenizer, PreTrainedTokenizerFast

from config import settings

# Configure logging
logging.basicConfig(
    level=getattr(logging, settings.log_level.upper()),
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


# ==============================================================================
# Pydantic Models for OpenAI-compatible API
# ==============================================================================

class ChatMessage(BaseModel):
    role: str = Field(..., description="Role: 'system', 'user', or 'assistant'")
    content: str = Field(..., description="Message content")


class ChatCompletionRequest(BaseModel):
    model: str = Field(default="lfm", description="Model identifier")
    messages: List[ChatMessage] = Field(..., description="Conversation messages")
    temperature: Optional[float] = Field(default=None, ge=0.0, le=2.0)
    top_p: Optional[float] = Field(default=None, ge=0.0, le=1.0)
    top_k: Optional[int] = Field(default=None, ge=0)
    max_tokens: Optional[int] = Field(default=None, ge=1)
    stream: bool = Field(default=False, description="Enable streaming response")
    stop: Optional[Union[str, List[str]]] = Field(default=None)


class CompletionRequest(BaseModel):
    model: str = Field(default="lfm", description="Model identifier")
    prompt: str = Field(..., description="Text prompt")
    temperature: Optional[float] = Field(default=None, ge=0.0, le=2.0)
    top_p: Optional[float] = Field(default=None, ge=0.0, le=1.0)
    top_k: Optional[int] = Field(default=None, ge=0)
    max_tokens: Optional[int] = Field(default=None, ge=1)
    stream: bool = Field(default=False, description="Enable streaming response")


class ChatCompletionChoice(BaseModel):
    index: int
    message: ChatMessage
    finish_reason: Optional[str] = None


class ChatCompletionResponse(BaseModel):
    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: List[ChatCompletionChoice]
    usage: Dict[str, int]


class CompletionChoice(BaseModel):
    index: int
    text: str
    finish_reason: Optional[str] = None


class CompletionResponse(BaseModel):
    id: str
    object: str = "text_completion"
    created: int
    model: str
    choices: List[CompletionChoice]
    usage: Dict[str, int]


class ModelInfo(BaseModel):
    id: str
    object: str = "model"
    created: int
    owned_by: str = "liquid-ai"


class ModelListResponse(BaseModel):
    object: str = "list"
    data: List[ModelInfo]


# ==============================================================================
# ONNX Model Manager
# ==============================================================================

# ONNX dtype mapping
ONNX_DTYPE = {
    "tensor(float)": np.float32,
    "tensor(float16)": np.float16,
    "tensor(int64)": np.int64
}


class ONNXModelManager:
    """Manages ONNX model with KV cache for efficient generation."""
    
    def __init__(self):
        self._session = None
        self._tokenizer = None
        self._cache_template = None
        self._use_position_ids = False
        self._lock = threading.Lock()
    
    @property
    def is_loaded(self) -> bool:
        return self._session is not None
    
    def download_model(self) -> str:
        """Download ONNX model files from HuggingFace."""
        model_id = settings.model_id
        variant = settings.model_variant
        
        logger.info(f"Downloading model: {model_id} (variant: {variant})")
        
        # Download main model file
        model_filename = f"onnx/model_{variant}.onnx"
        model_path = hf_hub_download(model_id, model_filename)
        
        # Download all data files for this variant
        for f in list_repo_files(model_id):
            if f.startswith(f"onnx/model_{variant}.onnx_data"):
                logger.info(f"Downloading: {f}")
                hf_hub_download(model_id, f)
        
        return model_path
    
    def load_model(self) -> None:
        """Load the ONNX model and tokenizer."""
        with self._lock:
            if self._session is not None:
                return
            
            logger.info("=" * 60)
            logger.info("Loading LFM2.5-1.2B-Instruct ONNX model...")
            logger.info(f"Model: {settings.model_id}")
            logger.info(f"Variant: {settings.model_variant} (Q8 = ~95% accuracy)")
            logger.info("=" * 60)
            
            start_time = time.time()
            
            # Download model
            model_path = self.download_model()
            
            # Configure ONNX Runtime for CPU
            sess_options = ort.SessionOptions()
            sess_options.intra_op_num_threads = settings.num_threads
            sess_options.inter_op_num_threads = settings.num_threads
            sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
            
            # Load ONNX session
            self._session = ort.InferenceSession(
                model_path,
                sess_options=sess_options,
                providers=['CPUExecutionProvider']
            )
            
            # Load tokenizer with fallback for models with invalid tokenizer_class
            try:
                self._tokenizer = AutoTokenizer.from_pretrained(
                    settings.model_id,
                    trust_remote_code=True
                )
            except ValueError as e:
                if "TokenizersBackend" in str(e):
                    # LFM models incorrectly specify TokenizersBackend as tokenizer_class
                    # Fallback to PreTrainedTokenizerFast which works with tokenizers backend
                    logger.warning(
                        "AutoTokenizer failed with TokenizersBackend error. "
                        "Falling back to PreTrainedTokenizerFast..."
                    )
                    self._tokenizer = PreTrainedTokenizerFast.from_pretrained(
                        settings.model_id,
                        trust_remote_code=True
                    )
                else:
                    raise
            
            # Initialize cache template
            self._init_cache_template()
            
            # Check if model uses position_ids
            input_names = {inp.name for inp in self._session.get_inputs()}
            self._use_position_ids = "position_ids" in input_names
            
            load_time = time.time() - start_time
            logger.info("=" * 60)
            logger.info(f"✓ Model loaded in {load_time:.2f}s")
            logger.info(f"  Threads: {settings.num_threads}")
            logger.info(f"  Provider: CPU")
            logger.info("=" * 60)
    
    def _init_cache_template(self) -> None:
        """Initialize KV cache template."""
        self._cache_template = {}
        for inp in self._session.get_inputs():
            if inp.name in {"input_ids", "attention_mask", "position_ids"}:
                continue
            
            shape = [d if isinstance(d, int) else 1 for d in inp.shape]
            for i, d in enumerate(inp.shape):
                if isinstance(d, str) and "sequence" in d.lower():
                    shape[i] = 0
            
            dtype = ONNX_DTYPE.get(inp.type, np.float32)
            self._cache_template[inp.name] = (shape, dtype)
    
    def _create_empty_cache(self) -> Dict[str, np.ndarray]:
        """Create a new empty KV cache."""
        return {
            name: np.zeros(shape, dtype=dtype)
            for name, (shape, dtype) in self._cache_template.items()
        }
    
    @property
    def session(self):
        if self._session is None:
            raise RuntimeError("Model not loaded")
        return self._session
    
    @property
    def tokenizer(self):
        if self._tokenizer is None:
            raise RuntimeError("Tokenizer not loaded")
        return self._tokenizer
    
    def generate(
        self,
        input_ids: np.ndarray,
        max_tokens: int = 512,
        temperature: float = 0.1,
        top_k: int = 50,
        top_p: float = 0.1,
        stop_tokens: Optional[List[int]] = None
    ) -> List[int]:
        """Generate tokens using ONNX model."""
        if stop_tokens is None:
            stop_tokens = [self._tokenizer.eos_token_id]
        
        cache = self._create_empty_cache()
        seq_len = input_ids.shape[1]
        generated_tokens = []
        
        for step in range(max_tokens):
            if step == 0:
                ids = input_ids
                pos = np.arange(seq_len, dtype=np.int64).reshape(1, -1)
            else:
                ids = np.array([[generated_tokens[-1]]], dtype=np.int64)
                pos = np.array([[seq_len + len(generated_tokens) - 1]], dtype=np.int64)
            
            attn_mask = np.ones((1, seq_len + len(generated_tokens)), dtype=np.int64)
            
            feed = {"input_ids": ids, "attention_mask": attn_mask, **cache}
            if self._use_position_ids:
                feed["position_ids"] = pos
            
            outputs = self._session.run(None, feed)
            
            # Get logits and apply temperature
            logits = outputs[0][0, -1]
            
            if temperature > 0:
                logits = logits / temperature
                
                # Apply top-k
                if top_k > 0:
                    indices_to_remove = np.argsort(logits)[:-top_k]
                    logits[indices_to_remove] = -np.inf
                
                # Apply top-p (nucleus sampling)
                if top_p < 1.0:
                    sorted_indices = np.argsort(logits)[::-1]
                    sorted_logits = logits[sorted_indices]
                    probs = np.exp(sorted_logits - np.max(sorted_logits))
                    probs = probs / probs.sum()
                    cumulative_probs = np.cumsum(probs)
                    sorted_indices_to_remove = cumulative_probs > top_p
                    sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].copy()
                    sorted_indices_to_remove[0] = False
                    indices_to_remove = sorted_indices[sorted_indices_to_remove]
                    logits[indices_to_remove] = -np.inf
                
                # Sample
                probs = np.exp(logits - np.max(logits))
                probs = probs / probs.sum()
                next_token = int(np.random.choice(len(probs), p=probs))
            else:
                next_token = int(np.argmax(logits))
            
            generated_tokens.append(next_token)
            
            # Update cache
            for i, out in enumerate(self._session.get_outputs()[1:], 1):
                name = out.name.replace("present_conv", "past_conv").replace("present.", "past_key_values.")
                if name in cache:
                    cache[name] = outputs[i]
            
            if next_token in stop_tokens:
                break
        
        return generated_tokens
    
    def generate_stream(
        self,
        input_ids: np.ndarray,
        max_tokens: int = 2000,
        temperature: float = 0.1,
        top_k: int = 50,
        top_p: float = 0.1,
        stop_tokens: Optional[List[int]] = None
    ):
        """Fixed and optimized streaming generation."""
        if stop_tokens is None:
            stop_tokens = [self._tokenizer.eos_token_id]
        
        cache = self._create_empty_cache()
        seq_len = input_ids.shape[1]
        
        # Pre-allocate inputs
        max_possible_len = seq_len + max_tokens
        attn_mask = np.ones((1, max_possible_len), dtype=np.int64)
        
        # Pre-compute flags
        use_temp = temperature > 0
        use_top_k = top_k > 0
        use_top_p = top_p < 1.0
        
        # Reuse this dict to avoid garbage collection overhead
        feed = {}
        
        # Initialize token storage
        generated_tokens = []
        
        for step in range(max_tokens):
            current_len = seq_len + step
            
            # Input Preparation
            if step == 0:
                ids = input_ids
                if self._use_position_ids:
                    pos = np.arange(seq_len, dtype=np.int64).reshape(1, -1)
            else:
                # FIX: Access list directly. O(1) speed, no UnboundLocalError.
                ids = np.array([[generated_tokens[-1]]], dtype=np.int64)
                if self._use_position_ids:
                    pos = np.array([[current_len - 1]], dtype=np.int64)
            
            # Update Feed Dict (In-place update is faster than creating new dict)
            feed.clear()
            feed["input_ids"] = ids
            feed["attention_mask"] = attn_mask[:, :current_len]
            if self._use_position_ids:
                feed["position_ids"] = pos
            feed.update(cache) # Merging cache is unavoidable
            
            # Inference
            outputs = self._session.run(None, feed)
            logits = outputs[0][0, -1]
            
            # --- Ultra-Fast Sampling ---
            if use_temp:
                logits /= temperature
                
                # 1. Top-K Selection (Partitioning is O(N))
                if use_top_k and top_k < len(logits):
                    # Moves largest k elements to the right; unordered
                    top_k_idx = np.argpartition(logits, -top_k)[-top_k:] 
                    # Mask everything else
                    mask = np.ones(logits.shape, dtype=bool)
                    mask[top_k_idx] = False
                    logits[mask] = -np.inf
                
                # 2. Top-P (Nucleus)
                if use_top_p:
                    valid_mask = logits > -np.inf
                    if valid_mask.any():
                        valid_logits = logits[valid_mask]
                        valid_indices = np.where(valid_mask)[0]
                        
                        # Sort only the valid candidates (small N)
                        sorted_indices = np.argsort(valid_logits)[::-1]
                        sorted_logits = valid_logits[sorted_indices]
                        
                        # Softmax on valid set
                        exp_logits = np.exp(sorted_logits - np.max(sorted_logits))
                        probs = exp_logits / exp_logits.sum()
                        
                        cumulative = np.cumsum(probs)
                        
                        # Find cutoff
                        cutoff = np.searchsorted(cumulative, top_p)
                        # Ensure we keep at least one token
                        cutoff = min(cutoff + 1, len(sorted_logits))
                        
                        # Filter indices
                        accepted_indices = sorted_indices[:cutoff]
                        accepted_probs = probs[:cutoff]
                        accepted_probs /= accepted_probs.sum() # Re-normalize
                        
                        # Fast Weighted Sample: Use searchsorted instead of np.random.choice
                        # This avoids Python overhead in np.random.choice
                        sample_idx = np.searchsorted(np.cumsum(accepted_probs), np.random.rand())
                        next_token = int(valid_indices[accepted_indices[sample_idx]])
                    else:
                        next_token = int(np.argmax(logits))
                else:
                    # Fallback if only Top-K was used
                    valid_mask = logits > -np.inf
                    valid_logits = logits[valid_mask]
                    valid_indices = np.where(valid_mask)[0]
                    exp_logits = np.exp(valid_logits - np.max(valid_logits))
                    probs = exp_logits / exp_logits.sum()
                    sample_idx = np.searchsorted(np.cumsum(probs), np.random.rand())
                    next_token = int(valid_indices[sample_idx])
            else:
                next_token = int(np.argmax(logits))
            
            # Storage
            generated_tokens.append(next_token)
            yield next_token
            
            if next_token in stop_tokens:
                break
                
            # Update Cache
            for i, out in enumerate(self._session.get_outputs()[1:], 1):
                name = out.name.replace("present_conv", "past_conv").replace("present.", "past_key_values.")
                if name in cache:
                    cache[name] = outputs[i]
    
    def unload(self) -> None:
        """Unload model from memory."""
        with self._lock:
            if self._session is not None:
                del self._session
                del self._tokenizer
                self._session = None
                self._tokenizer = None
                logger.info("Model unloaded")


# Global model manager
model_manager = ONNXModelManager()


# ==============================================================================
# Application Lifecycle
# ==============================================================================

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan handler."""
    logger.info("Starting LFM2.5 API Server (ONNX Runtime)...")
    
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, model_manager.load_model)
    
    yield
    
    logger.info("Shutting down...")
    model_manager.unload()


# ==============================================================================
# FastAPI Application
# ==============================================================================

app = FastAPI(
    title=settings.app_name,
    description="Fast CPU inference for LiquidAI LFM2.5-1.2B-Instruct using ONNX Runtime",
    version=settings.app_version,
    lifespan=lifespan,
    docs_url="/docs",
    redoc_url="/redoc",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all origins
    allow_credentials=False,  # Must be False when using wildcard origins
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"],  # Expose all headers for SSE
)


# Custom middleware to handle null origin (file:// protocol)
@app.middleware("http")
async def add_cors_for_null_origin(request: Request, call_next):
    """Handle CORS for null origin (when HTML is opened from file://)."""
    origin = request.headers.get("origin", "")
    response = await call_next(request)
    
    # If origin is null (file:// protocol), add explicit CORS headers
    if origin == "null" or not origin:
        response.headers["Access-Control-Allow-Origin"] = "*"
        response.headers["Access-Control-Allow-Methods"] = "GET, POST, PUT, DELETE, OPTIONS"
        response.headers["Access-Control-Allow-Headers"] = "*"
        response.headers["Access-Control-Expose-Headers"] = "*"
    
    return response


# ==============================================================================
# Helper Functions
# ==============================================================================

def generate_id() -> str:
    return f"chatcmpl-{uuid.uuid4().hex[:12]}"


async def stream_chat_completion(request: ChatCompletionRequest) -> AsyncGenerator[str, None]:
    """
    Optimized 'Zero-Latency' Streaming.
    Uses asyncio.Queue + call_soon_threadsafe to eliminate polling and blocking.
    """
    request_id = generate_id()
    created = int(time.time())
    
    # Capture the running event loop to bridge the background thread safely
    loop = asyncio.get_running_loop()
    # Async Queue allows 'await get()' which is non-blocking and instant
    async_queue = asyncio.Queue()
    
    tokenizer = model_manager.tokenizer
    
    # Prepare inputs
    messages = [{"role": m.role, "content": m.content} for m in request.messages]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    input_ids = np.array([tokenizer.encode(prompt, add_special_tokens=False)], dtype=np.int64)
    
    # Config
    max_tokens = request.max_tokens or settings.max_tokens
    temperature = request.temperature if request.temperature is not None else settings.temperature
    top_k = request.top_k if request.top_k is not None else settings.top_k
    top_p = request.top_p if request.top_p is not None else settings.top_p
    
    # Prepare stop tokens
    stop_tokens = [tokenizer.eos_token_id]
    if request.stop:
        if isinstance(request.stop, str):
            encoded = tokenizer.encode(request.stop, add_special_tokens=False)
            if encoded:
                stop_tokens.append(encoded[0])
        elif isinstance(request.stop, list):
            for stop_str in request.stop:
                encoded = tokenizer.encode(stop_str, add_special_tokens=False)
                if encoded:
                    stop_tokens.append(encoded[0])

    def generate_tokens():
        """
        Background Thread: Pushes data directly into the async loop.
        """
        try:
            # Use the optimized generate_stream from ONNXModelManager
            for token in model_manager.generate_stream(
                input_ids,
                max_tokens=max_tokens,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                stop_tokens=stop_tokens
            ):
                # CRITICAL: Schedule the 'put' on the main loop immediately
                # This wakes up the awaiter instantly—0ms latency overhead.
                loop.call_soon_threadsafe(async_queue.put_nowait, ("token", token))
        except Exception as e:
            logger.error(f"Stream generation error: {e}")
            loop.call_soon_threadsafe(async_queue.put_nowait, ("error", str(e)))
        finally:
            loop.call_soon_threadsafe(async_queue.put_nowait, ("done", None))

    # Start generation in background thread
    threading.Thread(target=generate_tokens, daemon=True).start()
    
    # Main Async Loop - No timeouts, no sleeps, pure event awaiting
    try:
        while True:
            # waits until data is pushed; yields control to other users while waiting
            msg_type, data = await async_queue.get()

            if msg_type == "token":
                text = tokenizer.decode([data], skip_special_tokens=True)
                if text:
                    chunk = {
                        "id": request_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": request.model,
                        "choices": [{
                            "index": 0,
                            "delta": {"content": text},
                            "finish_reason": None
                        }]
                    }
                    # Yield in the format expected by EventSourceResponse
                    yield {"data": json.dumps(chunk)}
            
            elif msg_type == "done":
                final = {
                    "id": request_id,
                    "object": "chat.completion.chunk",
                    "created": created,
                    "model": request.model,
                    "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
                }
                yield {"data": json.dumps(final)}
                yield {"data": "[DONE]"}
                break
                
            elif msg_type == "error":
                logger.error(f"Stream error: {data}")
                yield {"data": json.dumps({"error": {"message": data}})}
                break
                
    except asyncio.CancelledError:
        logger.info(f"Stream cancelled for request {request_id[:8]}")
        raise
    except Exception as e:
        logger.error(f"Streaming error: {e}")
        yield {"data": json.dumps({"error": {"message": str(e)}})}


# ==============================================================================
# API Endpoints
# ==============================================================================

@app.get("/", response_class=JSONResponse)
async def health_check():
    """Health check with model status."""
    return {
        "status": "ready" if model_manager.is_loaded else "loading",
        "model": {
            "id": settings.model_id,
            "variant": settings.model_variant,
            "loaded": model_manager.is_loaded,
            "backend": "ONNX Runtime"
        },
        "server": {
            "name": settings.app_name,
            "version": settings.app_version,
            "port": settings.port
        }
    }


@app.get("/health")
async def health():
    if not model_manager.is_loaded:
        raise HTTPException(status_code=503, detail="Model not loaded")
    return {"status": "healthy"}


@app.get("/v1/models", response_model=ModelListResponse)
async def list_models():
    return ModelListResponse(
        data=[
            ModelInfo(id="lfm", created=int(time.time())),
            ModelInfo(id="lfm-2.5-1.2b-instruct-onnx", created=int(time.time()))
        ]
    )


@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
    """OpenAI-compatible chat completion."""
    if not model_manager.is_loaded:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    if request.stream:
        return EventSourceResponse(
            stream_chat_completion(request),
            media_type="text/event-stream",
            ping=30000,  # 30 second keep-alive
            ping_message_factory=lambda: '{"type": "ping"}'
        )
    
    try:
        tokenizer = model_manager.tokenizer
        
        messages = [{"role": m.role, "content": m.content} for m in request.messages]
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        input_ids = np.array([tokenizer.encode(prompt, add_special_tokens=False)], dtype=np.int64)
        
        max_tokens = request.max_tokens or settings.max_tokens
        temperature = request.temperature if request.temperature is not None else settings.temperature
        top_k = request.top_k if request.top_k is not None else settings.top_k
        top_p = request.top_p if request.top_p is not None else settings.top_p
        
        start_time = time.time()
        
        loop = asyncio.get_event_loop()
        tokens = await loop.run_in_executor(
            None,
            lambda: model_manager.generate(
                input_ids,
                max_tokens=max_tokens,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p
            )
        )
        
        response_text = tokenizer.decode(tokens, skip_special_tokens=True)
        gen_time = time.time() - start_time
        
        logger.debug(f"Generated {len(tokens)} tokens in {gen_time:.2f}s")
        
        return ChatCompletionResponse(
            id=generate_id(),
            created=int(time.time()),
            model=request.model,
            choices=[
                ChatCompletionChoice(
                    index=0,
                    message=ChatMessage(role="assistant", content=response_text),
                    finish_reason="stop"
                )
            ],
            usage={
                "prompt_tokens": input_ids.shape[1],
                "completion_tokens": len(tokens),
                "total_tokens": input_ids.shape[1] + len(tokens)
            }
        )
        
    except Exception as e:
        logger.error(f"Chat completion error: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/v1/completions")
async def completions(request: CompletionRequest):
    """OpenAI-compatible text completion."""
    if not model_manager.is_loaded:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    try:
        tokenizer = model_manager.tokenizer
        input_ids = np.array([tokenizer.encode(request.prompt)], dtype=np.int64)
        
        max_tokens = request.max_tokens or settings.max_tokens
        temperature = request.temperature if request.temperature is not None else settings.temperature
        top_k = request.top_k if request.top_k is not None else settings.top_k
        top_p = request.top_p if request.top_p is not None else settings.top_p
        
        loop = asyncio.get_event_loop()
        tokens = await loop.run_in_executor(
            None,
            lambda: model_manager.generate(
                input_ids,
                max_tokens=max_tokens,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p
            )
        )
        
        response_text = tokenizer.decode(tokens, skip_special_tokens=True)
        
        return CompletionResponse(
            id=generate_id(),
            created=int(time.time()),
            model=request.model,
            choices=[
                CompletionChoice(index=0, text=response_text, finish_reason="stop")
            ],
            usage={
                "prompt_tokens": input_ids.shape[1],
                "completion_tokens": len(tokens),
                "total_tokens": input_ids.shape[1] + len(tokens)
            }
        )
        
    except Exception as e:
        logger.error(f"Completion error: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# ==============================================================================
# WebSocket Autocomplete Endpoint
# ==============================================================================

@app.websocket("/ws/autocomplete")
async def ws_autocomplete(websocket: WebSocket):
    """
    Persistent WebSocket endpoint for inline text predictions.

    Protocol:
      Client sends:  { "context": "last ~300 chars before cursor" }
      Server sends:  { "suggestion": "predicted next words" }
      Client sends:  { "type": "ping" }  →  Server sends: { "type": "pong" }

    Design decisions:
      - Persistent connection: avoids reconnect overhead per prediction
      - Low temperature (0.3): more deterministic for inline suggestions
      - Max 20 tokens: keeps predictions short and fast (~800ms)
      - Stop on sentence boundaries (., !, ?, newline): natural break points
      - Uses "raw completion" prompt (no chat template): faster, less overhead
    """
    await websocket.accept()
    logger.info("[ws/autocomplete] Client connected")

    try:
        while True:
            # Wait for a prediction request from the client
            raw = await websocket.receive_text()

            try:
                data = json.loads(raw)
            except json.JSONDecodeError:
                await websocket.send_text(json.dumps({"error": "Invalid JSON"}))
                continue

            # Heartbeat: respond to pings immediately
            if data.get("type") == "ping":
                await websocket.send_text(json.dumps({"type": "pong"}))
                continue

            context = data.get("context", "").strip()
            if not context:
                await websocket.send_text(json.dumps({"suggestion": ""}))
                continue

            if not model_manager.is_loaded:
                await websocket.send_text(json.dumps({"suggestion": ""}))
                continue

            # Generate prediction using the model
            try:
                tokenizer = model_manager.tokenizer
                max_tokens = min(data.get("max_tokens", 20), 30)  # Cap at 30

                # Use the chat template since this is an Instruct model.
                # Without it, the model repeats or hallucinates — it needs
                # the instruction format to understand it should CONTINUE text.
                messages = [
                    {
                        "role": "system",
                        "content": (
                            "You are a writing assistant. The user will give you text from a document. "
                            "Your job is to predict the next few words or sentence that naturally continues the text. "
                            "ONLY output the continuation — do NOT repeat any of the given text. "
                            "Keep it concise (1-2 short sentences max). "
                            "Match the tone, style, and language of the existing text."
                        )
                    },
                    {
                        "role": "user",
                        "content": f"Continue this text:\n\n{context}"
                    }
                ]

                prompt = tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
                input_ids = np.array(
                    [tokenizer.encode(prompt, add_special_tokens=False)],
                    dtype=np.int64
                )

                # Truncate input to last 512 tokens to keep inference fast
                if input_ids.shape[1] > 512:
                    input_ids = input_ids[:, -512:]

                # Generate in a background thread to keep the event loop free
                loop = asyncio.get_running_loop()
                tokens = await loop.run_in_executor(
                    None,
                    lambda: model_manager.generate(
                        input_ids,
                        max_tokens=max_tokens,
                        temperature=0.4,   # Slightly creative but still focused
                        top_k=40,
                        top_p=0.9,
                        stop_tokens=[
                            tokenizer.eos_token_id,
                            # Stop at paragraph boundary
                            *tokenizer.encode("\n", add_special_tokens=False),
                        ]
                    )
                )

                suggestion = tokenizer.decode(tokens, skip_special_tokens=True).strip()

                # Clean up: remove any accidental repetition of the context
                # (sometimes the model echoes the last few words)
                if suggestion and context:
                    # If suggestion starts with the end of context, trim the overlap
                    for overlap_len in range(min(len(suggestion), 30), 0, -1):
                        if context.endswith(suggestion[:overlap_len]):
                            suggestion = suggestion[overlap_len:].strip()
                            break

                await websocket.send_text(json.dumps({"suggestion": suggestion}))

            except Exception as e:
                logger.error(f"[ws/autocomplete] Prediction error: {e}")
                await websocket.send_text(json.dumps({"suggestion": ""}))

    except WebSocketDisconnect:
        logger.info("[ws/autocomplete] Client disconnected")
    except Exception as e:
        logger.error(f"[ws/autocomplete] Connection error: {e}")
        try:
            await websocket.close(code=1011, reason="Internal error")
        except Exception:
            pass


@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    logger.error(f"Unhandled exception: {exc}", exc_info=True)
    return JSONResponse(
        status_code=500,
        content={"error": {"message": "Internal server error", "type": "server_error"}}
    )


# ==============================================================================
# Main Entry Point
# ==============================================================================

if __name__ == "__main__":
    import uvicorn
    
    print(f"""
╔═══════════════════════════════════════════════════════════════╗
║           LFM2.5 FastAPI Backend (ONNX Runtime)               ║
╠═══════════════════════════════════════════════════════════════╣
║  Model:   LiquidAI/LFM2.5-1.2B-Instruct-ONNX                  ║
║  Variant: Q8 (~95% accuracy, fast CPU inference)              ║
║  Host:    {settings.host}:{settings.port}                                        ║
║  Docs:    http://{settings.host}:{settings.port}/docs                              ║
╚═══════════════════════════════════════════════════════════════╝
""")
    
    uvicorn.run(
        "app:app",
        host=settings.host,
        port=settings.port,
        log_level=settings.log_level,
        workers=1,
    )