Update main.py

main.py

@@ -30,9 +30,9 @@ DEFAULT_TEMP = float(os.getenv("DEFAULT_TEMPERATURE", "0.6"))
 DEFAULT_TOP_P = float(os.getenv("DEFAULT_TOP_P", "0.95"))
 DEFAULT_TOKENS = int(os.getenv("DEFAULT_MAX_TOKENS", "32000"))
 
-REQUEST_TIMEOUT = int(os.getenv("REQUEST_TIMEOUT", "120"))
+REQUEST_TIMEOUT = int(os.getenv("REQUEST_TIMEOUT", "180"))
 MAX_RETRIES = int(os.getenv("MAX_RETRIES", "3"))
-RETRY_BASE_DELAY = float(os.getenv("RETRY_BASE_DELAY", "1.5"))
+RETRY_BASE_DELAY = float(os.getenv("RETRY_BASE_DELAY", "1.7"))
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 log = logging.getLogger(__name__)
@@ -107,7 +107,7 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(
     title="FHR",
-    version="4.0.0",
+    version="4.1.0",
     lifespan=lifespan,
 )
 
@@ -123,7 +123,6 @@ app.add_middleware(
 # Utilities
 # ---------------------------------------------------------------------------
 
-
 def _content_str(m: Message) -> str:
     if isinstance(m.content, str):
         return m.content
@@ -160,24 +159,14 @@ def _build_prompt(messages: list[Message]) -> str:
 # ---------------------------------------------------------------------------
 
 def _extract_text(result: Any) -> str:
-    """
-    Robust extraction of assistant text from Gradio result.
-    Works with:
-    - tuple
-    - result.data
-    - dict["value"]
-    - direct list
-    """
-
     if hasattr(result, "data"):
         result = result.data
 
     if isinstance(result, tuple):
         result = list(result)
 
-    if isinstance(result, dict):
-        if "value" in result:
-            result = result["value"]
+    if isinstance(result, dict) and "value" in result:
+        result = result["value"]
 
     if isinstance(result, list) and result:
         last = result[-1]
@@ -197,7 +186,7 @@ def _extract_text(result: Any) -> str:
     if isinstance(result, str):
         return result.strip()
 
-    raise ValueError(f"Cannot extract text from result: {type(result)}")
+    raise ValueError("Unable to extract model response.")
 
 
 # ---------------------------------------------------------------------------
@@ -254,47 +243,56 @@ async def _call_falcon(prompt: str, req: ChatCompletionRequest) -> str:
 
 
 # ---------------------------------------------------------------------------
-# Real Streaming (if Space supports /stream)
+# SAFE STREAMING (HF Spaces stable)
 # ---------------------------------------------------------------------------
 
-async def _stream_real(prompt: str, req: ChatCompletionRequest) -> AsyncGenerator[str, None]:
-    client = await get_client()
-
-    settings = {
-        "model": req.model,
-        "temperature": req.temperature,
-        "max_new_tokens": req.max_tokens,
-        "top_p": req.top_p,
-    }
+async def _safe_stream(prompt: str, req: ChatCompletionRequest) -> AsyncGenerator[str, None]:
+    """
+    Stable streaming for HF Spaces:
+    1. Generate full response with retries
+    2. Stream chunks safely
+    """
 
-    await asyncio.to_thread(client.predict, api_name="/new_chat")
-
-    stream = await asyncio.to_thread(
-        client.submit,
-        input_value=prompt,
-        settings_form_value=settings,
-        api_name="/add_message",
-    )
+    text = await _call_falcon(prompt, req)
 
     cid = f"chatcmpl-{uuid.uuid4().hex}"
     created = int(time.time())
 
-    async for update in stream:
-        text = _extract_text(update)
-        chunk = {
+    try:
+        for i in range(0, len(text), 16):
+            chunk = {
+                "id": cid,
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": req.model,
+                "choices": [{
+                    "index": 0,
+                    "delta": {"content": text[i:i+16]},
+                    "finish_reason": None,
+                }],
+            }
+
+            yield f"data: {json.dumps(chunk)}\n\n"
+            await asyncio.sleep(0.02)
+
+        final = {
             "id": cid,
             "object": "chat.completion.chunk",
             "created": created,
             "model": req.model,
             "choices": [{
                 "index": 0,
-                "delta": {"content": text},
-                "finish_reason": None,
+                "delta": {},
+                "finish_reason": "stop",
             }],
         }
-        yield f"data: {json.dumps(chunk)}\n\n"
 
-    yield "data: [DONE]\n\n"
+        yield f"data: {json.dumps(final)}\n\n"
+        yield "data: [DONE]\n\n"
+
+    except Exception:
+        log.exception("Streaming crashed unexpectedly.")
+        yield "data: [DONE]\n\n"
 
 
 # ---------------------------------------------------------------------------
@@ -333,23 +331,20 @@ async def chat_completions(req: ChatCompletionRequest, _: None = Depends(verify_
 
     try:
         if req.stream:
-            try:
-                return StreamingResponse(
-                    _stream_real(prompt, req),
-                    media_type="text/event-stream",
-                )
-            except Exception:
-                log.warning("Real streaming failed, fallback to buffered.")
-                text = await _call_falcon(prompt, req)
-                return StreamingResponse(
-                    _fake_stream(text, req),
-                    media_type="text/event-stream",
-                )
+            return StreamingResponse(
+                _safe_stream(prompt, req),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "X-Accel-Buffering": "no",
+                },
+            )
 
         text = await _call_falcon(prompt, req)
         return JSONResponse(content=_make_response(text, req))
 
-    except Exception as e:
+    except Exception:
         log.exception("Final failure after retries.")
         raise HTTPException(
             status_code=502,