Model Routing

# pip install aroha
import asyncio
from aroha.routing import ModelSession, ModelBudgetMandate, TaskType
from aroha.routing.providers.anthropic import AnthropicProvider

async def main():
    provider = AnthropicProvider()

    mandate = ModelBudgetMandate(
        token_budget=4_000,
        cost_budget_usd=0.05,
        quality_threshold=0.85,
        task_type=TaskType.CODE,
    )

    session = ModelSession.create(provider=provider, mandate=mandate)
    response = await session.chat("Refactor this function to be async.")
    await session.close()

    print(response.content)
    print(response.model_id)       # e.g. "claude-sonnet-4-6"
    print(session.tracker.summary())

asyncio.run(main())

from aroha.routing import ModelBudgetMandate, TaskType

mandate = ModelBudgetMandate(
    # Required
    token_budget=8_000,            # hard cap: input + output tokens for the session
    cost_budget_usd=0.10,          # hard cap: USD spend for the session
    quality_threshold=0.85,        # minimum quality in [0, 1]
    task_type=TaskType.REASONING,  # governs benchmark selection & scoring

    # Optional — defaults shown
    context_window_required=4_096, # model must support at least this many tokens
    latency_budget_ms=30_000,      # max latency per LLM call in ms
    preferred_provider=None,       # "anthropic" | "openai" | "google" | None
    allow_cascade=True,            # allow escalation when quality < threshold
    max_cascade_steps=2,           # max escalation hops (not counting initial)
    expected_steps=1,              # hint to BudgetAllocator for multi-step sessions
)

# shadow_price — KKT dual variable: quality per dollar
# Cascade escalates only when  ΔQuality / ΔCost >= shadow_price
mandate.shadow_price   # -> quality_threshold / cost_budget_usd

# remaining() — narrows a mandate after consuming resources
narrowed = mandate.remaining(tokens_used=1200, cost_used=0.012)
# narrowed.token_budget     = 6_800
# narrowed.cost_budget_usd  = 0.088
# narrowed.expected_steps   = max(1, original - 1)

from aroha.routing import ModelSession, ModelBudgetMandate, TaskType, AllocatorMode
from aroha.routing.providers.anthropic import AnthropicProvider

session = ModelSession.create(
    provider=AnthropicProvider(),
    mandate=mandate,

    # Optional
    allocator_mode=AllocatorMode.THOMPSON,         # budget allocation strategy
    learner_path="/var/cache/aroha/learner.json",  # persist learner across restarts

    # Registry feedback loop (recommended in production)
    registry_url="https://registry.aroha-labs.com",
    registry_api_key="ak_...",
    feedback_flush_every=20,  # batch size before flushing to registry
)

response = await session.chat(
    prompt="Summarise the document above.",
    system="You are a technical writer.",  # optional system prompt
)

# ModelResponse fields:
response.content        # str — generated text
response.model_id       # str — model that produced the response
response.quality        # float — observed quality in [0, 1]
response.input_tokens   # int
response.output_tokens  # int
response.cost_usd       # float — cost of this specific call
response.n_escalations  # int — cascade hops used (0 = no escalation needed)
response.success        # bool — False when quality < threshold (best-effort)
response.total_tokens   # property: input + output

# warm_from_registry() — pull crowd-sourced priors before first call
await session.warm_from_registry()   # returns int: number of priors loaded

# close() — flush pending feedback signals and auto-save learner
await session.close()

# Always close() — especially important with registry_url set,
# to flush the last batch of quality signals.

mandate = ModelBudgetMandate(
    token_budget=12_000,
    cost_budget_usd=0.20,
    quality_threshold=0.85,
    task_type=TaskType.SUMMARIZATION,
    expected_steps=3,              # tell the allocator how many steps to expect
)

session = ModelSession.create(
    provider=provider,
    mandate=mandate,
    allocator_mode=AllocatorMode.IMPORTANCE_WEIGHTED,
)

# Each step() yields a narrowed ModelBudgetMandate for that slice of the budget.
async with session.step("retrieval", importance_weight=1.0, total_weight=4.0) as m:
    r1 = await session.chat("Extract key facts from this document.", step_mandate=m)

async with session.step("analysis", importance_weight=1.0, total_weight=4.0) as m:
    r2 = await session.chat("Identify the main argument.", step_mandate=m)

async with session.step("synthesis", importance_weight=2.0, total_weight=4.0) as m:
    r3 = await session.chat("Write a 200-word summary.", step_mandate=m)

await session.close()
print(session.tracker.summary())

tracker = session.tracker

# Running totals
tracker.tokens_used         # int — total tokens consumed so far
tracker.tokens_remaining    # int — remaining before token_budget is hit
tracker.cost_used           # float — USD spent so far
tracker.cost_remaining      # float — USD remaining
tracker.token_budget_exhausted  # bool
tracker.cost_budget_exhausted   # bool

# Per-call breakdown
for call in tracker.usage.calls:
    print(call.model_id, call.input_tokens, call.output_tokens, call.cost_usd)

# Summary dict
tracker.summary()
# {
#   "total_tokens": 1842,
#   "tokens_remaining": 6158,
#   "total_cost_usd": 0.01234567,
#   "cost_remaining": 0.08765433,
#   "n_calls": 3,
#   "models_used": ["claude-haiku-4-5", "claude-sonnet-4-6"],
# }

from aroha.routing.learner import OnlineLearner
from aroha.routing.catalog import BUILTIN_MODELS

# Share a learner across sessions for cross-session learning
learner = OnlineLearner(catalog=BUILTIN_MODELS)

session1 = ModelSession.create(provider, mandate1, learner=learner)
session2 = ModelSession.create(provider, mandate2, learner=learner)

# Persist to disk and reload
learner.auto_save("/var/cache/aroha/learner.json")
learner.load("/var/cache/aroha/learner.json")

# Inspect posteriors
posteriors = learner.posteriors()  # dict[str, dict] keyed by model_id
# {
#   "claude-sonnet-4-6": {
#     ("code", 2): {"alpha": 14.3, "beta": 1.7, "mean_quality": 0.894},
#     ...
#   },
#   ...
# }

# Full production setup with registry feedback loop
session = ModelSession.create(
    provider=AnthropicProvider(),
    mandate=mandate,
    registry_url="https://registry.aroha-labs.com",
    registry_api_key=os.environ["AROHA_API_KEY"],
    learner_path="/var/cache/aroha/learner.json",
    feedback_flush_every=20,
    allocator_mode=AllocatorMode.THOMPSON,
)

# Warm learner from registry before first call (optional but recommended)
priors_loaded = await session.warm_from_registry()
print(f"Warmed with {priors_loaded} priors")

response = await session.chat("Classify this support ticket: ...")

# close() flushes any buffered feedback signals
await session.close()

from aroha.routing.mandate import BudgetInfeasibleError

try:
    response = await session.chat("...")
except BudgetInfeasibleError as e:
    print(e)                      # human-readable message
    print(e.binding_constraint)   # e.g. "quality_threshold", "context_window", "cost_budget_usd"
except RuntimeError as e:
    # Raised when token_budget or cost_budget is already exhausted before the call
    print(e)

from aroha.routing.catalog import BUILTIN_MODELS, get_catalog, get_model

# All models
for m in BUILTIN_MODELS:
    print(m.model_id, m.quality_for(TaskType.CODE))

# Filtered view
anthropic_mid_plus = get_catalog(provider="anthropic", min_tier=ModelTier.MID)

# Lookup by ID
profile = get_model("claude-sonnet-4-6")
cost = profile.cost_estimate(input_tokens=500, output_tokens=200)