Update app.py

app.py

@@ -10,10 +10,10 @@ 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"
+HF_REPO = "paulsp94/Qwen3.5-2B-LiteRT-LM"
+HF_FILE = "model.litertlm"
 
-_SERVER_DIR   = os.path.dirname(os.path.abspath(__file__))
+_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.
@@ -21,19 +21,19 @@ _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:
-        pass
+    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_ID = "qwen3.5-2b"
 
 model_status = "loading"
-engine       = None
-_engine_ctx  = None
-engine_lock  = threading.Lock()
+engine = None
+_lm = None
+engine_lock = threading.Lock()
 
 app = Flask(__name__)
 CORS(app)
@@ -42,29 +42,45 @@ CORS(app)
 # ─── Model loading ─────────────────────────────────────────────────────────────
 
 def load_model():
-    global engine, model_status, _engine_ctx
+    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.set_min_log_severity(_lm.LogSeverity.SILENT)
+        import litert_lm as lm
+        _lm = lm
     except ImportError:
-        print("[INFO] litert_lm not installed — no model loaded", flush=True)
+        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:
-        _engine_ctx = _lm.Engine(
-            MODEL_PATH,
-            backend=_lm.interfaces.CPU(),
-            vision_backend=_lm.interfaces.CPU(),
-        )
-        engine      = _engine_ctx.__enter__()
+        # 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:
@@ -74,8 +90,8 @@ def load_model():
 
 # ─── OpenAI Request Parsing ────────────────────────────────────────────────────
 
-def parse_openai_messages(messages: list) -> tuple[str, bytes | None]:
-    """Parses OpenAI formatted messages into a flat text prompt and an optional image."""
+def parse_openai_messages(messages: list):
+    """Parses OpenAI formatted messages into a flat text prompt and optional image."""
     prompt_text = ""
     image_bytes = None
 
@@ -88,16 +104,17 @@ def parse_openai_messages(messages: list) -> tuple[str, bytes | None]:
         elif isinstance(content, list):
             prompt_text += f"{role}:\n"
             for part in content:
-                if part.get("type") == "text":
+                ptype = part.get("type")
+                if ptype == "text":
                     prompt_text += part.get("text", "") + "\n"
-                elif part.get("type") == "image_url":
+                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}")
+                            print(f"[WARN] Failed to decode base64 image: {e}", flush=True)
 
     prompt_text += "assistant: "
     return prompt_text.strip(), image_bytes
@@ -105,41 +122,65 @@ def parse_openai_messages(messages: list) -> tuple[str, bytes | None]:
 
 # ─── Inference Engine ──────────────────────────────────────────────────────────
 
-def _run_real_model_generator(ask: str, image_bytes: bytes | None):
+def _run_real_model_generator(ask: str, image_bytes):
     """Yields text chunks as they are generated by the model."""
-    import litert_lm
-    # engine_lock ensures only 1 request processes at a time to prevent RAM crashes
+    # Qwen 3.5 2B is text-only; image_bytes are ignored intentionally.
     with engine_lock:
-        with engine.create_conversation() as conv:
-            if image_bytes:
-                # NOTE: Qwen 3.5 2B is text-only. If a caller sends an image
-                # against this model, we ignore the image bytes rather than
-                # crash the engine, since the checkpoint has no vision tower.
-                msg = ask
-            else:
-                msg = ask
+        conv = None
+        try:
+            conv = engine.create_conversation()
 
-            for chunk in conv.send_message_async(msg):
-                for part in chunk.get("content", []):
-                    if part.get("type") == "text":
-                        text = part.get("text", "")
-                        if text:
-                            yield text
+            # 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] Received prompt. Vision included: {has_image}. Connect litert_lm for real output."
+    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.05)
+        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():
-    """OpenAI models endpoint."""
     return jsonify({
         "object": "list",
         "data": [{
@@ -150,23 +191,21 @@ def list_models():
         }]
     })
 
+
 @app.route("/v1/chat/completions", methods=["POST"])
 def chat_completions():
-    """OpenAI compatible chat completions endpoint."""
     data = request.get_json(silent=True) or {}
     messages = data.get("messages", [])
-    stream = data.get("stream", False)
-    
+    stream = bool(data.get("stream", False))
+
     if not messages:
-        return jsonify({"error": {"message": "Missing 'messages' array", "type": "invalid_request_error"}}), 400
+        return jsonify({
+            "error": {"message": "Missing 'messages' array", "type": "invalid_request_error"}
+        }), 400
 
     ask, image_bytes = parse_openai_messages(messages)
-    
-    # Determine which generator to use
-    if engine is None or model_status != "ready":
-        generator = _run_mock_generator(ask, bool(image_bytes))
-    else:
-        generator = _run_real_model_generator(ask, image_bytes)
+
+    use_mock = engine is None or model_status != "ready"
 
     req_model = data.get("model", MODEL_ID)
     cmpl_id = f"chatcmpl-{uuid.uuid4().hex}"
@@ -174,38 +213,46 @@ def chat_completions():
 
     if stream:
         def stream_response():
-            # 1. Initial chunk indicating role
             init_chunk = {
-                "id": cmpl_id, "object": "chat.completion.chunk", "created": created_time, "model": req_model,
+                "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"
 
-            # 2. Stream tokens
             try:
-                for text_chunk in generator:
+                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,
+                        "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 = {"error": str(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"
 
-            # 3. Final chunk indicating stop
             final_chunk = {
-                "id": cmpl_id, "object": "chat.completion.chunk", "created": created_time, "model": req_model,
+                "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:
-            full_text = "".join(list(generator))
+            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",
@@ -213,27 +260,26 @@ def chat_completions():
                 "model": req_model,
                 "choices": [{
                     "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": full_text
-                    },
+                    "message": {"role": "assistant", "content": full_text},
                     "finish_reason": "stop"
                 }],
                 "usage": {
-                    "prompt_tokens": 0,      # litert_lm token counting not implemented
-                    "completion_tokens": 0,  
+                    "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
+            return jsonify({
+                "error": {"message": f"Model error: {e}", "type": "server_error"}
+            }), 500
 
 
 # ─── Entry ─────────────────────────────────────────────────────────────────────
 
 if __name__ == "__main__":
-    port = int(os.environ.get("PORT", 5173))
+    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)
+    app.run(host="0.0.0.0", port=port, debug=False, threaded=True)