codex-lb 0.1.2__py3-none-any.whl

app/core/clients/usage.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+
+import aiohttp
+from aiohttp_retry import ExponentialRetry, RetryClient
+from pydantic import BaseModel, ConfigDict, ValidationError
+
+from app.core.clients.http import get_http_client
+from app.core.config.settings import get_settings
+from app.core.types import JsonObject
+from app.core.usage.models import UsagePayload
+from app.core.utils.request_id import get_request_id
+
+RETRYABLE_STATUS = {408, 429, 500, 502, 503, 504}
+RETRY_START_TIMEOUT = 0.5
+RETRY_MAX_TIMEOUT = 2.0
+
+logger = logging.getLogger(__name__)
+
+
+class UsageErrorDetail(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    message: str | None = None
+    error_description: str | None = None
+
+
+class UsageErrorEnvelope(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    error: UsageErrorDetail | str | None = None
+    error_description: str | None = None
+    message: str | None = None
+
+
+class UsageFetchError(Exception):
+    def __init__(self, status_code: int, message: str) -> None:
+        super().__init__(message)
+        self.status_code = status_code
+        self.message = message
+
+
+async def fetch_usage(
+    *,
+    access_token: str,
+    account_id: str | None,
+    base_url: str | None = None,
+    timeout_seconds: float | None = None,
+    max_retries: int | None = None,
+    client: RetryClient | None = None,
+) -> UsagePayload:
+    settings = get_settings()
+    usage_base = base_url or settings.upstream_base_url
+    url = _usage_url(usage_base)
+    timeout = aiohttp.ClientTimeout(total=timeout_seconds or settings.usage_fetch_timeout_seconds)
+    retries = max_retries if max_retries is not None else settings.usage_fetch_max_retries
+    headers = _usage_headers(access_token, account_id)
+    retry_client = client or get_http_client().retry_client
+    retry_options = _retry_options(retries + 1)
+
+    try:
+        async with retry_client.request(
+            "GET",
+            url,
+            headers=headers,
+            timeout=timeout,
+            retry_options=retry_options,
+        ) as resp:
+            data = await _safe_json(resp)
+            if resp.status >= 400:
+                message = _extract_error_message(data) or f"Usage fetch failed ({resp.status})"
+                logger.warning(
+                    "Usage fetch failed request_id=%s status=%s message=%s",
+                    get_request_id(),
+                    resp.status,
+                    message,
+                )
+                raise UsageFetchError(resp.status, message)
+            try:
+                return UsagePayload.model_validate(data)
+            except ValidationError as exc:
+                logger.warning(
+                    "Usage fetch invalid payload request_id=%s",
+                    get_request_id(),
+                )
+                raise UsageFetchError(502, "Invalid usage payload") from exc
+    except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
+        logger.warning(
+            "Usage fetch error request_id=%s error=%s",
+            get_request_id(),
+            exc,
+        )
+        raise UsageFetchError(0, f"Usage fetch failed: {exc}") from exc
+
+
+def _usage_url(base_url: str) -> str:
+    normalized = base_url.rstrip("/")
+    if "/backend-api" not in normalized:
+        normalized = f"{normalized}/backend-api"
+    return f"{normalized}/wham/usage"
+
+
+def _usage_headers(access_token: str, account_id: str | None) -> dict[str, str]:
+    headers = {"Authorization": f"Bearer {access_token}", "Accept": "application/json"}
+    request_id = get_request_id()
+    if request_id:
+        headers["x-request-id"] = request_id
+    if account_id and not account_id.startswith(("email_", "local_")):
+        headers["chatgpt-account-id"] = account_id
+    return headers
+
+
+async def _safe_json(resp: aiohttp.ClientResponse) -> JsonObject:
+    try:
+        data = await resp.json(content_type=None)
+    except Exception:
+        text = await resp.text()
+        return {"error": {"message": text.strip()}}
+    return data if isinstance(data, dict) else {"error": {"message": str(data)}}
+
+
+def _extract_error_message(payload: JsonObject) -> str | None:
+    envelope = UsageErrorEnvelope.model_validate(payload)
+    error = envelope.error
+    if isinstance(error, UsageErrorDetail):
+        return error.message or error.error_description
+    if isinstance(error, str):
+        return envelope.error_description or error
+    return envelope.message
+
+
+def _retry_options(attempts: int) -> ExponentialRetry:
+    return ExponentialRetry(
+        attempts=attempts,
+        start_timeout=RETRY_START_TIMEOUT,
+        max_timeout=RETRY_MAX_TIMEOUT,
+        factor=2.0,
+        statuses=RETRYABLE_STATUS,
+        exceptions={aiohttp.ClientError, asyncio.TimeoutError},
+        retry_all_server_errors=False,
+    )

app/core/config/settings.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from functools import lru_cache
+from pathlib import Path
+
+from pydantic import field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+BASE_DIR = Path(__file__).resolve().parents[3]
+
+DEFAULT_HOME_DIR = Path.home() / ".codex-lb"
+DEFAULT_DB_PATH = DEFAULT_HOME_DIR / "store.db"
+DEFAULT_ENCRYPTION_KEY_FILE = DEFAULT_HOME_DIR / "encryption.key"
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_prefix="CODEX_LB_",
+        env_file=(BASE_DIR / ".env", BASE_DIR / ".env.local"),
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+    database_url: str = f"sqlite+aiosqlite:///{DEFAULT_DB_PATH}"
+    upstream_base_url: str = "https://chatgpt.com/backend-api"
+    upstream_connect_timeout_seconds: float = 30.0
+    stream_idle_timeout_seconds: float = 300.0
+    auth_base_url: str = "https://auth.openai.com"
+    oauth_client_id: str = "app_EMoamEEZ73f0CkXaXp7hrann"
+    oauth_scope: str = "openid profile email"
+    oauth_timeout_seconds: float = 30.0
+    oauth_redirect_uri: str = "http://localhost:1455/auth/callback"
+    oauth_callback_host: str = "127.0.0.1"
+    oauth_callback_port: int = 1455  # Do not change the port. OpenAI dislikes changes.
+    token_refresh_timeout_seconds: float = 30.0
+    token_refresh_interval_days: int = 8
+    usage_fetch_timeout_seconds: float = 10.0
+    usage_fetch_max_retries: int = 2
+    usage_refresh_enabled: bool = True
+    usage_refresh_interval_seconds: int = 60
+    encryption_key_file: Path = DEFAULT_ENCRYPTION_KEY_FILE
+
+    @field_validator("database_url")
+    @classmethod
+    def _normalize_database_url(cls, value: str) -> str:
+        if not isinstance(value, str):
+            return value
+
+        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 value
+
+    @field_validator("encryption_key_file", mode="before")
+    @classmethod
+    def _normalize_encryption_key_file(cls, value: object) -> Path:
+        if isinstance(value, Path):
+            return value.expanduser()
+        if isinstance(value, str):
+            return Path(value).expanduser()
+        return value
+
+
+@lru_cache(maxsize=1)
+def get_settings() -> Settings:
+    return Settings()

app/core/crypto.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from cryptography.fernet import Fernet
+
+from app.core.config.settings import get_settings
+
+
+def _get_or_create_key(key_file: Path) -> bytes:
+    key_file.parent.mkdir(parents=True, exist_ok=True)
+    if key_file.exists():
+        return key_file.read_bytes()
+    key = Fernet.generate_key()
+    key_file.write_bytes(key)
+    key_file.chmod(0o600)
+    return key
+
+
+class TokenEncryptor:
+    def __init__(self, key: bytes | None = None, key_file: Path | None = None) -> None:
+        settings = get_settings()
+        resolved_file = key_file or settings.encryption_key_file
+        resolved_key = key or _get_or_create_key(resolved_file)
+        self._fernet = Fernet(resolved_key)
+
+    def encrypt(self, token: str) -> bytes:
+        return self._fernet.encrypt(token.encode())
+
+    def decrypt(self, encrypted: bytes) -> str:
+        return self._fernet.decrypt(encrypted).decode()
+
+
+def get_or_create_key(key_file: Path | None = None) -> bytes:
+    settings = get_settings()
+    resolved_file = key_file or settings.encryption_key_file
+    return _get_or_create_key(resolved_file)

app/core/errors.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import time
+from typing import Literal, TypedDict
+
+
+class OpenAIErrorDetail(TypedDict, total=False):
+    message: str
+    type: str
+    code: str
+    param: str
+    plan_type: str
+    resets_at: int | float
+    resets_in_seconds: int | float
+
+
+class OpenAIErrorEnvelope(TypedDict):
+    error: OpenAIErrorDetail
+
+
+class DashboardErrorDetail(TypedDict):
+    code: str
+    message: str
+
+
+class DashboardErrorEnvelope(TypedDict):
+    error: DashboardErrorDetail
+
+
+class ResponseFailedResponse(TypedDict, total=False):
+    id: str
+    object: str
+    created_at: int
+    status: str
+    error: OpenAIErrorDetail
+
+
+class ResponseFailedEvent(TypedDict):
+    type: Literal["response.failed"]
+    response: ResponseFailedResponse
+
+
+def openai_error(code: str, message: str, error_type: str = "server_error") -> OpenAIErrorEnvelope:
+    return {"error": {"message": message, "type": error_type, "code": code}}
+
+
+def dashboard_error(code: str, message: str) -> DashboardErrorEnvelope:
+    return {"error": {"code": code, "message": message}}
+
+
+def response_failed_event(
+    code: str,
+    message: str,
+    error_type: str = "server_error",
+    response_id: str | None = None,
+    created_at: int | None = None,
+    error_param: str | None = None,
+) -> ResponseFailedEvent:
+    error = openai_error(code, message, error_type)["error"]
+    if error_param:
+        error["param"] = error_param
+    if created_at is None:
+        created_at = int(time.time())
+    response: ResponseFailedResponse = {
+        "object": "response",
+        "status": "failed",
+        "error": error,
+    }
+    if response_id:
+        response["id"] = response_id
+    if created_at is not None:
+        response["created_at"] = created_at
+    return {"type": "response.failed", "response": response}

app/core/openai/models.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    StrictFloat,
+    StrictInt,
+    StrictStr,
+    ValidationError,
+    field_validator,
+)
+
+
+class OpenAIError(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    message: StrictStr | None = None
+    type: StrictStr | None = None
+    code: StrictStr | None = None
+    param: StrictStr | None = None
+    plan_type: StrictStr | None = None
+    resets_at: StrictInt | StrictFloat | None = None
+    resets_in_seconds: StrictInt | StrictFloat | None = None
+
+
+class OpenAIErrorEnvelope(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    error: OpenAIError | None = None
+
+
+class ResponseUsageDetails(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    cached_tokens: StrictInt | None = None
+    reasoning_tokens: StrictInt | None = None
+
+
+class ResponseUsage(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    input_tokens: StrictInt | None = None
+    output_tokens: StrictInt | None = None
+    total_tokens: StrictInt | None = None
+    input_tokens_details: ResponseUsageDetails | None = None
+    output_tokens_details: ResponseUsageDetails | None = None
+
+
+class OpenAIResponse(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    id: StrictStr | None = None
+    status: StrictStr | None = None
+    error: OpenAIError | None = None
+    usage: ResponseUsage | None = None
+
+    @field_validator("error", mode="before")
+    @classmethod
+    def _normalize_error(cls, value: object) -> OpenAIError | None:
+        if value is None:
+            return None
+        try:
+            return OpenAIError.model_validate(value)
+        except ValidationError:
+            return None
+
+    @field_validator("usage", mode="before")
+    @classmethod
+    def _normalize_usage(cls, value: object) -> ResponseUsage | None:
+        if value is None:
+            return None
+        try:
+            return ResponseUsage.model_validate(value)
+        except ValidationError:
+            return None
+
+
+class OpenAIEvent(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+
+    type: StrictStr
+    response: OpenAIResponse | None = None
+    error: OpenAIError | None = None
+
+    @field_validator("error", mode="before")
+    @classmethod
+    def _normalize_error(cls, value: object) -> OpenAIError | None:
+        if value is None:
+            return None
+        try:
+            return OpenAIError.model_validate(value)
+        except ValidationError:
+            return None
+
+
+class OpenAIResponsePayload(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    id: StrictStr | None = None
+    status: StrictStr | None = None
+    error: OpenAIError | None = None
+    usage: ResponseUsage | None = None
+
+    @field_validator("error", mode="before")
+    @classmethod
+    def _normalize_error(cls, value: object) -> OpenAIError | None:
+        if value is None:
+            return None
+        try:
+            return OpenAIError.model_validate(value)
+        except ValidationError:
+            return None
+
+    @field_validator("usage", mode="before")
+    @classmethod
+    def _normalize_usage(cls, value: object) -> ResponseUsage | None:
+        if value is None:
+            return None
+        try:
+            return ResponseUsage.model_validate(value)
+        except ValidationError:
+            return None

app/core/openai/parsing.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import json
+
+from pydantic import TypeAdapter, ValidationError
+
+from app.core.openai.models import OpenAIError, OpenAIErrorEnvelope, OpenAIEvent, OpenAIResponsePayload
+
+_EVENT_ADAPTER = TypeAdapter(OpenAIEvent)
+_ERROR_ADAPTER = TypeAdapter(OpenAIErrorEnvelope)
+_RESPONSE_ADAPTER = TypeAdapter(OpenAIResponsePayload)
+
+
+def parse_sse_event(line: str) -> OpenAIEvent | None:
+    data = None
+    if line.startswith("data:"):
+        data = line[5:].strip()
+    elif "\n" in line:
+        for part in line.splitlines():
+            if part.startswith("data:"):
+                data = part[5:].strip()
+                break
+    if data is None:
+        return None
+    if not data or data == "[DONE]":
+        return None
+    try:
+        payload = json.loads(data)
+    except json.JSONDecodeError:
+        return None
+    if not isinstance(payload, dict):
+        return None
+    try:
+        return _EVENT_ADAPTER.validate_python(payload)
+    except ValidationError:
+        return None
+
+
+def parse_error_payload(payload: object) -> OpenAIError | None:
+    if not isinstance(payload, dict):
+        return None
+    try:
+        envelope = _ERROR_ADAPTER.validate_python(payload)
+    except ValidationError:
+        return None
+    return envelope.error
+
+
+def parse_response_payload(payload: object) -> OpenAIResponsePayload | None:
+    if not isinstance(payload, dict):
+        return None
+    try:
+        return _RESPONSE_ADAPTER.validate_python(payload)
+    except ValidationError:
+        return None

app/core/openai/requests.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from app.core.types import JsonObject, JsonValue
+
+
+class ResponsesReasoning(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    effort: str | None = None
+    summary: str | None = None
+
+
+class ResponsesTextFormat(BaseModel):
+    model_config = ConfigDict(extra="allow", populate_by_name=True, serialize_by_alias=True)
+
+    type: str | None = None
+    strict: bool | None = None
+    schema_: JsonValue | None = Field(default=None, alias="schema")
+    name: str | None = None
+
+
+class ResponsesTextControls(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    verbosity: str | None = None
+    format: ResponsesTextFormat | None = None
+
+
+class ResponsesRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    model: str = Field(min_length=1)
+    instructions: str
+    input: list[JsonValue]
+    tools: list[JsonValue] = Field(default_factory=list)
+    tool_choice: str | None = None
+    parallel_tool_calls: bool | None = None
+    reasoning: ResponsesReasoning | None = None
+    store: bool | None = None
+    stream: bool | None = None
+    include: list[str] = Field(default_factory=list)
+    prompt_cache_key: str | None = None
+    text: ResponsesTextControls | None = None
+
+    def to_payload(self) -> JsonObject:
+        return self.model_dump(mode="json", exclude_none=True)
+
+
+class ResponsesCompactRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    model: str = Field(min_length=1)
+    instructions: str
+    input: list[JsonValue]
+
+    def to_payload(self) -> JsonObject:
+        return self.model_dump(mode="json", exclude_none=True)

app/core/types.py

@@ -0,0 +1,4 @@
+from __future__ import annotations
+
+type JsonValue = bool | int | float | str | None | list[JsonValue] | dict[str, JsonValue]
+type JsonObject = dict[str, JsonValue]