codex-lb 0.3.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

app/core/clients/proxy.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from typing import AsyncIterator, Mapping
+from typing import AsyncIterator, Mapping, Protocol, TypeAlias
 
 import aiohttp
 
@@ -28,6 +28,18 @@ class StreamIdleTimeoutError(Exception):
     pass
 
 
+class ErrorResponseProtocol(Protocol):
+    status: int
+    reason: str | None
+
+    async def json(self, *, content_type: str | None = None) -> object: ...
+
+    async def text(self, *, encoding: str | None = None, errors: str = "strict") -> str: ...
+
+
+ErrorResponse: TypeAlias = aiohttp.ClientResponse | ErrorResponseProtocol
+
+
 class ProxyResponseError(Exception):
     def __init__(self, status_code: int, payload: OpenAIErrorEnvelope) -> None:
         super().__init__(f"Proxy response error ({status_code})")
@@ -88,8 +100,10 @@ async def _iter_sse_lines(
         yield line
 
 
-async def _error_event_from_response(resp: aiohttp.ClientResponse) -> ResponseFailedEvent:
+async def _error_event_from_response(resp: ErrorResponse) -> ResponseFailedEvent:
     fallback_message = f"Upstream error: HTTP {resp.status}"
+    if resp.reason:
+        fallback_message += f" {resp.reason}"
     try:
         data = await resp.json(content_type=None)
     except Exception:
@@ -112,11 +126,16 @@ async def _error_event_from_response(resp: aiohttp.ClientResponse) -> ResponseFa
                 if key in payload:
                     event["response"]["error"][key] = payload[key]
             return event
+        message = _extract_upstream_message(data)
+        if message:
+            return response_failed_event("upstream_error", message, response_id=get_request_id())
     return response_failed_event("upstream_error", fallback_message, response_id=get_request_id())
 
 
-async def _error_payload_from_response(resp: aiohttp.ClientResponse) -> OpenAIErrorEnvelope:
+async def _error_payload_from_response(resp: ErrorResponse) -> OpenAIErrorEnvelope:
     fallback_message = f"Upstream error: HTTP {resp.status}"
+    if resp.reason:
+        fallback_message += f" {resp.reason}"
     try:
         data = await resp.json(content_type=None)
     except Exception:
@@ -128,9 +147,20 @@ async def _error_payload_from_response(resp: aiohttp.ClientResponse) -> OpenAIEr
         error = parse_error_payload(data)
         if error:
             return {"error": error.model_dump(exclude_none=True)}
+        message = _extract_upstream_message(data)
+        if message:
+            return openai_error("upstream_error", message)
     return openai_error("upstream_error", fallback_message)
 
 
+def _extract_upstream_message(data: dict) -> str | None:
+    for key in ("message", "detail", "error"):
+        value = data.get(key)
+        if isinstance(value, str) and value.strip():
+            return value
+    return None
+
+
 async def stream_responses(
     payload: ResponsesRequest,
     headers: Mapping[str, str],

app/core/config/settings.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from functools import lru_cache
 from pathlib import Path
 
-from pydantic import field_validator
+from pydantic import Field, field_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 BASE_DIR = Path(__file__).resolve().parents[3]
@@ -22,6 +22,9 @@ class Settings(BaseSettings):
     )
 
     database_url: str = f"sqlite+aiosqlite:///{DEFAULT_DB_PATH}"
+    database_pool_size: int = Field(default=15, gt=0)
+    database_max_overflow: int = Field(default=10, ge=0)
+    database_pool_timeout_seconds: float = Field(default=30.0, gt=0)
     upstream_base_url: str = "https://chatgpt.com/backend-api"
     upstream_connect_timeout_seconds: float = 30.0
     stream_idle_timeout_seconds: float = 300.0
@@ -42,24 +45,22 @@ class Settings(BaseSettings):
     database_migrations_fail_fast: bool = True
     log_proxy_request_shape: bool = False
     log_proxy_request_shape_raw_cache_key: bool = False
+    log_proxy_request_payload: bool = False
+    max_decompressed_body_bytes: int = Field(default=32 * 1024 * 1024, gt=0)
 
     @field_validator("database_url")
     @classmethod
-    def _normalize_database_url(cls, value: str) -> str:
-        if not isinstance(value, str):
-            return value
-
+    def _expand_database_url(cls, value: str) -> str:
         for prefix in ("sqlite+aiosqlite:///", "sqlite:///"):
             if value.startswith(prefix):
                 path = value[len(prefix) :]
                 if path.startswith("~"):
-                    expanded = str(Path(path).expanduser())
-                    return f"{prefix}{expanded}"
+                    return f"{prefix}{Path(path).expanduser()}"
         return value
 
     @field_validator("encryption_key_file", mode="before")
     @classmethod
-    def _normalize_encryption_key_file(cls, value: object) -> Path:
+    def _expand_encryption_key_file(cls, value: str | Path) -> Path:
         if isinstance(value, Path):
             return value.expanduser()
         if isinstance(value, str):

app/core/handlers/__init__.py

@@ -0,0 +1,3 @@
+from app.core.handlers.exceptions import add_exception_handlers
+
+__all__ = ["add_exception_handlers"]

app/core/handlers/exceptions.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from fastapi import FastAPI, Request
+from fastapi.exception_handlers import (
+    http_exception_handler,
+    request_validation_exception_handler,
+)
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse, Response
+from starlette.exceptions import HTTPException as StarletteHTTPException
+
+from app.core.errors import dashboard_error
+
+
+def add_exception_handlers(app: FastAPI) -> None:
+    @app.exception_handler(RequestValidationError)
+    async def validation_error_handler(
+        request: Request,
+        exc: RequestValidationError,
+    ) -> Response:
+        if request.url.path.startswith("/api/"):
+            return JSONResponse(
+                status_code=422,
+                content=dashboard_error("validation_error", "Invalid request payload"),
+            )
+        return await request_validation_exception_handler(request, exc)
+
+    @app.exception_handler(StarletteHTTPException)
+    async def http_error_handler(
+        request: Request,
+        exc: StarletteHTTPException,
+    ) -> Response:
+        if request.url.path.startswith("/api/"):
+            detail = exc.detail if isinstance(exc.detail, str) else "Request failed"
+            return JSONResponse(
+                status_code=exc.status_code,
+                content=dashboard_error(f"http_{exc.status_code}", detail),
+            )
+        return await http_exception_handler(request, exc)

app/core/middleware/__init__.py

@@ -0,0 +1,9 @@
+from app.core.middleware.api_errors import add_api_unhandled_error_middleware
+from app.core.middleware.request_decompression import add_request_decompression_middleware
+from app.core.middleware.request_id import add_request_id_middleware
+
+__all__ = [
+    "add_api_unhandled_error_middleware",
+    "add_request_decompression_middleware",
+    "add_request_id_middleware",
+]

app/core/middleware/api_errors.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Awaitable, Callable
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, Response
+
+from app.core.errors import dashboard_error
+from app.core.utils.request_id import get_request_id
+
+logger = logging.getLogger(__name__)
+
+
+def add_api_unhandled_error_middleware(app: FastAPI) -> None:
+    @app.middleware("http")
+    async def api_unhandled_error_middleware(
+        request: Request,
+        call_next: Callable[[Request], Awaitable[Response]],
+    ) -> Response:
+        try:
+            return await call_next(request)
+        except Exception:
+            if request.url.path.startswith("/api/"):
+                logger.exception(
+                    "Unhandled API error request_id=%s",
+                    get_request_id(),
+                )
+                return JSONResponse(
+                    status_code=500,
+                    content=dashboard_error("internal_error", "Unexpected error"),
+                )
+            raise

app/core/middleware/request_decompression.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import io
+from collections.abc import Awaitable, Callable
+from typing import Protocol
+
+import zstandard as zstd
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, Response
+
+from app.core.config.settings import get_settings
+from app.core.errors import dashboard_error
+
+
+class _DecompressedBodyTooLarge(Exception):
+    def __init__(self, max_size: int) -> None:
+        super().__init__(f"Decompressed body exceeded {max_size} bytes")
+        self.max_size = max_size
+
+
+class _Readable(Protocol):
+    def read(self, size: int = ...) -> bytes: ...
+
+
+def _read_limited(reader: _Readable, max_size: int) -> bytes:
+    buffer = bytearray()
+    total = 0
+    chunk_size = 64 * 1024
+    while True:
+        chunk = reader.read(chunk_size)
+        if not chunk:
+            break
+        total += len(chunk)
+        if total > max_size:
+            raise _DecompressedBodyTooLarge(max_size)
+        buffer.extend(chunk)
+    return bytes(buffer)
+
+
+def _replace_request_body(request: Request, body: bytes) -> None:
+    request._body = body
+    headers: list[tuple[bytes, bytes]] = []
+    for key, value in request.scope.get("headers", []):
+        if key.lower() in (b"content-encoding", b"content-length"):
+            continue
+        headers.append((key, value))
+    headers.append((b"content-length", str(len(body)).encode("ascii")))
+    request.scope["headers"] = headers
+    # Ensure subsequent request.headers reflects the updated scope headers.
+    request._headers = None
+
+
+def add_request_decompression_middleware(app: FastAPI) -> None:
+    @app.middleware("http")
+    async def request_decompression_middleware(
+        request: Request,
+        call_next: Callable[[Request], Awaitable[Response]],
+    ) -> Response:
+        content_encoding = request.headers.get("content-encoding")
+        if not content_encoding:
+            return await call_next(request)
+        encodings = [enc.strip().lower() for enc in content_encoding.split(",") if enc.strip()]
+        if encodings != ["zstd"]:
+            return await call_next(request)
+        body = await request.body()
+        settings = get_settings()
+        max_size = settings.max_decompressed_body_bytes
+        try:
+            decompressed = zstd.ZstdDecompressor().decompress(body, max_output_size=max_size)
+            if len(decompressed) > max_size:
+                raise _DecompressedBodyTooLarge(max_size)
+        except _DecompressedBodyTooLarge:
+            return JSONResponse(
+                status_code=413,
+                content=dashboard_error(
+                    "payload_too_large",
+                    "Request body exceeds the maximum allowed size",
+                ),
+            )
+        except Exception:
+            try:
+                with zstd.ZstdDecompressor().stream_reader(io.BytesIO(body)) as reader:
+                    decompressed = _read_limited(reader, max_size)
+            except _DecompressedBodyTooLarge:
+                return JSONResponse(
+                    status_code=413,
+                    content=dashboard_error(
+                        "payload_too_large",
+                        "Request body exceeds the maximum allowed size",
+                    ),
+                )
+            except Exception:
+                return JSONResponse(
+                    status_code=400,
+                    content=dashboard_error(
+                        "invalid_request",
+                        "Request body is zstd-compressed but could not be decompressed",
+                    ),
+                )
+        _replace_request_body(request, decompressed)
+        return await call_next(request)

app/core/middleware/request_id.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from uuid import uuid4
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+
+from app.core.utils.request_id import reset_request_id, set_request_id
+
+
+def add_request_id_middleware(app: FastAPI) -> None:
+    @app.middleware("http")
+    async def request_id_middleware(
+        request: Request,
+        call_next: Callable[[Request], Awaitable[JSONResponse]],
+    ) -> JSONResponse:
+        inbound_request_id = request.headers.get("x-request-id") or request.headers.get("request-id")
+        request_id = inbound_request_id or str(uuid4())
+        token = set_request_id(request_id)
+        try:
+            response = await call_next(request)
+        except Exception:
+            reset_request_id(token)
+            raise
+        response.headers.setdefault("x-request-id", request_id)
+        return response

app/core/openai/chat_requests.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import cast
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from app.core.openai.message_coercion import coerce_messages
+from app.core.openai.requests import ResponsesRequest, ResponsesTextControls, ResponsesTextFormat
+from app.core.types import JsonValue
+
+
+class ChatCompletionsRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    model: str = Field(min_length=1)
+    messages: list[dict[str, JsonValue]]
+    tools: list[JsonValue] = Field(default_factory=list)
+    tool_choice: str | dict[str, JsonValue] | None = None
+    parallel_tool_calls: bool | None = None
+    stream: bool | None = None
+    temperature: float | None = None
+    top_p: float | None = None
+    stop: str | list[str] | None = None
+    n: int | None = None
+    presence_penalty: float | None = None
+    frequency_penalty: float | None = None
+    logprobs: bool | None = None
+    top_logprobs: int | None = None
+    seed: int | None = None
+    response_format: JsonValue | None = None
+    max_tokens: int | None = None
+    max_completion_tokens: int | None = None
+    store: bool | None = None
+
+    @model_validator(mode="after")
+    def _validate_messages(self) -> "ChatCompletionsRequest":
+        if not self.messages:
+            raise ValueError("'messages' must be a non-empty list.")
+        return self
+
+    def to_responses_request(self) -> ResponsesRequest:
+        data = self.model_dump(mode="json", exclude_none=True)
+        messages = data.pop("messages")
+        data.pop("store", None)
+        data.pop("max_tokens", None)
+        data.pop("max_completion_tokens", None)
+        response_format = data.pop("response_format", None)
+        tools = _normalize_chat_tools(data.pop("tools", []))
+        tool_choice = _normalize_tool_choice(data.pop("tool_choice", None))
+        reasoning_effort = data.pop("reasoning_effort", None)
+        if reasoning_effort is not None and "reasoning" not in data:
+            data["reasoning"] = {"effort": reasoning_effort}
+        if response_format is not None:
+            _apply_response_format(data, response_format)
+        instructions, input_items = coerce_messages("", messages)
+        data["instructions"] = instructions
+        data["input"] = input_items
+        data["tools"] = tools
+        if tool_choice is not None:
+            data["tool_choice"] = tool_choice
+        return ResponsesRequest.model_validate(data)
+
+
+class ChatResponseFormatJsonSchema(BaseModel):
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
+
+    name: str | None = None
+    schema_: JsonValue | None = Field(default=None, alias="schema")
+    strict: bool | None = None
+
+
+class ChatResponseFormat(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    type: str = Field(min_length=1)
+    json_schema: ChatResponseFormatJsonSchema | None = None
+
+    @model_validator(mode="after")
+    def _validate_schema(self) -> "ChatResponseFormat":
+        if self.type == "json_schema" and self.json_schema is None:
+            raise ValueError("'response_format.json_schema' is required when type is 'json_schema'.")
+        return self
+
+
+def _normalize_chat_tools(tools: list[JsonValue]) -> list[JsonValue]:
+    normalized: list[JsonValue] = []
+    for tool in tools:
+        if not isinstance(tool, dict):
+            continue
+        tool_type = tool.get("type")
+        function = tool.get("function")
+        if isinstance(function, dict):
+            name = function.get("name")
+            if not isinstance(name, str) or not name:
+                continue
+            normalized.append(
+                {
+                    "type": tool_type or "function",
+                    "name": name,
+                    "description": function.get("description"),
+                    "parameters": function.get("parameters"),
+                }
+            )
+            continue
+        name = tool.get("name")
+        if isinstance(name, str) and name:
+            normalized.append(tool)
+    return normalized
+
+
+def _normalize_tool_choice(tool_choice: JsonValue | None) -> JsonValue | None:
+    if not isinstance(tool_choice, dict):
+        return tool_choice
+    tool_type = tool_choice.get("type")
+    function = tool_choice.get("function")
+    if isinstance(function, dict):
+        name = function.get("name")
+        if isinstance(name, str) and name:
+            return {"type": tool_type or "function", "name": name}
+    return tool_choice
+
+
+def _apply_response_format(data: dict[str, JsonValue], response_format: JsonValue) -> None:
+    text_controls = _parse_text_controls(data.get("text"))
+    if text_controls is None:
+        text_controls = ResponsesTextControls()
+    if text_controls.format is not None:
+        raise ValueError("Provide either 'response_format' or 'text.format', not both.")
+    text_controls.format = _response_format_to_text_format(response_format)
+    data["text"] = cast(JsonValue, text_controls.model_dump(mode="json", exclude_none=True))
+
+
+def _parse_text_controls(text: JsonValue | None) -> ResponsesTextControls | None:
+    if text is None:
+        return None
+    if not isinstance(text, Mapping):
+        raise ValueError("'text' must be an object when using 'response_format'.")
+    return ResponsesTextControls.model_validate(text)
+
+
+def _response_format_to_text_format(response_format: JsonValue) -> ResponsesTextFormat:
+    if isinstance(response_format, str):
+        return _text_format_from_type(response_format)
+    if isinstance(response_format, Mapping):
+        parsed = ChatResponseFormat.model_validate(response_format)
+        return _text_format_from_parsed(parsed)
+    raise ValueError("'response_format' must be a string or object.")
+
+
+def _text_format_from_type(format_type: str) -> ResponsesTextFormat:
+    if format_type in ("json_object", "text"):
+        return ResponsesTextFormat(type=format_type)
+    if format_type == "json_schema":
+        raise ValueError("'response_format' must include 'json_schema' when type is 'json_schema'.")
+    raise ValueError(f"Unsupported response_format.type: {format_type}")
+
+
+def _text_format_from_parsed(parsed: ChatResponseFormat) -> ResponsesTextFormat:
+    if parsed.type == "json_schema":
+        json_schema = parsed.json_schema
+        if json_schema is None:
+            raise ValueError("'response_format.json_schema' is required when type is 'json_schema'.")
+        return ResponsesTextFormat(
+            type=parsed.type,
+            schema_=json_schema.schema_,
+            name=json_schema.name,
+            strict=json_schema.strict,
+        )
+    if parsed.type in ("json_object", "text"):
+        return ResponsesTextFormat(type=parsed.type)
+    raise ValueError(f"Unsupported response_format.type: {parsed.type}")