gopidesupavan commented on code in PR #71946: URL: https://github.com/apache/airflow/pull/71946#discussion_r3834577798
########## providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py: ########## @@ -0,0 +1,317 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import logging +from abc import abstractmethod +from typing import TYPE_CHECKING, Any + +from pydantic_ai.exceptions import ModelRetry +from pydantic_ai.tools import ToolDefinition +from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool + +from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError +from airflow.providers.common.ai.utils.tool_definition import ( + build_args_validator, + return_schema_kwargs, + serialize_for_llm, +) +from airflow.providers.common.compat.sdk import Stats + +if TYPE_CHECKING: + from pydantic_ai._run_context import RunContext + +log = logging.getLogger(__name__) + +_PROMPT_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "The question or instruction to send to this agent.", + } + }, + "required": ["prompt"], +} + + +class BaseManagedAgentToolset(AbstractToolset[Any]): + """ + Base class exposing a vendor-managed agent as a single pydantic-ai tool. + + A managed agent runs its own reasoning loop on the vendor's infrastructure + (Snowflake Cortex Agents, Amazon Bedrock AgentCore, Azure AI Foundry hosted + agents, Vertex AI Agent Engine). Airflow submits one request and reads one + answer, so the Airflow-side agent features -- toolsets, human-in-the-loop + review, durable step replay -- apply to the *calling* agent and never reach + inside the managed agent. + + Subclasses implement :meth:`agent_ref` and :meth:`invoke`. Tool naming, + argument validation, result serialisation and logging are handled here so + every provider's implementation presents the same surface to the model. + + :param tool_name: Name the calling model sees, and the identifier it emits + when calling the tool. A verb phrase naming the specialist reads best, + e.g. ``ask_bookings_analyst``. + :param description: What this agent knows and when to consult it. Optional -- + it falls back to ``tool_name`` rendered as prose, matching how + ``HookToolset`` handles a method with no docstring. Worth writing anyway: + it is what tells the model to consult the agent rather than answer from + its own knowledge, and it is the only place to state a scope limit the + name cannot carry ("cannot see revenue figures"). Since the argument + schema is always a bare prompt, the name and this string are the whole + of what the model knows about the agent. + :param timeout: Seconds to wait for a single invocation. ``None`` defers to + the platform default, which subclasses supply -- a number chosen here + would silently disagree with the vendor operator's documented timeout + for the same service. + """ + + #: Whether a completed invocation may be replayed from the durable cache + #: instead of re-invoked. Off by default because a managed agent may act on + #: systems Airflow cannot observe, so replaying a cached answer could skip a + #: side effect. Read-only agents should opt in. + replayable: bool = False + + def __init__( + self, + *, + tool_name: str, + description: str | None = None, + timeout: float | None = None, + ) -> None: + if not tool_name: + raise ValueError("tool_name must be a non-empty string.") + self._tool_name = tool_name + # Same fallback as HookToolset uses for a method with no docstring. + self._description = (description or "").strip() or tool_name.replace("_", " ").capitalize() + self._timeout = timeout + + @property + @abstractmethod + def agent_ref(self) -> dict[str, str]: + """ + Normalised identity of the remote agent. + + Must contain ``platform`` and ``name``, e.g. + ``{"platform": "snowflake.cortex", "name": "ANALYTICS.REVENUE.BOOKINGS_ANALYST"}``. + Logged on every invocation, so the resolved remote identity behind a task + appears in that task's log even though the Dag only names a connection. + It is not pushed to XCom: ``FailoverManagedAgentToolset`` reports the + group rather than the responder, and the operational question -- how often + a standby is answering -- is carried by the ``managed_agent.served`` + counter instead. + """ + + @abstractmethod + async def invoke(self, prompt: str) -> Any: Review Comment: we might need invoke_sync eg: agentcore they need sync clients with boto3 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
