proxy.py

"""
Anthropic <-> OpenAI translation proxy.

Claude Code speaks Anthropic's /v1/messages schema. This proxy exposes
that same schema locally, translates each request into an OpenAI-style
/v1/chat/completions call against your real backend, and translates the
(streaming or non-streaming) response back into Anthropic's format.

Point Claude Code at this proxy via ANTHROPIC_BASE_URL=http://localhost:8317
and it never needs to know the real backend isn't Anthropic.
"""

import json
import os
import time
import uuid

import httpx
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse

app = FastAPI()

OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "http://localhost:8000")
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
PROXY_PORT = int(os.environ.get("PROXY_PORT", 8317))

client = httpx.AsyncClient(timeout=120.0)


# ---------- Anthropic request -> OpenAI request ----------

def anthropic_to_openai_request(body: dict) -> dict:
    messages = []

    system = body.get("system")
    if system:
        if isinstance(system, list):
            system_text = "\n".join(b.get("text", "") for b in system if b.get("type") == "text")
        else:
            system_text = system
        messages.append({"role": "system", "content": system_text})

    for msg in body.get("messages", []):
        role = msg["role"]
        content = msg["content"]

        if isinstance(content, str):
            messages.append({"role": role, "content": content})
            continue

        # content is a list of blocks (text, tool_use, tool_result, image)
        text_parts = []
        tool_calls = []
        for block in content:
            btype = block.get("type")
            if btype == "text":
                text_parts.append(block["text"])
            elif btype == "tool_use":
                tool_calls.append({
                    "id": block["id"],
                    "type": "function",
                    "function": {
                        "name": block["name"],
                        "arguments": json.dumps(block.get("input", {})),
                    },
                })
            elif btype == "tool_result":
                messages.append({
                    "role": "tool",
                    "tool_call_id": block["tool_use_id"],
                    "content": _flatten_tool_result(block.get("content", "")),
                })
            elif btype == "image":
                src = block.get("source", {})
                text_parts.append(f"[image omitted: {src.get('media_type', 'unknown')}]")

        entry = {"role": role, "content": "\n".join(text_parts) if text_parts else None}
        if tool_calls:
            entry["tool_calls"] = tool_calls
        if entry["content"] is not None or tool_calls:
            messages.append(entry)

    openai_body = {
        "model": os.environ.get("OPENAI_MODEL", body.get("model", "gpt-4")),
        "messages": messages,
        "max_tokens": body.get("max_tokens", 1024),
        "temperature": body.get("temperature", 1.0),
        "stream": body.get("stream", False),
    }

    if body.get("tools"):
        openai_body["tools"] = [
            {
                "type": "function",
                "function": {
                    "name": t["name"],
                    "description": t.get("description", ""),
                    "parameters": t.get("input_schema", {}),
                },
            }
            for t in body["tools"]
        ]

    return openai_body


def _flatten_tool_result(content) -> str:
    if isinstance(content, str):
        return content
    parts = []
    for block in content:
        if block.get("type") == "text":
            parts.append(block["text"])
    return "\n".join(parts)


# ---------- OpenAI response -> Anthropic response (non-streaming) ----------

def openai_to_anthropic_response(oa: dict, model: str) -> dict:
    choice = oa["choices"][0]
    message = choice["message"]
    content_blocks = []

    if message.get("content"):
        content_blocks.append({"type": "text", "text": message["content"]})

    for tc in message.get("tool_calls", []) or []:
        try:
            args = json.loads(tc["function"]["arguments"])
        except (json.JSONDecodeError, TypeError):
            args = {}
        content_blocks.append({
            "type": "tool_use",
            "id": tc["id"],
            "name": tc["function"]["name"],
            "input": args,
        })

    finish_map = {"stop": "end_turn", "length": "max_tokens", "tool_calls": "tool_use"}

    usage = oa.get("usage", {})

    return {
        "id": f"msg_{uuid.uuid4().hex[:24]}",
        "type": "message",
        "role": "assistant",
        "model": model,
        "content": content_blocks,
        "stop_reason": finish_map.get(choice.get("finish_reason"), "end_turn"),
        "stop_sequence": None,
        "usage": {
            "input_tokens": usage.get("prompt_tokens", 0),
            "output_tokens": usage.get("completion_tokens", 0),
        },
    }


# ---------- OpenAI SSE stream -> Anthropic SSE stream ----------

async def stream_openai_to_anthropic(oa_stream, model: str):
    message_id = f"msg_{uuid.uuid4().hex[:24]}"
    started = False
    block_open = False
    block_index = 0

    yield _sse("message_start", {
        "type": "message_start",
        "message": {
            "id": message_id, "type": "message", "role": "assistant",
            "model": model, "content": [], "stop_reason": None,
            "stop_sequence": None, "usage": {"input_tokens": 0, "output_tokens": 0},
        },
    })
    started = True

    async for line in oa_stream:
        if not line or not line.startswith("data: "):
            continue
        payload = line[len("data: "):].strip()
        if payload == "[DONE]":
            break
        try:
            chunk = json.loads(payload)
        except json.JSONDecodeError:
            continue

        delta = chunk["choices"][0].get("delta", {})

        if "content" in delta and delta["content"]:
            if not block_open:
                yield _sse("content_block_start", {
                    "type": "content_block_start", "index": block_index,
                    "content_block": {"type": "text", "text": ""},
                })
                block_open = True
            yield _sse("content_block_delta", {
                "type": "content_block_delta", "index": block_index,
                "delta": {"type": "text_delta", "text": delta["content"]},
            })

        if delta.get("tool_calls"):
            for tc in delta["tool_calls"]:
                if block_open:
                    yield _sse("content_block_stop", {"type": "content_block_stop", "index": block_index})
                    block_index += 1
                    block_open = False
                yield _sse("content_block_start", {
                    "type": "content_block_start", "index": block_index,
                    "content_block": {
                        "type": "tool_use",
                        "id": tc.get("id", f"toolu_{uuid.uuid4().hex[:16]}"),
                        "name": tc["function"]["name"],
                        "input": {},
                    },
                })
                args = tc["function"].get("arguments", "")
                if args:
                    yield _sse("content_block_delta", {
                        "type": "content_block_delta", "index": block_index,
                        "delta": {"type": "input_json_delta", "partial_json": args},
                    })
                yield _sse("content_block_stop", {"type": "content_block_stop", "index": block_index})
                block_index += 1

    if block_open:
        yield _sse("content_block_stop", {"type": "content_block_stop", "index": block_index})

    yield _sse("message_delta", {
        "type": "message_delta",
        "delta": {"stop_reason": "end_turn", "stop_sequence": None},
        "usage": {"output_tokens": 0},
    })
    yield _sse("message_stop", {"type": "message_stop"})


def _sse(event: str, data: dict) -> str:
    return f"event: {event}\ndata: {json.dumps(data)}\n\n"


# ---------- Route ----------

@app.post("/v1/messages")
async def messages(request: Request):
    body = await request.json()
    model = body.get("model", "claude-proxy")
    openai_body = anthropic_to_openai_request(body)

    headers = {"Content-Type": "application/json"}
    if OPENAI_API_KEY:
        headers["Authorization"] = f"Bearer {OPENAI_API_KEY}"

    if openai_body.get("stream"):
        async def event_gen():
            async with client.stream(
                "POST", f"{OPENAI_BASE_URL}/v1/chat/completions",
                json=openai_body, headers=headers,
            ) as resp:
                async for chunk in stream_openai_to_anthropic(resp.aiter_lines(), model):
                    yield chunk
        return StreamingResponse(event_gen(), media_type="text/event-stream")

    resp = await client.post(
        f"{OPENAI_BASE_URL}/v1/chat/completions",
        json=openai_body, headers=headers,
    )
    resp.raise_for_status()
    anthropic_resp = openai_to_anthropic_response(resp.json(), model)
    return JSONResponse(anthropic_resp)


@app.get("/health")
async def health():
    return {"status": "ok", "backend": OPENAI_BASE_URL}


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