Upload app.py with huggingface_hub

app.py

@@ -1,218 +1,319 @@
-import os
-import json
+"""
+Haiku API - OpenAI-compatible proxy for chatgpt.org/claude/chat
+Deploy to Hugging Face Spaces (Docker SDK)
+
+The chatgpt.org backend proxies through OpenRouter and returns
+OpenAI-compatible SSE chunks. We relay them directly.
+"""
+
 import asyncio
+import json
 import re
-import urllib.parse
-from typing import AsyncGenerator, Optional, List, Dict, Any
-from contextlib import asynccontextmanager
+import time
+from typing import Optional
+from urllib.parse import unquote
 
 import httpx
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import StreamingResponse, JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from sse_starlette.sse import EventSourceResponse
-from dotenv import load_dotenv
 
-load_dotenv()
+app = FastAPI(title="Haiku API", version="1.1.0")
 
-API_ENDPOINT = os.getenv("API_ENDPOINT", "")
-PROXY_URL = os.getenv("PROXY_URL", "")
+# ── Session State ────────────────────────────────────────────────
+class SessionState:
+    """Manages cookies and CSRF tokens for chatgpt.org."""
+    def __init__(self):
+        self.xsrf_token: Optional[str] = None
+        self.csrf_token: Optional[str] = None
+        self.cookies: Optional[httpx.Cookies] = None
+        self.last_refresh: float = 0
+        self.refresh_interval: float = 600  # 10 min
+        self._lock = asyncio.Lock()
 
-DEFAULT_MODELS = [
-    {"id": "openai/gpt-4o-mini", "slug": "chatgpt", "name": "ChatGPT"},
-    {"id": "anthropic/claude-haiku-4-5", "slug": "claude", "name": "Claude"},
-    {"id": "google/gemini-2.0-flash-001", "slug": "gemini", "name": "Gemini"},
-    {"id": "x-ai/grok-3-mini-beta", "slug": "grok", "name": "Grok"},
-    {"id": "deepseek/deepseek-chat-v3-0324", "slug": "deepseek", "name": "DeepSeek"},
-    {"id": "qwen/qwen-2.5-72b-instruct", "slug": "qwen", "name": "Qwen"},
-    {"id": "moonshotai/moonlight-16k", "slug": "kimi", "name": "Kimi"},
-    {"id": "perplexity/sonar", "slug": "perplexity", "name": "Perplexity"},
-]
+    async def refresh(self, client: httpx.AsyncClient):
+        """Visit chatgpt.org to obtain fresh session cookies + CSRF token."""
+        async with self._lock:
+            now = time.time()
+            if self.cookies and (now - self.last_refresh) < self.refresh_interval:
+                return
+
+            try:
+                resp = await client.get(
+                    "https://chatgpt.org/claude/chat",
+                    follow_redirects=True,
+                    headers={
+                        "User-Agent": (
+                            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+                            "AppleWebKit/537.36 (KHTML, like Gecko) "
+                            "Chrome/148.0.0.0 Safari/537.36"
+                        ),
+                        "Accept": (
+                            "text/html,application/xhtml+xml,"
+                            "application/xml;q=0.9,*/*;q=0.8"
+                        ),
+                    },
+                    timeout=30.0,
+                )
+
+                if resp.status_code != 200:
+                    print(f"[Session] GET returned {resp.status_code}")
+                    return
+
+                # Collect cookies
+                new_cookies = httpx.Cookies()
+                for name, value in resp.cookies.items():
+                    new_cookies.set(name, value, domain="chatgpt.org")
+                for header in resp.headers.get_list("set-cookie"):
+                    parts = header.split(";")[0]
+                    if "=" in parts:
+                        k, v = parts.split("=", 1)
+                        new_cookies.set(k.strip(), v.strip(), domain="chatgpt.org")
+
+                # XSRF from cookie
+                xsrf = new_cookies.get("XSRF-TOKEN", domain="chatgpt.org")
+                if xsrf:
+                    xsrf = unquote(xsrf)
+
+                # CSRF from HTML meta tag
+                csrf = None
+                m = re.search(
+                    r'<meta\s+name="csrf-token"\s+content="([^"]+)"', resp.text
+                )
+                if m:
+                    csrf = m.group(1)
+
+                self.cookies = new_cookies
+                self.xsrf_token = xsrf
+                self.csrf_token = csrf
+                self.last_refresh = now
+
+                print(
+                    f"[Session] OK — CSRF:{bool(csrf)} XSRF:{bool(xsrf)} "
+                    f"Cookies:{list(new_cookies.keys())}"
+                )
+
+            except Exception as e:
+                print(f"[Session] Refresh error: {e}")
 
-class Message(BaseModel):
-    role: str
-    content: str
 
-class ChatRequest(BaseModel):
-    model: str
-    messages: List[Message]
-    stream: bool = True
-    max_tokens: Optional[int] = None
-    temperature: Optional[float] = None
+session = SessionState()
 
-class AnthropicMessage(BaseModel):
+# ── HTTP Client ──────────────────────────────────────────────────
+http_client: Optional[httpx.AsyncClient] = None
+
+@app.on_event("startup")
+async def startup():
+    global http_client
+    http_client = httpx.AsyncClient(
+        verify=False,
+        timeout=httpx.Timeout(120.0, connect=10.0),
+    )
+    await session.refresh(http_client)
+
+
+@app.on_event("shutdown")
+async def shutdown():
+    if http_client:
+        await http_client.aclose()
+
+# ── Models ───────────────────────────────────────────────────────
+class Message(BaseModel):
     role: str
     content: str
 
-class AnthropicRequest(BaseModel):
-    model: str
-    messages: List[AnthropicMessage]
-    max_tokens: int = 1024
-    stream: bool = False
-
-def get_proxies():
-    if PROXY_URL:
-        return PROXY_URL
-    return None
-
-async def get_csrf_token(client: httpx.AsyncClient) -> str:
-    resp = await client.get(f"{API_ENDPOINT}/claude/chat", follow_redirects=True)
-    xsrf_cookie = resp.cookies.get("XSRF-TOKEN")
-    if xsrf_cookie:
-        return urllib.parse.unquote(xsrf_cookie)
-    match = re.search(r'XSRF-TOKEN=([^;]+)', resp.headers.get("set-cookie", ""))
-    if match:
-        return urllib.parse.unquote(match.group(1))
-    raise HTTPException(500, "Auth failed")
-
-async def stream_chat(client: httpx.AsyncClient, model: str, messages: List[Dict], csrf_token: str) -> AsyncGenerator[str, None]:
-    payload = {"model": model, "messages": [m.model_dump() for m in messages] if hasattr(messages[0], 'model_dump') else messages}
-    headers = {
+# ── Helpers ──────────────────────────────────────────────────────
+def _headers() -> dict:
+    h = {
+        "Accept": "*/*",
         "Content-Type": "application/json",
-        "X-CSRF-TOKEN": csrf_token,
-        "Accept": "text/event-stream",
-        "Origin": API_ENDPOINT,
-        "Referer": f"{API_ENDPOINT}/claude/chat",
+        "Origin": "https://chatgpt.org",
+        "Referer": "https://chatgpt.org/claude/chat",
+        "X-Requested-With": "XMLHttpRequest",
+        "User-Agent": (
+            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+            "AppleWebKit/537.36 (KHTML, like Gecko) "
+            "Chrome/148.0.0.0 Safari/537.36"
+        ),
     }
-    async with client.stream("POST", f"{API_ENDPOINT}/api/chat", json=payload, headers=headers, timeout=120.0) as resp:
-        if resp.status_code != 200:
-            error_body = await resp.aread()
-            raise HTTPException(resp.status_code, error_body.decode())
-        async for line in resp.aiter_lines():
-            if line.startswith("data: "):
-                data = line[6:]
-                if data == "[DONE]":
-                    yield f"data: [DONE]\n\n"
-                    break
-                try:
-                    parsed = json.loads(data)
-                    yield f"data: {json.dumps(parsed)}\n\n"
-                except json.JSONDecodeError:
-                    continue
-
-async def do_chat(model: str, messages: List[Dict], stream: bool = False):
-    proxy = get_proxies()
-    async with httpx.AsyncClient(proxy=proxy, cookies={}, http2=True, follow_redirects=True, timeout=120.0) as client:
-        csrf_token = await get_csrf_token(client)
-        
-        if stream:
-            return EventSourceResponse(stream_chat(client, model, messages, csrf_token))
-        else:
-            full_content = ""
-            usage = None
-            current_messages = [m.model_dump() if hasattr(m, 'model_dump') else m for m in messages]
-            max_continuations = 5
-            continuations = 0
-            
-            while continuations <= max_continuations:
-                if full_content:
-                    current_messages.append({"role": "assistant", "content": full_content})
-                    current_messages.append({"role": "user", "content": "Continue and complete from where you left off. Do not repeat, just continue."})
-                
-                chunk_count = 0
-                last_finish_reason = None
-                
-                async for chunk in stream_chat(client, model, current_messages, csrf_token):
-                    if chunk.startswith("data: ") and not chunk.startswith("data: [DONE]"):
-                        try:
-                            data = json.loads(chunk[6:-2])
-                            if "choices" in data and data["choices"]:
-                                delta = data["choices"][0].get("delta", {})
-                                content = delta.get("content", "")
-                                if content:
-                                    full_content += content
-                                    chunk_count += 1
-                                finish = data["choices"][0].get("finish_reason")
-                                if finish:
-                                    last_finish_reason = finish
-                            if "usage" in data:
-                                usage = data["usage"]
-                        except:
-                            pass
-                
-                if last_finish_reason == "stop" or chunk_count < 50:
-                    break
-                
-                continuations += 1
-            
-            return full_content, usage
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    yield
-
-app = FastAPI(title="Haiku API", description="Fast AI chat completions API", lifespan=lifespan)
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+    csrf = session.csrf_token or session.xsrf_token
+    if csrf:
+        h["X-CSRF-TOKEN"] = csrf
+    return h
 
-@app.get("/")
-async def root():
-    return {"status": "ok", "message": "Haiku API"}
 
-@app.get("/v1/models")
-async def list_models():
-    return {"object": "list", "data": [{"id": m["id"], "object": "model", "owned_by": m["slug"]} for m in DEFAULT_MODELS]}
+async def _chat(messages: list[dict], model: str):
+    """POST to chatgpt.org/api/chat; returns raw httpx Response."""
+    await session.refresh(http_client)
+
+    payload = {"model": model, "messages": messages}
+    resp = await http_client.post(
+        "https://chatgpt.org/api/chat",
+        json=payload,
+        headers=_headers(),
+        cookies=session.cookies,
+    )
+
+    # CSRF mismatch → refresh once
+    if resp.status_code == 419:
+        print("[Chat] 419 → refreshing session...")
+        session.last_refresh = 0
+        await session.refresh(http_client)
+        resp = await http_client.post(
+            "https://chatgpt.org/api/chat",
+            json=payload,
+            headers=_headers(),
+            cookies=session.cookies,
+        )
+
+    if resp.status_code == 429:
+        raise HTTPException(429, "Rate limited by upstream")
+    if resp.status_code != 200:
+        session.last_refresh = 0
+        raise HTTPException(resp.status_code, f"Upstream {resp.status_code}: {resp.text[:300]}")
+
+    return resp
+
+
+# ── SSE relay ────────────────────────────────────────────────────
+async def _relay_sse(resp):
+    """Relay the upstream SSE stream, filtering OPENROUTER PROCESSING lines."""
+    content_so_far = ""
+    async for raw_line in resp.aiter_lines():
+        line = raw_line.strip()
 
-@app.get("/v1/models/{model_id}")
-async def get_model(model_id: str):
-    for m in DEFAULT_MODELS:
-        if m["id"] == model_id:
-            return {"id": m["id"], "object": "model", "owned_by": m["slug"]}
-    raise HTTPException(404, "Model not found")
+        # Skip keep-alive / processing comments
+        if not line or line.startswith(":"):
+            continue
 
-# OpenAI format
+        if not line.startswith("data: "):
+            continue
+
+        payload = line[6:]  # strip "data: "
+
+        if payload.strip() == "[DONE]":
+            yield "data: [DONE]\n\n"
+            break
+
+        try:
+            chunk = json.loads(payload)
+        except json.JSONDecodeError:
+            continue
+
+        # The upstream already returns OpenAI-format chunks
+        # Just relay them as-is
+        # Collect content for non-streaming fallback
+        for choice in chunk.get("choices", []):
+            delta = choice.get("delta", {})
+            c = delta.get("content", "")
+            if c:
+                content_so_far += c
+
+        yield f"data: {payload}\n\n"
+
+    # If we got no content via streaming, send it as a single chunk
+    if content_so_far and not any(
+        c.get("finish_reason") == "stop"
+        for chunk_raw in [payload]
+        for c in json.loads(payload).get("choices", [])
+        if payload
+    ):
+        pass  # Content already streamed
+
+
+async def _collect_sse(resp):
+    """Collect all SSE chunks into a single text string."""
+    content = ""
+    async for raw_line in resp.aiter_lines():
+        line = raw_line.strip()
+        if not line or line.startswith(":"):
+            continue
+        if not line.startswith("data: "):
+            continue
+        payload = line[6:]
+        if payload.strip() == "[DONE]":
+            break
+        try:
+            chunk = json.loads(payload)
+            for choice in chunk.get("choices", []):
+                delta = choice.get("delta", {})
+                c = delta.get("content", "")
+                if c:
+                    content += c
+        except json.JSONDecodeError:
+            continue
+    return content
+
+
+# ── OpenAI-compatible endpoint ──────────────────────────────────
 @app.post("/v1/chat/completions")
-async def chat_completions(request: ChatRequest):
+async def chat_completions(request: Request):
     try:
-        result = await do_chat(request.model, request.messages, request.stream)
-        if request.stream:
-            return result
-        full_content, usage = result
-        return {
-            "id": f"chatcmpl-{hash(full_content) % 1000000}",
-            "object": "chat.completion",
-            "model": request.model,
-            "choices": [{"index": 0, "message": {"role": "assistant", "content": full_content}, "finish_reason": "stop"}],
-            "usage": usage
-        }
-    except Exception as e:
-        raise HTTPException(500, str(e))
-
-# Anthropic format
-@app.post("/v1/messages")
-async def anthropic_messages(request: AnthropicRequest):
-    try:
-        # Map model name
-        model_map = {
-            "claude-sonnet-4.5": "anthropic/claude-sonnet-4-5",
-            "claude-haiku-4.5": "anthropic/claude-haiku-4-5",
-            "claude-sonnet-4-5": "anthropic/claude-sonnet-4-5",
-            "claude-haiku-4-5": "anthropic/claude-haiku-4-5",
-        }
-        model = model_map.get(request.model, request.model)
-        
-        result = await do_chat(model, request.messages, request.stream)
-        if request.stream:
-            return result
-        full_content, usage = result
-        return {
-            "id": f"msg_{hash(full_content) % 1000000}",
-            "type": "message",
-            "role": "assistant",
-            "model": request.model,
-            "content": [{"type": "text", "text": full_content}],
-            "stop_reason": "end_turn",
-            "usage": {"input_tokens": usage.get("prompt_tokens", 0) if usage else 0, "output_tokens": usage.get("completion_tokens", 0) if usage else 0}
-        }
-    except Exception as e:
-        raise HTTPException(500, str(e))
+        body = await request.json()
+    except Exception:
+        raise HTTPException(400, "Invalid JSON")
+
+    model = body.get("model", "anthropic/claude-haiku-4-5")
+    messages = body.get("messages", [])
+    stream = body.get("stream", False)
+
+    if not messages:
+        raise HTTPException(400, "messages is empty")
+
+    resp = await _chat(messages, model)
+
+    if stream:
+        return StreamingResponse(
+            _relay_sse(resp), media_type="text/event-stream"
+        )
+
+    # Non-streaming: collect the full response
+    text = await _collect_sse(resp)
+    return JSONResponse({
+        "id": f"chatcmpl-{int(time.time())}",
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": model,
+        "choices": [{
+            "index": 0,
+            "message": {"role": "assistant", "content": text},
+            "finish_reason": "stop",
+        }],
+        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
+    })
+
+
+# ── Models / Health ─────────────────────────────────────────────
+@app.get("/v1/models")
+async def list_models():
+    return JSONResponse({
+        "object": "list",
+        "data": [
+            {"id": "anthropic/claude-haiku-4-5", "object": "model", "owned_by": "anthropic"},
+        ],
+    })
+
+
+@app.get("/")
+async def root():
+    return {"status": "ok", "endpoints": ["/v1/chat/completions", "/v1/models"]}
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "session_active": bool(session.cookies)}
+
+
+@app.get("/debug/session")
+async def debug_session():
+    return {
+        "has_cookies": bool(session.cookies),
+        "cookie_names": list(session.cookies.keys()) if session.cookies else [],
+        "has_csrf": bool(session.csrf_token),
+        "has_xsrf": bool(session.xsrf_token),
+        "last_refresh_ago": int(time.time() - session.last_refresh) if session.last_refresh else None,
+    }
+
 
 if __name__ == "__main__":
     import uvicorn