Update app.py

app.py

@@ -14,7 +14,7 @@ app = FastAPI()
 GEMMA_API_KEY = os.getenv("GEMMA_API_KEY")
 APP_API_KEY = os.getenv("APP_API_KEY")
 
-GEMMA_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemma-3-27b-it:generateContent"
+GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta"
 
 
 # -------- Models --------
@@ -27,50 +27,166 @@ class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
     stream: Optional[bool] = False
-    plain: Optional[bool] = False   # 👈 return plain text
+    plain: Optional[bool] = False
 
 
 # -------- Helpers --------
 def extract_text(messages):
     text = ""
+
     for msg in messages:
-        if isinstance(msg.content, list):
-            for item in msg.content:
+        content = msg.content
+
+        if isinstance(content, list):
+            for item in content:
                 if item.get("type") == "text":
                     text += item.get("text", "") + "\n"
-        else:
-            text += msg.content + "\n"
+                else:
+                    # ignore images and any unknown part types safely
+                    continue
+        elif isinstance(content, str):
+            text += content + "\n"
+
     return text.strip()
 
 
+def build_payload(prompt: str):
+    return {
+        "contents": [
+            {
+                "parts": [{"text": prompt}]
+            }
+        ]
+    }
+
+
+def get_stream_url(model_name: str) -> str:
+    return f"{GEMINI_BASE_URL}/models/{model_name}:streamGenerateContent?alt=sse&key={GEMMA_API_KEY}"
+
+
+def get_generate_url(model_name: str) -> str:
+    return f"{GEMINI_BASE_URL}/models/{model_name}:generateContent?key={GEMMA_API_KEY}"
+
+
+def parse_gemini_text(chunk_json: dict) -> str:
+    """
+    Gemini streaming chunks usually contain:
+    candidates[0].content.parts[0].text
+    """
+    try:
+        candidates = chunk_json.get("candidates", [])
+        if not candidates:
+            return ""
+        content = candidates[0].get("content", {})
+        parts = content.get("parts", [])
+        if not parts:
+            return ""
+        return parts[0].get("text", "") or ""
+    except Exception:
+        return ""
+
+
 # -------- Endpoint --------
 @app.post("/v1/chat/completions")
 def chat_completions(
     request: ChatRequest,
     authorization: Optional[str] = Header(None)
 ):
-    # ---- Auth ----
     if not authorization:
         raise HTTPException(status_code=401, detail="Missing Authorization header")
 
-    token = authorization.replace("Bearer ", "")
+    token = authorization.replace("Bearer ", "").strip()
     if token != APP_API_KEY:
         raise HTTPException(status_code=403, detail="Invalid API key")
 
+    if not GEMMA_API_KEY:
+        raise HTTPException(status_code=500, detail="GEMMA_API_KEY is not set")
+
+    model_name = request.model or "gemma-3-27b-it"
     prompt = extract_text(request.messages)
+    payload = build_payload(prompt)
 
-    payload = {
-        "contents": [
-            {
-                "parts": [{"text": prompt}]
-            }
-        ]
-    }
+    # -------- STREAM MODE --------
+    if request.stream:
+        def generate():
+            try:
+                url = get_stream_url(model_name)
+
+                with requests.post(
+                    url,
+                    json=payload,
+                    stream=True,
+                    timeout=120,
+                    headers={"Content-Type": "application/json"}
+                ) as res:
+                    res.raise_for_status()
+
+                    sent_role = False
+
+                    for raw_line in res.iter_lines(decode_unicode=True):
+                        if not raw_line:
+                            continue
+
+                        line = raw_line.strip()
+
+                        if line.startswith("data:"):
+                            line = line[5:].strip()
+
+                        if not line:
+                            continue
+
+                        # Some SSE implementations may send end markers
+                        if line == "[DONE]":
+                            break
+
+                        try:
+                            chunk_json = json.loads(line)
+                        except json.JSONDecodeError:
+                            continue
+
+                        text = parse_gemini_text(chunk_json)
+                        if not text:
+                            continue
+
+                        delta = {"content": text}
+                        if not sent_role:
+                            delta["role"] = "assistant"
+                            sent_role = True
+
+                        openai_chunk = {
+                            "id": "chatcmpl-gemma",
+                            "object": "chat.completion.chunk",
+                            "choices": [
+                                {
+                                    "index": 0,
+                                    "delta": delta,
+                                    "finish_reason": None
+                                }
+                            ]
+                        }
+
+                        yield f"data: {json.dumps(openai_chunk, ensure_ascii=False)}\n\n"
+
+                    yield "data: [DONE]\n\n"
+
+            except Exception as e:
+                error_chunk = {
+                    "error": str(e)
+                }
+                yield f"data: {json.dumps(error_chunk, ensure_ascii=False)}\n\n"
+                yield "data: [DONE]\n\n"
 
-    url = f"{GEMMA_URL}?key={GEMMA_API_KEY}"
+        return StreamingResponse(generate(), media_type="text/event-stream")
 
+    # -------- NON-STREAM --------
     try:
-        res = requests.post(url, json=payload)
+        url = get_generate_url(model_name)
+        res = requests.post(
+            url,
+            json=payload,
+            timeout=120,
+            headers={"Content-Type": "application/json"}
+        )
         res.raise_for_status()
         data = res.json()
 
@@ -79,33 +195,11 @@ def chat_completions(
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-    # -------- PLAIN TEXT (FASTEST) --------
+    # -------- PLAIN TEXT --------
     if request.plain:
         return PlainTextResponse(output)
 
-    # -------- STREAM (SIMULATED) --------
-    if request.stream:
-        def generate():
-            words = output.split(" ")
-            for word in words:
-                chunk = {
-                    "id": "chatcmpl-gemma",
-                    "object": "chat.completion.chunk",
-                    "choices": [
-                        {
-                            "delta": {"content": word + " "},
-                            "index": 0,
-                            "finish_reason": None
-                        }
-                    ]
-                }
-                yield f"data: {json.dumps(chunk)}\n\n"
-
-            yield "data: [DONE]\n\n"
-
-        return StreamingResponse(generate(), media_type="text/event-stream")
-
-    # -------- DEFAULT (OPENAI JSON) --------
+    # -------- OPENAI JSON --------
     return JSONResponse({
         "id": "chatcmpl-gemma",
         "object": "chat.completion",