agentbudget 0.1.0__py3-none-any.whl

agentbudget/__init__.py

@@ -0,0 +1,17 @@
+"""AgentBudget - Real-time cost enforcement for AI agent sessions."""
+
+__version__ = "0.1.0"
+
+from .budget import AgentBudget
+from .exceptions import AgentBudgetError, BudgetExhausted, InvalidBudget
+from .session import AsyncBudgetSession, BudgetSession, LoopDetected
+
+__all__ = [
+    "AgentBudget",
+    "AgentBudgetError",
+    "AsyncBudgetSession",
+    "BudgetExhausted",
+    "BudgetSession",
+    "InvalidBudget",
+    "LoopDetected",
+]

agentbudget/budget.py

@@ -0,0 +1,126 @@
+"""AgentBudget — top-level API for creating budget-enforced sessions."""
+
+from __future__ import annotations
+
+from typing import Callable, Optional, Union
+
+from .circuit_breaker import CircuitBreaker, LoopDetectorConfig
+from .exceptions import InvalidBudget
+from .ledger import Ledger
+from .session import AsyncBudgetSession, BudgetSession
+from .webhook import WebhookEmitter
+
+
+def _chain_callbacks(
+    user_cb: Optional[Callable], webhook_cb: Callable
+) -> Callable:
+    """Combine a user callback with a webhook callback."""
+    if user_cb is None:
+        return webhook_cb
+
+    def chained(report):
+        user_cb(report)
+        webhook_cb(report)
+
+    return chained
+
+
+def parse_budget(value: str | float | int) -> float:
+    """Parse a budget value into a float.
+
+    Accepts:
+        "$5.00", "$5", "5.00", "5", 5.0, 5
+    """
+    if isinstance(value, (int, float)):
+        if value <= 0:
+            raise InvalidBudget(str(value))
+        return float(value)
+
+    if isinstance(value, str):
+        cleaned = value.strip().lstrip("$").strip()
+        try:
+            amount = float(cleaned)
+        except ValueError:
+            raise InvalidBudget(value)
+        if amount <= 0:
+            raise InvalidBudget(value)
+        return amount
+
+    raise InvalidBudget(str(value))
+
+
+class AgentBudget:
+    """Create budget-enforced agent sessions.
+
+    Usage:
+        budget = AgentBudget(max_spend="$5.00")
+        with budget.session() as session:
+            response = session.wrap(llm_call(...))
+            session.track(tool_call(), cost=0.01)
+        print(session.report())
+    """
+
+    def __init__(
+        self,
+        max_spend: str | float | int,
+        soft_limit: float = 0.9,
+        max_repeated_calls: int = 10,
+        loop_window_seconds: float = 60.0,
+        on_soft_limit: Optional[Callable] = None,
+        on_hard_limit: Optional[Callable] = None,
+        on_loop_detected: Optional[Callable] = None,
+        webhook_url: Optional[str] = None,
+    ):
+        self._budget = parse_budget(max_spend)
+        self._soft_limit = soft_limit
+        self._loop_config = LoopDetectorConfig(
+            max_repeated_calls=max_repeated_calls,
+            time_window_seconds=loop_window_seconds,
+        )
+
+        # Wire up webhook emitter if URL is provided
+        if webhook_url:
+            emitter = WebhookEmitter(webhook_url)
+            self._on_soft_limit = _chain_callbacks(on_soft_limit, emitter.on_soft_limit)
+            self._on_hard_limit = _chain_callbacks(on_hard_limit, emitter.on_hard_limit)
+            self._on_loop_detected = _chain_callbacks(on_loop_detected, emitter.on_loop_detected)
+        else:
+            self._on_soft_limit = on_soft_limit
+            self._on_hard_limit = on_hard_limit
+            self._on_loop_detected = on_loop_detected
+
+    @property
+    def max_spend(self) -> float:
+        return self._budget
+
+    def session(self, session_id: Optional[str] = None) -> BudgetSession:
+        """Create a new budget session."""
+        ledger = Ledger(budget=self._budget)
+        circuit_breaker = CircuitBreaker(
+            soft_limit_fraction=self._soft_limit,
+            loop_config=self._loop_config,
+        )
+        return BudgetSession(
+            ledger=ledger,
+            session_id=session_id,
+            circuit_breaker=circuit_breaker,
+            on_soft_limit=self._on_soft_limit,
+            on_hard_limit=self._on_hard_limit,
+            on_loop_detected=self._on_loop_detected,
+        )
+
+    def async_session(self, session_id: Optional[str] = None) -> AsyncBudgetSession:
+        """Create a new async budget session."""
+        ledger = Ledger(budget=self._budget)
+        circuit_breaker = CircuitBreaker(
+            soft_limit_fraction=self._soft_limit,
+            loop_config=self._loop_config,
+        )
+        return AsyncBudgetSession(
+            ledger=ledger,
+            session_id=session_id,
+            circuit_breaker=circuit_breaker,
+            on_soft_limit=self._on_soft_limit,
+            on_hard_limit=self._on_hard_limit,
+            on_loop_detected=self._on_loop_detected,
+        )

agentbudget/circuit_breaker.py

@@ -0,0 +1,72 @@
+"""Circuit breaker — loop detection and budget threshold warnings."""
+
+from __future__ import annotations
+
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class LoopDetectorConfig:
+    """Configuration for loop detection."""
+
+    max_repeated_calls: int = 10
+    time_window_seconds: float = 60.0
+
+
+class LoopDetector:
+    """Detects when the same tool/model is called repeatedly in a short window."""
+
+    def __init__(self, config: Optional[LoopDetectorConfig] = None):
+        self._config = config or LoopDetectorConfig()
+        self._call_log: dict[str, list[float]] = defaultdict(list)
+
+    def record_call(self, key: str) -> bool:
+        """Record a call and return True if a loop is detected."""
+        now = time.time()
+        cutoff = now - self._config.time_window_seconds
+
+        # Prune old entries
+        self._call_log[key] = [
+            t for t in self._call_log[key] if t > cutoff
+        ]
+        self._call_log[key].append(now)
+
+        return len(self._call_log[key]) > self._config.max_repeated_calls
+
+    def reset(self) -> None:
+        """Clear all recorded calls."""
+        self._call_log.clear()
+
+
+class CircuitBreaker:
+    """Monitors budget usage and detects runaway loops."""
+
+    def __init__(
+        self,
+        soft_limit_fraction: float = 0.9,
+        loop_config: Optional[LoopDetectorConfig] = None,
+    ):
+        self._soft_limit_fraction = soft_limit_fraction
+        self._loop_detector = LoopDetector(loop_config)
+        self._soft_limit_triggered = False
+
+    @property
+    def soft_limit_triggered(self) -> bool:
+        return self._soft_limit_triggered
+
+    def check_budget(self, spent: float, budget: float) -> Optional[str]:
+        """Check budget thresholds. Returns warning message or None."""
+        if budget <= 0:
+            return None
+        fraction = spent / budget
+        if fraction >= self._soft_limit_fraction and not self._soft_limit_triggered:
+            self._soft_limit_triggered = True
+            return f"Soft limit reached: {fraction:.0%} of budget used (${spent:.4f} / ${budget:.2f})"
+        return None
+
+    def check_loop(self, key: str) -> bool:
+        """Record a call and return True if a loop is detected."""
+        return self._loop_detector.record_call(key)

agentbudget/exceptions.py

@@ -0,0 +1,24 @@
+"""Exceptions for AgentBudget."""
+
+
+class AgentBudgetError(Exception):
+    """Base exception for all AgentBudget errors."""
+
+
+class BudgetExhausted(AgentBudgetError):
+    """Raised when a session exceeds its allocated budget."""
+
+    def __init__(self, budget: float, spent: float):
+        self.budget = budget
+        self.spent = spent
+        super().__init__(
+            f"Budget exhausted: spent ${spent:.4f} of ${budget:.2f} budget"
+        )
+
+
+class InvalidBudget(AgentBudgetError):
+    """Raised when a budget value is invalid."""
+
+    def __init__(self, value: str):
+        self.value = value
+        super().__init__(f"Invalid budget value: {value!r}")

agentbudget/integrations/__init__.py

@@ -0,0 +1 @@
+"""Framework integrations for AgentBudget."""

agentbudget/integrations/crewai.py

@@ -0,0 +1,63 @@
+"""CrewAI integration for AgentBudget.
+
+Provides middleware that tracks costs for CrewAI agent runs.
+
+Usage:
+    from agentbudget.integrations.crewai import CrewAIBudgetMiddleware
+
+    middleware = CrewAIBudgetMiddleware(budget="$3.00")
+    # Use middleware.session to track costs in your CrewAI setup
+
+Requires: crewai (optional dependency)
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from ..budget import AgentBudget
+from ..session import BudgetSession
+
+
+class CrewAIBudgetMiddleware:
+    """Budget middleware for CrewAI agent runs.
+
+    Wraps a CrewAI execution with a budget session. Use the
+    `session` attribute to track costs within your CrewAI callbacks.
+    """
+
+    def __init__(
+        self,
+        budget: str | float | int,
+        session_id: Optional[str] = None,
+        on_soft_limit: Optional[Any] = None,
+        on_hard_limit: Optional[Any] = None,
+        on_loop_detected: Optional[Any] = None,
+    ):
+        self._agent_budget = AgentBudget(
+            max_spend=budget,
+            on_soft_limit=on_soft_limit,
+            on_hard_limit=on_hard_limit,
+            on_loop_detected=on_loop_detected,
+        )
+        self.session = self._agent_budget.session(session_id=session_id)
+
+    def __enter__(self) -> "CrewAIBudgetMiddleware":
+        self.session.__enter__()
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self.session.__exit__(*args)
+
+    def get_report(self) -> dict[str, Any]:
+        """Get the cost report for this middleware's session."""
+        return self.session.report()
+
+    def track(
+        self,
+        result: Any,
+        cost: float,
+        tool_name: Optional[str] = None,
+    ) -> Any:
+        """Track a cost within the CrewAI execution."""
+        return self.session.track(result, cost=cost, tool_name=tool_name)

agentbudget/integrations/langchain.py

@@ -0,0 +1,89 @@
+"""LangChain/LangGraph integration for AgentBudget.
+
+Provides a callback handler that tracks LLM costs automatically.
+
+Usage:
+    from agentbudget.integrations.langchain import LangChainBudgetCallback
+
+    callback = LangChainBudgetCallback(budget="$5.00")
+    agent.run(callbacks=[callback])
+
+    print(callback.session.report())
+
+Requires: langchain-core (optional dependency)
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from ..budget import AgentBudget, parse_budget
+from ..ledger import Ledger
+from ..pricing import calculate_llm_cost
+from ..session import BudgetSession
+from ..types import CostEvent, CostType
+
+try:
+    from langchain_core.callbacks import BaseCallbackHandler
+
+    _HAS_LANGCHAIN = True
+except ImportError:
+    _HAS_LANGCHAIN = False
+
+    # Provide a stub so the class definition doesn't fail at import
+    class BaseCallbackHandler:  # type: ignore[no-redef]
+        pass
+
+
+class LangChainBudgetCallback(BaseCallbackHandler):
+    """LangChain callback handler that enforces a per-run budget.
+
+    Tracks LLM call costs in real time and raises BudgetExhausted
+    when the budget is exceeded.
+    """
+
+    def __init__(
+        self,
+        budget: str | float | int,
+        session: Optional[BudgetSession] = None,
+        **kwargs: Any,
+    ):
+        if not _HAS_LANGCHAIN:
+            raise ImportError(
+                "langchain-core is required for LangChainBudgetCallback. "
+                "Install it with: pip install langchain-core"
+            )
+        super().__init__(**kwargs)
+        self._agent_budget = AgentBudget(max_spend=budget)
+        self.session = session or self._agent_budget.session()
+        self.session.__enter__()
+
+    def on_llm_end(self, response: Any, **kwargs: Any) -> None:
+        """Called when an LLM call finishes. Records the cost."""
+        llm_output = getattr(response, "llm_output", None) or {}
+        token_usage = llm_output.get("token_usage", {})
+        model_name = llm_output.get("model_name")
+
+        input_tokens = token_usage.get("prompt_tokens")
+        output_tokens = token_usage.get("completion_tokens")
+
+        if model_name and input_tokens is not None and output_tokens is not None:
+            cost = calculate_llm_cost(model_name, input_tokens, output_tokens)
+            if cost is not None:
+                event = CostEvent(
+                    cost=cost,
+                    cost_type=CostType.LLM,
+                    model=model_name,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                )
+                self.session._ledger.record(event)
+                self.session._check_after_record(call_key=model_name)
+
+    def on_tool_end(self, output: str, **kwargs: Any) -> None:
+        """Called when a tool finishes. Override to add cost tracking."""
+        pass
+
+    def get_report(self) -> dict[str, Any]:
+        """Get the cost report for this callback's session."""
+        return self.session.report()

agentbudget/ledger.py

@@ -0,0 +1,87 @@
+"""Budget ledger — tracks running totals and event history."""
+
+from __future__ import annotations
+
+import threading
+from typing import Any
+
+from .exceptions import BudgetExhausted
+from .types import CostEvent, CostType
+
+
+class Ledger:
+    """Thread-safe running balance tracker for a budget session."""
+
+    def __init__(self, budget: float):
+        self._budget = budget
+        self._spent = 0.0
+        self._events: list[CostEvent] = []
+        self._lock = threading.Lock()
+
+    @property
+    def budget(self) -> float:
+        return self._budget
+
+    @property
+    def spent(self) -> float:
+        with self._lock:
+            return self._spent
+
+    @property
+    def remaining(self) -> float:
+        with self._lock:
+            return self._budget - self._spent
+
+    @property
+    def events(self) -> list[CostEvent]:
+        with self._lock:
+            return list(self._events)
+
+    def record(self, event: CostEvent) -> None:
+        """Record a cost event. Raises BudgetExhausted if budget exceeded."""
+        with self._lock:
+            new_total = self._spent + event.cost
+            if new_total > self._budget:
+                raise BudgetExhausted(budget=self._budget, spent=new_total)
+            self._spent = new_total
+            self._events.append(event)
+
+    def would_exceed(self, cost: float) -> bool:
+        """Check if a cost would exceed the budget without recording it."""
+        with self._lock:
+            return (self._spent + cost) > self._budget
+
+    def breakdown(self) -> dict[str, Any]:
+        """Return a cost breakdown by type and model/tool."""
+        with self._lock:
+            llm_total = 0.0
+            llm_calls = 0
+            by_model: dict[str, float] = {}
+            tool_total = 0.0
+            tool_calls = 0
+            by_tool: dict[str, float] = {}
+
+            for event in self._events:
+                if event.cost_type == CostType.LLM:
+                    llm_total += event.cost
+                    llm_calls += 1
+                    if event.model:
+                        by_model[event.model] = by_model.get(event.model, 0.0) + event.cost
+                elif event.cost_type == CostType.TOOL:
+                    tool_total += event.cost
+                    tool_calls += 1
+                    if event.tool_name:
+                        by_tool[event.tool_name] = by_tool.get(event.tool_name, 0.0) + event.cost
+
+            return {
+                "llm": {
+                    "total": round(llm_total, 6),
+                    "calls": llm_calls,
+                    "by_model": {k: round(v, 6) for k, v in by_model.items()},
+                },
+                "tools": {
+                    "total": round(tool_total, 6),
+                    "calls": tool_calls,
+                    "by_tool": {k: round(v, 6) for k, v in by_tool.items()},
+                },
+            }

agentbudget/pricing.py

@@ -0,0 +1,78 @@
+"""Model pricing data for LLM cost calculation.
+
+Prices are per token in USD. Updated as of early 2025.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+# Mapping of model name -> (input_price_per_token, output_price_per_token)
+MODEL_PRICING: dict[str, tuple[float, float]] = {
+    # OpenAI
+    "gpt-4o": (2.50 / 1_000_000, 10.00 / 1_000_000),
+    "gpt-4o-2024-11-20": (2.50 / 1_000_000, 10.00 / 1_000_000),
+    "gpt-4o-2024-08-06": (2.50 / 1_000_000, 10.00 / 1_000_000),
+    "gpt-4o-mini": (0.15 / 1_000_000, 0.60 / 1_000_000),
+    "gpt-4o-mini-2024-07-18": (0.15 / 1_000_000, 0.60 / 1_000_000),
+    "gpt-4-turbo": (10.00 / 1_000_000, 30.00 / 1_000_000),
+    "gpt-4-turbo-2024-04-09": (10.00 / 1_000_000, 30.00 / 1_000_000),
+    "gpt-4": (30.00 / 1_000_000, 60.00 / 1_000_000),
+    "gpt-3.5-turbo": (0.50 / 1_000_000, 1.50 / 1_000_000),
+    "o1": (15.00 / 1_000_000, 60.00 / 1_000_000),
+    "o1-mini": (3.00 / 1_000_000, 12.00 / 1_000_000),
+    "o1-preview": (15.00 / 1_000_000, 60.00 / 1_000_000),
+    "o3-mini": (1.10 / 1_000_000, 4.40 / 1_000_000),
+    # Anthropic
+    "claude-3-5-sonnet-20241022": (3.00 / 1_000_000, 15.00 / 1_000_000),
+    "claude-3-5-sonnet-20240620": (3.00 / 1_000_000, 15.00 / 1_000_000),
+    "claude-3-5-haiku-20241022": (0.80 / 1_000_000, 4.00 / 1_000_000),
+    "claude-3-opus-20240229": (15.00 / 1_000_000, 75.00 / 1_000_000),
+    "claude-3-sonnet-20240229": (3.00 / 1_000_000, 15.00 / 1_000_000),
+    "claude-3-haiku-20240307": (0.25 / 1_000_000, 1.25 / 1_000_000),
+    # Google Gemini
+    "gemini-1.5-pro": (1.25 / 1_000_000, 5.00 / 1_000_000),
+    "gemini-1.5-pro-latest": (1.25 / 1_000_000, 5.00 / 1_000_000),
+    "gemini-1.5-flash": (0.075 / 1_000_000, 0.30 / 1_000_000),
+    "gemini-1.5-flash-latest": (0.075 / 1_000_000, 0.30 / 1_000_000),
+    "gemini-2.0-flash": (0.10 / 1_000_000, 0.40 / 1_000_000),
+    "gemini-2.0-flash-lite": (0.075 / 1_000_000, 0.30 / 1_000_000),
+    "gemini-1.0-pro": (0.50 / 1_000_000, 1.50 / 1_000_000),
+    # Mistral
+    "mistral-large-latest": (2.00 / 1_000_000, 6.00 / 1_000_000),
+    "mistral-large-2411": (2.00 / 1_000_000, 6.00 / 1_000_000),
+    "mistral-small-latest": (0.10 / 1_000_000, 0.30 / 1_000_000),
+    "mistral-small-2501": (0.10 / 1_000_000, 0.30 / 1_000_000),
+    "open-mistral-nemo": (0.15 / 1_000_000, 0.15 / 1_000_000),
+    "mistral-medium-latest": (2.70 / 1_000_000, 8.10 / 1_000_000),
+    "codestral-latest": (0.30 / 1_000_000, 0.90 / 1_000_000),
+    # Cohere
+    "command-r-plus": (2.50 / 1_000_000, 10.00 / 1_000_000),
+    "command-r": (0.15 / 1_000_000, 0.60 / 1_000_000),
+    "command-light": (0.30 / 1_000_000, 0.60 / 1_000_000),
+    "command": (1.00 / 1_000_000, 2.00 / 1_000_000),
+}
+
+
+def get_model_pricing(model: str) -> Optional[tuple[float, float]]:
+    """Look up per-token pricing for a model.
+
+    Returns (input_price_per_token, output_price_per_token) or None if unknown.
+    """
+    return MODEL_PRICING.get(model)
+
+
+def calculate_llm_cost(
+    model: str,
+    input_tokens: int,
+    output_tokens: int,
+) -> Optional[float]:
+    """Calculate the cost of an LLM call in USD.
+
+    Returns None if model pricing is not found.
+    """
+    pricing = get_model_pricing(model)
+    if pricing is None:
+        return None
+    input_price, output_price = pricing
+    return (input_tokens * input_price) + (output_tokens * output_price)