Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import json
 import requests
 from fastapi import FastAPI, Header, HTTPException
-from fastapi.responses import StreamingResponse, PlainTextResponse, JSONResponse
+from fastapi.responses import StreamingResponse, PlainTextResponse, Response
 from pydantic import BaseModel
 from typing import List, Optional, Union
 from dotenv import load_dotenv
@@ -13,11 +13,9 @@ app = FastAPI()
 
 GEMMA_API_KEY = os.getenv("GEMMA_API_KEY")
 APP_API_KEY = os.getenv("APP_API_KEY")
-
 GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta"
 
 
-# -------- Models --------
 class Message(BaseModel):
     role: str
     content: Union[str, List[dict]]
@@ -30,7 +28,17 @@ class ChatRequest(BaseModel):
     plain: Optional[bool] = False
 
 
-# -------- Helpers --------
+class UTF8JSONResponse(Response):
+    media_type = "application/json; charset=utf-8"
+
+    def render(self, content) -> bytes:
+        return json.dumps(
+            content,
+            ensure_ascii=False,
+            separators=(",", ":"),
+        ).encode("utf-8")
+
+
 def extract_text(messages):
     text = ""
 
@@ -42,7 +50,6 @@ def extract_text(messages):
                 if item.get("type") == "text":
                     text += item.get("text", "") + "\n"
                 else:
-                    # ignore images and any unknown part types safely
                     continue
         elif isinstance(content, str):
             text += content + "\n"
@@ -50,47 +57,49 @@ def extract_text(messages):
     return text.strip()
 
 
-def build_payload(prompt: str):
-    return {
-        "contents": [
-            {
-                "parts": [{"text": prompt}]
-            }
-        ]
-    }
+def fix_mojibake(text: str) -> str:
+    """
+    Repairs common UTF-8-as-latin1 mojibake such as:
+    It’s -> It’s
+    """
+    if not isinstance(text, str):
+        return text
 
+    suspicious = ("’", "“", "”", "‘", "–", "—", "Ã")
+    if any(s in text for s in suspicious):
+        try:
+            return text.encode("latin1").decode("utf-8")
+        except UnicodeError:
+            return text
 
-def get_stream_url(model_name: str) -> str:
-    return f"{GEMINI_BASE_URL}/models/{model_name}:streamGenerateContent?alt=sse&key={GEMMA_API_KEY}"
+    return text
 
 
-def get_generate_url(model_name: str) -> str:
+def gemini_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:
+def gemini_stream_url(model_name: str) -> str:
+    return f"{GEMINI_BASE_URL}/models/{model_name}:streamGenerateContent?alt=sse&key={GEMMA_API_KEY}"
+
+
+def extract_gemini_text(payload: dict) -> str:
+    candidates = payload.get("candidates") or []
+    if not candidates:
+        return ""
+
+    content = candidates[0].get("content") or {}
+    parts = content.get("parts") or []
+    if not parts:
         return ""
 
+    return parts[0].get("text", "") or ""
+
 
-# -------- Endpoint --------
 @app.post("/v1/chat/completions")
 def chat_completions(
     request: ChatRequest,
-    authorization: Optional[str] = Header(None)
+    authorization: Optional[str] = Header(None),
 ):
     if not authorization:
         raise HTTPException(status_code=401, detail="Missing Authorization header")
@@ -104,22 +113,30 @@ def chat_completions(
 
     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)
+                url = gemini_stream_url(model_name)
 
                 with requests.post(
                     url,
                     json=payload,
                     stream=True,
                     timeout=120,
-                    headers={"Content-Type": "application/json"}
+                    headers={"Content-Type": "application/json"},
                 ) as res:
                     res.raise_for_status()
+                    res.encoding = "utf-8"
 
                     sent_role = False
 
@@ -132,22 +149,20 @@ def chat_completions(
                         if line.startswith("data:"):
                             line = line[5:].strip()
 
-                        if not line:
+                        if not line or line == "[DONE]":
                             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)
+                        text = extract_gemini_text(chunk_json)
                         if not text:
                             continue
 
+                        text = fix_mojibake(text)
+
                         delta = {"content": text}
                         if not sent_role:
                             delta["role"] = "assistant"
@@ -160,9 +175,9 @@ def chat_completions(
                                 {
                                     "index": 0,
                                     "delta": delta,
-                                    "finish_reason": None
+                                    "finish_reason": None,
                                 }
-                            ]
+                            ],
                         }
 
                         yield f"data: {json.dumps(openai_chunk, ensure_ascii=False)}\n\n"
@@ -170,37 +185,44 @@ def chat_completions(
                     yield "data: [DONE]\n\n"
 
             except Exception as e:
-                error_chunk = {
-                    "error": str(e)
-                }
+                error_chunk = {"error": str(e)}
                 yield f"data: {json.dumps(error_chunk, ensure_ascii=False)}\n\n"
                 yield "data: [DONE]\n\n"
 
-        return StreamingResponse(generate(), media_type="text/event-stream")
+        return StreamingResponse(
+            generate(),
+            media_type="text/event-stream; charset=utf-8",
+            headers={"Cache-Control": "no-cache"},
+        )
 
     # -------- NON-STREAM --------
     try:
-        url = get_generate_url(model_name)
+        url = gemini_generate_url(model_name)
         res = requests.post(
             url,
             json=payload,
             timeout=120,
-            headers={"Content-Type": "application/json"}
+            headers={"Content-Type": "application/json"},
         )
         res.raise_for_status()
-        data = res.json()
+        res.encoding = "utf-8"
 
-        output = data["candidates"][0]["content"]["parts"][0]["text"]
+        data = res.json()
+        output = extract_gemini_text(data)
+        output = fix_mojibake(output)
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
     # -------- PLAIN TEXT --------
     if request.plain:
-        return PlainTextResponse(output)
+        return PlainTextResponse(
+            output,
+            media_type="text/plain; charset=utf-8",
+        )
 
     # -------- OPENAI JSON --------
-    return JSONResponse({
+    return UTF8JSONResponse({
         "id": "chatcmpl-gemma",
         "object": "chat.completion",
         "choices": [
@@ -208,9 +230,9 @@ def chat_completions(
                 "index": 0,
                 "message": {
                     "role": "assistant",
-                    "content": output
+                    "content": output,
                 },
-                "finish_reason": "stop"
+                "finish_reason": "stop",
             }
-        ]
+        ],
     })