codex-api-proxy 0.1.0__py3-none-any.whl

codex_api_proxy/main.py

@@ -0,0 +1,561 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import shutil
+import time
+import uuid
+from collections.abc import AsyncIterator, Awaitable, Callable
+from contextlib import asynccontextmanager
+
+from fastapi import Depends, FastAPI, Header, HTTPException
+from fastapi.responses import StreamingResponse
+
+from . import __version__
+from .app_server_runner import AppServerBusy, AppServerRunError, create_stdio_app_server_pool
+from .codex_runner import CodexRunError, run_codex_exec
+from .config import Settings, SettingsError
+from .prompt import messages_to_prompt
+from .schemas import ChatCompletionChoice, ChatCompletionRequest, ChatCompletionResponse, ChoiceMessage
+
+Runner = Callable[..., Awaitable[str]]
+AppServerRunner = Callable[..., AsyncIterator[str]]
+latency_logger = logging.getLogger("codex_api_proxy.latency")
+uvicorn_logger = logging.getLogger("uvicorn.error")
+
+
+def _elapsed_ms(started_at: float) -> float:
+    return round((time.perf_counter() - started_at) * 1000, 2)
+
+
+def _log_latency(payload: dict[str, object]) -> None:
+    message = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
+    latency_logger.info(message)
+    uvicorn_logger.info("codex_api_proxy.latency %s", message)
+
+
+def _sse_event(data: dict[str, object] | str) -> str:
+    if isinstance(data, str):
+        return f"data: {data}\n\n"
+    return f"data: {json.dumps(data, ensure_ascii=False, separators=(',', ':'))}\n\n"
+
+
+def _stream_chat_completion(
+    *,
+    completion_id: str,
+    created: int,
+    model: str,
+    content: str,
+    request_id: str,
+    request_started_at: float,
+):
+    base = {
+        "id": completion_id,
+        "object": "chat.completion.chunk",
+        "created": created,
+        "model": model,
+    }
+    _log_latency(
+        {
+            "event": "chat_completion_first_sse",
+            "request_id": request_id,
+            "time_to_first_sse_ms": _elapsed_ms(request_started_at),
+        }
+    )
+    yield _sse_event(
+        {
+            **base,
+            "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
+        }
+    )
+    yield _sse_event(
+        {
+            **base,
+            "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": None}],
+        }
+    )
+    yield _sse_event(
+        {
+            **base,
+            "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+        }
+    )
+    yield _sse_event("[DONE]")
+
+
+async def _stream_chat_completion_from_deltas(
+    *,
+    completion_id: str,
+    created: int,
+    model: str,
+    deltas: AsyncIterator[str],
+    request_id: str,
+    request_started_at: float,
+    phases_ms: dict[str, float],
+    on_complete: Callable[[], None] | None = None,
+    on_error: Callable[[], None] | None = None,
+):
+    base = {
+        "id": completion_id,
+        "object": "chat.completion.chunk",
+        "created": created,
+        "model": model,
+    }
+    _log_latency(
+        {
+            "event": "chat_completion_first_sse",
+            "request_id": request_id,
+            "time_to_first_sse_ms": _elapsed_ms(request_started_at),
+        }
+    )
+    yield _sse_event(
+        {
+            **base,
+            "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
+        }
+    )
+    first_content = True
+    started_at = time.perf_counter()
+    try:
+        try:
+            async for delta in deltas:
+                if first_content:
+                    first_content = False
+                    _log_latency(
+                        {
+                            "event": "chat_completion_first_content_sse",
+                            "request_id": request_id,
+                            "time_to_first_content_sse_ms": _elapsed_ms(request_started_at),
+                        }
+                    )
+                yield _sse_event(
+                    {
+                        **base,
+                        "choices": [{"index": 0, "delta": {"content": delta}, "finish_reason": None}],
+                    }
+                )
+        finally:
+            phases_ms["app_server_exec"] = _elapsed_ms(started_at)
+        yield _sse_event(
+            {
+                **base,
+                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+            }
+        )
+        yield _sse_event("[DONE]")
+    except Exception:
+        if on_error:
+            on_error()
+        raise
+    if on_complete:
+        on_complete()
+    _log_latency(
+        {
+            "event": "chat_completion_latency",
+            "request_id": request_id,
+            "status": "ok",
+            "stream": True,
+            "model": model,
+            "engine": "app-server",
+            "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+        }
+    )
+
+
+def create_app(
+    settings: Settings | None = None,
+    runner: Runner = run_codex_exec,
+    app_server_runner: AppServerRunner | None = None,
+) -> FastAPI:
+    active_settings = settings or Settings.from_env()
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        try:
+            yield
+        finally:
+            pool = app.state.app_server_pool
+            if pool is not None:
+                await pool.close()
+
+    app = FastAPI(title="Codex OpenAI-Compatible Proxy", version=__version__, lifespan=lifespan)
+    app.state.settings = active_settings
+    app.state.app_server_pool = None
+    app.state.started_at = time.time()
+    app.state.metrics = {
+        "requests_total": 0,
+        "requests_ok": 0,
+        "requests_error": 0,
+        "errors_by_status": {},
+    }
+    semaphore = asyncio.Semaphore(active_settings.max_concurrency)
+
+    def record_request_start() -> None:
+        app.state.metrics["requests_total"] += 1
+
+    def record_request_ok() -> None:
+        app.state.metrics["requests_ok"] += 1
+
+    def record_request_error(status_code: int) -> None:
+        app.state.metrics["requests_error"] += 1
+        errors_by_status = app.state.metrics["errors_by_status"]
+        key = str(status_code)
+        errors_by_status[key] = errors_by_status.get(key, 0) + 1
+
+    async def default_app_server_runner(**kwargs) -> AsyncIterator[str]:
+        pool = app.state.app_server_pool
+        if pool is None:
+            pool = create_stdio_app_server_pool(
+                codex_bin=active_settings.codex_bin,
+                proxy=active_settings.proxy,
+                codex_home=active_settings.app_server_codex_home,
+                codex_configs=active_settings.codex_configs or [],
+                workers=active_settings.workers,
+                max_queue_size=active_settings.max_queue_size,
+                queue_timeout_seconds=active_settings.queue_timeout_seconds,
+                timeout_seconds=active_settings.request_timeout_seconds,
+            )
+            app.state.app_server_pool = pool
+        async for chunk in pool.stream_completion(**kwargs):
+            yield chunk
+
+    active_app_server_runner = app_server_runner or default_app_server_runner
+
+    async def require_auth(authorization: str | None = Header(default=None)) -> None:
+        if not active_settings.api_key:
+            return
+        expected = f"Bearer {active_settings.api_key}"
+        if authorization != expected:
+            raise HTTPException(status_code=401, detail="Missing or invalid bearer token")
+
+    @app.get("/health")
+    async def health() -> dict[str, object]:
+        return {
+            "status": "ok",
+            "codex_available": shutil.which(active_settings.codex_bin) is not None,
+            "codex_bin": active_settings.codex_bin,
+        }
+
+    @app.get("/ready")
+    async def ready() -> dict[str, object]:
+        codex_available = shutil.which(active_settings.codex_bin) is not None
+        return {
+            "status": "ready" if codex_available else "not_ready",
+            "engine": active_settings.engine,
+            "codex_available": codex_available,
+            "codex_bin": active_settings.codex_bin,
+            "app_server_pool_started": app.state.app_server_pool is not None,
+        }
+
+    @app.get("/metrics")
+    async def metrics() -> dict[str, object]:
+        return {
+            **app.state.metrics,
+            "engine": active_settings.engine,
+            "uptime_seconds": round(time.time() - app.state.started_at, 3),
+            "app_server_pool_started": app.state.app_server_pool is not None,
+        }
+
+    @app.get("/v1/models")
+    async def models(_: None = Depends(require_auth)) -> dict[str, object]:
+        return {
+            "object": "list",
+            "data": [
+                {
+                    "id": "codex-local",
+                    "object": "model",
+                    "created": 0,
+                    "owned_by": "local",
+                }
+            ],
+        }
+
+    @app.post("/v1/chat/completions", response_model=None)
+    async def chat_completions(
+        request: ChatCompletionRequest,
+        _: None = Depends(require_auth),
+    ) -> ChatCompletionResponse | StreamingResponse:
+        record_request_start()
+        request_id = uuid.uuid4().hex
+        request_started_at = time.perf_counter()
+        phases_ms: dict[str, float] = {}
+        metadata = request.metadata or {}
+        cwd_started_at = time.perf_counter()
+        try:
+            cwd = active_settings.resolve_cwd(metadata.get("cwd"))
+        except SettingsError as exc:
+            phases_ms["cwd_resolve"] = _elapsed_ms(cwd_started_at)
+            _log_latency(
+                {
+                    "event": "chat_completion_latency",
+                    "request_id": request_id,
+                    "status": "error",
+                    "error_type": type(exc).__name__,
+                    "http_status": 400,
+                    "stream": request.stream,
+                    "model": request.model,
+                    "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+                }
+            )
+            record_request_error(400)
+            raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+        phases_ms["cwd_resolve"] = _elapsed_ms(cwd_started_at)
+        messages = [message.model_dump() for message in request.messages]
+        prompt_started_at = time.perf_counter()
+        prompt = messages_to_prompt(messages)
+        phases_ms["prompt_build"] = _elapsed_ms(prompt_started_at)
+        queue_started_at = time.perf_counter()
+
+        def record_codex_phase(name: str, elapsed_ms: float) -> None:
+            phases_ms[name] = round(elapsed_ms, 2)
+
+        if active_settings.engine == "app-server":
+            phases_ms["queue_wait"] = _elapsed_ms(queue_started_at)
+            response_started_at = time.perf_counter()
+            completion_id = f"chatcmpl-{request_id}"
+            created = int(time.time())
+            app_server_kwargs = {
+                "cwd": cwd,
+                "prompt": prompt,
+                "model": active_settings.model if request.model == "codex-local" else request.model,
+                "codex_configs": active_settings.codex_configs or [],
+                "ephemeral": active_settings.ephemeral,
+                "timeout_seconds": active_settings.request_timeout_seconds,
+                "latency_callback": record_codex_phase,
+            }
+            if request.stream:
+                phases_ms["response_build"] = _elapsed_ms(response_started_at)
+                return StreamingResponse(
+                    _stream_chat_completion_from_deltas(
+                        completion_id=completion_id,
+                        created=created,
+                        model=request.model,
+                        deltas=active_app_server_runner(**app_server_kwargs),
+                        request_id=request_id,
+                        request_started_at=request_started_at,
+                        phases_ms=phases_ms,
+                        on_complete=record_request_ok,
+                        on_error=lambda: record_request_error(502),
+                    ),
+                    media_type="text/event-stream",
+                )
+
+            app_server_started_at = time.perf_counter()
+            try:
+                chunks = [chunk async for chunk in active_app_server_runner(**app_server_kwargs)]
+            except AppServerBusy as exc:
+                phases_ms["app_server_exec"] = _elapsed_ms(app_server_started_at)
+                _log_latency(
+                    {
+                        "event": "chat_completion_latency",
+                        "request_id": request_id,
+                        "status": "error",
+                        "error_type": type(exc).__name__,
+                        "http_status": 429,
+                        "stream": request.stream,
+                        "model": request.model,
+                        "engine": "app-server",
+                        "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+                    }
+                )
+                record_request_error(429)
+                raise HTTPException(status_code=429, detail=str(exc)) from exc
+            except TimeoutError as exc:
+                phases_ms["app_server_exec"] = _elapsed_ms(app_server_started_at)
+                _log_latency(
+                    {
+                        "event": "chat_completion_latency",
+                        "request_id": request_id,
+                        "status": "error",
+                        "error_type": type(exc).__name__,
+                        "http_status": 504,
+                        "stream": request.stream,
+                        "model": request.model,
+                        "engine": "app-server",
+                        "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+                    }
+                )
+                record_request_error(504)
+                raise HTTPException(status_code=504, detail=str(exc)) from exc
+            except AppServerRunError as exc:
+                phases_ms["app_server_exec"] = _elapsed_ms(app_server_started_at)
+                _log_latency(
+                    {
+                        "event": "chat_completion_latency",
+                        "request_id": request_id,
+                        "status": "error",
+                        "error_type": type(exc).__name__,
+                        "http_status": 502,
+                        "stream": request.stream,
+                        "model": request.model,
+                        "engine": "app-server",
+                        "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+                    }
+                )
+                record_request_error(502)
+                raise HTTPException(status_code=502, detail=str(exc)) from exc
+            content = "".join(chunks)
+            phases_ms["app_server_exec"] = _elapsed_ms(app_server_started_at)
+            response = ChatCompletionResponse(
+                id=completion_id,
+                created=created,
+                model=request.model,
+                choices=[ChatCompletionChoice(message=ChoiceMessage(content=content))],
+            )
+            phases_ms["response_build"] = _elapsed_ms(response_started_at)
+            _log_latency(
+                {
+                    "event": "chat_completion_latency",
+                    "request_id": request_id,
+                    "status": "ok",
+                    "stream": False,
+                    "model": request.model,
+                    "engine": "app-server",
+                    "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+                }
+            )
+            record_request_ok()
+            return response
+
+        if active_settings.engine != "exec":
+            record_request_error(400)
+            raise HTTPException(status_code=400, detail=f"unsupported engine: {active_settings.engine}")
+
+        try:
+            async with semaphore:
+                phases_ms["queue_wait"] = _elapsed_ms(queue_started_at)
+                codex_started_at = time.perf_counter()
+                content = await runner(
+                    codex_bin=active_settings.codex_bin,
+                    cwd=cwd,
+                    prompt=prompt,
+                    timeout_seconds=active_settings.request_timeout_seconds,
+                    proxy=active_settings.proxy,
+                    model=active_settings.model,
+                    codex_configs=active_settings.codex_configs or [],
+                    ephemeral=active_settings.ephemeral,
+                    latency_callback=record_codex_phase,
+                )
+                phases_ms["codex_exec"] = _elapsed_ms(codex_started_at)
+        except SettingsError as exc:
+            phases_ms.setdefault("queue_wait", _elapsed_ms(queue_started_at))
+            _log_latency(
+                {
+                    "event": "chat_completion_latency",
+                    "request_id": request_id,
+                    "status": "error",
+                    "error_type": type(exc).__name__,
+                    "http_status": 400,
+                    "stream": request.stream,
+                    "model": request.model,
+                    "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+                }
+            )
+            record_request_error(400)
+            raise HTTPException(status_code=400, detail=str(exc)) from exc
+        except TimeoutError as exc:
+            phases_ms.setdefault("queue_wait", _elapsed_ms(queue_started_at))
+            _log_latency(
+                {
+                    "event": "chat_completion_latency",
+                    "request_id": request_id,
+                    "status": "error",
+                    "error_type": type(exc).__name__,
+                    "http_status": 504,
+                    "stream": request.stream,
+                    "model": request.model,
+                    "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+                }
+            )
+            record_request_error(504)
+            raise HTTPException(status_code=504, detail=str(exc)) from exc
+        except FileNotFoundError as exc:
+            detail = f"codex executable not found: {active_settings.codex_bin}"
+            phases_ms.setdefault("queue_wait", _elapsed_ms(queue_started_at))
+            _log_latency(
+                {
+                    "event": "chat_completion_latency",
+                    "request_id": request_id,
+                    "status": "error",
+                    "error_type": type(exc).__name__,
+                    "http_status": 503,
+                    "stream": request.stream,
+                    "model": request.model,
+                    "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+                }
+            )
+            record_request_error(503)
+            raise HTTPException(status_code=503, detail=detail) from exc
+        except CodexRunError as exc:
+            phases_ms.setdefault("queue_wait", _elapsed_ms(queue_started_at))
+            _log_latency(
+                {
+                    "event": "chat_completion_latency",
+                    "request_id": request_id,
+                    "status": "error",
+                    "error_type": type(exc).__name__,
+                    "http_status": 502,
+                    "stream": request.stream,
+                    "model": request.model,
+                    "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+                }
+            )
+            record_request_error(502)
+            raise HTTPException(status_code=502, detail=str(exc)) from exc
+
+        response_started_at = time.perf_counter()
+        completion_id = f"chatcmpl-{request_id}"
+        created = int(time.time())
+        if request.stream:
+            phases_ms["response_build"] = _elapsed_ms(response_started_at)
+            _log_latency(
+                {
+                    "event": "chat_completion_latency",
+                    "request_id": request_id,
+                    "status": "ok",
+                    "stream": True,
+                    "model": request.model,
+                    "engine": "exec",
+                    "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+                }
+            )
+            record_request_ok()
+            return StreamingResponse(
+                _stream_chat_completion(
+                    completion_id=completion_id,
+                    created=created,
+                    model=request.model,
+                    content=content,
+                    request_id=request_id,
+                    request_started_at=request_started_at,
+                ),
+                media_type="text/event-stream",
+            )
+
+        response = ChatCompletionResponse(
+            id=completion_id,
+            created=created,
+            model=request.model,
+            choices=[ChatCompletionChoice(message=ChoiceMessage(content=content))],
+        )
+        phases_ms["response_build"] = _elapsed_ms(response_started_at)
+        _log_latency(
+            {
+                "event": "chat_completion_latency",
+                "request_id": request_id,
+                "status": "ok",
+                "stream": False,
+                "model": request.model,
+                "engine": "exec",
+                "phases_ms": {**phases_ms, "total": _elapsed_ms(request_started_at)},
+            }
+        )
+        record_request_ok()
+        return response
+
+    return app
+
+
+app = create_app()

codex_api_proxy/prompt.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def _content_to_text(content: Any) -> str:
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts: list[str] = []
+        for item in content:
+            if isinstance(item, dict) and item.get("type") in {"text", "input_text"}:
+                parts.append(str(item.get("text", "")))
+            elif isinstance(item, dict):
+                item_type = item.get("type", "unknown")
+                parts.append(f"[{item_type} omitted by codex-api-proxy v1]")
+            else:
+                parts.append(str(item))
+        return "\n".join(part for part in parts if part)
+    return str(content)
+
+
+def messages_to_prompt(messages: list[dict[str, Any]]) -> str:
+    blocks: list[str] = []
+    for message in messages:
+        role = str(message.get("role", "user")).strip() or "user"
+        content = _content_to_text(message.get("content"))
+        blocks.append(f"[{role}]\n{content}")
+    return "\n\n".join(blocks) + "\n"

codex_api_proxy/schemas.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+
+class ChatMessage(BaseModel):
+    role: str
+    content: Any
+
+
+class ChatCompletionRequest(BaseModel):
+    model: str = "codex-local"
+    messages: list[ChatMessage] = Field(min_length=1)
+    stream: bool = False
+    temperature: float | None = None
+    top_p: float | None = None
+    max_tokens: int | None = None
+    presence_penalty: float | None = None
+    frequency_penalty: float | None = None
+    metadata: dict[str, Any] | None = None
+
+
+class ChoiceMessage(BaseModel):
+    role: Literal["assistant"] = "assistant"
+    content: str
+
+
+class ChatCompletionChoice(BaseModel):
+    index: int = 0
+    message: ChoiceMessage
+    finish_reason: str = "stop"
+
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: Literal["chat.completion"] = "chat.completion"
+    created: int
+    model: str
+    choices: list[ChatCompletionChoice]
+    usage: dict[str, int] = Field(
+        default_factory=lambda: {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "total_tokens": 0,
+        }
+    )