Connect Claude Code and Claude Desktop to CoreWeave Inference

from flask import Flask, request, jsonify
import requests
import os
import json
import logging

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# In-memory store linking session IDs to their generated titles
session_titles = {}

CW_API_TOKEN = os.environ["CW_API_TOKEN"]
CW_MODEL_NAME = os.environ.get("CW_MODEL_NAME")
if not CW_MODEL_NAME:
    raise SystemExit(
        "CW_MODEL_NAME is not set. List available models with "
        "'curl $CW_GATEWAY_ENDPOINT/v1/models', then export CW_MODEL_NAME "
        "before starting the proxy."
    )
CW_GATEWAY_ENDPOINT = os.environ["CW_GATEWAY_ENDPOINT"] + "/v1/chat/completions"
CW_USE_TOOLS = (
    os.environ.get("CW_USE_TOOLS", "true").lower()
    in ("true", "1", "yes")
)

CW_USER_AGENT = os.environ.get(
    "CW_USER_AGENT",
    "curl/8 claude-cw-proxy",
)

CW_OPENAI_PROJECT = os.environ.get("CW_OPENAI_PROJECT")

CW_ADVERTISED_MODEL = os.environ.get(
    "CW_ADVERTISED_MODEL",
    CW_MODEL_NAME,
)

logger.info("Tool use enabled: %s", CW_USE_TOOLS)


# Helpers
def _cw_headers():
    headers = {
        "Authorization": f"Bearer {CW_API_TOKEN}",
        "Content-Type": "application/json",
        "User-Agent": CW_USER_AGENT,
    }

    if CW_OPENAI_PROJECT:
        headers["OpenAI-Project"] = CW_OPENAI_PROJECT

    return headers


class GatewayError(Exception):
    def __init__(self, status_code, body):
        self.status_code = status_code
        self.body = body


def _cw_chat(payload):
    resp = requests.post(
        CW_GATEWAY_ENDPOINT,
        json=payload,
        headers=_cw_headers(),
    )
    if not resp.ok:
        raise GatewayError(resp.status_code, resp.text)
    return resp.json()


def extract_first_user_text(body):
    for msg in body.get("messages", []):
        if msg.get("role") == "user":
            content = msg.get("content", "")
            if isinstance(content, str):
                return content
            if isinstance(content, list):
                for b in content:
                    if b.get("type") == "text":
                        return b.get("text", "")
    return ""


# Format translation
def flatten_content(content):
    if isinstance(content, str):
        return content

    if isinstance(content, list):
        parts = []

        for b in content:
            if b.get("type") == "text":
                parts.append(b.get("text", ""))

            elif b.get("type") == "tool_result":
                result = b.get("content", "")

                if isinstance(result, list):
                    result = " ".join(
                        r.get("text", "")
                        for r in result
                        if r.get("type") == "text"
                    )

                parts.append(
                    f"[tool_result id={b.get('tool_use_id')}]: {result}"
                )

        return "\n".join(parts)

    return str(content)


def anthropic_tools_to_openai(tools):
    return [
        {
            "type": "function",
            "function": {
                "name": t["name"],
                "description": t.get("description", ""),
                "parameters": t.get("input_schema", {}),
            },
        }
        for t in tools
    ]


def anthropic_to_openai(body):
    messages = []

    system = body.get("system")
    if system:
        messages.append(
            {
                "role": "system",
                "content": flatten_content(system),
            }
        )

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

        if role == "assistant" and isinstance(content, list):
            text_parts = [
                b.get("text", "")
                for b in content
                if b.get("type") == "text"
            ]

            tool_calls = []

            if CW_USE_TOOLS:
                for b in content:
                    if b.get("type") == "tool_use":
                        tool_calls.append(
                            {
                                "id": b["id"],
                                "type": "function",
                                "function": {
                                    "name": b["name"],
                                    "arguments": json.dumps(
                                        b.get("input", {})
                                    ),
                                },
                            }
                        )

            oai_msg = {
                "role": "assistant",
                "content": "\n".join(text_parts) or None,
            }

            if tool_calls:
                oai_msg["tool_calls"] = tool_calls

            messages.append(oai_msg)
        elif role == "user" and isinstance(content, list):
            tool_results = []

            if CW_USE_TOOLS:
                tool_results = [
                    b for b in content
                    if b.get("type") == "tool_result"
                ]

            text_parts = [
                b for b in content
                if b.get("type") == "text"
            ]

            for tr in tool_results:
                result_content = tr.get("content", "")

                if isinstance(result_content, list):
                    result_content = "\n".join(
                        r.get("text", "")
                        for r in result_content
                        if r.get("type") == "text"
                    )

                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tr["tool_use_id"],
                        "content": result_content,
                    }
                )

            if text_parts:
                messages.append(
                    {
                        "role": "user",
                        "content": "\n".join(
                            b.get("text", "")
                            for b in text_parts
                        ),
                    }
                )

        else:
            messages.append(
                {
                    "role": role,
                    "content": flatten_content(content),
                }
            )

    payload = {
        "model": CW_MODEL_NAME,
        "messages": messages,
        "max_tokens": body.get("max_tokens", 4096),
    }

    tools = body.get("tools")

    if tools and CW_USE_TOOLS:
        payload["tools"] = anthropic_tools_to_openai(tools)
        payload["tool_choice"] = "auto"

    return payload


def openai_to_anthropic(oai):
    choice = oai["choices"][0]
    msg = choice["message"]
    usage = oai.get("usage", {})

    content = []

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

    for tc in msg.get("tool_calls") or []:
        fn = tc["function"]

        try:
            inp = json.loads(fn["arguments"])
        except Exception:
            inp = {"raw": fn["arguments"]}

        content.append(
            {
                "type": "tool_use",
                "id": tc["id"],
                "name": fn["name"],
                "input": inp,
            }
        )

    return {
        "id": oai.get("id", "proxy"),
        "type": "message",
        "role": "assistant",
        "model": oai.get("model") or CW_MODEL_NAME,
        "content": content,
        "stop_reason": (
            "tool_use"
            if msg.get("tool_calls")
            else "end_turn"
        ),
        "usage": {
            "input_tokens": usage.get("prompt_tokens", 0),
            "output_tokens": usage.get("completion_tokens", 0),
        },
    }


# Session title handling
def is_title_request(body):
    fmt = body.get("output_config", {}).get("format", {})

    if fmt.get("type") == "json_schema":
        if "title" in fmt.get("schema", {}).get("properties", {}):
            return True

    return False


def generate_session_title(user_text, session_id):
    payload = {
        "model": CW_MODEL_NAME,
        "messages": [
            {
                "role": "system",
                "content": (
                    "Generate a concise session title "
                    "(3-7 words, sentence case) "
                    "for a coding session based on the user's "
                    "first message. "
                    'Respond with ONLY valid JSON: {"title": "your title here"}'
                ),
            },
            {
                "role": "user",
                "content": user_text or "General coding session",
            },
        ],
        "max_tokens": 50,
    }

    try:
        raw = _cw_chat(payload)["choices"][0]["message"]["content"]
        raw = raw.strip().strip("```json").strip("```").strip()
        title = json.loads(raw).get("title", "Claude Code session")
    except Exception:
        title = "Claude Code session"

    session_titles[session_id] = title

    logger.info(
        "Session title generated: %s -> %s",
        session_id[:8] if session_id else "",
        title,
    )

    return title


# Inference
def run_inf(messages, tools, max_tokens, session_id=""):
    payload = {
        "model": CW_MODEL_NAME,
        "messages": messages,
        "max_tokens": max_tokens,
    }

    if tools:
        payload["tools"] = tools
        payload["tool_choice"] = "auto"

    logger.info(
        "Sending model: %s | session: %s",
        CW_MODEL_NAME,
        session_id[:8] if session_id else "",
    )

    return _cw_chat(payload)


# Routes
@app.route("/v1/messages", methods=["POST"])
def messages():
    body = request.get_json()

    try:
        session_id = json.loads(
            body.get("metadata", {}).get("user_id", "{}")
        ).get("session_id", "")
    except Exception:
        session_id = ""

    if is_title_request(body):
        logger.info("Title request detected - handling in proxy")

        user_text = extract_first_user_text(body)
        title = generate_session_title(user_text, session_id)

        return jsonify(
            {
                "id": "msg_title",
                "type": "message",
                "role": "assistant",
                "model": CW_MODEL_NAME,
                "content": [
                    {
                        "type": "text",
                        "text": json.dumps({"title": title}),
                    }
                ],
                "stop_reason": "end_turn",
                "usage": {
                    "input_tokens": 0,
                    "output_tokens": 0,
                },
            }
        )

    oai_payload = anthropic_to_openai(body)

    try:
        result = run_inf(
            oai_payload["messages"],
            oai_payload.get("tools", []),
            oai_payload["max_tokens"],
            session_id=session_id,
        )
    except GatewayError as e:
        return jsonify({"error": e.body}), e.status_code

    return jsonify(openai_to_anthropic(result))


def rough_token_count(text):
    if not text:
        return 0

    return max(1, len(text) // 4)


@app.route("/v1/messages/count_tokens", methods=["POST"])
def count_tokens():
    body = request.get_json(force=True)

    text_parts = []

    system = body.get("system")
    if system:
        text_parts.append(flatten_content(system))

    for msg in body.get("messages", []):
        text_parts.append(flatten_content(msg.get("content", "")))

    joined = "\n".join(text_parts)

    return jsonify({"input_tokens": rough_token_count(joined)})


@app.route("/health", methods=["GET"])
def health():
    return jsonify(
        {
            "status": "ok",
            "target_model": CW_MODEL_NAME,
            "advertised_model": CW_ADVERTISED_MODEL,
            "gateway": CW_GATEWAY_ENDPOINT,
            "tools_enabled": CW_USE_TOOLS,
            "project": CW_OPENAI_PROJECT,
        }
    )


@app.route("/v1/models", methods=["GET"])
def models():
    return jsonify(
        {
            "data": [
                {
                    "id": CW_ADVERTISED_MODEL,
                    "type": "model",
                    "display_name": CW_ADVERTISED_MODEL,
                    "created_at": "2026-01-01T00:00:00Z",
                    "max_input_tokens": 0,
                    "max_tokens": 0,
                    "capabilities": {
                        "batch": {"supported": False},
                        "citations": {"supported": False},
                        "code_execution": {"supported": False},
                        "context_management": {"supported": False},
                        "effort": {"supported": False},
                        "image_input": {"supported": True},
                        "pdf_input": {"supported": False},
                        "structured_outputs": {"supported": True},
                        "thinking": {"supported": True},
                    },
                }
            ],
            "first_id": CW_ADVERTISED_MODEL,
            "has_more": False,
            "last_id": CW_ADVERTISED_MODEL,
        }
    )


if __name__ == "__main__":
    for rule in app.url_map.iter_rules():
        print(rule)
    app.run(port=4000)