PyPI - codex-lb - Versions diffs - 0.1.2__py3-none-any.whl - Mend

codex-lb 0.1.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (80) hide show

app/__init__.py +5 -0
app/cli.py +24 -0
app/core/__init__.py +0 -0
app/core/auth/__init__.py +96 -0
app/core/auth/models.py +49 -0
app/core/auth/refresh.py +144 -0
app/core/balancer/__init__.py +19 -0
app/core/balancer/logic.py +140 -0
app/core/balancer/types.py +9 -0
app/core/clients/__init__.py +0 -0
app/core/clients/http.py +39 -0
app/core/clients/oauth.py +340 -0
app/core/clients/proxy.py +265 -0
app/core/clients/usage.py +143 -0
app/core/config/__init__.py +0 -0
app/core/config/settings.py +69 -0
app/core/crypto.py +37 -0
app/core/errors.py +73 -0
app/core/openai/__init__.py +0 -0
app/core/openai/models.py +122 -0
app/core/openai/parsing.py +55 -0
app/core/openai/requests.py +59 -0
app/core/types.py +4 -0
app/core/usage/__init__.py +185 -0
app/core/usage/logs.py +57 -0
app/core/usage/models.py +35 -0
app/core/usage/pricing.py +172 -0
app/core/usage/types.py +95 -0
app/core/utils/__init__.py +0 -0
app/core/utils/request_id.py +30 -0
app/core/utils/retry.py +16 -0
app/core/utils/sse.py +13 -0
app/core/utils/time.py +19 -0
app/db/__init__.py +0 -0
app/db/models.py +82 -0
app/db/session.py +44 -0
app/dependencies.py +123 -0
app/main.py +124 -0
app/modules/__init__.py +0 -0
app/modules/accounts/__init__.py +0 -0
app/modules/accounts/api.py +81 -0
app/modules/accounts/repository.py +80 -0
app/modules/accounts/schemas.py +66 -0
app/modules/accounts/service.py +211 -0
app/modules/health/__init__.py +0 -0
app/modules/health/api.py +10 -0
app/modules/oauth/__init__.py +0 -0
app/modules/oauth/api.py +57 -0
app/modules/oauth/schemas.py +32 -0
app/modules/oauth/service.py +356 -0
app/modules/oauth/templates/oauth_success.html +122 -0
app/modules/proxy/__init__.py +0 -0
app/modules/proxy/api.py +76 -0
app/modules/proxy/auth_manager.py +51 -0
app/modules/proxy/load_balancer.py +208 -0
app/modules/proxy/schemas.py +85 -0
app/modules/proxy/service.py +707 -0
app/modules/proxy/types.py +37 -0
app/modules/proxy/usage_updater.py +147 -0
app/modules/request_logs/__init__.py +0 -0
app/modules/request_logs/api.py +31 -0
app/modules/request_logs/repository.py +86 -0
app/modules/request_logs/schemas.py +25 -0
app/modules/request_logs/service.py +77 -0
app/modules/shared/__init__.py +0 -0
app/modules/shared/schemas.py +8 -0
app/modules/usage/__init__.py +0 -0
app/modules/usage/api.py +31 -0
app/modules/usage/repository.py +113 -0
app/modules/usage/schemas.py +62 -0
app/modules/usage/service.py +246 -0
app/static/7.css +1336 -0
app/static/index.css +543 -0
app/static/index.html +457 -0
app/static/index.js +1898 -0
codex_lb-0.1.2.dist-info/METADATA +108 -0
codex_lb-0.1.2.dist-info/RECORD +80 -0
codex_lb-0.1.2.dist-info/WHEEL +4 -0
codex_lb-0.1.2.dist-info/entry_points.txt +2 -0
codex_lb-0.1.2.dist-info/licenses/LICENSE +21 -0

app/modules/oauth/templates/oauth_success.html ADDED Viewed

@@ -0,0 +1,122 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
+  <title>Sign into Codex LB</title>
+  <link rel="icon"
+    href="data:image/svg+xml,%3Csvg xmlns=&quot;http://www.w3.org/2000/svg&quot; width=&quot;32&quot; height=&quot;32&quot; fill=&quot;none&quot; viewBox=&quot;0 0 32 32&quot;%3E%3Cpath stroke=&quot;%23000&quot; stroke-linecap=&quot;round&quot; stroke-width=&quot;2.484&quot; d=&quot;M22.356 19.797H17.17M9.662 12.29l1.979 3.576a.511.511 0 0 1-.005.504l-1.974 3.409M30.758 16c0 8.15-6.607 14.758-14.758 14.758-8.15 0-14.758-6.607-14.758-14.758C1.242 7.85 7.85 1.242 16 1.242c8.15 0 14.758 6.608 14.758 14.758Z&quot;/%3E%3C/svg%3E"
+    type="image/svg+xml">
+  <style>
+    .container {
+      margin: auto;
+      height: 100%;
+      display: flex;
+      align-items: center;
+      justify-content: center;
+      position: relative;
+      background: white;
+      font-family: system-ui, -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, 'Open Sans', 'Helvetica Neue', sans-serif;
+    }
+    .inner-container {
+      width: 400px;
+      flex-direction: column;
+      justify-content: flex-start;
+      align-items: center;
+      gap: 20px;
+      display: inline-flex;
+    }
+    .content {
+      align-self: stretch;
+      flex-direction: column;
+      justify-content: flex-start;
+      align-items: center;
+      gap: 20px;
+      display: flex;
+      margin-top: 15vh;
+    }
+    .title {
+      text-align: center;
+      color: var(--text-primary, #0D0D0D);
+      font-size: 32px;
+      font-weight: 400;
+      line-height: 40px;
+      word-wrap: break-word;
+    }
+    .setup-description {
+      align-self: stretch;
+      color: var(--text-secondary, #5D5D5D);
+      font-size: 14px;
+      font-weight: 400;
+      line-height: 20px;
+      word-wrap: break-word;
+    }
+    .logo {
+      display: flex;
+      align-items: center;
+      justify-content: center;
+      width: 4rem;
+      height: 4rem;
+      border-radius: 16px;
+      border: .5px solid rgba(0, 0, 0, 0.1);
+      box-shadow: rgba(0, 0, 0, 0.1) 0px 4px 16px 0px;
+      box-sizing: border-box;
+      background-color: rgb(255, 255, 255);
+    }
+  </style>
+</head>
+<body>
+  <div class="container">
+    <div class="inner-container">
+      <div class="content">
+        <div class="logo">
+          <svg xmlns="http://www.w3.org/2000/svg" width="32" height="32" fill="none" viewBox="0 0 32 32">
+            <path stroke="#000" stroke-linecap="round" stroke-width="2.484"
+              d="M22.356 19.797H17.17M9.662 12.29l1.979 3.576a.511.511 0 0 1-.005.504l-1.974 3.409M30.758 16c0 8.15-6.607 14.758-14.758 14.758-8.15 0-14.758-6.607-14.758-14.758C1.242 7.85 7.85 1.242 16 1.242c8.15 0 14.758 6.608 14.758 14.758Z">
+            </path>
+          </svg>
+        </div>
+        <div class="title">Signed in to Codex LB</div>
+      </div>
+      <div class="close-box" style="display: flex;">
+        <div class="setup-description" id="close-message">
+          Closing in <span id="close-countdown">3</span>s...
+        </div>
+      </div>
+    </div>
+  </div>
+  <script>
+    (() => {
+      const countdown = document.getElementById("close-countdown");
+      const message = document.getElementById("close-message");
+      let remaining = 3;
+      const tick = () => {
+        if (remaining <= 0) {
+          if (message) {
+            message.textContent = "Closing...";
+          }
+          window.close();
+          return;
+        }
+        if (countdown) {
+          countdown.textContent = String(remaining);
+        } else if (message) {
+          message.textContent = `Closing in ${remaining}s...`;
+        }
+        remaining -= 1;
+        window.setTimeout(tick, 1000);
+      };
+      tick();
+    })();
+  </script>
+</body>
+</html>

app/modules/proxy/__init__.py ADDED Viewed

File without changes

app/modules/proxy/api.py ADDED Viewed

@@ -0,0 +1,76 @@
+from __future__ import annotations
+from collections.abc import AsyncIterator
+from fastapi import APIRouter, Body, Depends, Request, Response
+from fastapi.responses import JSONResponse, StreamingResponse
+from app.core.clients.proxy import ProxyResponseError
+from app.core.errors import openai_error
+from app.core.openai.requests import ResponsesCompactRequest, ResponsesRequest
+from app.dependencies import ProxyContext, get_proxy_context
+from app.modules.proxy.schemas import RateLimitStatusPayload
+router = APIRouter(prefix="/backend-api/codex", tags=["proxy"])
+usage_router = APIRouter(tags=["proxy"])
+@router.post("/responses")
+async def responses(
+    request: Request,
+    payload: ResponsesRequest = Body(...),
+    context: ProxyContext = Depends(get_proxy_context),
+) -> Response:
+    rate_limit_headers = await context.service.rate_limit_headers()
+    stream = context.service.stream_responses(
+        payload,
+        request.headers,
+        propagate_http_errors=True,
+    )
+    try:
+        first = await stream.__anext__()
+    except StopAsyncIteration:
+        return StreamingResponse(
+            _prepend_first(None, stream),
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-cache", **rate_limit_headers},
+        )
+    except ProxyResponseError as exc:
+        return JSONResponse(status_code=exc.status_code, content=exc.payload, headers=rate_limit_headers)
+    return StreamingResponse(
+        _prepend_first(first, stream),
+        media_type="text/event-stream",
+        headers={"Cache-Control": "no-cache", **rate_limit_headers},
+    )
+@router.post("/responses/compact")
+async def responses_compact(
+    request: Request,
+    payload: ResponsesCompactRequest = Body(...),
+    context: ProxyContext = Depends(get_proxy_context),
+) -> JSONResponse:
+    rate_limit_headers = await context.service.rate_limit_headers()
+    try:
+        result = await context.service.compact_responses(payload, request.headers)
+    except NotImplementedError:
+        error = openai_error("not_implemented", "responses/compact is not implemented")
+        return JSONResponse(status_code=501, content=error, headers=rate_limit_headers)
+    except ProxyResponseError as exc:
+        return JSONResponse(status_code=exc.status_code, content=exc.payload, headers=rate_limit_headers)
+    return JSONResponse(content=result.model_dump(exclude_none=True), headers=rate_limit_headers)
+@usage_router.get("/api/codex/usage", response_model=RateLimitStatusPayload)
+async def codex_usage(
+    context: ProxyContext = Depends(get_proxy_context),
+) -> RateLimitStatusPayload:
+    payload = await context.service.get_rate_limit_payload()
+    return RateLimitStatusPayload.from_data(payload)
+async def _prepend_first(first: str | None, stream: AsyncIterator[str]) -> AsyncIterator[str]:
+    if first is not None:
+        yield first
+    async for line in stream:
+        yield line

app/modules/proxy/auth_manager.py ADDED Viewed

@@ -0,0 +1,51 @@
+from __future__ import annotations
+from app.core.auth.refresh import RefreshError, refresh_access_token, should_refresh
+from app.core.balancer import PERMANENT_FAILURE_CODES
+from app.core.crypto import TokenEncryptor
+from app.core.utils.time import utcnow
+from app.db.models import Account, AccountStatus
+from app.modules.accounts.repository import AccountsRepository
+class AuthManager:
+    def __init__(self, repo: AccountsRepository) -> None:
+        self._repo = repo
+        self._encryptor = TokenEncryptor()
+    async def ensure_fresh(self, account: Account, *, force: bool = False) -> Account:
+        if force or should_refresh(account.last_refresh):
+            return await self.refresh_account(account)
+        return account
+    async def refresh_account(self, account: Account) -> Account:
+        refresh_token = self._encryptor.decrypt(account.refresh_token_encrypted)
+        try:
+            result = await refresh_access_token(refresh_token)
+        except RefreshError as exc:
+            if exc.is_permanent:
+                reason = PERMANENT_FAILURE_CODES.get(exc.code, exc.message)
+                await self._repo.update_status(account.id, AccountStatus.DEACTIVATED, reason)
+                account.status = AccountStatus.DEACTIVATED
+                account.deactivation_reason = reason
+            raise
+        account.access_token_encrypted = self._encryptor.encrypt(result.access_token)
+        account.refresh_token_encrypted = self._encryptor.encrypt(result.refresh_token)
+        account.id_token_encrypted = self._encryptor.encrypt(result.id_token)
+        account.last_refresh = utcnow()
+        if result.plan_type:
+            account.plan_type = result.plan_type
+        if result.email:
+            account.email = result.email
+        await self._repo.update_tokens(
+            account.id,
+            access_token_encrypted=account.access_token_encrypted,
+            refresh_token_encrypted=account.refresh_token_encrypted,
+            id_token_encrypted=account.id_token_encrypted,
+            last_refresh=account.last_refresh,
+            plan_type=account.plan_type,
+            email=account.email,
+        )
+        return account

app/modules/proxy/load_balancer.py ADDED Viewed

@@ -0,0 +1,208 @@
+from __future__ import annotations
+import time
+from dataclasses import dataclass
+from typing import Iterable
+from app.core.balancer import (
+    AccountState,
+    handle_permanent_failure,
+    handle_quota_exceeded,
+    handle_rate_limit,
+    select_account,
+)
+from app.core.balancer.types import UpstreamError
+from app.db.models import Account, AccountStatus, UsageHistory
+from app.modules.accounts.repository import AccountsRepository
+from app.modules.proxy.usage_updater import UsageUpdater
+from app.modules.usage.repository import UsageRepository
+@dataclass
+class RuntimeState:
+    reset_at: int | None = None
+    last_error_at: float | None = None
+    last_selected_at: float | None = None
+    error_count: int = 0
+@dataclass
+class AccountSelection:
+    account: Account | None
+    error_message: str | None
+class LoadBalancer:
+    def __init__(self, accounts_repo: AccountsRepository, usage_repo: UsageRepository) -> None:
+        self._accounts_repo = accounts_repo
+        self._usage_repo = usage_repo
+        self._usage_updater = UsageUpdater(usage_repo, accounts_repo)
+        self._runtime: dict[str, RuntimeState] = {}
+    async def select_account(self) -> AccountSelection:
+        accounts = await self._accounts_repo.list_accounts()
+        latest_primary = await self._usage_repo.latest_by_account()
+        await self._usage_updater.refresh_accounts(accounts, latest_primary)
+        latest_primary = await self._usage_repo.latest_by_account()
+        latest_secondary = await self._usage_repo.latest_by_account(window="secondary")
+        states, account_map = _build_states(
+            accounts=accounts,
+            latest_primary=latest_primary,
+            latest_secondary=latest_secondary,
+            runtime=self._runtime,
+        )
+        result = select_account(states)
+        for state in states:
+            account = account_map.get(state.account_id)
+            if account:
+                await self._sync_state(account, state)
+        if result.account is None:
+            return AccountSelection(account=None, error_message=result.error_message)
+        selected = account_map.get(result.account.account_id)
+        if selected:
+            selected.status = result.account.status
+            selected.deactivation_reason = result.account.deactivation_reason
+            runtime = self._runtime.setdefault(selected.id, RuntimeState())
+            runtime.last_selected_at = time.time()
+        if selected is None:
+            return AccountSelection(account=None, error_message=result.error_message)
+        return AccountSelection(account=selected, error_message=None)
+    async def mark_rate_limit(self, account: Account, error: UpstreamError) -> None:
+        state = self._state_for(account)
+        handle_rate_limit(state, error)
+        await self._sync_state(account, state)
+    async def mark_quota_exceeded(self, account: Account, error: UpstreamError) -> None:
+        state = self._state_for(account)
+        handle_quota_exceeded(state, error)
+        await self._sync_state(account, state)
+    async def mark_permanent_failure(self, account: Account, error_code: str) -> None:
+        state = self._state_for(account)
+        handle_permanent_failure(state, error_code)
+        await self._sync_state(account, state)
+    async def record_error(self, account: Account) -> None:
+        state = self._state_for(account)
+        state.error_count += 1
+        state.last_error_at = time.time()
+        await self._sync_state(account, state)
+    def _state_for(self, account: Account) -> AccountState:
+        runtime = self._runtime.setdefault(account.id, RuntimeState())
+        return AccountState(
+            account_id=account.id,
+            status=account.status,
+            used_percent=None,
+            reset_at=runtime.reset_at,
+            last_error_at=runtime.last_error_at,
+            last_selected_at=runtime.last_selected_at,
+            error_count=runtime.error_count,
+            deactivation_reason=account.deactivation_reason,
+        )
+    async def _sync_state(self, account: Account, state: AccountState) -> None:
+        runtime = self._runtime.setdefault(account.id, RuntimeState())
+        runtime.reset_at = state.reset_at
+        runtime.last_error_at = state.last_error_at
+        runtime.error_count = state.error_count
+        if account.status != state.status or account.deactivation_reason != state.deactivation_reason:
+            await self._accounts_repo.update_status(
+                account.id,
+                state.status,
+                state.deactivation_reason,
+            )
+            account.status = state.status
+            account.deactivation_reason = state.deactivation_reason
+def _build_states(
+    *,
+    accounts: Iterable[Account],
+    latest_primary: dict[str, UsageHistory],
+    latest_secondary: dict[str, UsageHistory],
+    runtime: dict[str, RuntimeState],
+) -> tuple[list[AccountState], dict[str, Account]]:
+    states: list[AccountState] = []
+    account_map: dict[str, Account] = {}
+    for account in accounts:
+        state = _state_from_account(
+            account=account,
+            primary_entry=latest_primary.get(account.id),
+            secondary_entry=latest_secondary.get(account.id),
+            runtime=runtime.setdefault(account.id, RuntimeState()),
+        )
+        states.append(state)
+        account_map[account.id] = account
+    return states, account_map
+def _state_from_account(
+    *,
+    account: Account,
+    primary_entry: UsageHistory | None,
+    secondary_entry: UsageHistory | None,
+    runtime: RuntimeState,
+) -> AccountState:
+    primary_used = primary_entry.used_percent if primary_entry else None
+    secondary_used = secondary_entry.used_percent if secondary_entry else None
+    secondary_reset = secondary_entry.reset_at if secondary_entry else None
+    status, used_percent, reset_at = _apply_secondary_quota(
+        status=account.status,
+        primary_used=primary_used,
+        runtime_reset=runtime.reset_at,
+        secondary_used=secondary_used,
+        secondary_reset=secondary_reset,
+    )
+    return AccountState(
+        account_id=account.id,
+        status=status,
+        used_percent=used_percent,
+        reset_at=reset_at,
+        last_error_at=runtime.last_error_at,
+        last_selected_at=runtime.last_selected_at,
+        error_count=runtime.error_count,
+        deactivation_reason=account.deactivation_reason,
+    )
+def _apply_secondary_quota(
+    *,
+    status: AccountStatus,
+    primary_used: float | None,
+    runtime_reset: int | None,
+    secondary_used: float | None,
+    secondary_reset: int | None,
+) -> tuple[AccountStatus, float | None, int | None]:
+    used_percent = primary_used
+    reset_at = runtime_reset
+    if status in (AccountStatus.DEACTIVATED, AccountStatus.PAUSED):
+        return status, used_percent, reset_at
+    if secondary_used is None:
+        if status == AccountStatus.QUOTA_EXCEEDED and secondary_reset is not None:
+            reset_at = secondary_reset
+        return status, used_percent, reset_at
+    if secondary_used >= 100.0:
+        status = AccountStatus.QUOTA_EXCEEDED
+        used_percent = 100.0
+        if secondary_reset is not None:
+            reset_at = secondary_reset
+        return status, used_percent, reset_at
+    if status == AccountStatus.QUOTA_EXCEEDED:
+        status = AccountStatus.ACTIVE
+        reset_at = None
+    return status, used_percent, reset_at

app/modules/proxy/schemas.py ADDED Viewed

@@ -0,0 +1,85 @@
+from __future__ import annotations
+from pydantic import BaseModel, ConfigDict
+from app.core.types import JsonValue
+from app.modules.proxy.types import (
+    CreditStatusDetailsData,
+    RateLimitStatusDetailsData,
+    RateLimitStatusPayloadData,
+    RateLimitWindowSnapshotData,
+)
+class RateLimitWindowSnapshot(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    used_percent: int
+    limit_window_seconds: int
+    reset_after_seconds: int
+    reset_at: int
+    @classmethod
+    def from_data(cls, data: RateLimitWindowSnapshotData) -> "RateLimitWindowSnapshot":
+        return cls(
+            used_percent=data.used_percent,
+            limit_window_seconds=data.limit_window_seconds,
+            reset_after_seconds=data.reset_after_seconds,
+            reset_at=data.reset_at,
+        )
+class RateLimitStatusDetails(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    allowed: bool
+    limit_reached: bool
+    primary_window: RateLimitWindowSnapshot | None = None
+    secondary_window: RateLimitWindowSnapshot | None = None
+    @classmethod
+    def from_data(cls, data: RateLimitStatusDetailsData) -> "RateLimitStatusDetails":
+        return cls(
+            allowed=data.allowed,
+            limit_reached=data.limit_reached,
+            primary_window=RateLimitWindowSnapshot.from_data(data.primary_window) if data.primary_window else None,
+            secondary_window=RateLimitWindowSnapshot.from_data(data.secondary_window)
+            if data.secondary_window
+            else None,
+        )
+class CreditStatusDetails(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    has_credits: bool
+    unlimited: bool
+    balance: str | None = None
+    approx_local_messages: list[JsonValue] | None = None
+    approx_cloud_messages: list[JsonValue] | None = None
+    @classmethod
+    def from_data(cls, data: CreditStatusDetailsData) -> "CreditStatusDetails":
+        return cls(
+            has_credits=data.has_credits,
+            unlimited=data.unlimited,
+            balance=data.balance,
+            approx_local_messages=data.approx_local_messages,
+            approx_cloud_messages=data.approx_cloud_messages,
+        )
+class RateLimitStatusPayload(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    plan_type: str
+    rate_limit: RateLimitStatusDetails | None = None
+    credits: CreditStatusDetails | None = None
+    @classmethod
+    def from_data(cls, data: RateLimitStatusPayloadData) -> "RateLimitStatusPayload":
+        return cls(
+            plan_type=data.plan_type,
+            rate_limit=RateLimitStatusDetails.from_data(data.rate_limit) if data.rate_limit else None,
+            credits=CreditStatusDetails.from_data(data.credits) if data.credits else None,
+        )