Skip to content

Latest commit

 

History

History
264 lines (196 loc) · 6.63 KB

File metadata and controls

264 lines (196 loc) · 6.63 KB

UnitCause Python SDK

Real-time cost governance for AI agents. Set budgets. Stop loops. Ship confidently.

PyPI version Python 3.9+ License: MIT

Installation

pip install unitcause

With framework integrations:

pip install unitcause[langchain]
pip install unitcause[crewai]
pip install unitcause[autogen]
pip install unitcause[all]

Quick Start

import os
from unitcause import UnitCause

# Set your API key
os.environ["UNITCAUSE_API_KEY"] = "uc_live_..."

uc = UnitCause()

# Wrap your agent run in a session
with uc.session(agent_name="data-analyst", budget=5.00) as session:
    # Run your agent — UnitCause tracks every LLM call
    result = agent.run("Analyze Q4 revenue trends")
    
    print(f"Total cost: ${session.total_cost:.4f}")
    print(f"Tokens used: {session.total_tokens:,}")
    print(f"LLM calls: {session.call_count}")

Async Support

import asyncio
from unitcause import UnitCause

uc = UnitCause()

async def run_agent():
    async with uc.async_session(agent_name="researcher", budget=2.00) as session:
        result = await agent.arun("Find top ML papers from 2025")
        print(f"Cost: ${session.total_cost:.4f}")

asyncio.run(run_agent())

Manual Step Tracking

If you're not using auto-instrumentation, you can manually report each LLM call:

with uc.session(agent_name="my-agent", budget=1.00) as session:
    # After each LLM call, report it
    response = session.report_step(
        action="llm_call",
        model="gpt-4o",
        tokens_in=500,
        tokens_out=200,
        cost_usd=0.0035,
    )
    
    # The response tells you whether to continue
    if response.action == "kill":
        print(f"Session killed: {response.reason}")
        break

Enforcement & Callbacks

Budget Callbacks

def on_warning(session, utilization):
    print(f"⚠️ Budget {utilization:.0%} used")

def on_exceeded(session, cost):
    print(f"🛑 Budget exceeded: ${cost:.4f}")
    return False  # Return True to allow continuing

with uc.session(
    agent_name="analyst",
    budget=5.00,
    on_budget_warning=on_warning,
    on_budget_exceeded=on_exceeded,
) as session:
    result = agent.run("...")

Kill Callback

React when UnitCause kills a session (budget, loop detection, or manual kill):

def on_kill(session, reason):
    print(f"Session killed: {reason}")
    # Clean up resources, save state, etc.

with uc.session(
    agent_name="agent",
    budget=5.00,
    on_kill=on_kill,
) as session:
    agent.run("...")

Enforcement Actions

The SDK handles enforcement automatically:

  • warn — Logged, on_budget_warning fired
  • throttle — SDK auto-sleeps for the configured delay
  • killon_kill fired, then raises BudgetExceededError, LoopDetectedError, or SessionKilledError

Framework Integrations

LangChain

from langchain_openai import ChatOpenAI
from unitcause import UnitCause
from unitcause.integrations.langchain import UnitCauseCallbackHandler

uc = UnitCause()

with uc.session(agent_name="my-chain", budget=5.00) as session:
    handler = UnitCauseCallbackHandler(session)
    llm = ChatOpenAI(model="gpt-4o", callbacks=[handler])
    result = llm.invoke("Hello!")

CrewAI

from crewai import Agent, Task, Crew
from unitcause import UnitCause
from unitcause.integrations.crewai import unitcause_step_callback

uc = UnitCause()

with uc.session(agent_name="my-crew", budget=10.00) as session:
    callback = unitcause_step_callback(session)
    agent = Agent(role="Researcher", step_callback=callback, ...)
    crew = Crew(agents=[agent], tasks=[...])
    crew.kickoff()

AutoGen

from autogen import AssistantAgent, UserProxyAgent
from unitcause import UnitCause
from unitcause.integrations.autogen import UnitCauseHook

uc = UnitCause()

with uc.session(agent_name="autogen-chat", budget=5.00) as session:
    hook = UnitCauseHook(session)
    assistant = AssistantAgent("assistant", llm_config={...})
    hook.attach(assistant)
    
    user = UserProxyAgent("user")
    user.initiate_chat(assistant, message="Hello!")

Custom Pricing

Override or add model pricing for cost estimation:

from unitcause.integrations.pricing import register_pricing

register_pricing("my-fine-tuned-model", input_per_1k=0.01, output_per_1k=0.03)

Health Check

status = uc.health_check()
print(status)
# {'status': 'connected', 'latency_ms': 42, 'version': '0.2.0'}

Configuration

Environment Variables

Variable Default Description
UNITCAUSE_API_KEY Your API key
UNITCAUSE_BASE_URL https://api.unitcause.com API endpoint
UNITCAUSE_ENVIRONMENT production Environment label
UNITCAUSE_DISABLED false Disable all tracking
UNITCAUSE_LOG_LEVEL WARNING Logging level
UNITCAUSE_DEFAULT_BUDGET Default session budget (USD)

Programmatic Configuration

from unitcause import UnitCause, Config

uc = UnitCause(
    api_key="uc_live_...",
    base_url="https://api.unitcause.com",
    environment="staging",
    disabled=False,
    log_level="DEBUG",
    retry_config={
        "max_retries": 3,
        "backoff_factor": 0.5,
    },
)

Exception Handling

from unitcause import (
    BudgetExceededError,
    LoopDetectedError,
    SessionKilledError,
)

try:
    with uc.session(agent_name="agent", budget=1.00) as session:
        agent.run("...")
except BudgetExceededError as e:
    print(f"Budget exceeded: ${e.actual_cost:.4f} / ${e.budget_limit:.4f}")
except LoopDetectedError as e:
    print(f"Loop detected: {e.iteration_count} iterations")
except SessionKilledError as e:
    print(f"Session killed: {e.reason}")

What's New in 0.2.0

  • Enforcement engine — configurable warn / throttle / kill actions per policy
  • Loop detection — dual strategy (exact hash match + pattern cycle detection)
  • on_kill callback — react when sessions are killed
  • Throttle handling — SDK auto-sleeps on throttle responses
  • Framework integrations — LangChain, CrewAI, AutoGen adapters
  • Built-in cost estimation — 60+ models with customizable pricing
  • SessionKilledError — new exception for killed sessions

Documentation

Full documentation at unitcause.com/docs

License

MIT — see LICENSE