tetherai-python 0.1.0a0__tar.gz

tetherai_python-0.1.0a0/PKG-INFO

@@ -0,0 +1,35 @@
+Metadata-Version: 2.4
+Name: tetherai-python
+Version: 0.1.0a0
+Summary: AI budget guardrails for LLM applications
+Author-email: TetherAI Engineering <engineering@tetherai.com>
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/tetherai/tetherai-python
+Project-URL: Repository, https://github.com/tetherai/tetherai-python
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: tiktoken>=0.7.0
+Provides-Extra: crewai
+Requires-Dist: crewai>=1.0.0; extra == "crewai"
+Provides-Extra: litellm
+Requires-Dist: litellm>=1.40.0; extra == "litellm"
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0; extra == "dev"
+Requires-Dist: pytest-cov>=5.0; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.23; extra == "dev"
+Requires-Dist: pytest-mock>=3.12; extra == "dev"
+Requires-Dist: ruff>=0.4; extra == "dev"
+Requires-Dist: mypy>=1.10; extra == "dev"
+Requires-Dist: crewai>=1.0.0; extra == "dev"
+Requires-Dist: litellm>=1.40.0; extra == "dev"
+
+[![CI](https://github.com/tetherai/tetherai-python/actions/workflows/ci.yml/badge.svg)](https://github.com/tetherai/tetherai-python/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/tetherai)](https://pypi.org/project/tetherai/)

tetherai_python-0.1.0a0/README.md

@@ -0,0 +1,2 @@
+[![CI](https://github.com/tetherai/tetherai-python/actions/workflows/ci.yml/badge.svg)](https://github.com/tetherai/tetherai-python/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/tetherai)](https://pypi.org/project/tetherai/)

tetherai_python-0.1.0a0/pyproject.toml

@@ -0,0 +1,72 @@
+[project]
+name = "tetherai-python"
+version = "0.1.0-alpha"
+description = "AI budget guardrails for LLM applications"
+readme = "README.md"
+license = { text = "Apache-2.0" }
+requires-python = ">=3.10"
+authors = [
+    { name = "TetherAI Engineering", email = "engineering@tetherai.com" }
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+
+dependencies = [
+    "tiktoken>=0.7.0",
+]
+
+[project.optional-dependencies]
+crewai = ["crewai>=1.0.0"]
+litellm = ["litellm>=1.40.0"]
+dev = [
+    "pytest>=8.0",
+    "pytest-cov>=5.0",
+    "pytest-asyncio>=0.23",
+    "pytest-mock>=3.12",
+    "ruff>=0.4",
+    "mypy>=1.10",
+    "crewai>=1.0.0",
+    "litellm>=1.40.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/tetherai/tetherai-python"
+Repository = "https://github.com/tetherai/tetherai-python"
+
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+asyncio_mode = "auto"
+markers = [
+    "integration: integration tests that require more setup",
+    "slow: slow running tests",
+]
+
+[tool.ruff]
+target-version = "py310"
+line-length = 100
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I", "N", "UP", "B", "C4", "SIM"]
+ignore = ["E501"]
+
+[tool.mypy]
+python_version = "3.10"
+strict = true
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = true

tetherai_python-0.1.0a0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

tetherai_python-0.1.0a0/src/tetherai/__init__.py

@@ -0,0 +1,42 @@
+from typing import Any
+
+from tetherai._version import __version__
+from tetherai.circuit_breaker import enforce_budget
+from tetherai.config import TetherConfig, load_config
+from tetherai.exceptions import (
+    BudgetExceededError,
+    TetherError,
+    TokenCountError,
+    TurnLimitError,
+    UnknownModelError,
+)
+
+
+class Tether:
+    """TetherAI namespace class."""
+
+    enforce_budget = staticmethod(enforce_budget)
+
+
+tether = Tether
+
+
+def protect_crew(*args: Any, **kwargs: Any) -> Any:
+    from tetherai.crewai.integration import protect_crew as _protect_crew
+
+    return _protect_crew(*args, **kwargs)
+
+
+__all__ = [
+    "BudgetExceededError",
+    "TetherConfig",
+    "TetherError",
+    "TokenCountError",
+    "TurnLimitError",
+    "UnknownModelError",
+    "__version__",
+    "enforce_budget",
+    "load_config",
+    "protect_crew",
+    "tether",
+]

tetherai_python-0.1.0a0/src/tetherai/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0-alpha"

tetherai_python-0.1.0a0/src/tetherai/budget.py

@@ -0,0 +1,118 @@
+import threading
+from dataclasses import dataclass
+from typing import Any
+
+from tetherai.exceptions import BudgetExceededError, TurnLimitError
+
+
+@dataclass
+class CallRecord:
+    input_tokens: int
+    output_tokens: int
+    model: str
+    cost_usd: float
+    duration_ms: float
+
+
+class BudgetTracker:
+    def __init__(
+        self,
+        run_id: str,
+        max_usd: float,
+        max_turns: int | None = None,
+    ):
+        self.run_id = run_id
+        self.max_usd = max_usd
+        self.max_turns = max_turns
+        self._spent_usd = 0.0
+        self._turn_count = 0
+        self._calls: list[CallRecord] = []
+        self._lock = threading.Lock()
+
+    @property
+    def spent_usd(self) -> float:
+        with self._lock:
+            return self._spent_usd
+
+    @property
+    def remaining_usd(self) -> float:
+        with self._lock:
+            return max(0.0, self.max_usd - self._spent_usd)
+
+    @property
+    def turn_count(self) -> int:
+        with self._lock:
+            return self._turn_count
+
+    @property
+    def is_exceeded(self) -> bool:
+        with self._lock:
+            return self._spent_usd >= self.max_usd
+
+    def pre_check(self, estimated_input_cost: float) -> None:
+        with self._lock:
+            projected = self._spent_usd + estimated_input_cost
+            if projected > self.max_usd:
+                raise BudgetExceededError(
+                    message=f"Budget exceeded: ${projected:.2f} > ${self.max_usd:.2f}",
+                    run_id=self.run_id,
+                    budget_usd=self.max_usd,
+                    spent_usd=projected,
+                    last_model="unknown",
+                )
+
+    def record_call(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        model: str,
+        cost_usd: float,
+        duration_ms: float,
+    ) -> None:
+        if cost_usd < 0:
+            raise ValueError("cost_usd must be non-negative")
+
+        with self._lock:
+            if self.max_turns is not None and self._turn_count >= self.max_turns:
+                raise TurnLimitError(
+                    message=f"Turn limit exceeded: {self._turn_count} >= {self.max_turns}",
+                    run_id=self.run_id,
+                    max_turns=self.max_turns,
+                    current_turn=self._turn_count + 1,
+                )
+
+            self._spent_usd += cost_usd
+            self._turn_count += 1
+
+            self._calls.append(
+                CallRecord(
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                    model=model,
+                    cost_usd=cost_usd,
+                    duration_ms=duration_ms,
+                )
+            )
+
+            if self._spent_usd > self.max_usd:
+                self._spent_usd = self.max_usd
+
+    def get_summary(self) -> dict[str, Any]:
+        with self._lock:
+            return {
+                "run_id": self.run_id,
+                "budget_usd": self.max_usd,
+                "spent_usd": self._spent_usd,
+                "remaining_usd": max(0.0, self.max_usd - self._spent_usd),
+                "turn_count": self._turn_count,
+                "calls": [
+                    {
+                        "input_tokens": call.input_tokens,
+                        "output_tokens": call.output_tokens,
+                        "model": call.model,
+                        "cost_usd": call.cost_usd,
+                        "duration_ms": call.duration_ms,
+                    }
+                    for call in self._calls
+                ],
+            }

tetherai_python-0.1.0a0/src/tetherai/circuit_breaker.py

@@ -0,0 +1,132 @@
+import asyncio
+import functools
+import uuid
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from tetherai.budget import BudgetTracker
+from tetherai.config import TetherConfig
+from tetherai.exceptions import BudgetExceededError
+from tetherai.exporter import get_exporter
+from tetherai.interceptor import LLMInterceptor
+from tetherai.pricing import PricingRegistry
+from tetherai.token_counter import TokenCounter
+from tetherai.trace import TraceCollector
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def enforce_budget(
+    max_usd: float,
+    max_turns: int | None = None,
+    on_exceed: str = "raise",
+    trace_export: str | None = None,
+) -> Callable[[F], F]:
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return _run_with_budget(
+                func, max_usd, max_turns, on_exceed, trace_export, *args, **kwargs
+            )
+
+        @functools.wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            return await _run_with_budget_async(
+                func, max_usd, max_turns, on_exceed, trace_export, *args, **kwargs
+            )
+
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper  # type: ignore[return-value]
+        return wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def _run_with_budget(
+    func: Callable[..., Any],
+    max_usd: float,
+    max_turns: int | None,
+    on_exceed: str,
+    trace_export: str | None,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    run_id = f"run-{uuid.uuid4().hex[:8]}"
+    config = TetherConfig()
+
+    if trace_export is None:
+        trace_export = config.trace_export
+
+    budget_tracker = BudgetTracker(run_id=run_id, max_usd=max_usd, max_turns=max_turns)
+    token_counter = TokenCounter()
+    pricing = PricingRegistry()
+    trace_collector = TraceCollector()
+
+    trace_collector.start_trace(run_id, budget_tracker.get_summary())
+
+    interceptor = LLMInterceptor(
+        budget_tracker=budget_tracker,
+        token_counter=token_counter,
+        pricing=pricing,
+        trace_collector=trace_collector,
+    )
+
+    try:
+        interceptor.activate()
+        result = func(*args, **kwargs)
+        return result
+    except BudgetExceededError:
+        if on_exceed == "return_none":
+            return None
+        raise
+    finally:
+        interceptor.deactivate()
+        trace = trace_collector.end_trace()
+        if trace and trace_export != "none":
+            exporter = get_exporter(trace_export, config.trace_export_path)
+            exporter.export(trace)
+
+
+async def _run_with_budget_async(
+    func: Callable[..., Any],
+    max_usd: float,
+    max_turns: int | None,
+    on_exceed: str,
+    trace_export: str | None,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    run_id = f"run-{uuid.uuid4().hex[:8]}"
+    config = TetherConfig()
+
+    if trace_export is None:
+        trace_export = config.trace_export
+
+    budget_tracker = BudgetTracker(run_id=run_id, max_usd=max_usd, max_turns=max_turns)
+    token_counter = TokenCounter()
+    pricing = PricingRegistry()
+    trace_collector = TraceCollector()
+
+    trace_collector.start_trace(run_id, budget_tracker.get_summary())
+
+    interceptor = LLMInterceptor(
+        budget_tracker=budget_tracker,
+        token_counter=token_counter,
+        pricing=pricing,
+        trace_collector=trace_collector,
+    )
+
+    try:
+        interceptor.activate()
+        result = await func(*args, **kwargs)
+        return result
+    except BudgetExceededError:
+        if on_exceed == "return_none":
+            return None
+        raise
+    finally:
+        interceptor.deactivate()
+        trace = trace_collector.end_trace()
+        if trace and trace_export != "none":
+            exporter = get_exporter(trace_export, config.trace_export_path)
+            exporter.export(trace)

tetherai_python-0.1.0a0/src/tetherai/config.py

@@ -0,0 +1,91 @@
+import os
+from dataclasses import dataclass
+from typing import Any, Literal, cast
+
+TokenCounterBackend = Literal["tiktoken", "litellm", "auto"]
+PricingSource = Literal["bundled", "litellm"]
+TraceExport = Literal["console", "json", "none", "otlp"]
+
+
+@dataclass(frozen=True)
+class TetherConfig:
+    collector_url: str | None = None
+    default_budget_usd: float = 10.0
+    default_max_turns: int = 50
+    token_counter_backend: TokenCounterBackend = "auto"
+    pricing_source: PricingSource = "bundled"
+    log_level: str = "WARNING"
+    trace_export: TraceExport = "console"
+    trace_export_path: str = "./tetherai_traces/"
+
+    def __post_init__(self) -> None:
+        if self.default_budget_usd < 0:
+            raise ValueError("default_budget_usd must be non-negative")
+
+        if self.default_max_turns is not None and self.default_max_turns < 0:
+            raise ValueError("default_max_turns must be non-negative")
+
+        valid_backends = ("tiktoken", "litellm", "auto")
+        if self.token_counter_backend not in valid_backends:
+            raise ValueError(
+                f"Invalid token_counter_backend: {self.token_counter_backend}. "
+                f"Must be one of {valid_backends}"
+            )
+
+        valid_pricing = ("bundled", "litellm")
+        if self.pricing_source not in valid_pricing:
+            raise ValueError(
+                f"Invalid pricing_source: {self.pricing_source}. Must be one of {valid_pricing}"
+            )
+
+        valid_export = ("console", "json", "none", "otlp")
+        if self.trace_export not in valid_export:
+            raise ValueError(
+                f"Invalid trace_export: {self.trace_export}. Must be one of {valid_export}"
+            )
+
+    @classmethod
+    def from_env(cls) -> "TetherConfig":
+        return cls(
+            collector_url=os.getenv("TETHERAI_COLLECTOR_URL"),
+            default_budget_usd=float(os.getenv("TETHERAI_DEFAULT_BUDGET_USD", "10.0")),
+            default_max_turns=int(os.getenv("TETHERAI_DEFAULT_MAX_TURNS", "50")),
+            token_counter_backend=cls._resolve_backend(
+                os.getenv("TETHERAI_TOKEN_COUNTER_BACKEND", "auto")
+            ),
+            pricing_source=cast(
+                PricingSource, os.getenv("TETHERAI_PRICING_SOURCE", "bundled") or "bundled"
+            ),
+            log_level=os.getenv("TETHERAI_LOG_LEVEL", "WARNING"),
+            trace_export=cast(
+                TraceExport, os.getenv("TETHERAI_TRACE_EXPORT", "console") or "console"
+            ),
+            trace_export_path=os.getenv("TETHERAI_TRACE_EXPORT_PATH", "./tetherai_traces/"),
+        )
+
+    @staticmethod
+    def _resolve_backend(backend: str) -> TokenCounterBackend:
+        if backend == "auto":
+            try:
+                import litellm  # noqa: F401
+
+                return "litellm"
+            except ImportError:
+                return "tiktoken"
+        return backend  # type: ignore[return-value]
+
+    def resolve_backend(self) -> TokenCounterBackend:
+        return self._resolve_backend(self.token_counter_backend)
+
+
+def load_config(**kwargs: Any) -> TetherConfig:
+    env_config = TetherConfig.from_env()
+
+    config_dict = {}
+    for field_name in TetherConfig.__dataclass_fields__:
+        if field_name in kwargs:
+            config_dict[field_name] = kwargs[field_name]
+        else:
+            config_dict[field_name] = getattr(env_config, field_name)
+
+    return TetherConfig(**config_dict)

tetherai_python-0.1.0a0/src/tetherai/crewai/__init__.py

@@ -0,0 +1,3 @@
+from tetherai.crewai.integration import protect_crew, tether_step_callback, tether_task_callback
+
+__all__ = ["protect_crew", "tether_step_callback", "tether_task_callback"]

tetherai_python-0.1.0a0/src/tetherai/crewai/integration.py

@@ -0,0 +1,68 @@
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from crewai import Crew
+
+
+def _check_crewai_installed() -> None:
+    try:
+        import crewai  # noqa: F401
+    except ImportError:
+        raise ImportError(
+            "crewai is not installed. Install it with: pip install tetherai[crewai]"
+        ) from None
+
+
+def protect_crew(
+    crew: "Crew",
+    max_usd: float,
+    max_turns: int | None = None,
+) -> "Crew":
+    _check_crewai_installed()
+
+    from tetherai.circuit_breaker import enforce_budget
+
+    original_kickoff = crew.kickoff
+
+    @enforce_budget(max_usd=max_usd, max_turns=max_turns)
+    def wrapped_kickoff(*args: Any, **kwargs: Any) -> Any:
+        return original_kickoff(*args, **kwargs)
+
+    crew.kickoff = wrapped_kickoff  # type: ignore[method-assign]
+
+    for agent in crew.agents:
+        original_step_callback = getattr(agent, "step_callback", None)
+
+        def make_callback(original: Callable[..., Any] | None) -> Callable[..., Any]:
+            def callback(step_output: Any) -> None:
+                if original:
+                    original(step_output)
+
+            return callback
+
+        if original_step_callback is not None:
+            agent.step_callback = make_callback(original_step_callback)  # type: ignore[attr-defined]
+
+    for task in crew.tasks:
+        original_task_callback = getattr(task, "callback", None)
+
+        def make_task_callback(original: Callable[..., Any] | None) -> Callable[..., Any]:
+            def callback(task_output: Any) -> None:
+                if original:
+                    original(task_output)
+
+            return callback
+
+        if original_task_callback is not None:
+            task.callback = make_task_callback(original_task_callback)
+
+    return crew
+
+
+def tether_step_callback(step_output: Any) -> None:
+    pass
+
+
+def tether_task_callback(task_output: Any) -> None:
+    pass

tetherai_python-0.1.0a0/src/tetherai/exceptions.py

@@ -0,0 +1,92 @@
+class TetherError(Exception):
+    """Base exception for all TetherAI errors."""
+
+
+class BudgetExceededError(TetherError):
+    """Raised when a run's accumulated cost exceeds its budget."""
+
+    def __init__(
+        self,
+        message: str,
+        run_id: str,
+        budget_usd: float,
+        spent_usd: float,
+        last_model: str,
+        trace_url: str | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.run_id = run_id
+        self.budget_usd = budget_usd
+        self.spent_usd = spent_usd
+        self.last_model = last_model
+        self.trace_url = trace_url
+
+    def __str__(self) -> str:
+        return (
+            f"Budget exceeded: ${self.spent_usd:.2f} / ${self.budget_usd:.2f} on run {self.run_id}"
+        )
+
+    def __reduce__(self) -> tuple[type, tuple]:  # type: ignore[type-arg]
+        return (
+            self.__class__,
+            (
+                self.args[0],
+                self.run_id,
+                self.budget_usd,
+                self.spent_usd,
+                self.last_model,
+                self.trace_url,
+            ),
+        )
+
+
+class TurnLimitError(TetherError):
+    """Raised when an agent exceeds max allowed LLM calls."""
+
+    def __init__(
+        self,
+        message: str,
+        run_id: str,
+        max_turns: int,
+        current_turn: int,
+    ) -> None:
+        super().__init__(message)
+        self.run_id = run_id
+        self.max_turns = max_turns
+        self.current_turn = current_turn
+
+    def __str__(self) -> str:
+        return f"Turn limit exceeded: {self.current_turn} / {self.max_turns} on run {self.run_id}"
+
+    def __reduce__(self) -> tuple[type, tuple]:  # type: ignore[type-arg]
+        return (
+            self.__class__,
+            (
+                self.args[0],
+                self.run_id,
+                self.max_turns,
+                self.current_turn,
+            ),
+        )
+
+
+class TokenCountError(TetherError):
+    """Raised when token counting fails (e.g., unknown encoding)."""
+
+    def __init__(self, message: str, model: str | None = None) -> None:
+        super().__init__(message)
+        self.model = model
+
+    def __reduce__(self) -> tuple[type, tuple]:  # type: ignore[type-arg]
+        return (self.__class__, (self.args[0], self.model))
+
+
+class UnknownModelError(TetherError):
+    """Raised when a model is not found in the pricing registry."""
+
+    def __init__(self, message: str, model: str) -> None:
+        super().__init__(message)
+        self.model = model
+
+    def __reduce__(self) -> tuple[type, tuple]:  # type: ignore[type-arg]
+        return (self.__class__, (self.args[0], self.model))