Guardrails & Policies

# policies.yaml — load with niitaka.configure(policy_file="policies.yaml")
version: "1"

policies:
  # Warn at $0.30
  - agent_id: my-agent
    type: cost_limit
    priority: 5
    condition:
      cost_exceeded: 0.30
    action:
      type: warn

  # Hard abort at $0.50
  - agent_id: my-agent
    type: cost_limit
    priority: 10
    condition:
      cost_exceeded: 0.50
    action:
      type: abort

  # Cap the number of LLM + tool call steps
  - agent_id: my-agent
    type: step_limit
    priority: 10
    condition:
      steps_exceeded: 20
    action:
      type: abort

  # Retry transient errors with exponential backoff
  - agent_id: my-agent
    type: retry
    priority: 8
    condition:
      on_error: true
    action:
      max_retries: 3
      backoff_seconds: 2
      backoff: exponential
      on_errors:
        - RateLimitError
        - APITimeoutError
        - InternalServerError

  # Fall back to a cheaper model after retries are exhausted
  - agent_id: my-agent
    type: fallback
    priority: 7
    condition:
      on_error: true
    action:
      fallback_model: gpt-4o-mini
      on_errors:
        - RateLimitError
        - APITimeoutError

import os
import niitaka

niitaka.configure(
    api_key=os.getenv("NIITAKA_API_KEY"),
    policy_file="policies.yaml",      # file policies take priority over DB policies
)

import niitaka

# The SDK enforces guardrails automatically — no manual checks needed.
# To inspect active policy values from your agent code:
with niitaka.start_session(agent_id="my-agent"):
    result = niitaka.get_runtime_config()
    cfg = result.get("config", {})

    # Guardrail thresholds from the active shipped version / experiment variant:
    guardrails = cfg.get("guardrails", {})
    cost_limit = guardrails.get("cost_limit_usd")    # e.g. 0.50  (None if not set)
    max_steps  = guardrails.get("max_steps")          # e.g. 20    (None if not set)

    # Use for observability / logging — enforcement is handled by the SDK.
    print(f"Active cost limit: {cost_limit}, step limit: {max_steps}")

import niitaka

# A tool-calling loop makes several LLM calls for one logical task.
# Without policy_scope(), each start_session() resets cost/step accounting to
# zero, so cost_limit can never see the loop's cumulative spend — only ever
# one call at a time.
with niitaka.policy_scope(agent_id="my-agent", cost_limit_usd=0.50, max_steps=10):
    for step in range(10):
        with niitaka.start_session(agent_id="my-agent", goal=f"step {step}"):
            run_one_step()
    # cost_limit/step_limit fire against the SCOPE's running total across all
    # ten sessions above, not any single one of them.

# policy_scope() carries guardrail state across calls WITHIN one process.
# When the calls cross a process boundary (an upstream service proxying to a
# downstream one), two more pieces close that gap:

# Upstream — hand this scope's state to the downstream service:
with niitaka.policy_scope(agent_id="my-agent", cost_limit_usd=0.50):
    state = niitaka.get_active_scope_state()   # export the current cumulative total
    result = call_downstream_service(session_id=session.session_id, scope_state=state)
    niitaka.import_scope_state(result["scope_state"])  # merge downstream's post-call total back in

# Downstream — attach to the caller's session for visibility, import their
# scope state for accounting, and gate the call against the imported total:
with niitaka.attach_session(received_session_id, starting_sequence=received_sequence):
    with niitaka.policy_scope(imported_state=received_scope_state):
        niitaka.check_guardrails()   # raises AgentExecutionStopped if already over limit
        call_llm(...)
        return {"scope_state": niitaka.get_active_scope_state()}