Update main.py

main.py

@@ -1,5 +1,363 @@
 from __future__ import annotations
 
+import os, json, time, uuid, asyncio, logging
+from typing import Any, AsyncGenerator
+from contextlib import asynccontextmanager
+
+from dotenv import load_dotenv
+from fastapi import FastAPI, HTTPException, Request, Depends
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse, JSONResponse
+from pydantic import BaseModel
+from gradio_client import Client
+
+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"))
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+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.")
+    return _client
+
+# ---------------------------------------------------------------------------
+# Pydantic schemas
+# ---------------------------------------------------------------------------
+
+class Message(BaseModel):
+    role: str
+    content: str | list[dict] = ""
+    name: str | None = None
+
+class ChatCompletionRequest(BaseModel):
+    model: str = MODEL_ID
+    messages: list[Message]
+    temperature: float = DEFAULT_TEMP
+    top_p: float = DEFAULT_TOP_P
+    max_tokens: int = DEFAULT_TOKENS
+    stream: bool = False
+    frequency_penalty: float = 0
+    presence_penalty: float = 0
+    stop: str | list[str] | None = None
+    seed: int | None = None
+    user: str | None = None
+
+# ---------------------------------------------------------------------------
+# Auth
+# ---------------------------------------------------------------------------
+
+async def verify_key(request: Request) -> None:
+    if not API_KEY:
+        return
+    auth = request.headers.get("Authorization", "")
+    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)
+# ---------------------------------------------------------------------------
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup
+    log.info("Starting up - connecting to Gradio client...")
+    await get_client()
+    log.info("Startup complete.")
+    yield
+    # Shutdown (if needed)
+    log.info("Shutting down.")
+
+# ---------------------------------------------------------------------------
+# App
+# ---------------------------------------------------------------------------
+
+app = FastAPI(
+    title="Falcon H1R API",
+    version="3.2.0",
+    lifespan=lifespan,
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ---------------------------------------------------------------------------
+# Business logic
+# ---------------------------------------------------------------------------
+
+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")
+
+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)
+
+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.
+    """
+    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 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
+
+async def _call_falcon(prompt: str, req: ChatCompletionRequest) -> str:
+    """
+    Call Falcon H1R via Gradio client.
+    """
+    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...")
+    result = await asyncio.to_thread(
+        client.predict,
+        input_value=prompt,
+        settings_form_value=settings,
+        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()),
+        "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,
+        },
+    }
+
+async def _stream_sse(text: str, req: ChatCompletionRequest) -> AsyncGenerator[str, None]:
+    """Simulate streaming by chunking the full response."""
+    cid = f"chatcmpl-{uuid.uuid4().hex}"
+    created = int(time.time())
+    
+    # Stream in small chunks
+    for i in range(0, len(text), 6):
+        chunk = {
+            "id": cid,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": req.model,
+            "choices": [{
+                "index": 0,
+                "delta": {"role": "assistant", "content": text[i:i+6]},
+                "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
+# ---------------------------------------------------------------------------
+
+@app.get("/")
+async def root():
+    return {
+        "service": "FOC API",
+        "version": "3.2.0",
+        "endpoints": {
+            "health": "/health",
+            "models": "/v1/models",
+            "chat": "/v1/chat/completions",
+        },
+    }
+
+@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",
+    }]}
+
+@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:
+        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))
+    
+    
+    
+zobacz jak było w pierwszej powyższej wersji kodu, wtedy działało, a nastepnie wprowadz odpowiednie poprawki w ponizszym kodzie i podaj caly kod po modyfikacji: from __future__ import annotations
+
 import os
 import json
 import time