proxy.py

#!/usr/bin/env python3
import os
import json
import requests
from flask import Flask, request, Response, stream_with_context

app = Flask(__name__)

# All your keys
GEMINI_KEY_1 = os.environ.get("GEMINI_API_KEY", "")
GEMINI_KEY_2 = os.environ.get("GEMINI_CODER_API_KEY", "") or GEMINI_KEY_1
GEMINI_KEY_3 = os.environ.get("GEMINI_ARCHITECT_API_KEY", "") or GEMINI_KEY_1
GEMINI_KEY_4 = os.environ.get("GEMINI_BUGFIXER_API_KEY", "") or GEMINI_KEY_1
CEREBRAS_KEY = os.environ.get("CEREBRAS_API_KEY", "")

CEREBRAS_URL = "https://api.cerebras.ai/v1/chat/completions"
GEMINI_URL = "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions"

# Each Claude model name routes to a different free provider + key
# In Paperclip UI: pick the Claude model name for each agent
ROUTING = {
    # CEO — Cerebras, fast decisions, separate key
    "claude-haiku-4-5": {
        "url": CEREBRAS_URL,
        "model": "llama3.3-70b",
        "key": CEREBRAS_KEY,
        "label": "Cerebras - CEO"
    },
    # Researcher + DevOps — Gemini Key 1
    "claude-sonnet-4-6": {
        "url": GEMINI_URL,
        "model": "gemini-2.0-flash",
        "key": GEMINI_KEY_1,
        "label": "Gemini Flash - Researcher/DevOps"
    },
    # Coder — Gemini Key 2 (dedicated, 2.5 Pro)
    "claude-opus-4-6": {
        "url": GEMINI_URL,
        "model": "gemini-2.5-pro-exp-03-25",
        "key": GEMINI_KEY_2,
        "label": "Gemini Pro - Coder"
    },
    # Architect — Gemini Key 3 (dedicated, 2.5 Pro)
    "claude-sonnet-4-5": {
        "url": GEMINI_URL,
        "model": "gemini-2.5-pro-exp-03-25",
        "key": GEMINI_KEY_3,
        "label": "Gemini Pro - Architect"
    },
    # Bug Fixer — Gemini Key 4 (dedicated, 2.5 Pro)
    "claude-haiku-4-6": {
        "url": GEMINI_URL,
        "model": "gemini-2.5-pro-exp-03-25",
        "key": GEMINI_KEY_4,
        "label": "Gemini Pro - Bug Fixer"
    },
}

DEFAULT_ROUTE = ROUTING["claude-sonnet-4-6"]

def anthropic_to_openai(data, model):
    messages = []
    system = data.get("system", "")
    if system:
        if isinstance(system, list):
            system = " ".join([s.get("text", "") for s in system if isinstance(s, dict)])
        messages.append({"role": "system", "content": str(system)})
    for msg in data.get("messages", []):
        role = msg["role"]
        content = msg["content"]
        if isinstance(content, list):
            content = " ".join([
                c.get("text", "") for c in content
                if isinstance(c, dict) and c.get("type") == "text"
            ])
        messages.append({"role": role, "content": str(content)})
    return {
        "model": model,
        "messages": messages,
        "max_tokens": min(data.get("max_tokens", 8096), 8096),
        "stream": False,
        "temperature": 0.7
    }

@app.route("/v1/messages", methods=["POST"])
def messages():
    data = request.json
    requested = data.get("model", "claude-sonnet-4-6")
    route = ROUTING.get(requested, DEFAULT_ROUTE)

    print(f"[proxy] {requested} → {route['label']} ({route['model']})", flush=True)

    payload = anthropic_to_openai(data, route["model"])
    headers = {
        "Authorization": f"Bearer {route['key']}",
        "Content-Type": "application/json"
    }

    try:
        resp = requests.post(route["url"], json=payload, headers=headers, timeout=120)
        print(f"[proxy] status={resp.status_code}", flush=True)

        if resp.status_code != 200:
            print(f"[proxy] error: {resp.text[:300]}", flush=True)
            return {
                "type": "error",
                "error": {"type": "api_error", "message": resp.text}
            }, 500

        result = resp.json()
        text = result.get("choices", [{}])[0].get("message", {}).get("content", "")
        print(f"[proxy] got {len(text)} chars", flush=True)

        usage = result.get("usage", {})
        anthropic_resp = {
            "id": result.get("id", "msg_proxy"),
            "type": "message",
            "role": "assistant",
            "content": [{"type": "text", "text": text}],
            "model": requested,
            "stop_reason": "end_turn",
            "stop_sequence": None,
            "usage": {
                "input_tokens": usage.get("prompt_tokens", 0),
                "output_tokens": usage.get("completion_tokens", 0)
            }
        }

        if data.get("stream"):
            def generate():
                yield f'event: message_start\ndata: {json.dumps({"type":"message_start","message":anthropic_resp})}\n\n'
                yield f'event: content_block_start\ndata: {json.dumps({"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}})}\n\n'
                chunk_size = 200
                for i in range(0, len(text), chunk_size):
                    chunk = text[i:i+chunk_size]
                    yield f'event: content_block_delta\ndata: {json.dumps({"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":chunk}})}\n\n'
                yield f'event: content_block_stop\ndata: {json.dumps({"type":"content_block_stop","index":0})}\n\n'
                yield f'event: message_delta\ndata: {json.dumps({"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":None},"usage":{"output_tokens":usage.get("completion_tokens",0)}})}\n\n'
                yield f'event: message_stop\ndata: {json.dumps({"type":"message_stop"})}\n\n'
            return Response(stream_with_context(generate()), mimetype="text/event-stream")

        return anthropic_resp

    except Exception as e:
        print(f"[proxy] exception: {e}", flush=True)
        return {
            "type": "error",
            "error": {"type": "api_error", "message": str(e)}
        }, 500

@app.route("/v1/models", methods=["GET"])
def models():
    return {"data": [
        {"id": "claude-haiku-4-5",  "object": "model"},
        {"id": "claude-sonnet-4-6", "object": "model"},
        {"id": "claude-opus-4-6",   "object": "model"},
        {"id": "claude-sonnet-4-5", "object": "model"},
        {"id": "claude-haiku-4-6",  "object": "model"},
    ]}

@app.route("/", methods=["GET"])
@app.route("/health", methods=["GET"])
def health():
    return {
        "status": "ok",
        "cerebras": "set" if CEREBRAS_KEY else "missing",
        "gemini_1": "set" if GEMINI_KEY_1 else "missing",
        "gemini_2": "set" if GEMINI_KEY_2 else "missing",
        "gemini_3": "set" if GEMINI_KEY_3 else "missing",
        "gemini_4": "set" if GEMINI_KEY_4 else "missing",
    }

if __name__ == "__main__":
    print(f"[proxy] Cerebras={'SET' if CEREBRAS_KEY else 'MISSING'}", flush=True)
    print(f"[proxy] Gemini1={'SET' if GEMINI_KEY_1 else 'MISSING'}", flush=True)
    print(f"[proxy] Gemini2={'SET' if GEMINI_KEY_2 else 'MISSING'}", flush=True)
    print(f"[proxy] Gemini3={'SET' if GEMINI_KEY_3 else 'MISSING'}", flush=True)
    print(f"[proxy] Gemini4={'SET' if GEMINI_KEY_4 else 'MISSING'}", flush=True)
    app.run(host="0.0.0.0", port=8082, debug=False)