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

agentkeeper_runtime_sdk/azure_openai.py

@@ -0,0 +1,314 @@
+from __future__ import annotations
+
+import inspect
+from typing import Any, AsyncIterator, Iterator, Mapping, Optional
+
+from .client import (
+    AgentKeeperRuntimeClient,
+    compact,
+    create_agentkeeper_runtime_client,
+    get_value,
+    number_value,
+    safe_error_name,
+    string_value,
+    telemetry_metadata,
+)
+
+
+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": "azure_openai",
+        "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_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 _tool_names(request: Mapping[str, Any]) -> list[str]:
+    tools = request.get("tools")
+    if not isinstance(tools, list):
+        return []
+    names: list[str] = []
+    for tool in tools[:50]:
+        function = get_value(tool, "function")
+        name = string_value(get_value(function, "name")) or string_value(get_value(tool, "name"))
+        if name:
+            names.append(name)
+    return names
+
+
+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"))
+        or string_value(request.get("deployment"))
+        or string_value(request.get("deployment_name"))
+        or string_value(request.get("deploymentName"))
+    )
+
+
+def _input_type(value: Any) -> Optional[str]:
+    if value is None:
+        return None
+    if isinstance(value, list):
+        return "array"
+    if isinstance(value, Mapping):
+        return "object"
+    if isinstance(value, str):
+        return "string"
+    if isinstance(value, bool):
+        return "boolean"
+    if isinstance(value, (int, float)):
+        return "number"
+    return type(value).__name__
+
+
+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 = _tool_names(request)
+    response_format = request.get("response_format")
+    return compact({
+        **ids,
+        "operation": operation,
+        "model": _model_name(request),
+        "api_version": string_value(request.get("api_version")) or string_value(request.get("apiVersion")),
+        "message_count": len(request.get("messages")) if isinstance(request.get("messages"), list) else None,
+        "input_type": _input_type(request.get("input")),
+        "tool_count": len(tools) or None,
+        "tool_names": tools or None,
+        "response_format_type": string_value(get_value(response_format, "type")),
+        "max_tokens": number_value(request.get("max_tokens")) or number_value(request.get("maxTokens")),
+        "stream": request.get("stream") is True,
+    })
+
+
+def _usage_fields(response: Any) -> dict[str, Any]:
+    usage = get_value(response, "usage") or {}
+    return compact({
+        "token_input": (
+            number_value(get_value(usage, "prompt_tokens"))
+            or number_value(get_value(usage, "promptTokens"))
+            or number_value(get_value(usage, "input_tokens"))
+            or number_value(get_value(usage, "inputTokens"))
+        ),
+        "token_output": (
+            number_value(get_value(usage, "completion_tokens"))
+            or number_value(get_value(usage, "completionTokens"))
+            or number_value(get_value(usage, "output_tokens"))
+            or number_value(get_value(usage, "outputTokens"))
+        ),
+    })
+
+
+def _response_metadata(response: Any, metadata: Mapping[str, Any]) -> dict[str, Any]:
+    choices = get_value(response, "choices")
+    if not isinstance(choices, list):
+        choices = []
+    finish_reasons: list[str] = []
+    for choice in choices[:50]:
+        reason = string_value(get_value(choice, "finish_reason")) or string_value(get_value(choice, "finishReason"))
+        if reason:
+            finish_reasons.append(reason)
+    response_output = get_value(response, "output")
+    return compact({
+        **dict(metadata),
+        "response_id": string_value(get_value(response, "id")),
+        "response_object": string_value(get_value(response, "object")),
+        "choice_count": len(choices) or None,
+        "output_item_count": len(response_output) if isinstance(response_output, list) else None,
+        "finish_reasons": sorted(set(finish_reasons)) or None,
+        "stream_event_count": number_value(get_value(response, "event_count")),
+    })
+
+
+class _StreamProxy:
+    def __init__(self, stream: Any, track_once: Any, 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:
+        return getattr(self._stream, name)
+
+    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})
+
+    async def __aiter__(self) -> AsyncIterator[Any]:
+        event_count = 0
+        try:
+            async 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 _is_stream_like(result: Any, request: Mapping[str, Any]) -> bool:
+    if request.get("stream") is not True:
+        return False
+    return hasattr(result, "__aiter__") or hasattr(result, "__iter__")
+
+
+class _AzureOpenAIResourceProxy:
+    def __init__(self, resource: Any, ak: AgentKeeperRuntimeClient, options: Mapping[str, Any], path: tuple[str, ...] = ()) -> None:
+        self._resource = resource
+        self._ak = ak
+        self._options = dict(options)
+        self._path = path
+
+    def __getattr__(self, name: str) -> Any:
+        value = getattr(self._resource, name)
+        operation = ".".join((*self._path, name))
+        if operation in {"chat.completions.create", "responses.create", "embeddings.create"} and callable(value):
+            return self._wrap_method(operation, value)
+        if value is not None and not isinstance(value, (str, bytes, bytearray, int, float, bool)):
+            return _AzureOpenAIResourceProxy(value, self._ak, self._options, (*self._path, name))
+        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": "azure_openai",
+            "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"Azure OpenAI {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": "model_only",
+            "model_provider": "azure_openai",
+            "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"Azure OpenAI {operation} errored: {safe_error_name(error)}"
+                if error
+                else f"Azure OpenAI {operation} completed"
+            ),
+            "metadata": _response_metadata(response, metadata),
+        })
+
+    def _wrap_method(self, operation: str, method: Any) -> Any:
+        def wrapped(*args: Any, **kwargs: Any) -> Any:
+            request = _request_from_args(args, 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)
+
+            def wrap_or_track(result: Any) -> Any:
+                if _is_stream_like(result, request):
+                    return _StreamProxy(result, track_once, _model_name(request))
+                track_once(result)
+                return result
+
+            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
+                    return wrap_or_track(resolved)
+
+                return wait()
+
+            return wrap_or_track(result)
+
+        return wrapped
+
+
+def wrap_azure_openai_client(
+    azure_openai_client: Any,
+    client_or_options: AgentKeeperRuntimeClient | Mapping[str, Any],
+    **options: Any,
+) -> Any:
+    return _AzureOpenAIResourceProxy(azure_openai_client, _client_from(client_or_options), options)

agentkeeper_runtime_sdk/bedrock.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+import inspect
+import json
+from typing import Any, Mapping, Optional
+
+from .client import (
+    AgentKeeperRuntimeClient,
+    compact,
+    create_agentkeeper_runtime_client,
+    get_value,
+    number_value,
+    object_keys,
+    safe_error_name,
+    string_value,
+    telemetry_metadata,
+)
+
+
+def _client_from(value: AgentKeeperRuntimeClient | Mapping[str, Any]) -> AgentKeeperRuntimeClient:
+    if isinstance(value, AgentKeeperRuntimeClient):
+        return value
+    return create_agentkeeper_runtime_client(**dict(value))
+
+
+def _body_json(output: Any) -> Optional[dict[str, Any]]:
+    body = get_value(output, "body")
+    if body is None:
+        return None
+    if isinstance(body, str):
+        text = body
+    elif isinstance(body, (bytes, bytearray, memoryview)):
+        text = bytes(body).decode("utf-8", errors="replace")
+    else:
+        return None
+    try:
+        parsed = json.loads(text)
+    except json.JSONDecodeError:
+        return None
+    return parsed if isinstance(parsed, dict) else None
+
+
+def _usage_fields(output: Any) -> dict[str, Any]:
+    parsed = _body_json(output) or {}
+    usage = get_value(output, "usage") or parsed.get("usage") or parsed.get("amazonBedrockInvocationMetrics") or {}
+    return compact({
+        "token_input": (
+            number_value(get_value(usage, "inputTokens"))
+            or number_value(get_value(usage, "inputTokenCount"))
+            or number_value(get_value(usage, "input_tokens"))
+        ),
+        "token_output": (
+            number_value(get_value(usage, "outputTokens"))
+            or number_value(get_value(usage, "outputTokenCount"))
+            or number_value(get_value(usage, "output_tokens"))
+        ),
+    })
+
+
+def _model_name(params: Mapping[str, Any]) -> Optional[str]:
+    return string_value(params.get("modelId")) or string_value(params.get("modelIdentifier"))
+
+
+def _stop_reason(output: Any) -> Optional[str]:
+    parsed = _body_json(output) or {}
+    return (
+        string_value(get_value(output, "stopReason"))
+        or string_value(parsed.get("stopReason"))
+        or string_value(parsed.get("stop_reason"))
+    )
+
+
+class _BedrockRuntimeClientProxy:
+    def __init__(self, bedrock_client: Any, ak: AgentKeeperRuntimeClient, options: Optional[Mapping[str, Any]] = None) -> None:
+        self._bedrock_client = bedrock_client
+        self._ak = ak
+        self._options = dict(options or {})
+
+    def __getattr__(self, name: str) -> Any:
+        value = getattr(self._bedrock_client, name)
+        if name in {"converse", "invoke_model", "converse_stream", "invoke_model_with_response_stream"} and callable(value):
+            return self._wrap_method(name, value)
+        return value
+
+    def _base_event(self) -> dict[str, Any]:
+        return {
+            "runtime_integration": "bedrock",
+            "runtime_service": self._options.get("runtime_service") or self._options.get("runtimeService"),
+            "runtime_environment": self._options.get("runtime_environment") or self._options.get("runtimeEnvironment"),
+            **dict(self._options.get("event") or {}),
+        }
+
+    def _track_start(self, operation: str, params: Mapping[str, Any], metadata: Mapping[str, Any]) -> None:
+        self._ak.track({
+            **self._base_event(),
+            "event_kind": "model_call",
+            "capability": "model_only",
+            "model_provider": "aws_bedrock",
+            "model_name": _model_name(params),
+            "run_id": metadata.get("run_id"),
+            "trace_id": metadata.get("trace_id"),
+            "span_id": metadata.get("span_id"),
+            "verdict": "observed",
+            "evidence_summary": f"AWS Bedrock Runtime {operation} observed",
+            "metadata": metadata,
+        })
+
+    def _track_completion(
+        self,
+        operation: str,
+        params: Mapping[str, Any],
+        metadata: Mapping[str, Any],
+        output: Any,
+        error: Optional[BaseException] = None,
+    ) -> None:
+        self._ak.track({
+            **self._base_event(),
+            "event_kind": "model_call",
+            "capability": "model_only",
+            "model_provider": "aws_bedrock",
+            "model_name": _model_name(params),
+            **_usage_fields(output),
+            "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"AWS Bedrock Runtime {operation} errored: {safe_error_name(error)}"
+                if error
+                else f"AWS Bedrock Runtime {operation} completed"
+            ),
+            "metadata": compact({
+                **dict(metadata),
+                "operation": operation,
+                "request_keys": object_keys(params),
+                "stop_reason": _stop_reason(output),
+                "response_has_body": get_value(output, "body") is not None,
+            }),
+        })
+
+    def _wrap_method(self, operation: str, method: Any) -> Any:
+        def wrapped(*args: Any, **kwargs: Any) -> Any:
+            params = dict(kwargs)
+            if args and isinstance(args[0], Mapping):
+                params = {**dict(args[0]), **params}
+            metadata = compact({
+                **telemetry_metadata(
+                    runtime_service=self._options.get("runtime_service") or self._options.get("runtimeService"),
+                    runtime_environment=self._options.get("runtime_environment") or self._options.get("runtimeEnvironment"),
+                    run_id=self._options.get("run_id") or self._options.get("runId"),
+                    trace_id=self._options.get("trace_id") or self._options.get("traceId"),
+                    span_id=self._options.get("span_id") or self._options.get("spanId"),
+                ),
+                "operation": operation,
+                "request_keys": object_keys(params),
+            })
+            self._track_start(operation, params, metadata)
+            try:
+                result = method(*args, **kwargs)
+            except BaseException as error:
+                self._track_completion(operation, params, metadata, {}, error)
+                raise
+            if inspect.isawaitable(result):
+                async def wait() -> Any:
+                    try:
+                        resolved = await result
+                    except BaseException as error:
+                        self._track_completion(operation, params, metadata, {}, error)
+                        raise
+                    self._track_completion(operation, params, metadata, resolved)
+                    return resolved
+
+                return wait()
+            self._track_completion(operation, params, metadata, result)
+            return result
+
+        return wrapped
+
+
+def wrap_bedrock_runtime_client(
+    bedrock_client: Any,
+    client_or_options: AgentKeeperRuntimeClient | Mapping[str, Any],
+    options: Optional[Mapping[str, Any]] = None,
+) -> Any:
+    return _BedrockRuntimeClientProxy(bedrock_client, _client_from(client_or_options), options)
+