avenza 1.0.0__py3-none-any.whl

avenza/__init__.py

@@ -0,0 +1,30 @@
+"""
+Avenza Python SDK
+
+Instrument AI agents in one line. Cost, value, and SLO tracking — automatically.
+
+Quickstart:
+    from avenza import Agent
+
+    agent = Agent(name='Invoice Bot', risk_tier='T2')
+
+    with agent.run() as run:
+        result = process(data)
+        run.success = result.ok
+        run.log_value('task_completed', quantity=1, unit_value_usd=1.50)
+
+Auto-instrumentation is active by default — if you're using the official Anthropic,
+OpenAI, or Gemini client, token usage is captured without any additional code.
+"""
+from __future__ import annotations
+
+from .agent import Agent
+from .exceptions import AvenzaConfigError, AvenzaError
+
+try:
+    from importlib.metadata import version as _version
+    __version__: str = _version("avenza")
+except Exception:
+    __version__ = "dev"
+
+__all__ = ["Agent", "AvenzaError", "AvenzaConfigError", "__version__"]

avenza/_buffer.py

@@ -0,0 +1,52 @@
+"""
+Offline buffer — persists failed sends to .avenza_buffer.jsonl and retries
+on next SDK init. Designed for intermittent-connectivity environments.
+"""
+from __future__ import annotations
+
+import json
+import logging
+import os
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from .client import AvenzaClient
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_PATH = ".avenza_buffer.jsonl"
+
+
+class OfflineBuffer:
+    def __init__(self, path: str = _DEFAULT_PATH) -> None:
+        self._path = path
+
+    def save(self, path: str, payload: dict[str, Any]) -> None:
+        try:
+            with open(self._path, "a", encoding="utf-8") as f:
+                f.write(json.dumps({"path": path, "payload": payload}) + "\n")
+        except OSError:
+            pass  # Read-only FS or permission error — drop silently
+
+    def flush(self, client: "AvenzaClient") -> None:
+        if not os.path.exists(self._path):
+            return
+        try:
+            with open(self._path, encoding="utf-8") as f:
+                lines = f.readlines()
+            os.remove(self._path)
+        except OSError:
+            return
+
+        flushed = 0
+        for line in lines:
+            try:
+                entry = json.loads(line.strip())
+                if entry.get("path") and entry.get("payload") is not None:
+                    client.post_async(entry["path"], entry["payload"])
+                    flushed += 1
+            except (json.JSONDecodeError, KeyError):
+                pass
+
+        if flushed:
+            logger.info("avenza: flushed %d buffered run(s) from offline buffer", flushed)

avenza/_context.py

@@ -0,0 +1,30 @@
+"""
+contextvars-based active-run tracking.
+
+Each OS thread and each asyncio task gets its own isolated view of the current
+run, so token capture is always attributed to the correct run even under
+concurrent async tasks or multi-threaded agents.
+"""
+from __future__ import annotations
+
+import contextvars
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from .run import RunContext
+
+_current_run: contextvars.ContextVar[Optional["RunContext"]] = contextvars.ContextVar(
+    "avenza_current_run", default=None
+)
+
+
+def get_current_run() -> Optional["RunContext"]:
+    return _current_run.get()
+
+
+def set_current_run(run: "RunContext") -> contextvars.Token:
+    return _current_run.set(run)
+
+
+def reset_current_run(token: contextvars.Token) -> None:
+    _current_run.reset(token)

avenza/_instrument/__init__.py

@@ -0,0 +1,17 @@
+"""
+Auto-instrumentation dispatcher.
+
+patch_all() tries each provider's patch. If the provider's SDK is not
+installed, the patch is silently skipped — no error, no warning.
+"""
+from __future__ import annotations
+
+
+def patch_all() -> None:
+    from .anthropic_patch import patch_anthropic
+    from .openai_patch import patch_openai
+    from .gemini_patch import patch_gemini
+
+    patch_anthropic()
+    patch_openai()
+    patch_gemini()

avenza/_instrument/anthropic_patch.py

@@ -0,0 +1,77 @@
+"""
+Patches Anthropic's official Python client to capture token usage automatically.
+Wraps both the sync and async message creation methods.
+"""
+from __future__ import annotations
+
+import functools
+import logging
+
+from .._context import get_current_run
+
+logger = logging.getLogger(__name__)
+
+_patched = False
+
+
+def patch_anthropic() -> None:
+    global _patched
+    if _patched:
+        return
+
+    try:
+        import anthropic
+    except ImportError:
+        return  # SDK not installed — nothing to patch, no error
+
+    try:
+        _patch_sync(anthropic)
+        _patch_async(anthropic)
+        _patched = True
+        logger.debug("avenza: anthropic auto-instrumentation active")
+    except Exception as exc:
+        logger.debug("avenza: anthropic patch failed (non-fatal) — %s", exc)
+
+
+def _patch_sync(anthropic: object) -> None:
+    messages_cls = anthropic.resources.messages.Messages  # type: ignore[attr-defined]
+    original = messages_cls.create
+
+    @functools.wraps(original)
+    def wrapped(self: object, *args: object, **kwargs: object) -> object:
+        response = original(self, *args, **kwargs)
+        run = get_current_run()
+        if run is not None and hasattr(response, "usage"):
+            run._record_auto_tokens(
+                provider="anthropic",
+                model=str(kwargs.get("model", "unknown")),
+                input_tokens=getattr(response.usage, "input_tokens", 0),
+                output_tokens=getattr(response.usage, "output_tokens", 0),
+            )
+        return response
+
+    messages_cls.create = wrapped
+
+
+def _patch_async(anthropic: object) -> None:
+    try:
+        async_cls = anthropic.resources.messages.AsyncMessages  # type: ignore[attr-defined]
+    except AttributeError:
+        return  # older version — no async client
+
+    original_async = async_cls.create
+
+    @functools.wraps(original_async)
+    async def wrapped_async(self: object, *args: object, **kwargs: object) -> object:
+        response = await original_async(self, *args, **kwargs)
+        run = get_current_run()
+        if run is not None and hasattr(response, "usage"):
+            run._record_auto_tokens(
+                provider="anthropic",
+                model=str(kwargs.get("model", "unknown")),
+                input_tokens=getattr(response.usage, "input_tokens", 0),
+                output_tokens=getattr(response.usage, "output_tokens", 0),
+            )
+        return response
+
+    async_cls.create = wrapped_async

avenza/_instrument/gemini_patch.py

@@ -0,0 +1,76 @@
+"""
+Patches Google's Generative AI client to capture token usage automatically.
+Works with both google-generativeai (genai) and google-cloud-aiplatform clients.
+"""
+from __future__ import annotations
+
+import functools
+import logging
+
+from .._context import get_current_run
+
+logger = logging.getLogger(__name__)
+
+_patched = False
+
+
+def patch_gemini() -> None:
+    global _patched
+    if _patched:
+        return
+
+    try:
+        _patch_genai()
+        _patched = True
+    except ImportError:
+        pass  # Neither SDK installed — skip silently
+    except Exception as exc:
+        logger.debug("avenza: gemini patch failed (non-fatal) — %s", exc)
+
+
+def _patch_genai() -> None:
+    import google.generativeai as genai  # type: ignore[import]
+    model_cls = genai.GenerativeModel
+
+    original_sync = model_cls.generate_content
+
+    @functools.wraps(original_sync)
+    def wrapped_sync(self: object, *args: object, **kwargs: object) -> object:
+        response = original_sync(self, *args, **kwargs)
+        run = get_current_run()
+        if run is not None:
+            _capture_gemini_usage(run, response, getattr(self, "model_name", "gemini"))
+        return response
+
+    model_cls.generate_content = wrapped_sync
+    logger.debug("avenza: gemini (google-generativeai) auto-instrumentation active")
+
+    # Async variant
+    if hasattr(model_cls, "generate_content_async"):
+        original_async = model_cls.generate_content_async
+
+        @functools.wraps(original_async)
+        async def wrapped_async(self: object, *args: object, **kwargs: object) -> object:
+            response = await original_async(self, *args, **kwargs)
+            run = get_current_run()
+            if run is not None:
+                _capture_gemini_usage(run, response, getattr(self, "model_name", "gemini"))
+            return response
+
+        model_cls.generate_content_async = wrapped_async
+
+
+def _capture_gemini_usage(run: object, response: object, model_name: str) -> None:
+    # google-generativeai stores token counts in response.usage_metadata
+    usage = getattr(response, "usage_metadata", None)
+    if usage is None:
+        return
+    input_tokens  = getattr(usage, "prompt_token_count", 0) or 0
+    output_tokens = getattr(usage, "candidates_token_count", 0) or 0
+    if input_tokens or output_tokens:
+        run._record_auto_tokens(  # type: ignore[union-attr]
+            provider="google",
+            model=model_name,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+        )

avenza/_instrument/openai_patch.py

@@ -0,0 +1,77 @@
+"""
+Patches OpenAI's official Python client to capture token usage automatically.
+Handles both sync (Completions.create) and async (AsyncCompletions.create).
+"""
+from __future__ import annotations
+
+import functools
+import logging
+
+from .._context import get_current_run
+
+logger = logging.getLogger(__name__)
+
+_patched = False
+
+
+def patch_openai() -> None:
+    global _patched
+    if _patched:
+        return
+
+    try:
+        import openai
+    except ImportError:
+        return
+
+    try:
+        _patch_sync(openai)
+        _patch_async(openai)
+        _patched = True
+        logger.debug("avenza: openai auto-instrumentation active")
+    except Exception as exc:
+        logger.debug("avenza: openai patch failed (non-fatal) — %s", exc)
+
+
+def _patch_sync(openai: object) -> None:
+    completions_cls = openai.resources.chat.completions.Completions  # type: ignore[attr-defined]
+    original = completions_cls.create
+
+    @functools.wraps(original)
+    def wrapped(self: object, *args: object, **kwargs: object) -> object:
+        response = original(self, *args, **kwargs)
+        run = get_current_run()
+        if run is not None and hasattr(response, "usage") and response.usage is not None:
+            run._record_auto_tokens(
+                provider="openai",
+                model=str(kwargs.get("model", getattr(response, "model", "unknown"))),
+                input_tokens=getattr(response.usage, "prompt_tokens", 0),
+                output_tokens=getattr(response.usage, "completion_tokens", 0),
+            )
+        return response
+
+    completions_cls.create = wrapped
+
+
+def _patch_async(openai: object) -> None:
+    try:
+        async_cls = openai.resources.chat.completions.AsyncCompletions  # type: ignore[attr-defined]
+    except AttributeError:
+        return
+
+    original_async = async_cls.create
+
+    @functools.wraps(original_async)
+    async def wrapped_async(self: object, *args: object, **kwargs: object) -> object:
+        response = await original_async(self, *args, **kwargs)
+        run = get_current_run()
+        if run is not None and hasattr(response, "usage") and response.usage is not None:
+            run._record_auto_tokens(
+                provider="openai",
+                model=str(kwargs.get("model", getattr(response, "model", "unknown"))),
+                input_tokens=getattr(response.usage, "prompt_tokens", 0),
+                output_tokens=getattr(response.usage, "completion_tokens", 0),
+            )
+        return response
+
+    async_cls.create = wrapped_async

avenza/agent.py

@@ -0,0 +1,180 @@
+"""
+Agent — the entry point for instrumenting a single AI agent.
+
+agent = Agent(name='Invoice Bot', risk_tier='T2')
+
+with agent.run() as run:
+    result = do_work()
+    run.success = result.ok
+    run.log_value('task_completed', quantity=1, unit_value_usd=1.50)
+"""
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Optional
+
+from .client import AvenzaClient
+from .exceptions import AvenzaConfigError
+from .run import AgentRefGetter, RunContext
+
+logger = logging.getLogger(__name__)
+
+
+class Agent:
+    """
+    Represents a single AI agent registered in Avenza.
+
+    Parameters
+    ----------
+    name            : Agent display name — used to find-or-create on the platform.
+    risk_tier       : 'T1' (autonomous), 'T2' (approve-first), or 'T3' (assist-only).
+    api_key         : Bearer token with sdk scope. Falls back to AVENZA_API_KEY env var.
+    model           : LLM model identifier — used for cost table lookup.
+    owner           : Email of the accountable human. Defaults to API key's owning user.
+    auto_instrument : Patch official LLM client libraries automatically (default True).
+    offline_buffer  : Queue failed sends to local disk and retry on next init (default True).
+    base_url        : Override for self-hosted deployments.
+    language        : Runtime language tag sent in registration heartbeat.
+    """
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        risk_tier: str = "T1",
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        owner: Optional[str] = None,
+        auto_instrument: bool = True,
+        offline_buffer: bool = True,
+        base_url: Optional[str] = None,
+        language: str = "python",
+    ) -> None:
+        resolved_key = api_key or os.environ.get("AVENZA_API_KEY")
+        if not resolved_key:
+            raise AvenzaConfigError(
+                "No API key provided. Set AVENZA_API_KEY in your environment "
+                "or pass api_key= to Agent()."
+            )
+
+        self.name = name
+        self.risk_tier = risk_tier
+        self._model = model
+
+        import importlib.metadata as _meta
+        try:
+            sdk_version = _meta.version("avenza")
+        except Exception:
+            sdk_version = "dev"
+
+        self._client = AvenzaClient(
+            api_key=resolved_key,
+            base_url=base_url or os.environ.get("AVENZA_URL"),
+            offline_buffer=offline_buffer,
+        )
+
+        self._ref_getter = AgentRefGetter()
+
+        # Kick off agent registration on the background thread
+        self._client.post_async(
+            "/agents/register",
+            {
+                "agent_ref":   None,  # server assigns; name is the match key
+                "name":        name,
+                "risk_tier":   risk_tier,
+                "sdk_version": sdk_version,
+                "language":    language,
+                **({"model": model} if model else {}),
+                **({"owner": owner} if owner else {}),
+            },
+        )
+        # Registration response will never come back through post_async.
+        # We do a synchronous registration call to get the agent_ref.
+        # This is the one-time network cost per Agent() instantiation.
+        self._register_sync(name, risk_tier, model, owner, sdk_version, language)
+
+        if auto_instrument:
+            from ._instrument import patch_all
+            patch_all()
+
+    def run(self, run_ref: Optional[str] = None) -> RunContext:
+        """Return a RunContext to wrap a single agent execution."""
+        return RunContext(
+            client=self._client,
+            agent_ref_getter=self._ref_getter,
+            run_ref=run_ref,
+        )
+
+    def task(
+        self,
+        *,
+        value_type: Optional[str] = None,
+        unit_value_usd: float = 0.0,
+        quantity: float = 1.0,
+    ) -> Any:
+        """Decorator that wraps a function in a run automatically."""
+        import functools
+
+        def decorator(fn: Any) -> Any:
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                with self.run() as run:
+                    try:
+                        result = fn(*args, **kwargs)
+                        run.success = bool(result)
+                        if value_type:
+                            run.log_value(value_type, quantity=quantity, unit_value_usd=unit_value_usd)
+                        return result
+                    except Exception:
+                        run.success = False
+                        raise
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                async with self.run() as run:
+                    try:
+                        result = await fn(*args, **kwargs)
+                        run.success = bool(result)
+                        if value_type:
+                            run.log_value(value_type, quantity=quantity, unit_value_usd=unit_value_usd)
+                        return result
+                    except Exception:
+                        run.success = False
+                        raise
+
+            import asyncio as _asyncio
+            if _asyncio.iscoroutinefunction(fn):
+                return async_wrapper
+            return sync_wrapper
+
+        return decorator
+
+    def _register_sync(
+        self,
+        name: str,
+        risk_tier: str,
+        model: Optional[str],
+        owner: Optional[str],
+        sdk_version: str,
+        language: str,
+    ) -> None:
+        try:
+            data = self._client.post_sync(
+                "/agents/register",
+                {
+                    "agent_ref":   name,  # used as a hint; server may assign differently
+                    "name":        name,
+                    "risk_tier":   risk_tier,
+                    "sdk_version": sdk_version,
+                    "language":    language,
+                    **({"model": model} if model else {}),
+                    **({"owner": owner} if owner else {}),
+                },
+            )
+            agent_ref: str = data.get("agent_ref", name)
+            self._ref_getter.set(agent_ref)
+            logger.debug("avenza: registered agent=%s ref=%s", name, agent_ref)
+        except Exception as exc:
+            logger.debug("avenza: registration failed, will retry — %s", exc)
+            self._ref_getter.set(name.upper().replace(" ", "-")[:20])