Skip to content

UnitCause is a real-time cost governance platform for AI agents. It provides budget enforcement, loop detection, cost analytics, and policy management via a web dashboard, REST API, and Python SDK. Set budgets, stop runaway agents, and monitor spending with ease.

License

Notifications You must be signed in to change notification settings

unitcause/unitcause-python

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 

History

2 Commits
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

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
  • kill β€” on_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

About

UnitCause is a real-time cost governance platform for AI agents. It provides budget enforcement, loop detection, cost analytics, and policy management via a web dashboard, REST API, and Python SDK. Set budgets, stop runaway agents, and monitor spending with ease.

Topics

Resources

License

Stars

Watchers

Forks

Packages

 
 
 

Contributors

Languages