Update main.py

main.py

@@ -1,6 +1,11 @@
 from __future__ import annotations
 
-import os, json, time, uuid, asyncio, logging
+import os
+import json
+import time
+import uuid
+import asyncio
+import logging
 from typing import Any, AsyncGenerator
 from contextlib import asynccontextmanager
 
@@ -16,12 +21,18 @@ load_dotenv()
 # ---------------------------------------------------------------------------
 # Config
 # ---------------------------------------------------------------------------
-API_KEY         = os.getenv("API_KEY", "")
-HF_SPACE_URL    = os.getenv("HF_SPACE_URL", "")
-MODEL_ID        = os.getenv("MODEL_ID", "")
-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", "1024"))
+
+API_KEY = os.getenv("API_KEY", "")
+HF_SPACE_URL = os.getenv("HF_SPACE_URL", "")
+MODEL_ID = os.getenv("MODEL_ID", "")
+
+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"))
+MAX_RETRIES = int(os.getenv("MAX_RETRIES", "3"))
+RETRY_BASE_DELAY = float(os.getenv("RETRY_BASE_DELAY", "1.5"))
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 log = logging.getLogger(__name__)
@@ -29,18 +40,21 @@ log = logging.getLogger(__name__)
 # ---------------------------------------------------------------------------
 # Gradio client (singleton)
 # ---------------------------------------------------------------------------
+
 _client: Client | None = None
 
+
 async def get_client() -> Client:
     global _client
     if _client is None:
         log.info("Connecting to %s", HF_SPACE_URL)
         _client = await asyncio.to_thread(Client, HF_SPACE_URL)
-        log.info("Connected.")
+        log.info("Connected to Space.")
     return _client
 
+
 # ---------------------------------------------------------------------------
-# Pydantic schemas
+# Schemas
 # ---------------------------------------------------------------------------
 
 class Message(BaseModel):
@@ -48,6 +62,7 @@ class Message(BaseModel):
     content: str | list[dict] = ""
     name: str | None = None
 
+
 class ChatCompletionRequest(BaseModel):
     model: str = MODEL_ID
     messages: list[Message]
@@ -61,6 +76,7 @@ class ChatCompletionRequest(BaseModel):
     seed: int | None = None
     user: str | None = None
 
+
 # ---------------------------------------------------------------------------
 # Auth
 # ---------------------------------------------------------------------------
@@ -72,27 +88,26 @@ async def verify_key(request: Request) -> None:
     if not auth.startswith("Bearer ") or auth[7:] != API_KEY:
         raise HTTPException(status_code=401, detail="Invalid or missing API key")
 
+
 # ---------------------------------------------------------------------------
-# Lifespan context manager (modern FastAPI pattern)
+# Lifespan
 # ---------------------------------------------------------------------------
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Startup
-    log.info("Starting up - connecting to Gradio client...")
+    log.info("Startup: connecting to Gradio client...")
     await get_client()
-    log.info("Startup complete.")
     yield
-    # Shutdown (if needed)
-    log.info("Shutting down.")
+    log.info("Shutdown.")
+
 
 # ---------------------------------------------------------------------------
 # App
 # ---------------------------------------------------------------------------
 
 app = FastAPI(
-    title="Falcon H1R API",
-    version="3.2.0",
+    title="FHR",
+    version="4.0.0",
     lifespan=lifespan,
 )
 
@@ -105,177 +120,167 @@ app.add_middleware(
 )
 
 # ---------------------------------------------------------------------------
-# Business logic
+# Utilities
 # ---------------------------------------------------------------------------
 
+
 def _content_str(m: Message) -> str:
     if isinstance(m.content, str):
         return m.content
-    return "".join(p.get("text", "") for p in m.content if p.get("type") == "text")
+    return "".join(
+        p.get("text", "") or p.get("content", "")
+        for p in m.content
+        if isinstance(p, dict)
+    )
+
 
 def _build_prompt(messages: list[Message]) -> str:
-    """Flatten messages into a single prompt string."""
     system, parts = [], []
     for m in messages:
-        c = _content_str(m)
-        if m.role == "system":    system.append(c)
-        elif m.role == "user":    parts.append(c)
-        elif m.role == "assistant": parts.append(f"[ASSISTANT]\n{c}")
-    prefix = "[SYSTEM]\n" + "\n".join(system) + "\n[/SYSTEM]\n" if system else ""
-    return prefix + "\n".join(parts)
+        c = _content_str(m).strip()
+        if not c:
+            continue
+
+        if m.role == "system":
+            system.append(c)
+        elif m.role == "assistant":
+            parts.append(f"[ASSISTANT]\n{c}")
+        else:
+            parts.append(c)
+
+    prefix = ""
+    if system:
+        prefix = "[SYSTEM]\n" + "\n".join(system) + "\n[/SYSTEM]\n\n"
+
+    return prefix + "\n\n".join(parts)
+
+
+# ---------------------------------------------------------------------------
+# Robust Extraction
+# ---------------------------------------------------------------------------
 
 def _extract_text(result: Any) -> str:
     """
-    Extract assistant reply from Gradio client.predict() result.
-    
-    gradio_client returns either:
-      - tuple: (output1, output2, ...) where one element is the chatbot data
-      - object with .data attribute containing a list
-    
-    We need to find the conversation list and extract the last message.
+    Robust extraction of assistant text from Gradio result.
+    Works with:
+    - tuple
+    - result.data
+    - dict["value"]
+    - direct list
     """
-    try:
-        # Handle both tuple and object with .data
-        if isinstance(result, tuple):
-            data = result
-        elif hasattr(result, 'data'):
-            data = result.data
-        else:
-            data = [result]
-        
-        log.info("Raw result type: %s, length: %s", type(data).__name__, len(data) if hasattr(data, '__len__') else 'N/A')
-        
-        # Search through all returned values for the conversation
-        conversation = None
-        for idx, item in enumerate(data):
-            log.debug("Item %d type: %s", idx, type(item).__name__)
-            
-            # Check if this item is a dict with 'value' key (chatbot component)
-            if isinstance(item, dict) and "value" in item:
-                val = item["value"]
-                if isinstance(val, list) and val:
-                    conversation = val
-                    log.info("Found conversation in dict at index %d, length: %d", idx, len(val))
-                    break
-            
-            # Check if item itself is a list (direct conversation)
-            elif isinstance(item, list) and item:
-                # Verify it looks like a conversation (list of message dicts/tuples)
-                first = item[0]
-                if isinstance(first, (dict, list, tuple)):
-                    conversation = item
-                    log.info("Found conversation as list at index %d, length: %d", idx, len(item))
-                    break
-        
-        if conversation is None:
-            raise ValueError(f"Cannot find conversation in result. Data structure: {json.dumps(str(data)[:500])}")
-        
-        # Extract last message
-        last = conversation[-1]
-        log.info("Last message type: %s, value: %s", type(last).__name__, str(last)[:200])
-        
-        # Handle different message formats
-        content = None
-        
-        # Format 1: dict with 'content' key (Gradio 4.x)
+
+    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, list) and result:
+        last = result[-1]
+
         if isinstance(last, dict):
-            content = last.get("content", "")
-        
-        # Format 2: tuple/list [user_msg, assistant_msg] (Gradio 3.x)
-        elif isinstance(last, (list, tuple)) and len(last) >= 2:
-            content = last[1] or ""
-        
-        # Format 3: plain string
-        elif isinstance(last, str):
-            content = last
-        
-        if content is None:
-            raise ValueError(f"Cannot extract content from last message: {last}")
-        
-        # If content is a list of content blocks, extract text
-        if isinstance(content, list):
-            text_parts = []
-            for block in content:
-                if isinstance(block, dict):
-                    if block.get("type") == "text":
-                        text_parts.append(block.get("content", block.get("text", "")).strip())
-                elif isinstance(block, str):
-                    text_parts.append(block.strip())
-            return "".join(text_parts)
-        
-        return str(content).strip()
-        
-    except Exception as e:
-        log.error("_extract_text failed: %s", e, exc_info=True)
-        log.error("Raw result dump: %s", str(result)[:1000])
-        raise ValueError(f"Failed to extract text: {e}") from e
+            if "content" in last:
+                return str(last["content"]).strip()
+            if "value" in last:
+                return str(last["value"]).strip()
 
-async def _call_falcon(prompt: str, req: ChatCompletionRequest) -> str:
-    """
-    Call Falcon H1R via Gradio client.
-    """
+        if isinstance(last, (list, tuple)) and len(last) >= 2:
+            return str(last[1]).strip()
+
+        if isinstance(last, str):
+            return last.strip()
+
+    if isinstance(result, str):
+        return result.strip()
+
+    raise ValueError(f"Cannot extract text from result: {type(result)}")
+
+
+# ---------------------------------------------------------------------------
+# Retry Wrapper
+# ---------------------------------------------------------------------------
+
+async def _call_with_retries(func, *args, **kwargs):
+    for attempt in range(1, MAX_RETRIES + 1):
+        try:
+            return await asyncio.wait_for(func(*args, **kwargs), timeout=REQUEST_TIMEOUT)
+        except Exception as e:
+            if attempt >= MAX_RETRIES:
+                log.error("All retries failed.")
+                raise
+
+            delay = RETRY_BASE_DELAY ** attempt
+            log.warning(
+                "Attempt %d failed: %s | retrying in %.2fs",
+                attempt,
+                str(e),
+                delay,
+            )
+            await asyncio.sleep(delay)
+
+
+# ---------------------------------------------------------------------------
+# Falcon Call
+# ---------------------------------------------------------------------------
+
+async def _call_falcon_once(prompt: str, req: ChatCompletionRequest) -> str:
     client = await get_client()
-    
+
     settings = {
         "model": req.model,
         "temperature": req.temperature,
         "max_new_tokens": req.max_tokens,
         "top_p": req.top_p,
     }
-    
-    # Step 1: Reset chat
-    log.info("Resetting chat session...")
-    await asyncio.to_thread(
-        client.predict,
-        api_name="/new_chat"
-    )
-    
-    # Step 2: Send message
-    log.info("Sending message to Falcon...")
+
+    await asyncio.to_thread(client.predict, api_name="/new_chat")
+
     result = await asyncio.to_thread(
         client.predict,
         input_value=prompt,
         settings_form_value=settings,
-        api_name="/add_message"
+        api_name="/add_message",
     )
-    
-    log.info("Received result, extracting text...")
+
     return _extract_text(result)
 
-def _make_response(text: str, req: ChatCompletionRequest) -> dict:
-    pt = sum(len(_content_str(m)) for m in req.messages) // 4
-    ct = len(text) // 4
-    return {
-        "id": f"chatcmpl-{uuid.uuid4().hex}",
-        "object": "chat.completion",
-        "created": int(time.time()),
+
+async def _call_falcon(prompt: str, req: ChatCompletionRequest) -> str:
+    return await _call_with_retries(_call_falcon_once, prompt, req)
+
+
+# ---------------------------------------------------------------------------
+# Real Streaming (if Space supports /stream)
+# ---------------------------------------------------------------------------
+
+async def _stream_real(prompt: str, req: ChatCompletionRequest) -> AsyncGenerator[str, None]:
+    client = await get_client()
+
+    settings = {
         "model": req.model,
-        "system_fingerprint": f"fp_{uuid.uuid4().hex[:8]}",
-        "choices": [{
-            "index": 0,
-            "message": {
-                "role": "assistant",
-                "content": text,
-                "tool_calls": None,
-                "function_call": None,
-            },
-            "finish_reason": "stop",
-            "logprobs": None,
-        }],
-        "usage": {
-            "prompt_tokens": pt,
-            "completion_tokens": ct,
-            "total_tokens": pt + ct,
-        },
+        "temperature": req.temperature,
+        "max_new_tokens": req.max_tokens,
+        "top_p": req.top_p,
     }
 
-async def _stream_sse(text: str, req: ChatCompletionRequest) -> AsyncGenerator[str, None]:
-    """Simulate streaming by chunking the full response."""
+    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",
+    )
+
     cid = f"chatcmpl-{uuid.uuid4().hex}"
     created = int(time.time())
-    
-    # Stream in small chunks
-    for i in range(0, len(text), 6):
+
+    async for update in stream:
+        text = _extract_text(update)
         chunk = {
             "id": cid,
             "object": "chat.completion.chunk",
@@ -283,73 +288,70 @@ async def _stream_sse(text: str, req: ChatCompletionRequest) -> AsyncGenerator[s
             "model": req.model,
             "choices": [{
                 "index": 0,
-                "delta": {"role": "assistant", "content": text[i:i+6]},
+                "delta": {"content": text},
                 "finish_reason": None,
             }],
         }
         yield f"data: {json.dumps(chunk)}\n\n"
-        await asyncio.sleep(0.01)
-    
-    # Final chunk
-    pt = sum(len(_content_str(m)) for m in req.messages) // 4
-    ct = len(text) // 4
-    final = {
-        "id": cid,
-        "object": "chat.completion.chunk",
-        "created": created,
-        "model": req.model,
-        "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
-        "usage": {"prompt_tokens": pt, "completion_tokens": ct, "total_tokens": pt + ct},
-    }
-    yield f"data: {json.dumps(final)}\n\n"
+
     yield "data: [DONE]\n\n"
 
+
 # ---------------------------------------------------------------------------
-# Routes
+# OpenAI Response Builder
 # ---------------------------------------------------------------------------
 
-@app.get("/")
-async def root():
+def _make_response(text: str, req: ChatCompletionRequest) -> dict:
+    pt = sum(len(_content_str(m)) for m in req.messages) // 4
+    ct = len(text) // 4
+
     return {
-        "service": "FOC API",
-        "version": "3.2.0",
-        "endpoints": {
-            "health": "/health",
-            "models": "/v1/models",
-            "chat": "/v1/chat/completions",
+        "id": f"chatcmpl-{uuid.uuid4().hex}",
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": req.model,
+        "choices": [{
+            "index": 0,
+            "message": {"role": "assistant", "content": text},
+            "finish_reason": "stop",
+        }],
+        "usage": {
+            "prompt_tokens": pt,
+            "completion_tokens": ct,
+            "total_tokens": pt + ct,
         },
     }
 
-@app.get("/health")
-async def health():
-    return {"status": "ok", "model": MODEL_ID, "space": HF_SPACE_URL}
 
-@app.get("/v1/models")
-async def list_models(_: None = Depends(verify_key)):
-    return {"object": "list", "data": [{
-        "id": MODEL_ID,
-        "object": "model",
-        "created": 1710000000,
-        "owned_by": "tiiuae",
-    }]}
+# ---------------------------------------------------------------------------
+# Routes
+# ---------------------------------------------------------------------------
 
 @app.post("/v1/chat/completions")
 async def chat_completions(req: ChatCompletionRequest, _: None = Depends(verify_key)):
     prompt = _build_prompt(req.messages)
-    log.info("Request | model=%s temp=%.2f tokens=%d stream=%s",
-             req.model, req.temperature, req.max_tokens, req.stream)
-    
+
     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",
+                )
+
         text = await _call_falcon(prompt, req)
-    except Exception as exc:
-        log.exception("Falcon call failed")
-        raise HTTPException(status_code=502, detail=f"Upstream error: {exc}") from exc
-    
-    if req.stream:
-        return StreamingResponse(
-            _stream_sse(text, req),
-            media_type="text/event-stream",
-            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
-        )
-    
-    return JSONResponse(content=_make_response(text, req))
+        return JSONResponse(content=_make_response(text, req))
+
+    except Exception as e:
+        log.exception("Final failure after retries.")
+        raise HTTPException(
+            status_code=502,
+            detail="Model temporarily unavailable. Please try again.",
+        )