@macaw-ai/client 0.0.0-stage → 0.9.11

package/payload/secureAI/macaw_adapters/pydantic_ai/__init__.py

@@ -0,0 +1,36 @@
+"""
+MACAW Adapter for Pydantic AI
+
+Drop-in replacement for pydantic_ai.Agent. Model calls become
+tool:<app>/generate and tool calls become tool:<app>/<tool_name>, so both are
+governed by the MAPL policies you already write.
+
+Usage:
+    from macaw_adapters.pydantic_ai import SecureAgent
+    from pydantic_ai.models.openai import OpenAIChatModel
+
+    agent = SecureAgent(
+        OpenAIChatModel("gpt-4o"),
+        app_name="acme",
+        tools=[query_catalog, get_lineage],
+    )
+    result = await agent.run("which tables contain PII?")
+
+Multi-user:
+    alice_agent = agent.bind_to_user(alice_client)
+
+Prerequisites:
+    - MACAW Client Library: https://console.macawsecurity.ai
+    - pip install macaw-adapters[pydantic-ai]
+"""
+
+from macaw_adapters.pydantic_ai.agent import Agent, SecureAgent
+from macaw_adapters.pydantic_ai.model import SecureModel
+from macaw_adapters.pydantic_ai.toolset import SecureToolset
+
+__all__ = [
+    "SecureAgent",
+    "Agent",  # drop-in alias
+    "SecureModel",
+    "SecureToolset",
+]

package/payload/secureAI/macaw_adapters/pydantic_ai/_core.py

@@ -0,0 +1,257 @@
+"""
+Shared MACAW binding for the Pydantic AI adapter.
+
+MacawBinding holds the MACAWClient, the set of registered resource names, and
+the models reachable by the generate handler. SecureModel and SecureToolset are
+constructed with a binding and call through it.
+
+    MacawBinding(app_name=...)          registers the app; issues as the app
+    binding.for_user(user_client)       shares that registration; issues as the user
+
+    binding.register_tool(name, fn)     -> tool:<app>/<name>, dispatched by MACAW
+    binding.register_ack(name)          -> tool:<app>/<name>, authorization only
+    binding.invoke(resource, params)    -> invoke_tool(tool:<app>/<resource>, ...)
+"""
+
+import asyncio
+import inspect
+import logging
+from typing import Any, Callable, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+# Carries ModelRequestParameters through invoke_tool. invoke_tool has no channel
+# other than `parameters`, so this rides alongside the policy inputs.
+REQUEST_PARAMETERS_KEY = "_request_parameters"
+
+try:
+    from macaw_client import MACAWClient, PermissionDenied
+except ImportError:  # pragma: no cover - surfaced at construction time
+    MACAWClient = None
+    PermissionDenied = None
+
+
+def is_permission_denied(exc: BaseException) -> bool:
+    """
+    True if `exc` is a MACAW policy denial.
+
+    macaw_client exports one PermissionDenied, and a second class of the same
+    name is raised from the SDK core. Match on the name so a denial is never
+    mistaken for a transport error.
+    """
+    if PermissionDenied is not None and isinstance(exc, PermissionDenied):
+        return True
+    return type(exc).__name__ == "PermissionDenied"
+
+
+def as_handler(fn: Callable) -> Callable[[Dict[str, Any]], Any]:
+    """
+    Adapt a tool function into a MACAW handler.
+
+    Handlers receive a params dict and run on MACAW's handler pool after the PEP
+    has allowed the call. Async tools are driven to completion there - the pool
+    thread has no event loop of its own.
+    """
+
+    def handler(params: Dict[str, Any]) -> Any:
+        result = fn(**params)
+        if inspect.isawaitable(result):
+            result = asyncio.run(result)
+        return result
+
+    return handler
+
+
+def _acknowledge(params: Dict[str, Any]) -> Dict[str, Any]:
+    """Handler for resources MACAW authorizes but does not execute."""
+    return {"authorized": True}
+
+
+class MacawBinding:
+    """Owns the MACAWClient and turns every seam into an invoke_tool call."""
+
+    def __init__(
+        self,
+        app_name: Optional[str] = None,
+        intent_policy: Optional[Dict[str, Any]] = None,
+        jwt_token: Optional[str] = None,
+        user_name: Optional[str] = None,
+    ):
+        if MACAWClient is None:
+            raise ImportError(
+                "MACAWClient not installed. Download from https://console.macawsecurity.ai"
+            )
+
+        self.app_name = app_name or "secure-pydantic-app"
+        self._mode = "user" if jwt_token else "service"
+        self._registered = set()
+        self._models = {}  # model_name -> wrapped pydantic_ai Model
+
+        # The model seam. Declaring "messages" as a prompt is what makes
+        # invoke_tool create authenticated prompts for the conversation.
+        self.tools = {
+            f"tool:{self.app_name}/generate": {
+                "handler": self._handle_generate,
+                "prompts": ["messages"],
+            }
+        }
+
+        identity = (
+            {"user_name": user_name, "iam_token": jwt_token, "agent_type": "user"}
+            if jwt_token
+            else {}
+        )
+
+        self.macaw_client = MACAWClient(
+            app_name=self.app_name,
+            app_version="1.0.0",
+            intent_policy=intent_policy or {},
+            tools=self.tools,
+            **identity,
+        )
+
+        if self.macaw_client.register():
+            self.server_id = self.macaw_client.agent_id
+            logger.info(
+                "SecureAgent registered as %s (mode: %s)", self.server_id, self._mode
+            )
+        else:
+            raise RuntimeError("Failed to register with MACAW LocalAgent")
+
+    # ------------------------------------------------------------------ users
+
+    def for_user(self, user_client: Any) -> "UserBinding":
+        """Return a binding that issues invocations as `user_client`."""
+        if self._mode != "service":
+            raise ValueError("bind_to_user() is only valid on a service-mode agent")
+        if not hasattr(user_client, "agent_id") or not hasattr(user_client, "invoke_tool"):
+            raise ValueError("bind_to_user() requires a valid MACAWClient instance")
+        if not getattr(user_client, "registered", False):
+            raise ValueError(
+                "bind_to_user() requires a registered MACAWClient. Call register() first."
+            )
+
+        agent_type = getattr(user_client, "agent_type", None)
+        if agent_type and agent_type != "user":
+            logger.warning(
+                "bind_to_user() called with agent_type=%r (expected 'user'). "
+                "User identity and policy enforcement may not work as expected.",
+                agent_type,
+            )
+
+        return UserBinding(self, user_client)
+
+    # -------------------------------------------------------------- resources
+
+    def register_tool(self, name: str, fn: Callable) -> str:
+        """Register a callable MACAW will dispatch."""
+        return self._register(name, as_handler(fn))
+
+    def register_ack(self, name: str) -> str:
+        """Register a resource MACAW authorizes but does not execute."""
+        return self._register(name, _acknowledge)
+
+    def register_model(self, wrapped) -> None:
+        """Make `wrapped` reachable by the tool:<app>/generate handler."""
+        self._models[wrapped.model_name] = wrapped
+
+    def _register(self, name: str, handler: Callable) -> str:
+        mapl_name = f"tool:{self.app_name}/{name}"
+        if mapl_name not in self._registered:
+            self.macaw_client.register_tool(mapl_name, handler)
+            self._registered.add(mapl_name)
+            logger.info("Registered resource: %s", mapl_name)
+        return mapl_name
+
+    # ------------------------------------------------------------- invocation
+
+    def invoke(self, resource: str, params: Dict[str, Any], **kwargs) -> Any:
+        """
+        Invoke a resource through MACAW.
+
+        Blocking: policy, signing, and audit all happen inside this call. Callers
+        on an event loop must run it off-thread.
+        """
+        return self.issuer.invoke_tool(
+            tool_name=f"tool:{self.app_name}/{resource}",
+            parameters=params,
+            target_agent=self.server_id,
+            **kwargs,
+        )
+
+    @property
+    def issuer(self):
+        """Client that issues invocations. UserBinding overrides this."""
+        return self.macaw_client
+
+    # ------------------------------------------------------------ model seam
+
+    def _handle_generate(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Handler for tool:<app>/generate.
+
+        Runs on MACAW's handler pool once the PEP allows the call, so the
+        provider request and its response both pass through invoke_tool - the
+        same path an owned tool call takes.
+
+        Streaming authorizes here and streams from the wrapped model in the
+        caller, so no stream event is re-serialized on the way back.
+        """
+        from pydantic import TypeAdapter
+        from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelResponse
+        from pydantic_ai.models import ModelRequestParameters
+
+        p = dict(params)
+        if p.get("stream"):
+            return {"authorized": True}
+
+        model_name = p.pop("model")
+        messages = ModelMessagesTypeAdapter.validate_python(p.pop("messages"))
+        request_parameters = TypeAdapter(ModelRequestParameters).validate_python(
+            p.pop(REQUEST_PARAMETERS_KEY)
+        )
+        p.pop("native_tools", None)  # policy input; not a model setting
+
+        wrapped = self._models.get(model_name)
+        if wrapped is None:
+            raise ValueError(f"No model registered under {model_name!r}")
+
+        response = asyncio.run(wrapped.request(messages, p or None, request_parameters))
+        return TypeAdapter(ModelResponse).dump_python(response, mode="json")
+
+
+class UserBinding:
+    """A MacawBinding that issues invocations as a specific user."""
+
+    def __init__(self, base: MacawBinding, user_client: Any):
+        self._base = base
+        self._user_client = user_client
+
+    @property
+    def issuer(self):
+        return self._user_client
+
+    @property
+    def app_name(self) -> str:
+        return self._base.app_name
+
+    @property
+    def server_id(self) -> str:
+        return self._base.server_id
+
+    def register_tool(self, name: str, fn: Callable) -> str:
+        return self._base.register_tool(name, fn)
+
+    def register_ack(self, name: str) -> str:
+        return self._base.register_ack(name)
+
+    def register_model(self, wrapped) -> None:
+        self._base.register_model(wrapped)
+
+    def invoke(self, resource: str, params: Dict[str, Any], **kwargs) -> Any:
+        return self._user_client.invoke_tool(
+            tool_name=f"tool:{self.app_name}/{resource}",
+            parameters=params,
+            target_agent=self.server_id,
+            **kwargs,
+        )

package/payload/secureAI/macaw_adapters/pydantic_ai/agent.py

@@ -0,0 +1,218 @@
+"""
+SecureAgent - drop-in replacement for pydantic_ai.Agent.
+
+    # Before
+    from pydantic_ai import Agent
+    agent = Agent(OpenAIChatModel("gpt-4o"), tools=[query_catalog])
+
+    # After
+    from macaw_adapters.pydantic_ai import SecureAgent
+    agent = SecureAgent(OpenAIChatModel("gpt-4o"), app_name="acme",
+                        tools=[query_catalog])
+
+Every way Pydantic AI takes tools is governed:
+
+    tools=[fn]                       registered with MACAW
+    @agent.tool / @agent.tool_plain  registered with MACAW
+    toolsets=[FunctionToolset(...)]  functions registered with MACAW
+    toolsets=[MCPToolset(...)]       authorized here, executed by the server it
+                                     connects to (governed there by SecureMCP or
+                                     SecureMCPProxy)
+
+run(), run_sync(), run_stream(), and override() wrap any model or toolsets
+passed to them, so a per-call override cannot route around the adapter.
+
+Multi-user:
+
+    alice_agent = agent.bind_to_user(alice_client)
+"""
+
+import logging
+from contextlib import contextmanager
+from typing import Any, Callable, List, Optional, Sequence
+
+from pydantic_ai import Agent as _Agent
+from pydantic_ai.toolsets import AbstractToolset, FunctionToolset
+
+from macaw_adapters.pydantic_ai._core import MacawBinding
+from macaw_adapters.pydantic_ai.model import SecureModel
+from macaw_adapters.pydantic_ai.toolset import SecureToolset
+
+logger = logging.getLogger(__name__)
+
+
+def _functions_of(toolset: AbstractToolset) -> Optional[List[Callable]]:
+    """
+    The plain functions behind a toolset, or None if MACAW cannot hold them.
+
+    FunctionToolset exposes `.tools` as {name: Tool}, with Tool.function being
+    the original callable.
+    """
+    if not isinstance(toolset, FunctionToolset):
+        return None
+    return [t.function for t in toolset.tools.values() if hasattr(t, "function")]
+
+
+class SecureAgent(_Agent):
+    """A pydantic_ai.Agent whose model and tool calls route through MACAW."""
+
+    def __init__(
+        self,
+        model=None,
+        *,
+        app_name: Optional[str] = None,
+        intent_policy: Optional[dict] = None,
+        jwt_token: Optional[str] = None,
+        user_name: Optional[str] = None,
+        tools: Sequence[Callable] = (),
+        toolsets: Optional[Sequence[AbstractToolset]] = None,
+        _binding: Optional[MacawBinding] = None,
+        **kwargs,
+    ):
+        self._macaw = _binding or MacawBinding(
+            app_name=app_name,
+            intent_policy=intent_policy,
+            jwt_token=jwt_token,
+            user_name=user_name,
+        )
+
+        # Kept so bind_to_user() can rebuild an identical agent on a user binding.
+        self._raw_model = model
+        self._raw_tools = list(tools)
+        self._raw_toolsets = list(toolsets or [])
+        self._raw_kwargs = dict(kwargs)
+
+        # tools= and the @agent.tool decorators share one governed toolset.
+        functions = list(tools)
+        passthrough = []
+        for toolset in self._raw_toolsets:
+            fns = _functions_of(toolset)
+            if fns is None:
+                passthrough.append(toolset)
+            else:
+                functions.extend(fns)
+
+        self._tools = SecureToolset(
+            self._macaw, FunctionToolset(tools=functions), owned=functions
+        )
+
+        super().__init__(
+            self._secure_model(model),
+            toolsets=[self._tools] + [self._secure_toolset(t) for t in passthrough],
+            **kwargs,
+        )
+
+    # --------------------------------------------------------------- wrapping
+
+    def _secure_model(self, model):
+        if model is None or isinstance(model, SecureModel):
+            return model
+        if isinstance(model, str):
+            from pydantic_ai.models import infer_model
+
+            model = infer_model(model)
+
+        # A router such as FallbackModel keeps its candidates in .models and
+        # reports a composite model_name. Secure the candidates instead, so
+        # policy always sees the model that actually calls a provider.
+        members = getattr(model, "models", None)
+        if members:
+            model.models = [self._secure_model(m) for m in members]
+            return model
+
+        return SecureModel(self._macaw, model)
+
+    def _secure_toolset(self, toolset: AbstractToolset):
+        if isinstance(toolset, SecureToolset):
+            return toolset
+        fns = _functions_of(toolset)
+        if fns is None:
+            return SecureToolset(self._macaw, toolset)
+        return SecureToolset(self._macaw, toolset, owned=fns)
+
+    def _secure_toolsets(self, toolsets):
+        if toolsets is None:
+            return None
+        return [self._secure_toolset(t) for t in toolsets]
+
+    # ------------------------------------------------------------- decorators
+
+    def tool(self, func=None, /, **kwargs):
+        """Register a tool that takes RunContext. Governed like any other."""
+        return self._tool_decorator(func, takes_ctx=True, **kwargs)
+
+    def tool_plain(self, func=None, /, **kwargs):
+        """Register a tool that does not take RunContext."""
+        return self._tool_decorator(func, takes_ctx=False, **kwargs)
+
+    def _tool_decorator(self, func, *, takes_ctx: bool, **kwargs):
+        def register(fn):
+            self._tools.wrapped.add_function(fn, takes_ctx=takes_ctx, **kwargs)
+            name = kwargs.get("name") or fn.__name__
+            if takes_ctx:
+                # A tool taking RunContext cannot be dispatched by MACAW, which
+                # passes only a params dict, so it is authorized instead.
+                self._macaw.register_ack(name)
+            else:
+                self._tools.own(fn, name=name)
+            return fn
+
+        return register if func is None else register(func)
+
+    # ------------------------------------------------------------------ seams
+    # Per-call model/toolsets replace the agent's own, so they are wrapped too.
+
+    async def run(self, *args, model=None, toolsets=None, **kwargs):
+        return await super().run(
+            *args,
+            model=self._secure_model(model),
+            toolsets=self._secure_toolsets(toolsets),
+            **kwargs,
+        )
+
+    def run_sync(self, *args, model=None, toolsets=None, **kwargs):
+        return super().run_sync(
+            *args,
+            model=self._secure_model(model),
+            toolsets=self._secure_toolsets(toolsets),
+            **kwargs,
+        )
+
+    def run_stream(self, *args, model=None, toolsets=None, **kwargs):
+        return super().run_stream(
+            *args,
+            model=self._secure_model(model),
+            toolsets=self._secure_toolsets(toolsets),
+            **kwargs,
+        )
+
+    @contextmanager
+    def override(self, *, model=None, toolsets=None, **kwargs):
+        overrides = dict(kwargs)
+        if model is not None:
+            overrides["model"] = self._secure_model(model)
+        if toolsets is not None:
+            overrides["toolsets"] = self._secure_toolsets(toolsets)
+        with super().override(**overrides):
+            yield
+
+    # ------------------------------------------------------------------ users
+
+    def bind_to_user(self, user_client: Any) -> "SecureAgent":
+        """
+        Return an agent that issues invocations as `user_client`.
+
+        Shares this agent's registration and tools; policy is evaluated for the
+        user on every model and tool call.
+        """
+        return SecureAgent(
+            self._raw_model,
+            tools=self._raw_tools,
+            toolsets=self._raw_toolsets,
+            _binding=self._macaw.for_user(user_client),
+            **self._raw_kwargs,
+        )
+
+
+# Drop-in alias (same class, native naming)
+Agent = SecureAgent

package/payload/secureAI/macaw_adapters/pydantic_ai/model.py

@@ -0,0 +1,99 @@
+"""
+SecureModel - the model seam.
+
+Wraps any Pydantic AI model so every request is invoked as tool:<app>/generate.
+MACAW dispatches to the wrapped model once the PEP allows the call, so the
+request and the response both travel through invoke_tool - the same path a tool
+call takes.
+
+What policy sees per call:
+
+    model          the model id, e.g. "anthropic.claude-haiku-4-5"
+    messages       the conversation, declared as a prompt so invoke_tool creates
+                   authenticated prompts from it
+    native_tools   provider-side tools offered on this call
+    <settings>     max_tokens, temperature, and anything else the caller set
+"""
+
+import logging
+from contextlib import asynccontextmanager
+from typing import Any, Dict, List, Optional
+
+import anyio
+from pydantic import TypeAdapter
+from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelResponse
+from pydantic_ai.models import ModelRequestParameters
+from pydantic_ai.models.wrapper import WrapperModel
+
+from macaw_adapters.pydantic_ai._core import REQUEST_PARAMETERS_KEY
+
+logger = logging.getLogger(__name__)
+
+_RESPONSE = TypeAdapter(ModelResponse)
+_REQUEST_PARAMETERS = TypeAdapter(ModelRequestParameters)
+
+
+class SecureModel(WrapperModel):
+    """A Pydantic AI model whose every request goes through the MACAW PEP."""
+
+    def __init__(self, binding, wrapped):
+        super().__init__(wrapped)
+        self._macaw = binding
+        binding.register_model(self.wrapped)
+
+    @staticmethod
+    def _native_tool_names(request_parameters) -> List[str]:
+        """
+        Names of the provider-side tools offered on this request.
+
+        Read from the request rather than from what the application declared, so
+        policy sees what the provider was actually asked for.
+        """
+        native = getattr(request_parameters, "native_tools", None) or []
+        return sorted(
+            getattr(tool, "kind", None) or type(tool).__name__ for tool in native
+        )
+
+    def _params(
+        self,
+        messages: List[Any],
+        model_settings: Optional[Dict[str, Any]],
+        request_parameters: ModelRequestParameters,
+    ) -> Dict[str, Any]:
+        return {
+            "model": self.wrapped.model_name,
+            "messages": ModelMessagesTypeAdapter.dump_python(messages, mode="json"),
+            "native_tools": self._native_tool_names(request_parameters),
+            REQUEST_PARAMETERS_KEY: _REQUEST_PARAMETERS.dump_python(
+                request_parameters, mode="json"
+            ),
+            **(model_settings or {}),
+        }
+
+    async def request(self, messages, model_settings, model_request_parameters):
+        params = self._params(messages, model_settings, model_request_parameters)
+        result = await anyio.to_thread.run_sync(
+            lambda: self._macaw.invoke("generate", params)
+        )
+        return _RESPONSE.validate_python(result)
+
+    @asynccontextmanager
+    async def request_stream(
+        self, messages, model_settings, model_request_parameters, run_context=None
+    ):
+        """
+        Streaming path.
+
+        The call is authorized through invoke_tool and the events are then
+        streamed from the wrapped model directly, so no stream event is
+        re-serialized on the way to the caller.
+        """
+        params = self._params(messages, model_settings, model_request_parameters)
+        params["stream"] = True
+        await anyio.to_thread.run_sync(
+            lambda: self._macaw.invoke("generate", dict(params, _authorize_only=True))
+        )
+        async with self.wrapped.request_stream(
+            messages, model_settings, model_request_parameters, run_context
+        ) as stream:
+            yield stream

package/payload/secureAI/macaw_adapters/pydantic_ai/toolset.py

@@ -0,0 +1,86 @@
+"""
+SecureToolset - the tool seam.
+
+Wraps a Pydantic AI toolset so every call is invoked as tool:<app>/<tool_name>.
+
+Two cases, decided by whether MACAW can hold the callable:
+
+    owned       plain functions (tools=, @agent.tool_plain, FunctionToolset).
+                Registered with MACAW, which dispatches them - request and
+                response both travel through invoke_tool.
+
+    not owned   toolsets whose execution MACAW cannot hold, such as MCPToolset
+                or a tool taking RunContext. The call is authorized through
+                invoke_tool and then executed by the wrapped toolset.
+
+Either way no tool runs until the PEP allows it.
+"""
+
+import logging
+from typing import Any, Callable, Dict, Iterable, Optional
+
+import anyio
+from pydantic_ai.exceptions import ToolFailed
+from pydantic_ai.toolsets import AbstractToolset, WrapperToolset
+
+from macaw_adapters.pydantic_ai._core import is_permission_denied
+
+logger = logging.getLogger(__name__)
+
+
+class SecureToolset(WrapperToolset):
+    """A Pydantic AI toolset whose every call goes through the MACAW PEP."""
+
+    def __init__(
+        self,
+        binding,
+        wrapped: AbstractToolset,
+        owned: Iterable[Callable] = (),
+    ):
+        super().__init__(wrapped)
+        self._macaw = binding
+        self._owned = set()
+        for fn in owned:
+            self.own(fn)
+
+    def own(self, fn: Callable, name: Optional[str] = None) -> Callable:
+        """Register `fn` with MACAW so MACAW dispatches it."""
+        tool_name = name or fn.__name__
+        self._macaw.register_tool(tool_name, fn)
+        self._owned.add(tool_name)
+        return fn
+
+    async def get_tools(self, ctx) -> Dict[str, Any]:
+        """
+        Resolve tools from the wrapped toolset and make each one a MACAW
+        resource. Schemas and validation stay Pydantic AI's.
+        """
+        tools = await super().get_tools(ctx)
+        for name in tools:
+            if name not in self._owned:
+                self._macaw.register_ack(name)
+        return tools
+
+    async def call_tool(self, name: str, tool_args: Dict[str, Any], ctx, tool) -> Any:
+        """
+        Invoke through MACAW. Owned tools are dispatched by MACAW; the rest are
+        authorized here and executed by the wrapped toolset.
+
+        A denial is re-raised as ToolFailed so the model sees it and adapts.
+        Retrying would not change the outcome, which is why this is ToolFailed
+        rather than ModelRetry.
+        """
+        try:
+            result = await anyio.to_thread.run_sync(
+                lambda: self._macaw.invoke(name, tool_args)
+            )
+        except Exception as exc:
+            if is_permission_denied(exc):
+                raise ToolFailed(
+                    f"Access denied by security policy: {name}. {exc}"
+                ) from exc
+            raise
+
+        if name in self._owned:
+            return result
+        return await self.wrapped.call_tool(name, tool_args, ctx, tool)