main.py

import logging
import os
import time
from contextlib import asynccontextmanager

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

from config import (
    UPSTREAM_URL, SUPPORTED_MODELS, DEFAULT_MODEL,
    UPSTREAM_TIMEOUT_SECS, UPSTREAM_CONNECT_TIMEOUT_SECS,
    UPSTREAM_MAX_CONNECTIONS, UPSTREAM_MAX_KEEPALIVE_CONNECTIONS,
    UPSTREAM_KEEPALIVE_EXPIRY_SECS, UPSTREAM_PROXY_URL,
)
from response_cache import get_cache_service
from translate import (
    pick_model, build_payload,
    build_upstream_messages_anthropic, build_upstream_messages_openai,
)
from upstream import create_upstream_client, close_upstream_client, fetch_completion_artifact
from render import (
    render_anthropic_json_from_artifact, render_openai_json_from_artifact,
    anthropic_stream_from_artifact, openai_stream_from_artifact,
    anthropic_aggregate, openai_aggregate,
)
from stream import (
    anthropic_stream_plain, anthropic_stream_with_tools,
    openai_stream_plain, openai_stream_with_tools,
    build_cache_headers, build_stream_headers,
    wait_for_inflight_artifact, finalize_stream_cache,
)
from upstream import LiveArtifactCapture

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


@asynccontextmanager
async def lifespan(_: FastAPI):
    client = create_upstream_client()
    logger.info("upstream client created (proxy=%s)", bool(UPSTREAM_PROXY_URL))
    try:
        yield
    finally:
        await close_upstream_client()
        logger.info("upstream client closed")


app = FastAPI(title="p5js.ai 2 API (Anthropic + OpenAI)", lifespan=lifespan)


@app.get("/")
@app.get("/health")
async def health():
    cache_service = get_cache_service()
    return {
        "status": "ok",
        "service": "p5js.ai-2api",
        "endpoints": ["/v1/messages", "/v1/chat/completions", "/v1/models"],
        "models": sorted(SUPPORTED_MODELS),
        "tool_use": "pseudo (XML-based, upstream does not support native tool_use)",
        "cache": cache_service.describe(),
        "upstream": {
            "url": UPSTREAM_URL,
            "proxy_configured": bool(UPSTREAM_PROXY_URL or os.getenv("HTTPS_PROXY") or os.getenv("HTTP_PROXY")),
            "timeout_secs": UPSTREAM_TIMEOUT_SECS,
            "connect_timeout_secs": UPSTREAM_CONNECT_TIMEOUT_SECS,
            "max_connections": UPSTREAM_MAX_CONNECTIONS,
            "max_keepalive_connections": UPSTREAM_MAX_KEEPALIVE_CONNECTIONS,
            "keepalive_expiry_secs": UPSTREAM_KEEPALIVE_EXPIRY_SECS,
        },
    }


@app.get("/v1/models")
async def list_models():
    now = int(time.time())
    return {
        "data": [
            {"id": m, "object": "model", "created": now, "owned_by": "anthropic"}
            for m in sorted(SUPPORTED_MODELS)
        ],
        "object": "list",
    }


@app.post("/v1/messages")
async def anthropic_messages(request: Request):
    try:
        body = await request.json()
    except Exception:
        raise HTTPException(status_code=400, detail="invalid json body")

    model = pick_model(body.get("model"))
    stream = bool(body.get("stream", False))
    tools = body.get("tools") or None
    has_tools = bool(tools)
    upstream_messages = build_upstream_messages_anthropic(body.get("system"), body.get("messages", []), tools)
    payload = build_payload(model, upstream_messages)
    cache_service = get_cache_service()
    bypass_cache = cache_service.should_bypass(request.headers)
    if bypass_cache:
        await cache_service.record_bypass()
    live_status = "MISS" if cache_service.config.enabled and not bypass_cache else ("BYPASS" if bypass_cache else "DISABLED")
    cache_key = cache_service.build_key(
        protocol_family="anthropic",
        resolved_model=model,
        upstream_messages=upstream_messages,
        auth_scope=cache_service.auth_scope_from_headers(request.headers),
        has_tools=has_tools,
    )
    ttl_secs = cache_service.config.ttl_for(has_tools)

    if cache_service.config.enabled and not bypass_cache:
        lookup = await cache_service.get(cache_key)
        if lookup.artifact is not None:
            logger.debug("anthropic cache HIT key=%s source=%s", cache_key[:24], lookup.source)
            if stream:
                return StreamingResponse(
                    anthropic_stream_from_artifact(lookup.artifact, has_tools),
                    media_type="text/event-stream",
                    headers=build_stream_headers("HIT", lookup.source),
                )
            return JSONResponse(
                render_anthropic_json_from_artifact(lookup.artifact, has_tools),
                headers=build_cache_headers("HIT", lookup.source),
            )

        is_leader, future = await cache_service.inflight.start(cache_key)
        if not is_leader:
            await cache_service.record_inflight_wait()
            artifact = await wait_for_inflight_artifact(future)
            if artifact is not None:
                if stream:
                    return StreamingResponse(
                        anthropic_stream_from_artifact(artifact, has_tools),
                        media_type="text/event-stream",
                        headers=build_stream_headers("HIT", "inflight"),
                    )
                return JSONResponse(
                    render_anthropic_json_from_artifact(artifact, has_tools),
                    headers=build_cache_headers("HIT", "inflight"),
                )
        else:
            if stream:
                capture = LiveArtifactCapture(model)
                live_stream = anthropic_stream_with_tools(payload, capture) if has_tools else anthropic_stream_plain(payload, capture)
                return StreamingResponse(
                    finalize_stream_cache(live_stream, capture, cache_key, ttl_secs),
                    media_type="text/event-stream",
                    headers=build_stream_headers("MISS", "live"),
                )
            try:
                artifact = await fetch_completion_artifact(payload)
                await cache_service.set(cache_key, artifact, ttl_secs)
                await cache_service.inflight.resolve(cache_key, artifact)
            except Exception as exc:
                await cache_service.inflight.reject(cache_key, exc)
                raise
            return JSONResponse(
                render_anthropic_json_from_artifact(artifact, has_tools),
                headers=build_cache_headers("MISS", "live"),
            )

    if stream:
        gen = anthropic_stream_with_tools(payload) if has_tools else anthropic_stream_plain(payload)
        return StreamingResponse(
            gen,
            media_type="text/event-stream",
            headers=build_stream_headers(live_status, "live"),
        )
    return JSONResponse(
        await anthropic_aggregate(payload, has_tools),
        headers=build_cache_headers(live_status, "live"),
    )


@app.post("/v1/chat/completions")
async def openai_chat_completions(request: Request):
    try:
        body = await request.json()
    except Exception:
        raise HTTPException(status_code=400, detail="invalid json body")

    requested_model = body.get("model") or DEFAULT_MODEL
    model = pick_model(requested_model)
    stream = bool(body.get("stream", False))
    tools = body.get("tools") or None
    has_tools = bool(tools)
    upstream_messages = build_upstream_messages_openai(body.get("messages", []), tools)
    payload = build_payload(model, upstream_messages)
    cache_service = get_cache_service()
    bypass_cache = cache_service.should_bypass(request.headers)
    if bypass_cache:
        await cache_service.record_bypass()
    live_status = "MISS" if cache_service.config.enabled and not bypass_cache else ("BYPASS" if bypass_cache else "DISABLED")
    cache_key = cache_service.build_key(
        protocol_family="openai",
        resolved_model=model,
        upstream_messages=upstream_messages,
        auth_scope=cache_service.auth_scope_from_headers(request.headers),
        has_tools=has_tools,
    )
    ttl_secs = cache_service.config.ttl_for(has_tools)

    if cache_service.config.enabled and not bypass_cache:
        lookup = await cache_service.get(cache_key)
        if lookup.artifact is not None:
            logger.debug("openai cache HIT key=%s source=%s", cache_key[:24], lookup.source)
            if stream:
                return StreamingResponse(
                    openai_stream_from_artifact(lookup.artifact, has_tools),
                    media_type="text/event-stream",
                    headers=build_stream_headers("HIT", lookup.source),
                )
            return JSONResponse(
                render_openai_json_from_artifact(lookup.artifact, has_tools),
                headers=build_cache_headers("HIT", lookup.source),
            )

        is_leader, future = await cache_service.inflight.start(cache_key)
        if not is_leader:
            await cache_service.record_inflight_wait()
            artifact = await wait_for_inflight_artifact(future)
            if artifact is not None:
                if stream:
                    return StreamingResponse(
                        openai_stream_from_artifact(artifact, has_tools),
                        media_type="text/event-stream",
                        headers=build_stream_headers("HIT", "inflight"),
                    )
                return JSONResponse(
                    render_openai_json_from_artifact(artifact, has_tools),
                    headers=build_cache_headers("HIT", "inflight"),
                )
        else:
            if stream:
                capture = LiveArtifactCapture(model)
                live_stream = openai_stream_with_tools(payload, requested_model, capture) if has_tools else openai_stream_plain(payload, requested_model, capture)
                return StreamingResponse(
                    finalize_stream_cache(live_stream, capture, cache_key, ttl_secs),
                    media_type="text/event-stream",
                    headers=build_stream_headers("MISS", "live"),
                )
            try:
                artifact = await fetch_completion_artifact(payload)
                await cache_service.set(cache_key, artifact, ttl_secs)
                await cache_service.inflight.resolve(cache_key, artifact)
            except Exception as exc:
                await cache_service.inflight.reject(cache_key, exc)
                raise
            return JSONResponse(
                render_openai_json_from_artifact(artifact, has_tools),
                headers=build_cache_headers("MISS", "live"),
            )

    if stream:
        gen = openai_stream_with_tools(payload, requested_model) if has_tools else openai_stream_plain(payload, requested_model)
        return StreamingResponse(
            gen,
            media_type="text/event-stream",
            headers=build_stream_headers(live_status, "live"),
        )
    return JSONResponse(
        await openai_aggregate(payload, requested_model, has_tools),
        headers=build_cache_headers(live_status, "live"),
    )


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