agent_os_kernel 3.1.0__py3-none-any.whl

agent_os/integrations/openai_adapter.py

@@ -0,0 +1,816 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+OpenAI Assistants Integration
+
+Wraps OpenAI Assistants API with Agent OS governance.
+
+Usage:
+    from agent_os.integrations import OpenAIKernel
+    from openai import OpenAI
+
+    client = OpenAI()
+    kernel = OpenAIKernel(policy="strict")
+
+    # Create assistant as normal
+    assistant = client.beta.assistants.create(
+        name="Trading Bot",
+        instructions="You analyze market data",
+        model="gpt-4-turbo"
+    )
+
+    # Wrap for governance
+    governed_assistant = kernel.wrap(assistant, client)
+
+    # All runs are now governed!
+    thread = governed_assistant.create_thread()
+    governed_assistant.add_message(thread.id, "Analyze AAPL")
+    run = governed_assistant.run(thread.id)  # Governed execution
+
+Features:
+- Pre-execution policy checks
+- Tool call interception and validation
+- Real-time run monitoring
+- SIGKILL support (cancel run on violation)
+- Full audit trail
+"""
+
+import json
+import logging
+import random
+import time
+from collections.abc import Generator
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Callable, Optional
+
+from .base import BaseIntegration, ExecutionContext, GovernancePolicy
+
+logger = logging.getLogger("agent_os.openai")
+
+
+@dataclass
+class AssistantContext(ExecutionContext):
+    """Extended execution context for OpenAI Assistants.
+
+    Tracks assistant-specific state including thread IDs, run IDs,
+    function call history, and cumulative token usage for governance
+    enforcement.
+
+    Attributes:
+        assistant_id: The OpenAI assistant identifier.
+        thread_ids: List of thread IDs created during this session.
+        run_ids: List of run IDs executed during this session.
+        function_calls: History of function/tool calls made by the assistant.
+        prompt_tokens: Cumulative prompt tokens consumed across all runs.
+        completion_tokens: Cumulative completion tokens consumed across all runs.
+    """
+
+    assistant_id: str = ""
+    thread_ids: list[str] = field(default_factory=list)
+    run_ids: list[str] = field(default_factory=list)
+    function_calls: list[dict] = field(default_factory=list)
+
+    # Token tracking
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+
+
+# Transient error base classes for retry detection
+_TRANSIENT_ERROR_NAMES = ("RateLimitError", "APIConnectionError", "Timeout", "APITimeoutError")
+
+
+def _is_transient(exc: Exception) -> bool:
+    """Return True if the exception is a transient OpenAI error."""
+    return type(exc).__name__ in _TRANSIENT_ERROR_NAMES
+
+
+def retry_with_backoff(
+    fn: Callable[..., Any],
+    *args: Any,
+    max_retries: int = 3,
+    base_delay: float = 1.0,
+    max_delay: float = 30.0,
+    **kwargs: Any,
+) -> Any:
+    """Call *fn* with exponential backoff + jitter on transient errors.
+
+    Args:
+        fn: Callable to invoke.
+        max_retries: Number of retry attempts after the initial call.
+        base_delay: Base delay in seconds for backoff calculation.
+        max_delay: Upper bound for the computed delay.
+
+    Returns:
+        The return value of *fn*.
+
+    Raises:
+        The last caught exception if all retries are exhausted.
+    """
+    last_exc: Optional[Exception] = None
+    for attempt in range(max_retries + 1):
+        try:
+            return fn(*args, **kwargs)
+        except Exception as exc:
+            if not _is_transient(exc) or attempt == max_retries:
+                raise
+            last_exc = exc
+            delay = min(base_delay * (2 ** attempt) + random.uniform(0, 1), max_delay)
+            logger.warning(
+                "Retry %d/%d for %s after %s (delay=%.2fs)",
+                attempt + 1,
+                max_retries,
+                fn.__name__ if hasattr(fn, "__name__") else str(fn),
+                type(exc).__name__,
+                delay,
+            )
+            time.sleep(delay)
+    raise last_exc  # type: ignore[misc]
+
+
+class OpenAIKernel(BaseIntegration):
+    """
+    OpenAI Assistants adapter for Agent OS.
+
+    Provides governance for:
+    - Assistant creation/modification
+    - Thread management
+    - Run execution
+    - Tool/function calls
+    - File operations
+
+    Example:
+        kernel = OpenAIKernel(policy=GovernancePolicy(
+            max_tokens=10000,
+            allowed_tools=["code_interpreter", "retrieval"],
+            blocked_patterns=["password", "api_key", "secret"]
+        ))
+
+        governed = kernel.wrap(assistant, client)
+        result = governed.run(thread_id)
+    """
+
+    def __init__(
+        self,
+        policy: Optional[GovernancePolicy] = None,
+        max_retries: int = 3,
+        timeout_seconds: float = 300.0,
+    ):
+        """Initialise the OpenAI governance kernel.
+
+        Args:
+            policy: Governance policy to enforce. When ``None`` the default
+                ``GovernancePolicy`` is used.
+            max_retries: Maximum number of retry attempts for transient
+                OpenAI errors (default 3).
+            timeout_seconds: Default timeout in seconds for operations
+                (default 300).
+        """
+        super().__init__(policy)
+        self.max_retries = max_retries
+        self.timeout_seconds = timeout_seconds
+        self._wrapped_assistants: dict[str, Any] = {}  # assistant_id -> original
+        self._clients: dict[str, Any] = {}  # assistant_id -> client
+        self._cancelled_runs: set[str] = set()
+        self._start_time = time.monotonic()
+        self._last_error: Optional[str] = None
+
+    def wrap(self, agent: Any, client: Any = None) -> "GovernedAssistant":
+        """Wrap an OpenAI Assistant with governance.
+
+        This is the primary wrapping method, consistent with all other
+        adapters.  OpenAI Assistants require both an assistant object
+        **and** a client, so ``client`` must be provided.
+
+        Args:
+            agent: OpenAI Assistant object.
+            client: OpenAI client instance (required).
+
+        Returns:
+            GovernedAssistant with full governance.
+
+        Raises:
+            TypeError: If *client* is not provided.
+        """
+        if client is None:
+            raise TypeError(
+                "OpenAIKernel.wrap() requires a 'client' argument: "
+                "kernel.wrap(assistant, client)"
+            )
+        assistant_id = agent.id
+        ctx = AssistantContext(
+            agent_id=assistant_id,
+            session_id=f"oai-{int(time.time())}",
+            policy=self.policy,
+            assistant_id=assistant_id
+        )
+        self.contexts[assistant_id] = ctx
+        self._wrapped_assistants[assistant_id] = agent
+        self._clients[assistant_id] = client
+
+        return GovernedAssistant(
+            assistant=agent,
+            client=client,
+            kernel=self,
+            ctx=ctx
+        )
+
+    def wrap_assistant(self, assistant: Any, client: Any) -> "GovernedAssistant":
+        """Wrap an OpenAI Assistant with governance.
+
+        .. deprecated::
+            Use :meth:`wrap` instead::
+
+                governed = kernel.wrap(assistant, client)
+
+        Args:
+            assistant: OpenAI Assistant object.
+            client: OpenAI client instance.
+
+        Returns:
+            GovernedAssistant with full governance.
+        """
+        import warnings
+        warnings.warn(
+            "wrap_assistant() is deprecated, use wrap(assistant, client) instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.wrap(assistant, client)
+
+    def unwrap(self, governed_agent: Any) -> Any:
+        """Retrieve the original unwrapped assistant.
+
+        Args:
+            governed_agent: A ``GovernedAssistant`` or any object.
+
+        Returns:
+            The original OpenAI assistant object if *governed_agent* is a
+            ``GovernedAssistant``; otherwise returns *governed_agent* as-is.
+        """
+        if isinstance(governed_agent, GovernedAssistant):
+            return governed_agent._assistant
+        return governed_agent
+
+    def cancel_run(self, thread_id: str, run_id: str, client: Any):
+        """Cancel a run (SIGKILL equivalent).
+
+        Immediately marks the run as cancelled locally and issues a cancel
+        request to the OpenAI API.  If the API call fails (e.g. the run
+        has already completed), the error is silently ignored.
+
+        Args:
+            thread_id: The thread the run belongs to.
+            run_id: The run to cancel.
+            client: OpenAI client used to issue the cancellation.
+        """
+        self._cancelled_runs.add(run_id)
+        try:
+            client.beta.threads.runs.cancel(
+                thread_id=thread_id,
+                run_id=run_id
+            )
+        except Exception:  # noqa: BLE001 — best-effort cancel, run may already be complete
+            logger.warning("Run cancel failed (may already be complete): thread=%s run=%s", thread_id, run_id, exc_info=True)
+
+    def is_cancelled(self, run_id: str) -> bool:
+        """Check whether a run has been cancelled via :meth:`cancel_run`.
+
+        Args:
+            run_id: The run identifier to check.
+
+        Returns:
+            ``True`` if the run was previously cancelled.
+        """
+        return run_id in self._cancelled_runs
+
+    def health_check(self) -> dict[str, Any]:
+        """Return adapter health status.
+
+        Returns:
+            A dict with ``status``, ``backend``, ``last_error``, and
+            ``uptime_seconds`` keys.
+        """
+        uptime = time.monotonic() - self._start_time
+        has_clients = bool(self._clients)
+        if self._last_error:
+            status = "degraded"
+        elif not has_clients:
+            status = "healthy"
+        else:
+            status = "healthy"
+        return {
+            "status": status,
+            "backend": "openai",
+            "backend_connected": has_clients,
+            "last_error": self._last_error,
+            "uptime_seconds": round(uptime, 2),
+        }
+
+
+class GovernedAssistant:
+    """
+    OpenAI Assistant wrapped with Agent OS governance.
+
+    All API calls are intercepted for policy enforcement.
+    """
+
+    def __init__(
+        self,
+        assistant: Any,
+        client: Any,
+        kernel: OpenAIKernel,
+        ctx: AssistantContext
+    ):
+        self._assistant = assistant
+        self._client = client
+        self._kernel = kernel
+        self._ctx = ctx
+        self._tool_registry: dict[str, Callable] = {}
+
+    def register_tool(self, name: str, func: Callable) -> None:
+        """Register a tool function for automatic execution."""
+        self._tool_registry[name] = func
+
+    @property
+    def id(self) -> str:
+        """Assistant ID"""
+        return self._assistant.id
+
+    @property
+    def name(self) -> str:
+        """Assistant name"""
+        return self._assistant.name
+
+    # =========================================================================
+    # Thread Management
+    # =========================================================================
+
+    def create_thread(self, **kwargs) -> Any:
+        """Create a new conversation thread.
+
+        The thread ID is automatically recorded in the execution context
+        for audit purposes.
+
+        Args:
+            **kwargs: Forwarded to ``client.beta.threads.create()``.
+
+        Returns:
+            The newly created OpenAI thread object.
+        """
+        thread = self._client.beta.threads.create(**kwargs)
+        self._ctx.thread_ids.append(thread.id)
+        return thread
+
+    def get_thread(self, thread_id: str) -> Any:
+        """Retrieve an existing thread by ID.
+
+        Args:
+            thread_id: The thread to retrieve.
+
+        Returns:
+            The OpenAI thread object.
+        """
+        return self._client.beta.threads.retrieve(thread_id)
+
+    def delete_thread(self, thread_id: str) -> bool:
+        """Delete a thread and remove it from the execution context.
+
+        Args:
+            thread_id: The thread to delete.
+
+        Returns:
+            ``True`` if the thread was successfully deleted.
+        """
+        result = self._client.beta.threads.delete(thread_id)
+        if thread_id in self._ctx.thread_ids:
+            self._ctx.thread_ids.remove(thread_id)
+        return result.deleted
+
+    # =========================================================================
+    # Message Management
+    # =========================================================================
+
+    def add_message(
+        self,
+        thread_id: str,
+        content: str,
+        role: str = "user",
+        **kwargs
+    ) -> Any:
+        """Add a message to a thread with pre-execution policy checks.
+
+        The message content is validated against ``blocked_patterns`` before
+        being sent to the OpenAI API.
+
+        Args:
+            thread_id: Target thread.
+            content: Message text.
+            role: Message role (default ``"user"``).
+            **kwargs: Additional parameters forwarded to the API.
+
+        Returns:
+            The created OpenAI message object.
+
+        Raises:
+            PolicyViolationError: If the content matches a blocked pattern.
+        """
+        # Pre-check: blocked patterns
+        allowed, reason = self._kernel.pre_execute(self._ctx, content)
+        if not allowed:
+            raise PolicyViolationError(f"Message blocked: {reason}")
+
+        message = self._client.beta.threads.messages.create(
+            thread_id=thread_id,
+            role=role,
+            content=content,
+            **kwargs
+        )
+        return message
+
+    def list_messages(self, thread_id: str, **kwargs) -> list:
+        """List messages in a thread.
+
+        Args:
+            thread_id: The thread whose messages to list.
+            **kwargs: Additional parameters (e.g. ``limit``, ``order``).
+
+        Returns:
+            A list of message objects in the thread.
+        """
+        return self._client.beta.threads.messages.list(
+            thread_id=thread_id,
+            **kwargs
+        )
+
+    # =========================================================================
+    # Run Execution (Core Governance)
+    # =========================================================================
+
+    def run(
+        self,
+        thread_id: str,
+        instructions: Optional[str] = None,
+        tools: Optional[list] = None,
+        poll_interval: float = 1.0,
+        **kwargs
+    ) -> Any:
+        """
+        Execute a governed run.
+
+        This is the primary method for executing the assistant.
+        All tool calls and outputs are validated against policy.
+
+        Args:
+            thread_id: Thread to run on
+            instructions: Optional override instructions
+            tools: Optional tools to enable
+            poll_interval: How often to check run status
+            **kwargs: Additional run parameters
+
+        Returns:
+            Completed run object
+
+        Raises:
+            PolicyViolationError: If policy is violated
+            RunCancelledException: If run was SIGKILL'd
+        """
+        # Pre-check
+        if instructions:
+            allowed, reason = self._kernel.pre_execute(self._ctx, instructions)
+            if not allowed:
+                raise PolicyViolationError(f"Instructions blocked: {reason}")
+
+        # Validate tools against policy
+        if tools:
+            self._validate_tools(tools)
+
+        # Create run
+        run_kwargs = {
+            "thread_id": thread_id,
+            "assistant_id": self._assistant.id,
+            **kwargs
+        }
+        if instructions:
+            run_kwargs["instructions"] = instructions
+        if tools:
+            run_kwargs["tools"] = tools
+
+        run = self._client.beta.threads.runs.create(**run_kwargs)
+        self._ctx.run_ids.append(run.id)
+
+        # Poll until complete (with governance checks)
+        return self._poll_run(thread_id, run.id, poll_interval)
+
+    def run_stream(
+        self,
+        thread_id: str,
+        instructions: Optional[str] = None,
+        **kwargs
+    ) -> Generator:
+        """
+        Stream a governed run.
+
+        Yields events as they arrive, with real-time policy checks.
+        """
+        # Pre-check
+        if instructions:
+            allowed, reason = self._kernel.pre_execute(self._ctx, instructions)
+            if not allowed:
+                raise PolicyViolationError(f"Instructions blocked: {reason}")
+
+        # Create streaming run
+        with self._client.beta.threads.runs.stream(
+            thread_id=thread_id,
+            assistant_id=self._assistant.id,
+            instructions=instructions,
+            **kwargs
+        ) as stream:
+            for event in stream:
+                # Check for cancellation
+                if hasattr(event, 'data') and hasattr(event.data, 'id'):
+                    if self._kernel.is_cancelled(event.data.id):
+                        raise RunCancelledException("Run was cancelled (SIGKILL)")
+
+                # Yield event
+                yield event
+
+    def _poll_run(
+        self,
+        thread_id: str,
+        run_id: str,
+        poll_interval: float
+    ) -> Any:
+        """
+        Poll run status with governance checks.
+        """
+        while True:
+            # Check for SIGKILL
+            if self._kernel.is_cancelled(run_id):
+                raise RunCancelledException("Run was cancelled (SIGKILL)")
+
+            run = self._client.beta.threads.runs.retrieve(
+                thread_id=thread_id,
+                run_id=run_id
+            )
+
+            # Update token counts
+            if hasattr(run, 'usage') and run.usage:
+                self._ctx.prompt_tokens += run.usage.prompt_tokens or 0
+                self._ctx.completion_tokens += run.usage.completion_tokens or 0
+
+                # Check token limit
+                total = self._ctx.prompt_tokens + self._ctx.completion_tokens
+                if total > self._kernel.policy.max_tokens:
+                    self._kernel.cancel_run(thread_id, run_id, self._client)
+                    raise PolicyViolationError(
+                        f"Token limit exceeded: {total} > {self._kernel.policy.max_tokens}"
+                    )
+
+            # Handle different statuses
+            if run.status == "completed":
+                self._kernel.post_execute(self._ctx, run)
+                return run
+
+            elif run.status == "requires_action":
+                # Tool calls need approval
+                run = self._handle_tool_calls(thread_id, run)
+
+            elif run.status in ["failed", "cancelled", "expired"]:
+                return run
+
+            elif run.status in ["queued", "in_progress"]:
+                time.sleep(poll_interval)
+
+            else:
+                # Unknown status
+                time.sleep(poll_interval)
+
+    def _handle_tool_calls(self, thread_id: str, run: Any) -> Any:
+        """
+        Handle tool calls with policy validation.
+        """
+        tool_calls = run.required_action.submit_tool_outputs.tool_calls
+        tool_outputs = []
+
+        for tool_call in tool_calls:
+            # Record tool call
+            call_info = {
+                "id": tool_call.id,
+                "type": tool_call.type,
+                "function": tool_call.function.name if hasattr(tool_call, 'function') else None,
+                "arguments": tool_call.function.arguments if hasattr(tool_call, 'function') else None,
+                "timestamp": datetime.now().isoformat()
+            }
+            self._ctx.function_calls.append(call_info)
+            self._ctx.tool_calls.append(call_info)
+
+            # Check tool call count
+            if len(self._ctx.tool_calls) > self._kernel.policy.max_tool_calls:
+                self._kernel.cancel_run(thread_id, run.id, self._client)
+                raise PolicyViolationError(
+                    f"Tool call limit exceeded: {len(self._ctx.tool_calls)} > {self._kernel.policy.max_tool_calls}"
+                )
+
+            # Validate function name
+            if hasattr(tool_call, 'function'):
+                func_name = tool_call.function.name
+                if self._kernel.policy.allowed_tools:
+                    if func_name not in self._kernel.policy.allowed_tools:
+                        self._kernel.cancel_run(thread_id, run.id, self._client)
+                        raise PolicyViolationError(
+                            f"Tool not allowed: {func_name}"
+                        )
+
+            # Check human approval requirement
+            if self._kernel.policy.require_human_approval:
+                tool_outputs.append({
+                    "tool_call_id": tool_call.id,
+                    "output": json.dumps({
+                        "status": "requires_approval",
+                        "function": func_name if hasattr(tool_call, 'function') else "unknown",
+                        "message": "Tool execution requires human approval per governance policy"
+                    })
+                })
+                continue
+
+            # Execute via tool registry if available
+            output = None
+            if hasattr(self, '_tool_registry') and self._tool_registry:
+                func_name_exec = tool_call.function.name if hasattr(tool_call, 'function') else None
+                if func_name_exec and func_name_exec in self._tool_registry:
+                    try:
+                        args = json.loads(tool_call.function.arguments) if hasattr(tool_call, 'function') else {}
+                        result = self._tool_registry[func_name_exec](**args)
+                        output = json.dumps(result) if not isinstance(result, str) else result
+                    except Exception as e:
+                        logger.warning("Tool execution failed for %s", func_name_exec, exc_info=True)
+                        output = json.dumps({"status": "error", "message": str(e)})
+
+            if output is None:
+                output = json.dumps({
+                    "status": "no_executor",
+                    "function": tool_call.function.name if hasattr(tool_call, 'function') else "unknown",
+                    "message": "No tool executor registered for this function"
+                })
+
+            tool_outputs.append({
+                "tool_call_id": tool_call.id,
+                "output": output
+            })
+
+        # Submit outputs
+        return self._client.beta.threads.runs.submit_tool_outputs(
+            thread_id=thread_id,
+            run_id=run.id,
+            tool_outputs=tool_outputs
+        )
+
+    def _validate_tools(self, tools: list):
+        """Validate tools against policy"""
+        if not self._kernel.policy.allowed_tools:
+            return  # No restrictions
+
+        for tool in tools:
+            tool_type = tool.get("type") if isinstance(tool, dict) else getattr(tool, "type", None)
+            if tool_type and tool_type not in self._kernel.policy.allowed_tools:
+                raise PolicyViolationError(f"Tool type not allowed: {tool_type}")
+
+    # =========================================================================
+    # Signal Handling
+    # =========================================================================
+
+    def sigkill(self, thread_id: str, run_id: str):
+        """Send SIGKILL to a running assistant — immediately cancels the run.
+
+        This is the primary mechanism for forcibly stopping a governed
+        assistant that has violated policy or needs emergency termination.
+
+        Args:
+            thread_id: Thread containing the run.
+            run_id: The run to kill.
+
+        Example:
+            >>> governed.sigkill(thread_id="thread_abc", run_id="run_xyz")
+        """
+        self._kernel.cancel_run(thread_id, run_id, self._client)
+
+    def sigstop(self, thread_id: str, run_id: str):
+        """Send SIGSTOP to a running assistant.
+
+        .. note::
+
+            The OpenAI Assistants API does not support pausing a run, so
+            this behaves identically to :meth:`sigkill` (cancels the run).
+
+        Args:
+            thread_id: Thread containing the run.
+            run_id: The run to stop.
+        """
+        self._kernel.cancel_run(thread_id, run_id, self._client)
+
+    # =========================================================================
+    # Utility
+    # =========================================================================
+
+    def get_context(self) -> AssistantContext:
+        """Return the execution context containing the full audit trail.
+
+        Returns:
+            The ``AssistantContext`` for this governed assistant, including
+            thread IDs, run IDs, function call history, and token usage.
+        """
+        return self._ctx
+
+    def get_token_usage(self) -> dict:
+        """Return cumulative token usage statistics.
+
+        Returns:
+            A dict with keys ``prompt_tokens``, ``completion_tokens``,
+            ``total_tokens``, and ``limit``.
+
+        Example:
+            >>> governed.get_token_usage()
+            {'prompt_tokens': 120, 'completion_tokens': 80, 'total_tokens': 200, 'limit': 10000}
+        """
+        return {
+            "prompt_tokens": self._ctx.prompt_tokens,
+            "completion_tokens": self._ctx.completion_tokens,
+            "total_tokens": self._ctx.prompt_tokens + self._ctx.completion_tokens,
+            "limit": self._kernel.policy.max_tokens
+        }
+
+    def __getattr__(self, name):
+        """Proxy attribute access to the underlying OpenAI assistant.
+
+        Allows transparent access to assistant properties (e.g. ``model``,
+        ``instructions``) that are not explicitly overridden by this wrapper.
+        """
+        return getattr(self._assistant, name)
+
+
+class PolicyViolationError(Exception):
+    """Raised when an assistant action violates governance policy.
+
+    Contains a human-readable reason describing which policy was violated.
+    """
+
+    pass
+
+
+class RunCancelledException(Exception):
+    """Raised when a run is forcibly cancelled via SIGKILL.
+
+    Indicates that ``cancel_run`` (or ``sigkill``) was invoked, either
+    directly or automatically by the governance layer (e.g. token limit
+    exceeded, disallowed tool call).
+    """
+
+    pass
+
+
+# ============================================================================
+# Convenience Functions
+# ============================================================================
+
+def wrap(
+    assistant: Any,
+    client: Any,
+    policy: Optional[GovernancePolicy] = None,
+    max_retries: int = 3,
+    timeout_seconds: float = 300.0,
+) -> GovernedAssistant:
+    """Quick wrapper for OpenAI Assistants.
+
+    Example::
+
+        from agent_os.integrations.openai_adapter import wrap
+
+        governed = wrap(my_assistant, openai_client)
+        result = governed.run(thread_id)
+    """
+    return OpenAIKernel(
+        policy, max_retries=max_retries, timeout_seconds=timeout_seconds
+    ).wrap(assistant, client)
+
+
+def wrap_assistant(
+    assistant: Any,
+    client: Any,
+    policy: Optional[GovernancePolicy] = None,
+    max_retries: int = 3,
+    timeout_seconds: float = 300.0,
+) -> GovernedAssistant:
+    """Quick wrapper for OpenAI Assistants.
+
+    .. deprecated::
+        Use :func:`wrap` instead.
+    """
+    import warnings
+    warnings.warn(
+        "wrap_assistant() is deprecated, use wrap() instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return wrap(assistant, client, policy=policy, max_retries=max_retries,
+                timeout_seconds=timeout_seconds)