tetherai-python 0.1.0a0__py3-none-any.whl

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/_version.py

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

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/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/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/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/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/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))

tetherai/exporter.py

@@ -0,0 +1,60 @@
+import json
+import sys
+from pathlib import Path
+from typing import Protocol, runtime_checkable
+
+from tetherai.trace import Trace
+
+
+@runtime_checkable
+class TraceExporter(Protocol):
+    def export(self, trace: Trace) -> None: ...
+
+
+class ConsoleExporter:
+    def export(self, trace: Trace) -> None:
+        print(f"=== TetherAI Trace: {trace.run_id} ===", file=sys.stderr)
+        print(f"Total Cost: ${trace.total_cost:.4f}", file=sys.stderr)
+        print(f"Input Tokens: {trace.total_input_tokens}", file=sys.stderr)
+        print(f"Output Tokens: {trace.total_output_tokens}", file=sys.stderr)
+        print(f"Spans: {len(trace.spans)}", file=sys.stderr)
+        print(file=sys.stderr)
+
+        for i, span in enumerate(trace.spans):
+            print(f"  [{i + 1}] {span.span_type}: {span.model or 'N/A'}", file=sys.stderr)
+            if span.cost_usd is not None:
+                print(f"      Cost: ${span.cost_usd:.6f}", file=sys.stderr)
+            if span.input_tokens:
+                print(f"      Input: {span.input_tokens} tokens", file=sys.stderr)
+            if span.output_tokens:
+                print(f"      Output: {span.output_tokens} tokens", file=sys.stderr)
+            print(file=sys.stderr)
+
+
+class JSONFileExporter:
+    def __init__(self, output_dir: str = "./tetherai_traces/"):
+        self.output_dir = Path(output_dir)
+
+    def export(self, trace: Trace) -> None:
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        filename = f"{trace.run_id}.json"
+        filepath = self.output_dir / filename
+
+        with open(filepath, "w") as f:
+            json.dump(trace.to_dict(), f, indent=2)
+
+
+class NoopExporter:
+    def export(self, trace: Trace) -> None:
+        pass
+
+
+def get_exporter(exporter_type: str, output_dir: str = "./tetherai_traces/") -> TraceExporter:
+    if exporter_type == "console":
+        return ConsoleExporter()
+    elif exporter_type == "json":
+        return JSONFileExporter(output_dir=output_dir)
+    elif exporter_type == "none" or exporter_type == "noop":
+        return NoopExporter()
+    else:
+        raise ValueError(f"Unknown exporter type: {exporter_type}")

tetherai/interceptor.py

@@ -0,0 +1,258 @@
+import time
+from collections.abc import Callable
+from typing import Any
+
+from tetherai.budget import BudgetTracker
+from tetherai.exceptions import BudgetExceededError, TetherError
+from tetherai.pricing import PricingRegistry
+from tetherai.token_counter import TokenCounter
+from tetherai.trace import Span, TraceCollector
+
+
+class LLMInterceptor:
+    def __init__(
+        self,
+        budget_tracker: BudgetTracker,
+        token_counter: TokenCounter,
+        pricing: PricingRegistry,
+        trace_collector: TraceCollector,
+    ):
+        self.budget_tracker = budget_tracker
+        self.token_counter = token_counter
+        self.pricing = pricing
+        self.trace_collector = trace_collector
+
+        self._original_completion: Callable[..., Any] | None = None
+        self._original_acompletion: Callable[..., Any] | None = None
+        self._active = False
+
+    def activate(self) -> None:
+        if self._active:
+            raise TetherError("Interceptor is already active")
+
+        try:
+            import litellm
+        except ImportError:
+            return
+
+        self._original_completion = litellm.completion
+        self._original_acompletion = litellm.acompletion
+
+        litellm.completion = self._patched_completion
+        litellm.acompletion = self._patched_acompletion
+        self._active = True
+
+    def deactivate(self) -> None:
+        if not self._active:
+            return
+
+        try:
+            import litellm
+        except ImportError:
+            self._active = False
+            return
+
+        if self._original_completion:
+            litellm.completion = self._original_completion
+        if self._original_acompletion:
+            litellm.acompletion = self._original_acompletion
+
+        self._active = False
+
+    def __enter__(self) -> "LLMInterceptor":
+        self.activate()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.deactivate()
+
+    def _patched_completion(self, *args: Any, **kwargs: Any) -> Any:
+        return self._intercept_call(self._original_completion, *args, **kwargs)
+
+    async def _patched_acompletion(self, *args: Any, **kwargs: Any) -> Any:
+        return await self._intercept_call_async(self._original_acompletion, *args, **kwargs)
+
+    def _intercept_call(
+        self, original_fn: Callable[..., Any] | None, *args: Any, **kwargs: Any
+    ) -> Any:
+        if original_fn is None:
+            raise TetherError("Interceptor not properly activated")
+        model = kwargs.get("model", args[0] if args else "unknown")
+        messages = kwargs.get("messages", [])
+
+        start_time = time.time()
+
+        try:
+            input_tokens = self.token_counter.count_messages(messages, model)
+        except Exception:
+            input_tokens = 0
+
+        try:
+            estimated_input_cost = self.pricing.get_input_cost(model) * input_tokens
+        except Exception:
+            estimated_input_cost = 0
+
+        try:
+            self.budget_tracker.pre_check(estimated_input_cost)
+        except BudgetExceededError:
+            raise
+
+        span = Span(
+            run_id=self.budget_tracker.run_id,
+            span_type="llm_call",
+            model=model,
+            input_tokens=input_tokens,
+            input_preview=messages[0].get("content", "")[:200] if messages else None,
+        )
+        self.trace_collector.add_span(span)
+
+        try:
+            response = original_fn(*args, **kwargs)
+        except Exception:
+            span.status = "error"
+            span.duration_ms = (time.time() - start_time) * 1000
+            raise
+
+        duration_ms = (time.time() - start_time) * 1000
+
+        try:
+            usage = getattr(response, "usage", None)
+            if usage:
+                output_tokens = getattr(usage, "completion_tokens", 0)
+                actual_input_tokens = getattr(usage, "prompt_tokens", input_tokens)
+            else:
+                output_tokens = 0
+                actual_input_tokens = input_tokens
+
+            cost_usd = self.pricing.estimate_call_cost(model, actual_input_tokens, output_tokens)
+
+            span.output_tokens = output_tokens
+            span.input_tokens = actual_input_tokens
+            span.cost_usd = cost_usd
+            span.duration_ms = duration_ms
+            span.status = "ok"
+
+            try:
+                content = response.choices[0].message.content if response.choices else ""
+                span.output_preview = content[:200] if content else None
+            except Exception:
+                pass
+
+            self.budget_tracker.record_call(
+                actual_input_tokens,
+                output_tokens,
+                model,
+                cost_usd,
+                duration_ms,
+            )
+
+        except Exception:
+            pass
+
+        return response
+
+    async def _intercept_call_async(
+        self, original_fn: Callable[..., Any] | None, *args: Any, **kwargs: Any
+    ) -> Any:
+        if original_fn is None:
+            raise TetherError("Interceptor not properly activated")
+        model = kwargs.get("model", args[0] if args else "unknown")
+        messages = kwargs.get("messages", [])
+
+        start_time = time.time()
+
+        try:
+            input_tokens = self.token_counter.count_messages(messages, model)
+        except Exception:
+            input_tokens = 0
+
+        try:
+            estimated_input_cost = self.pricing.get_input_cost(model) * input_tokens
+        except Exception:
+            estimated_input_cost = 0
+
+        try:
+            self.budget_tracker.pre_check(estimated_input_cost)
+        except BudgetExceededError:
+            raise
+
+        span = Span(
+            run_id=self.budget_tracker.run_id,
+            span_type="llm_call",
+            model=model,
+            input_tokens=input_tokens,
+            input_preview=messages[0].get("content", "")[:200] if messages else None,
+        )
+        self.trace_collector.add_span(span)
+
+        try:
+            response = await original_fn(*args, **kwargs)
+        except Exception:
+            span.status = "error"
+            span.duration_ms = (time.time() - start_time) * 1000
+            raise
+
+        duration_ms = (time.time() - start_time) * 1000
+
+        try:
+            usage = getattr(response, "usage", None)
+            if usage:
+                output_tokens = getattr(usage, "completion_tokens", 0)
+                actual_input_tokens = getattr(usage, "prompt_tokens", input_tokens)
+            else:
+                output_tokens = 0
+                actual_input_tokens = input_tokens
+
+            cost_usd = self.pricing.estimate_call_cost(model, actual_input_tokens, output_tokens)
+
+            span.output_tokens = output_tokens
+            span.input_tokens = actual_input_tokens
+            span.cost_usd = cost_usd
+            span.duration_ms = duration_ms
+            span.status = "ok"
+
+            try:
+                content = response.choices[0].message.content if response.choices else ""
+                span.output_preview = content[:200] if content else None
+            except Exception:
+                pass
+
+            self.budget_tracker.record_call(
+                actual_input_tokens,
+                output_tokens,
+                model,
+                cost_usd,
+                duration_ms,
+            )
+
+        except Exception:
+            pass
+
+        return response
+
+    def track_call(
+        self,
+        model: str,
+        input_tokens: int,
+        output_tokens: int,
+    ) -> None:
+        cost_usd = self.pricing.estimate_call_cost(model, input_tokens, output_tokens)
+        self.budget_tracker.pre_check(self.pricing.get_input_cost(model) * input_tokens)
+        self.budget_tracker.record_call(
+            input_tokens,
+            output_tokens,
+            model,
+            cost_usd,
+            0.0,
+        )
+
+        span = Span(
+            run_id=self.budget_tracker.run_id,
+            span_type="llm_call",
+            model=model,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cost_usd=cost_usd,
+            status="ok",
+        )
+        self.trace_collector.add_span(span)

tetherai/pricing.py

@@ -0,0 +1,99 @@
+from tetherai.exceptions import UnknownModelError
+
+BUNDLED_PRICING: dict[str, tuple[float, float]] = {
+    "gpt-4o": (0.0025, 0.01),
+    "gpt-4o-mini": (0.00015, 0.0006),
+    "gpt-4-turbo": (0.01, 0.03),
+    "gpt-4": (0.03, 0.06),
+    "gpt-3.5-turbo": (0.0005, 0.002),
+    "claude-3-5-sonnet-20241022": (0.003, 0.015),
+    "claude-3-5-sonnet": (0.003, 0.015),
+    "claude-3-opus-20240229": (0.015, 0.075),
+    "claude-3-opus": (0.015, 0.075),
+    "claude-3-sonnet-20240229": (0.003, 0.015),
+    "claude-3-sonnet": (0.003, 0.015),
+    "claude-3-haiku-20240307": (0.00025, 0.00125),
+    "claude-3-haiku": (0.00025, 0.00125),
+    "gemini-1.5-pro": (0.00125, 0.005),
+    "gemini-1.5-flash": (0.000075, 0.0003),
+    "gemini-1.5-flash-8b": (0.0000375, 0.00015),
+    "llama-3-70b": (0.0008, 0.0008),
+    "llama-3-8b": (0.0002, 0.0002),
+    "mixtral-8x7b": (0.00024, 0.00024),
+    "mistral-small": (0.001, 0.003),
+    "mistral-medium": (0.0024, 0.0072),
+    "mistral-large": (0.004, 0.012),
+}
+
+MODEL_ALIASES: dict[str, str] = {
+    "gpt4o": "gpt-4o",
+    "gpt-4o": "gpt-4o",
+    "gpt4o-mini": "gpt-4o-mini",
+    "gpt-4-turbo": "gpt-4-turbo",
+    "gpt4": "gpt-4",
+    "gpt-4": "gpt-4",
+    "gpt-3.5-turbo": "gpt-3.5-turbo",
+    "claude-sonnet": "claude-3-5-sonnet-20241022",
+    "claude-3.5-sonnet": "claude-3-5-sonnet-20241022",
+    "claude-opus": "claude-3-opus-20240229",
+    "claude-3-opus": "claude-3-opus-20240229",
+    "claude-sonnet-20240229": "claude-3-sonnet-20240229",
+    "claude-3-sonnet-20240229": "claude-3-sonnet-20240229",
+    "claude-haiku": "claude-3-haiku-20240307",
+    "claude-3-haiku": "claude-3-haiku-20240307",
+}
+
+
+class PricingRegistry:
+    def __init__(self, source: str = "bundled"):
+        self._source = source
+        self._custom_models: dict[str, tuple[float, float]] = {}
+        self._bundled = BUNDLED_PRICING.copy()
+
+    def get_input_cost(self, model: str) -> float:
+        resolved = self.resolve_model_alias(model)
+        if resolved in self._custom_models:
+            return self._custom_models[resolved][0]
+        if resolved in self._bundled:
+            return self._bundled[resolved][0]
+        if self._source == "litellm":
+            return self._get_litellm_cost(model, "input")
+        raise UnknownModelError(f"Unknown model: {model}", model)
+
+    def get_output_cost(self, model: str) -> float:
+        resolved = self.resolve_model_alias(model)
+        if resolved in self._custom_models:
+            return self._custom_models[resolved][1]
+        if resolved in self._bundled:
+            return self._bundled[resolved][1]
+        if self._source == "litellm":
+            return self._get_litellm_cost(model, "output")
+        raise UnknownModelError(f"Unknown model: {model}", model)
+
+    def estimate_call_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
+        input_cost = self.get_input_cost(model) * input_tokens
+        output_cost = self.get_output_cost(model) * output_tokens
+        return input_cost + output_cost
+
+    def resolve_model_alias(self, model: str) -> str:
+        normalized = model.lower().strip()
+        return MODEL_ALIASES.get(normalized, model)
+
+    def register_custom_model(self, model: str, input_cost: float, output_cost: float) -> None:
+        self._custom_models[model] = (input_cost, output_cost)
+
+    def _get_litellm_cost(self, model: str, direction: str) -> float:
+        try:
+            import litellm
+        except ImportError:
+            raise UnknownModelError(
+                f"Unknown model: {model} (litellm not installed)", model
+            ) from None
+        cost = litellm.cost_per_token(model, direction)  # type: ignore[arg-type,attr-defined]
+        if isinstance(cost, tuple):
+            return cost[0] if direction == "input" else cost[1]
+        return cost
+
+
+def get_pricing_registry(source: str = "bundled") -> PricingRegistry:
+    return PricingRegistry(source=source)

tetherai/token_counter.py

@@ -0,0 +1,150 @@
+import logging
+from typing import Any
+
+from tetherai.exceptions import TokenCountError
+
+logger = logging.getLogger(__name__)
+
+TOKENIZER_CACHE: dict[str, Any] = {}
+
+CHATML_FORMATTING = {
+    "system": {"prefix": "<|im_start|>system\n", "suffix": "<|im_end|>\n"},
+    "user": {"prefix": "<|im_start|>user\n", "suffix": "<|im_end|>\n"},
+    "assistant": {"prefix": "<|im_start|>assistant\n", "suffix": "<|im_end|>\n"},
+    "tool": {"prefix": "<|im_start|>tool\n", "suffix": "<|im_end|>\n"},
+}
+
+
+def _get_tiktoken_encoder(encoding_name: str = "cl100k_base") -> Any:
+    if encoding_name in TOKENIZER_CACHE:
+        return TOKENIZER_CACHE[encoding_name]
+
+    try:
+        import tiktoken
+
+        encoder = tiktoken.get_encoding(encoding_name)
+        TOKENIZER_CACHE[encoding_name] = encoder
+        return encoder
+    except Exception as e:
+        raise TokenCountError(f"Failed to load tiktoken: {e}") from e
+
+
+def _get_litellm_tokenizer(model: str) -> Any:
+    try:
+        import litellm
+
+        return litellm.token_counter  # type: ignore[attr-defined]
+    except ImportError as e:
+        raise TokenCountError("litellm not installed", model) from e
+
+
+class TokenCounter:
+    def __init__(self, backend: str = "auto"):
+        self._backend = backend
+        self._tiktoken_encoder = None
+        self._litellm_tokenizer = None
+
+        if backend == "auto":
+            try:
+                import litellm  # noqa: F401
+
+                self._backend = "litellm"
+            except ImportError:
+                self._backend = "tiktoken"
+
+        if self._backend == "tiktoken":
+            self._tiktoken_encoder = _get_tiktoken_encoder()
+
+    def count_tokens(self, text: str, model: str = "gpt-4o") -> int:
+        if not text:
+            return 0
+
+        if self._backend == "tiktoken":
+            return self._count_with_tiktoken(text, model)
+        elif self._backend == "litellm":
+            return self._count_with_litellm(text, model)
+        else:
+            raise TokenCountError(f"Unknown backend: {self._backend}")
+
+    def count_messages(self, messages: list[dict[str, str]], model: str = "gpt-4o") -> int:
+        if not messages:
+            return 0
+
+        if self._backend == "tiktoken":
+            return self._count_messages_with_tiktoken(messages, model)
+        elif self._backend == "litellm":
+            return self._count_messages_with_litellm(messages, model)
+        else:
+            raise TokenCountError(f"Unknown backend: {self._backend}")
+
+    def _count_with_tiktoken(self, text: str, model: str) -> int:
+        if model.startswith("claude-"):
+            logger.warning(
+                f"Using tiktoken for Claude model {model}. "
+                f"Token counts may be inaccurate (up to 12% error)."
+            )
+
+        encoder = self._tiktoken_encoder
+        if encoder is None:
+            encoder = _get_tiktoken_encoder()
+
+        return len(encoder.encode(text))
+
+    def _count_with_litellm(self, text: str, model: str) -> int:
+        if self._litellm_tokenizer is None:
+            self._litellm_tokenizer = _get_litellm_tokenizer(model)
+
+        if model.startswith("claude-"):
+            try:
+                return self._litellm_tokenizer(model=model, text=text)  # type: ignore[no-any-return,misc]
+            except Exception:
+                logger.warning(
+                    f"litellm token_counter failed for {model}, falling back to tiktoken"
+                )
+                return self._count_with_tiktoken(text, model)
+
+        return self._litellm_tokenizer(model=model, text=text)  # type: ignore[no-any-return,misc]
+
+    def _count_messages_with_tiktoken(self, messages: list[dict[str, str]], model: str) -> int:
+        encoder = self._tiktoken_encoder
+        if encoder is None:
+            encoder = _get_tiktoken_encoder()
+
+        if model.startswith("claude-"):
+            logger.warning(
+                f"Using tiktoken for Claude model {model}. Token counts may be inaccurate."
+            )
+
+        total_tokens = 0
+        for message in messages:
+            role = message.get("role", "user")
+            content = message.get("content", "")
+
+            formatting = CHATML_FORMATTING.get(role, CHATML_FORMATTING["user"])
+            formatted = f"{formatting['prefix']}{content}{formatting['suffix']}"
+            total_tokens += len(encoder.encode(formatted))
+
+        total_tokens += 3
+        return total_tokens
+
+    def _count_messages_with_litellm(self, messages: list[dict[str, str]], model: str) -> int:
+        if self._litellm_tokenizer is None:
+            self._litellm_tokenizer = _get_litellm_tokenizer(model)
+
+        try:
+            return self._litellm_tokenizer(model=model, messages=messages)  # type: ignore[no-any-return,misc]
+        except Exception:
+            logger.warning(f"litellm token_counter failed for {model}, falling back to tiktoken")
+            return self._count_messages_with_tiktoken(messages, model)
+
+
+def count_tokens(text: str, model: str = "gpt-4o", backend: str = "auto") -> int:
+    counter = TokenCounter(backend=backend)
+    return counter.count_tokens(text, model)
+
+
+def count_messages(
+    messages: list[dict[str, str]], model: str = "gpt-4o", backend: str = "auto"
+) -> int:
+    counter = TokenCounter(backend=backend)
+    return counter.count_messages(messages, model)

tetherai/trace.py

@@ -0,0 +1,117 @@
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+MAX_PREVIEW_LENGTH = 200
+
+
+def generate_id() -> str:
+    import uuid
+
+    return uuid.uuid4().hex[:16]
+
+
+@dataclass
+class Span:
+    span_id: str = field(default_factory=generate_id)
+    parent_span_id: str | None = None
+    run_id: str = ""
+    timestamp: datetime = field(default_factory=datetime.now)
+    duration_ms: float = 0.0
+    span_type: str = "llm_call"
+    model: str | None = None
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+    cost_usd: float | None = None
+    status: str = "ok"
+    metadata: dict[str, Any] = field(default_factory=dict)
+    input_preview: str | None = None
+    output_preview: str | None = None
+
+    def __post_init__(self) -> None:
+        if self.input_preview and len(self.input_preview) > MAX_PREVIEW_LENGTH:
+            self.input_preview = self.input_preview[:MAX_PREVIEW_LENGTH] + "..."
+
+        if self.output_preview and len(self.output_preview) > MAX_PREVIEW_LENGTH:
+            self.output_preview = self.output_preview[:MAX_PREVIEW_LENGTH] + "..."
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "span_id": self.span_id,
+            "parent_span_id": self.parent_span_id,
+            "run_id": self.run_id,
+            "timestamp": self.timestamp.isoformat(),
+            "duration_ms": self.duration_ms,
+            "span_type": self.span_type,
+            "model": self.model,
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "cost_usd": self.cost_usd,
+            "status": self.status,
+            "metadata": self.metadata,
+            "input_preview": self.input_preview,
+            "output_preview": self.output_preview,
+        }
+
+
+@dataclass
+class Trace:
+    run_id: str
+    spans: list[Span] = field(default_factory=list)
+    budget_summary: dict[str, Any] = field(default_factory=dict)
+    start_time: datetime = field(default_factory=datetime.now)
+    end_time: datetime | None = None
+
+    def add_span(self, span: Span) -> None:
+        self.spans.append(span)
+
+    @property
+    def total_cost(self) -> float:
+        return sum(span.cost_usd or 0 for span in self.spans)
+
+    @property
+    def total_input_tokens(self) -> int:
+        return sum(span.input_tokens or 0 for span in self.spans)
+
+    @property
+    def total_output_tokens(self) -> int:
+        return sum(span.output_tokens or 0 for span in self.spans)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "run_id": self.run_id,
+            "spans": [span.to_dict() for span in self.spans],
+            "budget_summary": self.budget_summary,
+            "start_time": self.start_time.isoformat(),
+            "end_time": self.end_time.isoformat() if self.end_time else None,
+            "total_cost": self.total_cost,
+            "total_input_tokens": self.total_input_tokens,
+            "total_output_tokens": self.total_output_tokens,
+        }
+
+
+class TraceCollector:
+    def __init__(self) -> None:
+        self._current_trace: Trace | None = None
+
+    def start_trace(self, run_id: str, budget_summary: dict[str, Any] | None = None) -> Trace:
+        self._current_trace = Trace(
+            run_id=run_id,
+            budget_summary=budget_summary or {},
+        )
+        return self._current_trace
+
+    def end_trace(self) -> Trace | None:
+        if self._current_trace:
+            self._current_trace.end_time = datetime.now()
+            trace = self._current_trace
+            self._current_trace = None
+            return trace
+        return None
+
+    def add_span(self, span: Span) -> None:
+        if self._current_trace:
+            self._current_trace.add_span(span)
+
+    def get_current_trace(self) -> Trace | None:
+        return self._current_trace

tetherai_python-0.1.0a0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,17 @@
+tetherai/__init__.py,sha256=OTKdW-yPFkSNB_1U5FbKJ7EUuH3NT-IlvSOpJrgSY1U,852
+tetherai/_version.py,sha256=04Pi7kDrn0FVD-1c5mk_o8ryOaXTkkiLyp6TW-qxcaE,28
+tetherai/budget.py,sha256=UUY481RFsiTPAwmaq2cYEx5Kdb7ifsTvYbTblmUvUfo,3552
+tetherai/circuit_breaker.py,sha256=FLIfVhJ2WnGRRXaofqH5w6JUlZdR2Kl7mvi_a9JiuBo,3878
+tetherai/config.py,sha256=lAo9hPj2fpU9V1AJw0PAOFPwJhloxxxD0UVpl814Rm4,3432
+tetherai/exceptions.py,sha256=kx21E9QV7oQJ1-bDfxVfmrokrHg7LaHgu7dI-W5s-Do,2643
+tetherai/exporter.py,sha256=tsyk1ol-n_w_hrUjRT3zdfwmcWFHu6DvmpR8f1CuewU,2137
+tetherai/interceptor.py,sha256=HOYSLoduxPM-6kt-cg06RFCJi4uBllXHg5g2BFTIemE,8163
+tetherai/pricing.py,sha256=Cdj6QDRvZOoYlO6aQORrVsQFHcXYto1lOxd7C6EMifg,3913
+tetherai/token_counter.py,sha256=DkvL2qpHpKJzp0ODNs9xZ-bzzwEJhExmE8g64CeaHME,5326
+tetherai/trace.py,sha256=blJ4liFo_mdzzDX8VeAu5_OlmgUxZQnZ-1VvFHoqUBE,3821
+tetherai/crewai/__init__.py,sha256=LDjmhTX2DfeN4z6-YY2wODCT60NlK551xA4pB6VI9Rw,173
+tetherai/crewai/integration.py,sha256=nDWk4DQhbGZmzxnEK7Ra_x8C_VbmvEjdn0Zagg032WY,1918
+tetherai_python-0.1.0a0.dist-info/METADATA,sha256=i6Wr5UmNn89mMoJvq8KB174yev1EcccODgF9IHlAyQ4,1605
+tetherai_python-0.1.0a0.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91
+tetherai_python-0.1.0a0.dist-info/top_level.txt,sha256=kKKtbtmezQqlgMwH9lmXDCY_GTSz1Z36J_VdCqUcxJM,9
+tetherai_python-0.1.0a0.dist-info/RECORD,,

tetherai_python-0.1.0a0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

tetherai_python-0.1.0a0.dist-info/top_level.txt

@@ -0,0 +1 @@
+tetherai