hexgate 0.2.8__tar.gz → 0.2.9__tar.gz

{hexgate-0.2.8 → hexgate-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hexgate
-Version: 0.2.8
+Version: 0.2.9
 Summary: Hexgate — authorization infrastructure for AI agents (agent runtime + cloud client).
 License-Expression: MIT
 Requires-Python: >=3.13
@@ -23,7 +23,7 @@ Requires-Dist: openai-agents>=0.0.10
 Requires-Dist: langgraph>=0.2
 Requires-Dist: nest_asyncio>=1.6
 Requires-Dist: openinference-instrumentation-openai-agents>=0.1
-Requires-Dist: google-adk>=1.0
+Requires-Dist: google-adk>=1.14
 Requires-Dist: google-genai>=1.0
 Requires-Dist: litellm>=1.50
 Requires-Dist: openinference-instrumentation-google-adk>=0.1.11

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/google/runner.py

@@ -9,14 +9,16 @@ from typing import Any, AsyncGenerator, Generator
 
 import nest_asyncio
 from google.adk.agents import BaseAgent
+from google.adk.apps import App
 from google.adk.runners import Runner
 from google.adk.sessions import BaseSessionService
 from google.genai import types
 from langfuse import get_client, propagate_attributes
 from openinference.instrumentation.google_adk import GoogleADKInstrumentor
 
+from hexgate.adapters.google.usage import HexgateUsagePlugin
 from hexgate.adapters.google.wrapper import wrap_google_agent
-from hexgate.agents.factory import ApprovalHandler
+from hexgate.approvals import ApprovalHandler
 from hexgate.cloud.client import HexgateClient, HexgateConfig
 from hexgate.config.env import resolve_api_key
 from hexgate.runtime import User
@@ -51,9 +53,11 @@ class HexgateRunner:
             approval_handler=approval_handler,
             client=client,
         )
+        plugins = list(runner_kwargs.pop("plugins", None) or [])
+        plugins.append(HexgateUsagePlugin(api_key=self.api_key))
+        app = App(name=app_name, root_agent=self._wrapped_agent, plugins=plugins)
         self._runner = Runner(
-            agent=self._wrapped_agent,
-            app_name=app_name,
+            app=app,
             session_service=session_service,
             **runner_kwargs,
         )

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/google/tools.py

@@ -20,7 +20,7 @@ from google.adk.tools.function_tool import FunctionTool
 from google.adk.tools.tool_context import ToolContext
 
 from hexgate.agents.approvals import resolve_approval_async
-from hexgate.agents.factory import ApprovalHandler
+from hexgate.approvals import ApprovalHandler
 from hexgate.security.decision import DecisionOutcome
 from hexgate.security.enforcer import PolicyEnforcer
 

hexgate-0.2.9/hexgate/adapters/google/usage.py

@@ -0,0 +1,44 @@
+"""Google ADK per-call token usage capture via a ``BasePlugin``.
+
+``after_model_callback`` fires once per underlying model call, so a single
+run with several turns (tool-calling loops, sub-agent handoffs) can emit
+more than one usage event. ``callback_context.agent_name`` is read per-call
+rather than fixed at construction, since one ``Runner`` can drive several
+named sub-agents.
+"""
+
+from __future__ import annotations
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
+
+from hexgate.tracing.usage import emit_llm_usage
+
+
+class HexgateUsagePlugin(BasePlugin):
+    """Emits one :class:`~hexgate.tracing.usage.LlmUsageEvent` per
+    ``after_model_callback`` callback. Never rewrites the response — always
+    returns ``None`` so the real model output reaches the agent unchanged."""
+
+    def __init__(self, *, api_key: str) -> None:
+        super().__init__(name="hexgate_usage")
+        self._api_key = api_key
+
+    async def after_model_callback(
+        self,
+        *,
+        callback_context: CallbackContext,
+        llm_response: LlmResponse,
+    ) -> LlmResponse | None:
+        usage = llm_response.usage_metadata
+        if usage is None:
+            return None
+        emit_llm_usage(
+            callback_context.agent_name,
+            llm_response.model_version or "",
+            usage.prompt_token_count or 0,
+            usage.candidates_token_count or 0,
+            api_key=self._api_key,
+        )
+        return None

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/google/wrapper.py

@@ -15,7 +15,7 @@ from typing import TYPE_CHECKING
 from google.adk.agents import BaseAgent
 
 from hexgate.adapters.google.tools import wrap_tools
-from hexgate.agents.factory import ApprovalHandler
+from hexgate.approvals import ApprovalHandler
 from hexgate.security.binding import PolicyBinding, resolve_policy
 from hexgate.security.enforcer import build_enforcer
 

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/langchain/agent.py

@@ -9,6 +9,7 @@ from langfuse import get_client, propagate_attributes
 from langfuse.langchain import CallbackHandler
 from langgraph.graph.state import CompiledStateGraph
 
+from hexgate.adapters.langchain.usage import HexgateUsageCallbackHandler
 from hexgate.runtime import User
 
 if TYPE_CHECKING:
@@ -33,6 +34,7 @@ class HexgateLangchainAgent:
         agent: CompiledStateGraph,
         api_key: str,
         tool_names: list[str],
+        agent_name: str = "default",
         binding: PolicyBinding | None = None,
         ban_gate: BanGate | None = None,
     ) -> None:
@@ -43,6 +45,9 @@ class HexgateLangchainAgent:
         self._tool_names = tool_names
         self._langfuse = get_client()
         self._callback_handler = CallbackHandler()
+        self._usage_handler = HexgateUsageCallbackHandler(
+            agent_name=agent_name, api_key=api_key
+        )
 
     async def _refresh_async(self) -> None:
         """Refresh the policy binding, if attached (async entry points)."""
@@ -72,11 +77,12 @@ class HexgateLangchainAgent:
         }
 
     def _with_callbacks(self, config: RunnableConfig | None) -> RunnableConfig:
-        """Append the Hexgate callback handler to ``config['callbacks']``."""
+        """Append the Hexgate callback handlers to ``config['callbacks']``."""
         merged: RunnableConfig = dict(config) if config else {}
         callbacks = list(merged.get("callbacks") or [])
-        if self._callback_handler not in callbacks:
-            callbacks.append(self._callback_handler)
+        for handler in (self._callback_handler, self._usage_handler):
+            if handler not in callbacks:
+                callbacks.append(handler)
         merged["callbacks"] = callbacks
         return merged
 

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/langchain/tools.py

@@ -26,7 +26,7 @@ from hexgate.agents.approvals import (
 from hexgate.agents.approvals import (
     resolve_approval_sync as _resolve_approval_sync,
 )
-from hexgate.agents.factory import ApprovalHandler
+from hexgate.approvals import ApprovalHandler
 from hexgate.security.decision import DecisionOutcome
 from hexgate.security.enforcer import PolicyEnforcer
 from hexgate.tools.decorators import TOOL_METADATA_ATTR

hexgate-0.2.9/hexgate/adapters/langchain/usage.py

@@ -0,0 +1,84 @@
+"""LangChain per-call token usage capture via a ``BaseCallbackHandler``.
+
+LangGraph propagates ``config["callbacks"]`` down through every node, so a
+handler placed there (see ``HexgateLangchainAgent._with_callbacks``) has its
+``on_llm_end`` invoked by the framework itself on each underlying chat-model
+call — not something this module drives directly. A single ``.invoke()`` can
+therefore emit more than one usage event, once per LLM turn in the run.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from langchain_core.callbacks import BaseCallbackHandler
+from langchain_core.outputs import LLMResult
+
+from hexgate.tracing.usage import emit_llm_usage
+
+
+class HexgateUsageCallbackHandler(BaseCallbackHandler):
+    """Emits one :class:`~hexgate.tracing.usage.LlmUsageEvent` per
+    ``on_llm_end`` callback.
+
+    Reads the standardized ``UsageMetadata`` off the response message when
+    the provider populates it, falling back to ``llm_output["token_usage"]``
+    for providers that only fill the legacy field. Does nothing when
+    neither is present — a provider that reports no usage must not
+    synthesize a zeroed event.
+    """
+
+    def __init__(self, *, agent_name: str, api_key: str | None = None) -> None:
+        self._agent_name = agent_name
+        self._api_key = api_key
+
+    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        # Async so LangChain's AsyncCallbackManager awaits this inline on the
+        # real event loop (asyncio.iscoroutinefunction check in
+        # _ahandle_event_for_handler) instead of dispatching it to a thread
+        # pool executor — a plain sync def here runs off-loop, and
+        # emit_llm_usage's sender never gets a valid loop to schedule its
+        # HTTP send on, silently dropping every event.
+        usage = _extract_usage(response)
+        if usage is None:
+            return
+        input_tokens, output_tokens = usage
+        model = _model_name(response)
+        emit_llm_usage(
+            self._agent_name,
+            model,
+            input_tokens,
+            output_tokens,
+            api_key=self._api_key,
+        )
+
+
+def _model_name(response: LLMResult) -> str:
+    """Read the model name off the response.
+
+    ``llm_output`` is ``None`` for a streamed response (LangChain's default
+    ``_combine_llm_outputs`` — and even providers that override it, like
+    ChatOpenAI, only combine per-chunk outputs that are non-``None``, which
+    streaming chunks typically aren't) — confirmed empirically against a
+    real streaming ``ChatOpenAI`` call, not assumed. ``response_metadata``
+    on the message is populated in both the streaming and non-streaming
+    case, so it's the primary source; ``llm_output`` stays as a fallback
+    for providers that only populate the legacy field.
+    """
+    message = response.generations[0][0].message  # type: ignore[union-attr]
+    response_metadata = getattr(message, "response_metadata", None) or {}
+    return response_metadata.get("model_name") or (response.llm_output or {}).get(
+        "model_name", ""
+    )
+
+
+def _extract_usage(response: LLMResult) -> tuple[int, int] | None:
+    """(input_tokens, output_tokens), or None when the provider reported no
+    usage at all."""
+    usage_metadata = response.generations[0][0].message.usage_metadata  # type: ignore[union-attr]
+    if usage_metadata:
+        return usage_metadata["input_tokens"], usage_metadata["output_tokens"]
+    token_usage = (response.llm_output or {}).get("token_usage")
+    if not token_usage:
+        return None
+    return token_usage.get("prompt_tokens", 0), token_usage.get("completion_tokens", 0)

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/langchain/wrapper.py

@@ -59,6 +59,7 @@ def wrap_langchain_agent(
     return HexgateLangchainAgent(
         agent=agent,
         api_key=resolved_key,
+        agent_name=agent_name,
         tool_names=tool_names,
         binding=PolicyBinding(enforcer, resolved.source),
         ban_gate=resolve_ban_gate(agent_name, api_key=resolved_key, client=client),

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/openai/runner.py

@@ -14,6 +14,7 @@ import nest_asyncio
 from agents import (
     Agent,
     RunConfig,
+    RunHooks,
     Runner,
     RunResult,
     RunResultStreaming,
@@ -21,11 +22,13 @@ from agents import (
     TContext,
     TResponseInputItem,
 )
+from agents.lifecycle import RunHooksBase
 from langfuse import get_client, propagate_attributes
 from openinference.instrumentation.openai_agents import OpenAIAgentsInstrumentor
 
+from hexgate.adapters.openai.usage import HexgateUsageHooks
 from hexgate.adapters.openai.wrapper import wrap_openai_agent
-from hexgate.agents.factory import ApprovalHandler
+from hexgate.approvals import ApprovalHandler
 from hexgate.cloud.client import HexgateClient, HexgateConfig
 from hexgate.config.env import resolve_api_key
 from hexgate.runtime import User
@@ -34,6 +37,47 @@ from hexgate.security.binding import PolicyBinding, resolve_policy
 from hexgate.security.enforcer import build_enforcer
 
 
+class _CompositeRunHooks(RunHooks):
+    """Fan a run's lifecycle callbacks out to multiple ``RunHooks``.
+
+    ``Runner.run*`` accepts exactly one ``hooks=`` object; when the caller
+    already passed one, ``HexgateUsageHooks`` must not replace it — this
+    composes both and forwards every ``RunHooksBase`` callback to each in
+    turn.
+    """
+
+    def __init__(self, hooks: list[RunHooksBase]) -> None:
+        self._hooks = hooks
+
+    async def on_llm_start(self, context, agent, system_prompt, input_items) -> None:
+        for hook in self._hooks:
+            await hook.on_llm_start(context, agent, system_prompt, input_items)
+
+    async def on_llm_end(self, context, agent, response) -> None:
+        for hook in self._hooks:
+            await hook.on_llm_end(context, agent, response)
+
+    async def on_agent_start(self, context, agent) -> None:
+        for hook in self._hooks:
+            await hook.on_agent_start(context, agent)
+
+    async def on_agent_end(self, context, agent, output) -> None:
+        for hook in self._hooks:
+            await hook.on_agent_end(context, agent, output)
+
+    async def on_handoff(self, context, from_agent, to_agent) -> None:
+        for hook in self._hooks:
+            await hook.on_handoff(context, from_agent, to_agent)
+
+    async def on_tool_start(self, context, agent, tool) -> None:
+        for hook in self._hooks:
+            await hook.on_tool_start(context, agent, tool)
+
+    async def on_tool_end(self, context, agent, tool, result) -> None:
+        for hook in self._hooks:
+            await hook.on_tool_end(context, agent, tool, result)
+
+
 class HexgateRunner:
     """Runner for OpenAI agents with Hexgate tool policy and observability."""
 
@@ -109,12 +153,21 @@ class HexgateRunner:
         with propagate_attributes(**kwargs):
             yield
 
+    def _merge_hooks(self, hooks: RunHooks | None) -> RunHooks:
+        """Compose caller-supplied ``hooks`` with the usage hook — never
+        clobber a hooks object the caller already passed."""
+        usage_hooks = HexgateUsageHooks(api_key=self.api_key)
+        if hooks is None:
+            return usage_hooks
+        return _CompositeRunHooks([hooks, usage_hooks])
+
     async def run(
         self,
         agent: Agent,
         input: str | list[TResponseInputItem] | RunState[TContext],
         user: User,
         run_config: RunConfig | None = None,
+        hooks: RunHooks | None = None,
         **kwargs,
     ) -> RunResult:
         """Run the OpenAI agent asynchronously inside a User scope."""
@@ -132,7 +185,11 @@ class HexgateRunner:
         async with user:
             with self._propagate(user, agent.name):
                 return await Runner.run(
-                    wrapped_agent, input, run_config=run_config, **kwargs
+                    wrapped_agent,
+                    input,
+                    run_config=run_config,
+                    hooks=self._merge_hooks(hooks),
+                    **kwargs,
                 )
 
     def run_sync(
@@ -141,6 +198,7 @@ class HexgateRunner:
         input: str | list[TResponseInputItem] | RunState[TContext],
         user: User,
         run_config: RunConfig | None = None,
+        hooks: RunHooks | None = None,
         **kwargs,
     ) -> RunResult:
         """Run the OpenAI agent synchronously inside a User scope."""
@@ -158,7 +216,11 @@ class HexgateRunner:
         with user.sync_scope():
             with self._propagate(user, agent.name):
                 return Runner.run_sync(
-                    wrapped_agent, input, run_config=run_config, **kwargs
+                    wrapped_agent,
+                    input,
+                    run_config=run_config,
+                    hooks=self._merge_hooks(hooks),
+                    **kwargs,
                 )
 
     def run_streamed(
@@ -167,6 +229,7 @@ class HexgateRunner:
         input: str | list[TResponseInputItem] | RunState[TContext],
         user: User,
         run_config: RunConfig | None = None,
+        hooks: RunHooks | None = None,
         **kwargs,
     ) -> RunResultStreaming:
         """Stream the OpenAI agent inside a User scope.
@@ -193,7 +256,11 @@ class HexgateRunner:
         with user.sync_scope():
             with self._propagate(user, agent.name):
                 result = Runner.run_streamed(
-                    wrapped_agent, input, run_config=run_config, **kwargs
+                    wrapped_agent,
+                    input,
+                    run_config=run_config,
+                    hooks=self._merge_hooks(hooks),
+                    **kwargs,
                 )
 
         original_stream_events = result.stream_events

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/openai/tools.py

@@ -19,7 +19,7 @@ from agents import FunctionTool
 from agents.tool import ToolContext
 
 from hexgate.agents.approvals import resolve_approval_async
-from hexgate.agents.factory import ApprovalHandler
+from hexgate.approvals import ApprovalHandler
 from hexgate.security.decision import DecisionOutcome
 from hexgate.security.enforcer import PolicyEnforcer
 

hexgate-0.2.9/hexgate/adapters/openai/usage.py

@@ -0,0 +1,39 @@
+"""OpenAI Agents SDK per-call token usage capture via ``RunHooks``.
+
+``Runner.run``/``run_sync``/``run_streamed`` invoke ``on_llm_end`` once per
+underlying model call, so a single run with several turns (tool-calling
+loops, handoffs) can emit more than one usage event.
+"""
+
+from __future__ import annotations
+
+from agents import Agent, RunContextWrapper
+from agents.items import ModelResponse
+from agents.lifecycle import RunHooks
+
+from hexgate.tracing.usage import emit_llm_usage
+
+
+class HexgateUsageHooks(RunHooks):
+    """Emits one :class:`~hexgate.tracing.usage.LlmUsageEvent` per
+    ``on_llm_end`` callback."""
+
+    def __init__(self, *, api_key: str) -> None:
+        self._api_key = api_key
+
+    async def on_llm_end(
+        self,
+        context: RunContextWrapper,
+        agent: Agent,
+        response: ModelResponse,
+    ) -> None:
+        # agent.model is `str | Model | None` — only the str case gives a
+        # clean name; a Model implementation has no guaranteed name field.
+        model = agent.model if isinstance(agent.model, str) else ""
+        emit_llm_usage(
+            agent.name,
+            model,
+            response.usage.input_tokens,
+            response.usage.output_tokens,
+            api_key=self._api_key,
+        )

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/openai/wrapper.py

@@ -15,7 +15,7 @@ import dataclasses
 from agents import Agent
 
 from hexgate.adapters.openai.tools import wrap_tools
-from hexgate.agents.factory import ApprovalHandler
+from hexgate.approvals import ApprovalHandler
 from hexgate.security.enforcer import PolicyEnforcer
 
 

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/pydantic_ai/agent.py

@@ -10,6 +10,7 @@ from pydantic_ai import Agent
 from pydantic_ai.agent import AgentRun, AgentRunResult
 from pydantic_ai.result import StreamedRunResult
 
+from hexgate.adapters.pydantic_ai.usage import emit_run_usage
 from hexgate.runtime import User
 
 if TYPE_CHECKING:
@@ -100,7 +101,9 @@ class HexgatePydanticAgent:
         await self._refresh_async()
         await self._check_ban_async(user)
         async with self._abind(user, "run"):
-            return await self._agent.run(*args, **kwargs)
+            result = await self._agent.run(*args, **kwargs)
+            emit_run_usage(self._agent_name, self._agent, result, api_key=self._api_key)
+            return result
 
     def run_sync(
         self,
@@ -112,7 +115,9 @@ class HexgatePydanticAgent:
         self._refresh()
         self._check_ban(user)
         with self._bind(user, "run_sync"):
-            return self._agent.run_sync(*args, **kwargs)
+            result = self._agent.run_sync(*args, **kwargs)
+            emit_run_usage(self._agent_name, self._agent, result, api_key=self._api_key)
+            return result
 
     @asynccontextmanager
     async def run_stream(
@@ -127,6 +132,15 @@ class HexgatePydanticAgent:
         async with self._abind(user, "run_stream"):
             async with self._agent.run_stream(*args, **kwargs) as result:
                 yield result
+                # Emit usage only if the run completed, not if the caller aborted mid-stream,
+                # because the usage counts from pydantic's side are 0 until the run completes.
+                # This can happen if a user cancels a LLM request mid-response: we will never
+                # know the total number of input / output tokens, however they are still charged by the LLM provider.
+                # This is a known limitation of pydantic_ai's usage reporting, and we will not be able to report usage in this case.
+                if result.is_complete:
+                    emit_run_usage(
+                        self._agent_name, self._agent, result, api_key=self._api_key
+                    )
 
     @asynccontextmanager
     async def iter(
@@ -141,6 +155,10 @@ class HexgatePydanticAgent:
         async with self._abind(user, "iter"):
             async with self._agent.iter(*args, **kwargs) as run:
                 yield run
+                if run.result is not None:
+                    emit_run_usage(
+                        self._agent_name, self._agent, run, api_key=self._api_key
+                    )
 
     def __getattr__(self, name: str) -> Any:
         """Delegate unknown attributes to the wrapped agent.

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/pydantic_ai/tools.py

@@ -20,7 +20,7 @@ from pydantic_ai.exceptions import ModelRetry
 from pydantic_ai.tools import Tool
 
 from hexgate.agents.approvals import resolve_approval_async
-from hexgate.agents.factory import ApprovalHandler
+from hexgate.approvals import ApprovalHandler
 from hexgate.security.decision import DecisionOutcome
 from hexgate.security.enforcer import PolicyEnforcer
 

hexgate-0.2.9/hexgate/adapters/pydantic_ai/usage.py

@@ -0,0 +1,38 @@
+"""Pydantic AI has no per-call usage hook — usage is read from the run
+result after the call completes and reported as one aggregate event per
+agent run (not per LLM call), a documented limitation vs. the other three
+adapters.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic_ai import Agent
+
+from hexgate.manifest.pydantic_ai import extract_model
+from hexgate.tracing.usage import emit_llm_usage
+
+
+def emit_run_usage(agent_name: str, agent: Agent, result: Any, *, api_key: str) -> None:
+    """Emit one aggregate LlmUsageEvent for a completed pydantic_ai run.
+
+    ``result`` is anything exposing ``.usage()`` and ``.response`` —
+    ``AgentRunResult``, ``StreamedRunResult``, and ``AgentRun`` (from
+    ``run``/``run_sync``, ``run_stream``, and ``iter`` respectively) all
+    qualify. Model name comes from the actual run's response when
+    available (pydantic_ai supports per-call model overrides), else the
+    agent's statically configured model.
+    """
+    usage = result.usage()
+    response = getattr(result, "response", None)
+    model_name = (
+        getattr(response, "model_name", None) or extract_model(agent.model) or ""
+    )
+    emit_llm_usage(
+        agent_name,
+        model_name,
+        usage.input_tokens,
+        usage.output_tokens,
+        api_key=api_key,
+    )

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/adapters/pydantic_ai/wrapper.py

@@ -17,7 +17,7 @@ from pydantic_ai.tools import Tool
 
 from hexgate.adapters.pydantic_ai.agent import HexgatePydanticAgent
 from hexgate.adapters.pydantic_ai.tools import wrap_tools
-from hexgate.agents.factory import ApprovalHandler
+from hexgate.approvals import ApprovalHandler
 from hexgate.cloud.client import HexgateClient, HexgateConfig
 from hexgate.config.env import resolve_api_key
 from hexgate.security.bans import resolve_ban_gate

{hexgate-0.2.8 → hexgate-0.2.9}/hexgate/agents/approvals.py

@@ -12,7 +12,7 @@ from __future__ import annotations
 from inspect import isawaitable
 from typing import Any
 
-from hexgate.agents.factory import ApprovalHandler
+from hexgate.approvals import ApprovalHandler
 from hexgate.security.decision import Decision