Skip to content

[v2][budget] SDK: send per-run budget to agents via execution_params #953

Description

@aix-ahmet

[v2][budget] SDK: send per-run budget to agents via execution_params

Goal

Let SDK users attach a budget (cost / duration / iteration caps) to a single agent
run. The budget feature already exists and is enforced in the aixplain-agents
execution service (BudgetConfig + BudgetController). This task only makes the
Python SDK produce the budget payload. It is run-time only — not persisted on
the agent.

Budget contract — SHARED, DO NOT DEVIATE

This exact shape is agreed across three repos (SDK ← this task, platform-backend,
aixplain-agents). If any key differs, budget is silently ignored downstream.

Budget rides inside executionParams as a nested budget object:

"executionParams": {
  "maxTokens": 2048,
  "maxIterations": 5,
  "budget": {
    "maxCost": 1.0,            // float, USD spending cap for the run
    "maxDurationSeconds": 300, // float, wall-clock seconds
    "maxIterations": 50        // int, node-execution cap (budget guardrail)
  }
}
  • All wire keys are camelCase. budget and every field inside it are optional.
  • Omitting budget ⇒ no budget enforcement (current behavior — must stay backward compatible).
  • warnAtPercent is intentionally NOT sent; the backend BudgetConfig default 0.8 applies.
  • Python user-facing API is snake_case (max_cost, max_duration_seconds, max_iterations)
    and serializes to the camelCase keys above.

Scope (this repo only)

File: aixplain/v2/agent.py

  1. Add a small serializable Budget type (dataclass, snake_case fields
    max_cost: Optional[float], max_duration_seconds: Optional[float],
    max_iterations: Optional[int]) that serializes to the camelCase keys
    maxCost / maxDurationSeconds / maxIterations, dropping None fields.
    Follow the existing @dataclass_json + metadata=config(field_name=...) pattern
    used elsewhere in this file.
  2. Add budget: NotRequired[Optional[Union[Dict, Budget]]] to AgentRunParams
    (run-time params, ~line 218). Accept either a Budget instance or a plain dict.
  3. In build_run_payload (~line 1229): pop budget from kwargs, normalize it to the
    nested camelCase dict, and set execution_params["budget"] only when provided.
    Do not add a default — absence must produce no budget key.
    Reuse / extend _EXEC_PARAMS_MAP if a snake→camel mapping for budget is needed.
  4. Keep max_cost/max_duration_seconds/max_iterations snake→camel conversion local
    to the Budget serialization so user-supplied dicts in either case work.

Out of scope

  • No creation-time / persisted budget field on Agent.
  • No platform-backend or agents-service changes (separate tasks).
  • Do not expose warnAtPercent.

Acceptance criteria

  • agent.run("q", budget=Budget(max_cost=1.0, max_duration_seconds=300, max_iterations=50))
    produces a run payload whose executionParams.budget == {"maxCost":1.0,"maxDurationSeconds":300,"maxIterations":50}.
  • Passing budget as a plain dict (snake_case or camelCase) yields the same camelCase payload.
  • Partial budgets work: Budget(max_cost=2.0)executionParams.budget == {"maxCost":2.0} (no null keys).
  • No budget passed ⇒ no budget key in executionParams (byte-for-byte unchanged from today).
  • New unit test file tests/unit/v2/test_agent_budget.py covering the above.

Verification

python -m pytest tests/unit

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions