main.py

"""
Dual-Protocol API Server
========================
Supports BOTH:
  • OpenAI Python library  →  POST /v1/chat/completions
  • Claude Code / Anthropic SDK  →  POST /v1/messages

Usage as base URL
-----------------
OpenAI Python:
    client = OpenAI(base_url="http://localhost:7860/v1", api_key="any")

Anthropic SDK / Claude Code:
    export ANTHROPIC_BASE_URL=http://localhost:7860
    export ANTHROPIC_API_KEY=any-value
    claude          # or use the Python SDK with base_url=...
"""

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware

import requests
import json
import uvicorn
import uuid
import traceback
import time

app = FastAPI(title="Dual-Protocol LLM Proxy")

# =========================================================
# CORS
# =========================================================

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# =========================================================
# CONFIG
# =========================================================

HF_SPACE_URL    = "https://akhaliq-ling-2-6-1t.hf.space/stream_chat"
HF_SPACE_ORIGIN = "https://akhaliq-ling-2-6-1t.hf.space"
MODEL_NAME      = "ling-2.6-1t"

# How many times to retry the upstream on timeout/5xx before giving up
UPSTREAM_RETRIES      = 3
# Per-attempt connect timeout (seconds) — allows HF cold-start wake-up
UPSTREAM_CONNECT_TIMEOUT = 30
# Per-attempt read timeout (seconds)
UPSTREAM_READ_TIMEOUT    = 300   # 5 min; HF spaces can be slow

UPSTREAM_HEADERS = {
    "accept":       "*/*",
    "content-type": "application/json",
    "origin":       HF_SPACE_ORIGIN,
    "referer":      HF_SPACE_ORIGIN + "/",
    "user-agent":   "Mozilla/5.0",
}

# =========================================================
# HELPERS — content normalisation
# =========================================================

def normalize_content(content) -> str:
    """
    Accept any of:
      - plain string
      - OpenAI/Anthropic multimodal list  [{"type": "text", "text": "..."}, ...]
      - None
    Returns a plain string.
    """
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict):
                t = item.get("type", "")
                if t in ("text", "input_text"):
                    parts.append(item.get("text", ""))
                # image / document blocks are silently dropped
            elif isinstance(item, str):
                parts.append(item)
        return "\n".join(parts)
    if isinstance(content, str):
        return content
    if content is None:
        return ""
    return str(content)


# =========================================================
# HELPERS — upstream warmup + call
# =========================================================

def warmup_upstream() -> dict:
    """
    Send a tiny probe request to wake the HF Space from sleep.
    Returns {"ok": True, "elapsed": N} or {"ok": False, "error": "..."}.
    """
    probe = {
        "messages":      [{"role": "user", "content": "hi"}],
        "system_prompt": "",
    }
    t0 = time.time()
    try:
        r = requests.post(
            HF_SPACE_URL,
            headers=UPSTREAM_HEADERS,
            json=probe,
            stream=True,
            timeout=(UPSTREAM_CONNECT_TIMEOUT, UPSTREAM_READ_TIMEOUT),
        )
        # drain just the first line so the connection is confirmed alive
        for _ in r.iter_lines():
            break
        return {"ok": True, "elapsed": round(time.time() - t0, 2)}
    except Exception as e:
        return {"ok": False, "error": str(e), "elapsed": round(time.time() - t0, 2)}


def call_upstream(system_prompt: str, messages: list, stream: bool):
    """
    Forward to the HF-space backend with automatic retry on timeout / 5xx.
    Returns a requests.Response object (always opened in stream mode).
    Raises RuntimeError if all retries are exhausted.
    """
    payload = {
        "messages":      messages,
        "system_prompt": system_prompt,
    }
    last_exc = None
    for attempt in range(1, UPSTREAM_RETRIES + 1):
        try:
            resp = requests.post(
                HF_SPACE_URL,
                headers=UPSTREAM_HEADERS,
                json=payload,
                stream=True,
                timeout=(UPSTREAM_CONNECT_TIMEOUT, UPSTREAM_READ_TIMEOUT),
            )
            if resp.status_code >= 500:
                raise RuntimeError(f"upstream returned {resp.status_code}")
            return resp
        except (requests.exceptions.Timeout,
                requests.exceptions.ConnectionError,
                RuntimeError) as e:
            last_exc = e
            wait = 2 ** (attempt - 1)   # 1s, 2s, 4s …
            print(f"[upstream] attempt {attempt}/{UPSTREAM_RETRIES} failed: {e}  — retry in {wait}s")
            time.sleep(wait)

    raise RuntimeError(f"upstream unreachable after {UPSTREAM_RETRIES} attempts: {last_exc}")


def collect_full_text(response) -> str:
    """Drain an upstream streaming response and return the concatenated text."""
    full = ""
    for line in response.iter_lines():
        if not line:
            continue
        try:
            decoded = line.decode("utf-8")
            if decoded.startswith("data:"):
                data_str = decoded[len("data:"):].strip()
                if data_str == "[DONE]":
                    break
                parsed = json.loads(data_str)
                full += parsed.get("token", "")
        except Exception:
            continue
    return full


# =========================================================
# SHARED message converter (OpenAI → internal)
# =========================================================

def convert_openai_messages(messages: list):
    """
    Split OpenAI-style messages (which may include role='system' anywhere)
    into (system_prompt_str, user_assistant_messages).
    """
    system_parts = []
    converted = []
    for msg in messages:
        role    = msg.get("role", "user")
        content = normalize_content(msg.get("content", ""))
        if role in ("system", "developer"):
            system_parts.append(content)
        else:
            converted.append({"role": role, "content": content})
    return "\n".join(system_parts), converted


# =========================================================
# ROOT / HEALTH / MODELS
# =========================================================

@app.get("/")
async def root():
    return {
        "status":  "ok",
        "message": "Dual-Protocol API Server (OpenAI + Anthropic/Claude-Code)",
        "model":   MODEL_NAME,
        "endpoints": [
            "POST /v1/chat/completions  ← OpenAI Python library",
            "POST /v1/messages          ← Anthropic SDK / Claude Code",
        ],
    }


@app.get("/health")
async def health():
    return {"status": "healthy"}


@app.get("/warmup")
@app.post("/warmup")
async def warmup():
    """
    Wake the upstream HF Space from sleep.
    Call this once before running tests or the first real request.
    Returns when the upstream has responded to a probe message.
    """
    result = warmup_upstream()
    status = 200 if result["ok"] else 503
    return JSONResponse(content={
        "upstream": HF_SPACE_URL,
        **result,
    }, status_code=status)


@app.get("/upstream-status")
async def upstream_status():
    """Quick reachability check for the upstream HF Space (no LLM call)."""
    try:
        r = requests.get(HF_SPACE_ORIGIN, timeout=(10, 10), allow_redirects=True)
        return {"reachable": True, "http_status": r.status_code, "url": HF_SPACE_ORIGIN}
    except Exception as e:
        return JSONResponse(
            content={"reachable": False, "error": str(e), "url": HF_SPACE_ORIGIN},
            status_code=503,
        )


# ----- OpenAI-style model list -----
@app.get("/v1/models")
async def openai_models():
    return {
        "object": "list",
        "data": [{
            "id":       MODEL_NAME,
            "object":   "model",
            "created":  int(time.time()),
            "owned_by": "custom",
        }],
    }


# ----- Anthropic-style model list (Claude Code may hit this) -----
@app.get("/v1/models", include_in_schema=False)   # duplicate handled below
async def _noop(): pass                           # shadowed; keep for completeness


# =========================================================
# ENDPOINT 1 — OpenAI  /v1/chat/completions
# =========================================================

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    try:
        body     = await request.json()
        messages = body.get("messages", [])
        stream   = body.get("stream", False)

        system_prompt, converted_messages = convert_openai_messages(messages)

        upstream = call_upstream(system_prompt, converted_messages, stream)

        # ── streaming ──────────────────────────────────────────────────
        if stream:
            async def openai_stream():
                cid = f"chatcmpl-{uuid.uuid4().hex}"
                for line in upstream.iter_lines():
                    if not line:
                        continue
                    try:
                        decoded = line.decode("utf-8")
                        if not decoded.startswith("data:"):
                            continue
                        data_str = decoded[len("data:"):].strip()
                        if data_str == "[DONE]":
                            yield "data: [DONE]\n\n"
                            break
                        parsed = json.loads(data_str)
                        token  = parsed.get("token", "")
                        chunk  = {
                            "id":      cid,
                            "object":  "chat.completion.chunk",
                            "created": int(time.time()),
                            "model":   MODEL_NAME,
                            "choices": [{
                                "index":         0,
                                "delta":         {"role": "assistant", "content": token},
                                "finish_reason": None,
                            }],
                        }
                        yield f"data: {json.dumps(chunk)}\n\n"
                    except Exception:
                        continue

                # emit final chunk with finish_reason
                cid = f"chatcmpl-{uuid.uuid4().hex}"
                final_chunk = {
                    "id":      cid,
                    "object":  "chat.completion.chunk",
                    "created": int(time.time()),
                    "model":   MODEL_NAME,
                    "choices": [{
                        "index":         0,
                        "delta":         {},
                        "finish_reason": "stop",
                    }],
                }
                yield f"data: {json.dumps(final_chunk)}\n\n"

            return StreamingResponse(openai_stream(), media_type="text/event-stream")

        # ── non-streaming ──────────────────────────────────────────────
        full_text = collect_full_text(upstream)
        return JSONResponse({
            "id":      f"chatcmpl-{uuid.uuid4().hex}",
            "object":  "chat.completion",
            "created": int(time.time()),
            "model":   MODEL_NAME,
            "choices": [{
                "index":         0,
                "message":       {"role": "assistant", "content": full_text},
                "finish_reason": "stop",
            }],
            "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
        })

    except Exception as e:
        traceback.print_exc()
        return JSONResponse(
            status_code=500,
            content={"error": {"message": str(e), "type": "server_error"}},
        )


# =========================================================
# ENDPOINT 2 — Anthropic / Claude Code  /v1/messages
#
# Request shape (what Claude Code sends):
# {
#   "model": "...",
#   "max_tokens": 4096,
#   "stream": true,
#   "system": "...",           ← top-level, NOT inside messages
#   "messages": [
#     {"role": "user", "content": "..."},
#     ...
#   ],
#   "tools": [...],            ← optional, silently accepted
# }
#
# Response shape expected by Claude Code:
# Non-stream:
#   { "id", "type": "message", "role": "assistant",
#     "content": [{"type": "text", "text": "..."}],
#     "model", "stop_reason": "end_turn", "stop_sequence": null,
#     "usage": {"input_tokens": N, "output_tokens": N} }
#
# Streaming SSE:
#   event: message_start
#   data: {"type":"message_start","message":{...}}
#
#   event: content_block_start
#   data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
#
#   event: ping
#   data: {"type":"ping"}
#
#   event: content_block_delta
#   data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
#
#   event: content_block_stop
#   data: {"type":"content_block_stop","index":0}
#
#   event: message_delta
#   data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":N}}
#
#   event: message_stop
#   data: {"type":"message_stop"}
# =========================================================

def _sse_event(event_name: str, data: dict) -> str:
    """Format a single SSE event in Anthropic's named-event style."""
    return f"event: {event_name}\ndata: {json.dumps(data)}\n\n"


@app.post("/v1/messages")
async def anthropic_messages(request: Request):
    try:
        body       = await request.json()
        stream     = body.get("stream", False)
        # Anthropic puts system at the top level
        system_raw = body.get("system", "")
        system_prompt = normalize_content(system_raw)
        messages   = body.get("messages", [])

        # Normalise content in each message (may be multimodal list)
        converted_messages = [
            {"role": m.get("role", "user"), "content": normalize_content(m.get("content", ""))}
            for m in messages
            if m.get("role") not in ("system",)   # safety: drop any stray system roles
        ]

        msg_id   = f"msg_{uuid.uuid4().hex[:24]}"
        upstream = call_upstream(system_prompt, converted_messages, stream)

        # ── streaming (Anthropic SSE format) ──────────────────────────
        if stream:
            async def anthropic_stream():
                # 1. message_start
                yield _sse_event("message_start", {
                    "type": "message_start",
                    "message": {
                        "id":            msg_id,
                        "type":          "message",
                        "role":          "assistant",
                        "content":       [],
                        "model":         MODEL_NAME,
                        "stop_reason":   None,
                        "stop_sequence": None,
                        "usage":         {"input_tokens": 0, "output_tokens": 1},
                    },
                })

                # 2. content_block_start
                yield _sse_event("content_block_start", {
                    "type":          "content_block_start",
                    "index":         0,
                    "content_block": {"type": "text", "text": ""},
                })

                # 3. ping
                yield _sse_event("ping", {"type": "ping"})

                # 4. stream tokens as content_block_delta events
                output_tokens = 0
                for line in upstream.iter_lines():
                    if not line:
                        continue
                    try:
                        decoded = line.decode("utf-8")
                        if not decoded.startswith("data:"):
                            continue
                        data_str = decoded[len("data:"):].strip()
                        if data_str == "[DONE]":
                            break
                        parsed = json.loads(data_str)
                        token  = parsed.get("token", "")
                        if not token:
                            continue
                        output_tokens += len(token.split())   # rough estimate
                        yield _sse_event("content_block_delta", {
                            "type":  "content_block_delta",
                            "index": 0,
                            "delta": {"type": "text_delta", "text": token},
                        })
                    except Exception:
                        continue

                # 5. content_block_stop
                yield _sse_event("content_block_stop", {
                    "type":  "content_block_stop",
                    "index": 0,
                })

                # 6. message_delta (carries stop_reason)
                yield _sse_event("message_delta", {
                    "type":  "message_delta",
                    "delta": {"stop_reason": "end_turn", "stop_sequence": None},
                    "usage": {"output_tokens": output_tokens},
                })

                # 7. message_stop
                yield _sse_event("message_stop", {"type": "message_stop"})

            return StreamingResponse(
                anthropic_stream(),
                media_type="text/event-stream",
                headers={
                    # Claude Code checks for this header
                    "anthropic-version": "2023-06-01",
                    "x-request-id":      msg_id,
                },
            )

        # ── non-streaming ──────────────────────────────────────────────
        full_text = collect_full_text(upstream)
        return JSONResponse(
            content={
                "id":            msg_id,
                "type":          "message",
                "role":          "assistant",
                "content":       [{"type": "text", "text": full_text}],
                "model":         MODEL_NAME,
                "stop_reason":   "end_turn",
                "stop_sequence": None,
                "usage":         {"input_tokens": 0, "output_tokens": 0},
            },
            headers={
                "anthropic-version": "2023-06-01",
                "x-request-id":      msg_id,
            },
        )

    except Exception as e:
        traceback.print_exc()
        # Return an Anthropic-shaped error
        return JSONResponse(
            status_code=500,
            content={
                "type":  "error",
                "error": {"type": "api_error", "message": str(e)},
            },
        )


# =========================================================
# ENDPOINT 3 — OpenAI Responses API  /v1/responses
# (some older OpenAI clients / Claude Code builds use this)
# =========================================================

@app.post("/v1/responses")
async def responses_api(request: Request):
    try:
        body       = await request.json()
        stream     = body.get("stream", False)
        input_data = body.get("input", "")

        messages = []
        if isinstance(input_data, list):
            for item in input_data:
                role    = item.get("role", "user")
                content = normalize_content(item.get("content", ""))
                messages.append({"role": role, "content": content})
        else:
            messages.append({"role": "user", "content": str(input_data)})

        # Delegate to the OpenAI chat completions handler
        class _FakeRequest:
            async def json(self):
                return {"messages": messages, "stream": stream}

        return await chat_completions(_FakeRequest())

    except Exception as e:
        traceback.print_exc()
        return JSONResponse(
            status_code=500,
            content={"error": {"message": str(e), "type": "server_error"}},
        )


# =========================================================
# START
# =========================================================

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)