Update server.py

server.py

@@ -5,10 +5,11 @@ import ctypes
 import threading
 import json
 import time
+import uuid
 from flask import Flask, request, jsonify, Response
 from flask_cors import CORS
 
-# --- KEY CHANGE: Balik sa 2B Model ---
+# --- Model Configuration ---
 HF_REPO    = "litert-community/gemma-4-E2B-it-litert-lm"
 HF_FILE    = "gemma-4-E2B-it.litertlm"
 
@@ -27,6 +28,7 @@ if os.path.exists(_vk_stub):
 os.environ.setdefault("GLOG_minloglevel", "3")
 
 MODEL_PATH = os.environ.get("GEMMA_MODEL_PATH", _DEFAULT_PATH).strip()
+MODEL_ID   = "gemma-4-e2b"
 
 model_status = "loading"
 engine       = None
@@ -70,27 +72,38 @@ def load_model():
         model_status = "error"
 
 
-# ─── Fallback functions (No Model) ──────────────────────────────────────────────
+# ─── OpenAI Request Parsing ────────────────────────────────────────────────────
 
-def _analyze_image(image_bytes: bytes) -> dict:
-    from PIL import Image
-    img = Image.open(io.BytesIO(image_bytes))
-    w, h   = img.size
-    fmt    = (img.format or "image").lower()
-    return dict(w=w, h=h, fmt=fmt)
+def parse_openai_messages(messages: list) -> tuple[str, bytes | None]:
+    """Parses OpenAI formatted messages into a flat text prompt and an optional image."""
+    prompt_text = ""
+    image_bytes = None
 
-def _describe_image(ask: str, image_bytes: bytes) -> str:
-    try:
-        i = _analyze_image(image_bytes)
-        return f"[MOCK] This is a {i['fmt']} image ({i['w']}×{i['h']} px). Connect the real model for full vision analysis."
-    except Exception as e:
-        return f"Could not analyze image: {e}"
+    for msg in messages:
+        role = msg.get("role", "user")
+        content = msg.get("content", "")
 
-def _text_reply(ask: str) -> str:
-    return f"[MOCK] Hello! I received: '{ask}'. Connect the Gemma model to see real answers."
+        if isinstance(content, str):
+            prompt_text += f"{role}: {content}\n"
+        elif isinstance(content, list):
+            prompt_text += f"{role}:\n"
+            for part in content:
+                if part.get("type") == "text":
+                    prompt_text += part.get("text", "") + "\n"
+                elif part.get("type") == "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}")
 
+    prompt_text += "assistant: "
+    return prompt_text.strip(), image_bytes
 
-# ─── Real model inference (litert_lm Generator) ────────────────────────────────
+
+# ─── Inference Engine ──────────────────────────────────────────────────────────
 
 def _run_real_model_generator(ask: str, image_bytes: bytes | None):
     """Yields text chunks as they are generated by the model."""
@@ -114,122 +127,107 @@ def _run_real_model_generator(ask: str, image_bytes: bytes | None):
                             yield text
 
 
-# ─── Request extraction ────────────────────────────────────────────────────────
-
-def extract_request() -> tuple[str, bytes | None, bool]:
-    # Check query params first (e.g. ?stream=true)
-    stream = str(request.args.get("stream", "")).lower() in ["true", "1", "yes"]
-
-    if request.method == "GET":
-        return request.args.get("ask", "").strip(), None, stream
-
-    ct = request.content_type or ""
-    
-    # Handle Form Data (File Uploads)
-    if "multipart/form-data" in ct:
-        ask = request.form.get("ask", "").strip()
-        f   = request.files.get("image")
-        if "stream" in request.form:
-            stream = str(request.form.get("stream")).lower() in ["true", "1", "yes"]
-        return ask, (f.read() if f else None), stream
-
-    # Handle JSON
-    data        = request.get_json(silent=True) or {}
-    ask         = data.get("ask", "").strip()
-    image_bytes = None
-    raw         = data.get("image", "")
-    
-    if "stream" in data:
-        stream = str(data.get("stream")).lower() in ["true", "1", "yes"]
-
-    if raw:
-        try:
-            if "," in raw:
-                raw = raw.split(",", 1)[1]
-            image_bytes = base64.b64decode(raw)
-        except Exception:
-            pass
-            
-    return ask, image_bytes, stream
+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."
+    for word in msg.split():
+        yield word + " "
+        time.sleep(0.05)
 
 
 # ─── Routes ────────────────────────────────────────────────────────────────────
 
-@app.route("/")
-def index():
+@app.route("/v1/models", methods=["GET"])
+def list_models():
+    """OpenAI models endpoint."""
     return jsonify({
-        "service": "Gemma 4 API (2B Model)",
-        "model_status": model_status,
-        "guide": {
-            "1. Text Only Chat": {
-                "GET_Example": "/gemma?ask=Hello&stream=false",
-                "POST_JSON": {"ask": "What is AI?", "stream": True}
-            },
-            "2. Image with Text": {
-                "POST_JSON": {
-                    "ask": "Describe this image",
-                    "image": "<base64_string_here>",
-                    "stream": False
-                },
-                "POST_FormData": {
-                    "ask": "What color is this?",
-                    "image": "@file.jpg (File Upload)",
-                    "stream": "true"
-                }
-            },
-            "Streaming Info": "Set 'stream': true to receive Server-Sent Events (SSE) by token. Set false to wait for 1 full JSON response."
-        }
+        "object": "list",
+        "data": [{
+            "id": MODEL_ID,
+            "object": "model",
+            "created": int(time.time()),
+            "owned_by": "litert-community"
+        }]
     })
 
-
-@app.route("/health")
-def health():
-    return jsonify({"status": "ok", "model_status": model_status})
-
-
-@app.route("/gemma", methods=["GET", "POST"])
-def gemma():
-    ask, image_bytes, stream = extract_request()
+@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)
     
-    if not ask:
-        return jsonify({"error": "Missing 'ask' parameter"}), 400
+    if not messages:
+        return jsonify({"error": {"message": "Missing 'messages' array", "type": "invalid_request_error"}}), 400
 
-    # Fallback if model isn't ready
+    ask, image_bytes = parse_openai_messages(messages)
+    
+    # Determine which generator to use
     if engine is None or model_status != "ready":
-        fallback_msg = _describe_image(ask, image_bytes) if image_bytes else _text_reply(ask)
-        
-        if stream:
-            def mock_stream():
-                for word in fallback_msg.split():
-                    yield f"data: {json.dumps({'text': word + ' '})}\n\n"
-                    time.sleep(0.05)
-                yield "data: [DONE]\n\n"
-            return Response(mock_stream(), mimetype="text/event-stream")
-        else:
-            return jsonify({"ask": ask, "response": fallback_msg, "has_image": bool(image_bytes)})
-
-    # Real Model Logic
+        generator = _run_mock_generator(ask, bool(image_bytes))
+    else:
+        generator = _run_real_model_generator(ask, image_bytes)
+
+    req_model = data.get("model", MODEL_ID)
+    cmpl_id = f"chatcmpl-{uuid.uuid4().hex}"
+    created_time = int(time.time())
+
     if stream:
-        def generate_stream():
+        def stream_response():
+            # 1. Initial chunk indicating role
+            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"
+
+            # 2. Stream tokens
             try:
-                for text_chunk in _run_real_model_generator(ask, image_bytes):
-                    yield f"data: {json.dumps({'text': text_chunk})}\n\n"
-                yield "data: [DONE]\n\n"
+                for text_chunk in generator:
+                    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:
-                yield f"data: {json.dumps({'error': str(e)})}\n\n"
-                
-        return Response(generate_stream(), mimetype="text/event-stream")
+                err_chunk = {"error": str(e)}
+                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,
+                "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(_run_real_model_generator(ask, image_bytes)))
-            return jsonify({
-                "ask": ask,
-                "response": full_text,
-                "has_image": image_bytes is not None,
-                "model_status": model_status
-            })
+            full_text = "".join(list(generator))
+            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,      # litert_lm token counting not implemented
+                    "completion_tokens": 0,  
+                    "total_tokens": 0
+                }
+            }
+            return jsonify(response)
         except Exception as e:
-            return jsonify({"error": f"Model error: {e}"}), 500
+            return jsonify({"error": {"message": f"Model error: {e}", "type": "server_error"}}), 500
 
 
 # ─── Entry ─────────────────────────────────────────────────────────────────────
@@ -237,5 +235,5 @@ def gemma():
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 5173))
     threading.Thread(target=load_model, daemon=True).start()
-    print(f"[INFO] Gemma API on :{port}", flush=True)
+    print(f"[INFO] Gemma OpenAI-Compatible API listening on :{port}", flush=True)
     app.run(host="0.0.0.0", port=port, debug=False)