148 lines
4.9 KiB
Python
148 lines
4.9 KiB
Python
"""Billing-mode abstraction — the single ``claude_invoke`` seam (design §3.1).
|
|
|
|
Every Claude-calling node imports :func:`claude_invoke` from here. The seam
|
|
selects the Claude auth/billing path from config:
|
|
|
|
* ``SUBSCRIPTION`` — OAuth token (the R720 default; headless Agent SDK),
|
|
* ``API`` — metered ``ANTHROPIC_API_KEY``,
|
|
* ``BEDROCK`` — cross-account Bedrock (the rare cross-family tiebreak).
|
|
|
|
Switching modes is a config flip, not a code change. In ``SUBSCRIPTION`` mode
|
|
the seam pops/unsets any stray ``ANTHROPIC_API_KEY`` from the environment
|
|
before invoking, so an inherited key cannot silently override OAuth (§3.1).
|
|
|
|
This module is the contract leaf builders import verbatim; the actual SDK call
|
|
is delegated to an injectable ``_invoker`` so the seam stays testable and the
|
|
transport/SDK wiring lives in the leaves.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from typing import Any, Callable, Mapping
|
|
|
|
__all__ = [
|
|
"BillingMode",
|
|
"ClaudeResult",
|
|
"claude_invoke",
|
|
"resolve_mode",
|
|
"set_invoker",
|
|
]
|
|
|
|
# Environment variable that carries the metered API key. Popped in
|
|
# subscription mode so OAuth cannot be silently overridden.
|
|
_API_KEY_ENV = "ANTHROPIC_API_KEY"
|
|
|
|
# Config key (env or mapping) naming the desired billing mode.
|
|
_MODE_ENV = "AGENT_TEAM_BILLING_MODE"
|
|
|
|
|
|
class BillingMode(Enum):
|
|
"""Claude auth/billing path selector (§3.1)."""
|
|
|
|
SUBSCRIPTION = "subscription"
|
|
API = "api"
|
|
BEDROCK = "bedrock"
|
|
|
|
|
|
@dataclass
|
|
class ClaudeResult:
|
|
"""Result of a :func:`claude_invoke` call.
|
|
|
|
``text`` is the model's response text. ``mode`` records which billing path
|
|
served the call. ``usage`` carries token/cost accounting for the budget
|
|
ledger (§6.6); ``raw`` is the untouched provider response for callers that
|
|
need more.
|
|
"""
|
|
|
|
text: str
|
|
mode: BillingMode
|
|
usage: dict[str, Any] = field(default_factory=dict)
|
|
raw: Any = None
|
|
|
|
|
|
# Pluggable invoker: signature (prompt, mode, **kw) -> ClaudeResult. The
|
|
# default raises so an un-wired environment fails loudly rather than silently
|
|
# returning nothing; leaves call set_invoker() to bind the real SDK path.
|
|
Invoker = Callable[..., ClaudeResult]
|
|
|
|
|
|
def _unconfigured_invoker(prompt: str, *, mode: BillingMode, **kw: Any) -> ClaudeResult:
|
|
raise RuntimeError(
|
|
"claude_invoke has no invoker bound; call billing.set_invoker(fn) to "
|
|
"wire the Claude SDK path (subscription OAuth / API / Bedrock)."
|
|
)
|
|
|
|
|
|
_invoker: Invoker = _unconfigured_invoker
|
|
|
|
|
|
def set_invoker(invoker: Invoker) -> None:
|
|
"""Bind the function that performs the actual Claude SDK call.
|
|
|
|
Leaves call this once at startup with an implementation that honours the
|
|
resolved :class:`BillingMode`. Keeping the SDK call injectable keeps this
|
|
seam dependency-free and unit-testable.
|
|
"""
|
|
global _invoker
|
|
_invoker = invoker
|
|
|
|
|
|
def resolve_mode(config: Mapping[str, Any] | None) -> BillingMode:
|
|
"""Resolve the billing mode from ``config`` (falling back to env).
|
|
|
|
Precedence: an explicit ``billing_mode`` in ``config`` (a
|
|
:class:`BillingMode` or its string value), then the
|
|
``AGENT_TEAM_BILLING_MODE`` env var, then the ``SUBSCRIPTION`` default.
|
|
"""
|
|
raw: Any = None
|
|
if config is not None:
|
|
raw = config.get("billing_mode")
|
|
if raw is None:
|
|
raw = os.environ.get(_MODE_ENV)
|
|
if raw is None:
|
|
return BillingMode.SUBSCRIPTION
|
|
if isinstance(raw, BillingMode):
|
|
return raw
|
|
try:
|
|
return BillingMode(str(raw).strip().lower())
|
|
except ValueError as exc:
|
|
valid = ", ".join(m.value for m in BillingMode)
|
|
raise ValueError(
|
|
f"unknown billing mode {raw!r}; expected one of: {valid}"
|
|
) from exc
|
|
|
|
|
|
def claude_invoke(
|
|
prompt: str,
|
|
*,
|
|
mode: BillingMode | None = None,
|
|
config: Mapping[str, Any] | None = None,
|
|
**kw: Any,
|
|
) -> ClaudeResult:
|
|
"""Invoke Claude through the configured billing path (§3.1).
|
|
|
|
``mode`` overrides config when given; otherwise it is resolved via
|
|
:func:`resolve_mode`. In ``SUBSCRIPTION`` mode any stray
|
|
``ANTHROPIC_API_KEY`` is popped from ``os.environ`` for the duration of the
|
|
call so OAuth cannot be silently overridden, then restored afterward.
|
|
|
|
The actual SDK call is delegated to the bound invoker (see
|
|
:func:`set_invoker`); this function owns only mode selection and the
|
|
subscription-mode env hygiene that the design mandates.
|
|
"""
|
|
effective = mode if mode is not None else resolve_mode(config)
|
|
|
|
if effective is BillingMode.SUBSCRIPTION:
|
|
# Pop the stray key for the duration of the call; restore on exit so we
|
|
# don't mutate the caller's environment permanently.
|
|
stashed = os.environ.pop(_API_KEY_ENV, None)
|
|
try:
|
|
return _invoker(prompt, mode=effective, **kw)
|
|
finally:
|
|
if stashed is not None:
|
|
os.environ[_API_KEY_ENV] = stashed
|
|
|
|
return _invoker(prompt, mode=effective, **kw)
|