agentkeeper-runtime-sdk 0.1.0b1__py3-none-any.whl

agentkeeper_runtime_sdk/__init__.py

@@ -0,0 +1,45 @@
+from .anthropic import (
+    wrap_anthropic_client,
+    wrap_anthropic_runnable_tool,
+    wrap_anthropic_tool,
+)
+from .bedrock import wrap_bedrock_runtime_client
+from .azure_openai import wrap_azure_openai_client
+from .claude_managed_agents import wrap_claude_managed_agents_client
+from .client import (
+    AgentKeeperRuntimeClient,
+    TrackResult,
+    create_agentkeeper_runtime_client,
+    find_raw_payload_keys,
+)
+from .langchain import (
+    create_langchain_callback_handler,
+    create_langgraph_callback_handler,
+    wrap_langchain_tool,
+    wrap_langgraph_tool,
+)
+from .openai_agents import (
+    wrap_openai_agents_run,
+    wrap_openai_agents_sdk,
+    wrap_openai_agents_tool,
+)
+
+__all__ = [
+    "AgentKeeperRuntimeClient",
+    "TrackResult",
+    "create_agentkeeper_runtime_client",
+    "find_raw_payload_keys",
+    "create_langchain_callback_handler",
+    "create_langgraph_callback_handler",
+    "wrap_langchain_tool",
+    "wrap_langgraph_tool",
+    "wrap_bedrock_runtime_client",
+    "wrap_azure_openai_client",
+    "wrap_claude_managed_agents_client",
+    "wrap_openai_agents_run",
+    "wrap_openai_agents_sdk",
+    "wrap_openai_agents_tool",
+    "wrap_anthropic_client",
+    "wrap_anthropic_runnable_tool",
+    "wrap_anthropic_tool",
+]

agentkeeper_runtime_sdk/anthropic.py

@@ -0,0 +1,441 @@
+from __future__ import annotations
+
+import inspect
+from typing import Any, Callable, Iterator, Mapping, Optional
+
+from .client import (
+    AgentKeeperRuntimeClient,
+    ToolPolicyDecision,
+    compact,
+    create_agentkeeper_runtime_client,
+    get_value,
+    number_value,
+    object_keys,
+    safe_error_name,
+    string_value,
+    telemetry_metadata,
+)
+
+_RUNNABLE_TOOL_WRAPPED = "_agentkeeper_anthropic_runnable_tool_wrapped"
+
+
+def _client_from(value: AgentKeeperRuntimeClient | Mapping[str, Any]) -> AgentKeeperRuntimeClient:
+    if isinstance(value, AgentKeeperRuntimeClient):
+        return value
+    return create_agentkeeper_runtime_client(**dict(value))
+
+
+def _base_event(options: Mapping[str, Any]) -> dict[str, Any]:
+    return {
+        "runtime_integration": "anthropic_sdk",
+        "runtime_service": options.get("runtime_service") or options.get("runtimeService"),
+        "runtime_environment": options.get("runtime_environment") or options.get("runtimeEnvironment"),
+        **dict(options.get("event") or {}),
+    }
+
+
+def _request_metadata(request: Mapping[str, Any], operation: str, options: Mapping[str, Any]) -> dict[str, Any]:
+    ids = telemetry_metadata(
+        runtime_service=options.get("runtime_service") or options.get("runtimeService"),
+        runtime_environment=options.get("runtime_environment") or options.get("runtimeEnvironment"),
+        run_id=options.get("run_id") or options.get("runId"),
+        trace_id=options.get("trace_id") or options.get("traceId"),
+        span_id=options.get("span_id") or options.get("spanId"),
+        user_id=options.get("user_id") or options.get("userId"),
+        convo_id=options.get("convo_id") or options.get("convoId"),
+    )
+    tools = request.get("tools")
+    tool_choice = request.get("tool_choice") if isinstance(request.get("tool_choice"), Mapping) else {}
+    return compact({
+        **ids,
+        "operation": operation,
+        "model": string_value(request.get("model")),
+        "max_tokens": number_value(request.get("max_tokens")),
+        "message_count": len(request.get("messages")) if isinstance(request.get("messages"), list) else None,
+        "tool_count": len(tools) if isinstance(tools, list) else None,
+        "tool_names": [
+            name for name in (string_value(get_value(tool, "name")) for tool in tools or []) if name
+        ][:50] or None,
+        "tool_choice_type": string_value(tool_choice.get("type")) if isinstance(tool_choice, Mapping) else None,
+        "tool_choice_name": string_value(tool_choice.get("name")) if isinstance(tool_choice, Mapping) else None,
+        "stream": request.get("stream") is True or operation.endswith(".stream"),
+        "has_system": bool(request.get("system")),
+    })
+
+
+def _usage_fields(response: Any) -> dict[str, Any]:
+    usage = get_value(response, "usage") or {}
+    return compact({
+        "token_input": number_value(get_value(usage, "input_tokens")) or number_value(get_value(usage, "inputTokens")),
+        "token_output": number_value(get_value(usage, "output_tokens")) or number_value(get_value(usage, "outputTokens")),
+    })
+
+
+def _model_name(request: Mapping[str, Any], response: Any = None) -> Optional[str]:
+    return string_value(get_value(response, "model")) or string_value(request.get("model"))
+
+
+def _response_metadata(response: Any, metadata: Mapping[str, Any]) -> dict[str, Any]:
+    content = get_value(response, "content")
+    if not isinstance(content, list):
+        content = []
+    tool_names = []
+    for block in content:
+        if string_value(get_value(block, "type")) == "tool_use":
+            name = string_value(get_value(block, "name"))
+            if name:
+                tool_names.append(name)
+    return compact({
+        **dict(metadata),
+        "response_id": string_value(get_value(response, "id")),
+        "role": string_value(get_value(response, "role")),
+        "stop_reason": string_value(get_value(response, "stop_reason") or get_value(response, "stopReason")),
+        "stop_sequence_present": get_value(response, "stop_sequence") is not None,
+        "content_block_count": len(content) or None,
+        "tool_use_count": len(tool_names) or None,
+        "tool_use_names": tool_names[:50] or None,
+        "stream_event_count": number_value(get_value(response, "event_count")),
+    })
+
+
+class _StreamProxy:
+    def __init__(
+        self,
+        stream: Any,
+        track_once: Callable[[Any, Optional[BaseException]], None],
+        model: Optional[str],
+    ) -> None:
+        self._stream = stream
+        self._track_once = track_once
+        self._model = model
+        self._tracked = False
+
+    def _track(self, response: Any, error: Optional[BaseException] = None) -> None:
+        if self._tracked:
+            return
+        self._tracked = True
+        self._track_once(response, error)
+
+    def __getattr__(self, name: str) -> Any:
+        value = getattr(self._stream, name)
+        if name in {
+            "get_final_message",
+            "final_message",
+            "finalMessage",
+            "done",
+            "run_until_done",
+            "runUntilDone",
+        } and callable(value):
+            def wrapped(*args: Any, **kwargs: Any) -> Any:
+                try:
+                    result = value(*args, **kwargs)
+                except BaseException as error:
+                    self._track({"model": self._model}, error)
+                    raise
+                if inspect.isawaitable(result):
+                    async def wait() -> Any:
+                        try:
+                            resolved = await result
+                        except BaseException as error:
+                            self._track({"model": self._model}, error)
+                            raise
+                        self._track(resolved)
+                        return resolved
+
+                    return wait()
+                self._track(result)
+                return result
+
+            return wrapped
+        return value
+
+    def __iter__(self) -> Iterator[Any]:
+        event_count = 0
+        try:
+            for event in self._stream:
+                event_count += 1
+                yield event
+        except BaseException as error:
+            self._track({"model": self._model, "event_count": event_count}, error)
+            raise
+        self._track({"model": self._model, "event_count": event_count})
+
+    def __enter__(self) -> Any:
+        entered = self._stream.__enter__()
+        if entered is self._stream:
+            return self
+        self._stream = entered
+        return self
+
+    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any:
+        return self._stream.__exit__(exc_type, exc, tb)
+
+    async def __aenter__(self) -> Any:
+        entered = await self._stream.__aenter__()
+        if entered is not self._stream:
+            self._stream = entered
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> Any:
+        return await self._stream.__aexit__(exc_type, exc, tb)
+
+
+def _request_from_args(args: tuple[Any, ...], kwargs: Mapping[str, Any]) -> dict[str, Any]:
+    if args and isinstance(args[0], Mapping):
+        return {**dict(args[0]), **dict(kwargs)}
+    return dict(kwargs)
+
+
+def _prepare_tool_runner_request(request: dict[str, Any], ak: AgentKeeperRuntimeClient, options: Mapping[str, Any]) -> dict[str, Any]:
+    tools = request.get("tools")
+    if not isinstance(tools, list):
+        return request
+    return {
+        **request,
+        "tools": [
+            wrap_anthropic_runnable_tool(tool, ak, **options)
+            if (isinstance(tool, Mapping) and callable(tool.get("run"))) or callable(getattr(tool, "run", None))
+            else tool
+            for tool in tools
+        ],
+    }
+
+
+class _AnthropicResourceProxy:
+    def __init__(self, resource: Any, ak: AgentKeeperRuntimeClient, options: Mapping[str, Any], prefix: str = "messages") -> None:
+        self._resource = resource
+        self._ak = ak
+        self._options = dict(options)
+        self._prefix = prefix
+
+    def __getattr__(self, name: str) -> Any:
+        value = getattr(self._resource, name)
+        if name in {"create", "stream", "tool_runner", "toolRunner"} and callable(value):
+            return self._wrap_method(name, value)
+        return value
+
+    def _track_start(self, request: Mapping[str, Any], operation: str, metadata: Mapping[str, Any]) -> None:
+        self._ak.track({
+            **_base_event(self._options),
+            "event_kind": "model_call",
+            "capability": "model_only",
+            "model_provider": "anthropic",
+            "model_name": _model_name(request),
+            "run_id": metadata.get("run_id"),
+            "trace_id": metadata.get("trace_id"),
+            "span_id": metadata.get("span_id"),
+            "verdict": "observed",
+            "evidence_summary": f"Anthropic SDK {operation} observed",
+            "metadata": metadata,
+        })
+
+    def _track_completion(
+        self,
+        request: Mapping[str, Any],
+        operation: str,
+        metadata: Mapping[str, Any],
+        response: Any,
+        error: Optional[BaseException] = None,
+    ) -> None:
+        self._ak.track({
+            **_base_event(self._options),
+            "event_kind": "model_call",
+            "capability": "post_run",
+            "model_provider": "anthropic",
+            "model_name": _model_name(request, response),
+            **_usage_fields(response),
+            "run_id": metadata.get("run_id"),
+            "trace_id": metadata.get("trace_id"),
+            "span_id": metadata.get("span_id"),
+            "verdict": "observed" if error else "passed",
+            "severity": "medium" if error else "low",
+            "evidence_summary": (
+                f"Anthropic SDK {operation} errored: {safe_error_name(error)}"
+                if error
+                else f"Anthropic SDK {operation} completed"
+            ),
+            "metadata": _response_metadata(response, metadata),
+        })
+
+    def _wrap_method(self, name: str, method: Callable[..., Any]) -> Callable[..., Any]:
+        operation = f"{self._prefix}.{name}"
+
+        def wrapped(*args: Any, **kwargs: Any) -> Any:
+            request = _request_from_args(args, kwargs)
+            if name in {"tool_runner", "toolRunner"}:
+                request = _prepare_tool_runner_request(request, self._ak, self._options)
+                args = (request,) if args and isinstance(args[0], Mapping) else args
+                kwargs = request if not args else kwargs
+            metadata = _request_metadata(request, operation, self._options)
+            self._track_start(request, operation, metadata)
+            tracked = False
+
+            def track_once(response: Any, error: Optional[BaseException] = None) -> None:
+                nonlocal tracked
+                if tracked:
+                    return
+                tracked = True
+                self._track_completion(request, operation, metadata, response, error)
+
+            try:
+                result = method(*args, **kwargs)
+            except BaseException as error:
+                track_once({}, error)
+                raise
+
+            if inspect.isawaitable(result):
+                async def wait() -> Any:
+                    try:
+                        resolved = await result
+                    except BaseException as error:
+                        track_once({}, error)
+                        raise
+                    if name in {"stream", "tool_runner", "toolRunner"}:
+                        return _StreamProxy(resolved, track_once, _model_name(request))
+                    track_once(resolved)
+                    return resolved
+
+                return wait()
+
+            if name == "stream" or request.get("stream") is True or name in {"tool_runner", "toolRunner"}:
+                return _StreamProxy(result, track_once, _model_name(request))
+
+            track_once(result)
+            return result
+
+        return wrapped
+
+
+class _AnthropicClientProxy:
+    def __init__(self, client: Any, ak: AgentKeeperRuntimeClient, options: Mapping[str, Any]) -> None:
+        self._client = client
+        self._ak = ak
+        self._options = dict(options)
+
+    def __getattr__(self, name: str) -> Any:
+        value = getattr(self._client, name)
+        if name == "messages" and value is not None:
+            return _AnthropicResourceProxy(value, self._ak, self._options, "messages")
+        if name == "beta" and value is not None:
+            return _BetaProxy(value, self._ak, self._options)
+        return value
+
+
+class _BetaProxy:
+    def __init__(self, beta: Any, ak: AgentKeeperRuntimeClient, options: Mapping[str, Any]) -> None:
+        self._beta = beta
+        self._ak = ak
+        self._options = dict(options)
+
+    def __getattr__(self, name: str) -> Any:
+        value = getattr(self._beta, name)
+        if name == "messages" and value is not None:
+            return _AnthropicResourceProxy(value, self._ak, self._options, "beta.messages")
+        return value
+
+
+def wrap_anthropic_client(
+    anthropic_client: Any,
+    client_or_options: AgentKeeperRuntimeClient | Mapping[str, Any],
+    **options: Any,
+) -> Any:
+    return _AnthropicClientProxy(anthropic_client, _client_from(client_or_options), options)
+
+
+def _tool_name(tool: Any, explicit: Optional[str]) -> str:
+    return (
+        explicit
+        or string_value(get_value(tool, "name"))
+        or string_value(get_value(tool, "type"))
+        or string_value(getattr(tool.__class__, "__name__", None))
+        or "anthropic_sdk_tool"
+    )
+
+
+class _RunnableToolProxy:
+    def __init__(self, tool: Any, run: Callable[..., Any]) -> None:
+        self._tool = tool
+        self.run = run
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._tool, name)
+
+
+def wrap_anthropic_runnable_tool(
+    tool: Any,
+    client_or_options: AgentKeeperRuntimeClient | Mapping[str, Any],
+    *,
+    tool_name: Optional[str] = None,
+    evaluate: Optional[Callable[[Any], ToolPolicyDecision | Any]] = None,
+    evaluate_tool: Optional[Callable[[str, Any], ToolPolicyDecision | Any]] = None,
+    **options: Any,
+) -> Any:
+    if isinstance(tool, Mapping) and tool.get(_RUNNABLE_TOOL_WRAPPED):
+        return tool
+    if getattr(tool, _RUNNABLE_TOOL_WRAPPED, False):
+        return tool
+
+    ak = _client_from(client_or_options)
+    name = _tool_name(tool, tool_name)
+    event = {
+        **_base_event(options),
+        "runtime_integration": "anthropic_sdk",
+        "tool_type": "anthropic_sdk",
+        "metadata": {"tool_runner": True, **dict(options.get("metadata") or {})},
+    }
+
+    def evaluator(tool_args: Any) -> ToolPolicyDecision | Any:
+        if evaluate:
+            return evaluate(tool_args)
+        if evaluate_tool:
+            return evaluate_tool(name, tool_args)
+        return {"verdict": "passed"}
+
+    def original_run(*args: Any, **kwargs: Any) -> Any:
+        run = tool.get("run") if isinstance(tool, Mapping) else getattr(tool, "run")
+        return run(*args, **kwargs)
+
+    def wrapped_run(*args: Any, **kwargs: Any) -> Any:
+        tool_args = args[0] if args else dict(kwargs)
+        return ak.run_guarded_tool(name, tool_args, lambda: original_run(*args, **kwargs), evaluate=evaluator, event=event)
+
+    if isinstance(tool, Mapping):
+        return {**dict(tool), "run": wrapped_run, _RUNNABLE_TOOL_WRAPPED: True}
+    wrapped = _RunnableToolProxy(tool, wrapped_run)
+    setattr(wrapped, _RUNNABLE_TOOL_WRAPPED, True)
+    return wrapped
+
+
+def wrap_anthropic_tool(
+    tool_name: str,
+    execute: Callable[[Any], Any],
+    client_or_options: AgentKeeperRuntimeClient | Mapping[str, Any],
+    *,
+    evaluate: Optional[Callable[[Any], ToolPolicyDecision | Any]] = None,
+    evaluate_tool: Optional[Callable[[str, Any], ToolPolicyDecision | Any]] = None,
+    **options: Any,
+) -> Callable[[Any], Any]:
+    ak = _client_from(client_or_options)
+
+    def evaluator(tool_args: Any) -> ToolPolicyDecision | Any:
+        if evaluate:
+            return evaluate(tool_args)
+        if evaluate_tool:
+            return evaluate_tool(tool_name, tool_args)
+        return {"verdict": "passed"}
+
+    def wrapped(tool_args: Any) -> Any:
+        return ak.run_guarded_tool(
+            tool_name,
+            tool_args,
+            lambda: execute(tool_args),
+            evaluate=evaluator,
+            event={
+                **_base_event(options),
+                "runtime_integration": "anthropic_sdk",
+                "tool_type": "anthropic_sdk",
+            },
+        )
+
+    wrapped.__name__ = getattr(execute, "__name__", tool_name)
+    return wrapped