app.py

import os
import io
import base64
import ctypes
import threading
import json
import time
import uuid
from flask import Flask, request, jsonify, Response
from flask_cors import CORS

# --- Model Configuration ---
HF_REPO = "paulsp94/Qwen3.5-2B-LiteRT-LM"
HF_FILE = "model.litertlm"

_SERVER_DIR = os.path.dirname(os.path.abspath(__file__))
_DEFAULT_PATH = os.path.join(_SERVER_DIR, "models", "qwen", HF_FILE)

# litert_lm links against libvulkan.so.1 even on CPU-only runs.
_vk_stub = os.path.join(_SERVER_DIR, "libvulkan.so.1")
if os.path.exists(_vk_stub):
    try:
        ctypes.CDLL(_vk_stub, mode=ctypes.RTLD_GLOBAL)
    except OSError as e:
        print(f"[WARN] Could not preload vulkan stub: {e}", flush=True)

# Suppress verbose C++ logs from litert_lm
os.environ.setdefault("GLOG_minloglevel", "3")

MODEL_PATH = os.environ.get("GEMMA_MODEL_PATH", _DEFAULT_PATH).strip()
MODEL_ID = "qwen3.5-2b"

model_status = "loading"
engine = None
_lm = None
engine_lock = threading.Lock()

app = Flask(__name__)
CORS(app)


# ─── Model loading ─────────────────────────────────────────────────────────────

def load_model():
    global engine, model_status, _lm
    if not MODEL_PATH:
        print("[INFO] GEMMA_MODEL_PATH not set — no model loaded", flush=True)
        model_status = "no_model_path"
        return

    try:
        import litert_lm as lm
        _lm = lm
    except ImportError:
        print("[INFO] litert_lm not installed — running in mock mode", flush=True)
        model_status = "no_litert_lm"
        return

    # Try to silence logs if the API exists
    try:
        _lm.set_min_log_severity(_lm.LogSeverity.SILENT)
    except Exception:
        pass

    if not os.path.exists(MODEL_PATH):
        print(f"[WARN] Model file not found: {MODEL_PATH}", flush=True)
        model_status = "model_file_missing"
        return

    try:
        # The litert_lm Engine API. Build args defensively depending
        # on what the installed version exposes.
        try:
            cpu_backend = _lm.interfaces.CPU()
            engine = _lm.Engine(
                MODEL_PATH,
                backend=cpu_backend,
                vision_backend=cpu_backend,
            )
        except (AttributeError, TypeError):
            # Fallback: simpler constructor signature
            engine = _lm.Engine(MODEL_PATH)

        model_status = "ready"
        print(f"[INFO] Model ready → {MODEL_PATH}", flush=True)
    except Exception as e:
        print(f"[ERROR] Failed to load model: {e}", flush=True)
        model_status = "error"


# ─── OpenAI Request Parsing ────────────────────────────────────────────────────

def parse_openai_messages(messages: list):
    """Parses OpenAI formatted messages into a flat text prompt and optional image."""
    prompt_text = ""
    image_bytes = None

    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")

        if isinstance(content, str):
            prompt_text += f"{role}: {content}\n"
        elif isinstance(content, list):
            prompt_text += f"{role}:\n"
            for part in content:
                ptype = part.get("type")
                if ptype == "text":
                    prompt_text += part.get("text", "") + "\n"
                elif ptype == "image_url":
                    url = part.get("image_url", {}).get("url", "")
                    if url.startswith("data:image"):
                        try:
                            b64_data = url.split(",", 1)[1]
                            image_bytes = base64.b64decode(b64_data)
                        except Exception as e:
                            print(f"[WARN] Failed to decode base64 image: {e}", flush=True)

    prompt_text += "assistant: "
    return prompt_text.strip(), image_bytes


# ─── Inference Engine ──────────────────────────────────────────────────────────

def _run_real_model_generator(ask: str, image_bytes):
    """Yields text chunks as they are generated by the model."""
    # Qwen 3.5 2B is text-only; image_bytes are ignored intentionally.
    with engine_lock:
        conv = None
        try:
            conv = engine.create_conversation()

            # Support both context-manager and plain object styles
            if hasattr(conv, "__enter__"):
                conv_obj = conv.__enter__()
            else:
                conv_obj = conv

            stream = conv_obj.send_message_async(ask)

            for chunk in stream:
                # Chunk may be a plain string or a structured dict
                if isinstance(chunk, str):
                    if chunk:
                        yield chunk
                elif isinstance(chunk, dict):
                    for part in chunk.get("content", []):
                        if part.get("type") == "text":
                            text = part.get("text", "")
                            if text:
                                yield text
                else:
                    # Try common attribute names
                    text = getattr(chunk, "text", None)
                    if text:
                        yield text
        finally:
            if conv is not None and hasattr(conv, "__exit__"):
                try:
                    conv.__exit__(None, None, None)
                except Exception:
                    pass


def _run_mock_generator(ask: str, has_image: bool):
    """Fallback generator when the model is missing/loading."""
    msg = (f"[MOCK] Model status: {model_status}. "
           f"Vision included: {has_image}. "
           f"Connect litert_lm + model file for real output.")
    for word in msg.split():
        yield word + " "
        time.sleep(0.02)


# ─── Routes ────────────────────────────────────────────────────────────────────

@app.route("/health", methods=["GET"])
def health():
    return jsonify({"status": model_status}), 200


@app.route("/v1/models", methods=["GET"])
def list_models():
    return jsonify({
        "object": "list",
        "data": [{
            "id": MODEL_ID,
            "object": "model",
            "created": int(time.time()),
            "owned_by": "paulsp94"
        }]
    })


@app.route("/v1/chat/completions", methods=["POST"])
def chat_completions():
    data = request.get_json(silent=True) or {}
    messages = data.get("messages", [])
    stream = bool(data.get("stream", False))

    if not messages:
        return jsonify({
            "error": {"message": "Missing 'messages' array", "type": "invalid_request_error"}
        }), 400

    ask, image_bytes = parse_openai_messages(messages)

    use_mock = engine is None or model_status != "ready"

    req_model = data.get("model", MODEL_ID)
    cmpl_id = f"chatcmpl-{uuid.uuid4().hex}"
    created_time = int(time.time())

    if stream:
        def stream_response():
            init_chunk = {
                "id": cmpl_id, "object": "chat.completion.chunk",
                "created": created_time, "model": req_model,
                "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]
            }
            yield f"data: {json.dumps(init_chunk)}\n\n"

            try:
                gen = (_run_mock_generator(ask, bool(image_bytes)) if use_mock
                       else _run_real_model_generator(ask, image_bytes))
                for text_chunk in gen:
                    chunk = {
                        "id": cmpl_id, "object": "chat.completion.chunk",
                        "created": created_time, "model": req_model,
                        "choices": [{"index": 0, "delta": {"content": text_chunk}, "finish_reason": None}]
                    }
                    yield f"data: {json.dumps(chunk)}\n\n"
            except Exception as e:
                err_chunk = {
                    "id": cmpl_id, "object": "chat.completion.chunk",
                    "created": created_time, "model": req_model,
                    "choices": [{"index": 0, "delta": {"content": f"[ERROR] {e}"}, "finish_reason": "stop"}]
                }
                yield f"data: {json.dumps(err_chunk)}\n\n"

            final_chunk = {
                "id": cmpl_id, "object": "chat.completion.chunk",
                "created": created_time, "model": req_model,
                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
            }
            yield f"data: {json.dumps(final_chunk)}\n\n"
            yield "data: [DONE]\n\n"

        return Response(stream_response(), mimetype="text/event-stream")

    else:
        try:
            gen = (_run_mock_generator(ask, bool(image_bytes)) if use_mock
                   else _run_real_model_generator(ask, image_bytes))
            full_text = "".join(gen)
            response = {
                "id": cmpl_id,
                "object": "chat.completion",
                "created": created_time,
                "model": req_model,
                "choices": [{
                    "index": 0,
                    "message": {"role": "assistant", "content": full_text},
                    "finish_reason": "stop"
                }],
                "usage": {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0
                }
            }
            return jsonify(response)
        except Exception as e:
            return jsonify({
                "error": {"message": f"Model error: {e}", "type": "server_error"}
            }), 500


# ─── Entry ─────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    threading.Thread(target=load_model, daemon=True).start()
    print(f"[INFO] Qwen 3.5 2B OpenAI-Compatible API listening on :{port}", flush=True)
    app.run(host="0.0.0.0", port=port, debug=False, threaded=True)