gopidesupavan commented on code in PR #67438:
URL: https://github.com/apache/airflow/pull/67438#discussion_r3365532224
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -172,41 +196,195 @@ def _provider_factory(pname: str) -> Any:
self._model = infer_model(model_name)
return self._model
- @overload
- def create_agent(
- self, output_type: type[OutputT], *, instructions: str, **agent_kwargs
- ) -> Agent[None, OutputT]: ...
+ # ------------------------------------------------------------------
+ # BaseAIHook abstract interface
+ # ------------------------------------------------------------------
- @overload
- def create_agent(self, *, instructions: str, **agent_kwargs) ->
Agent[None, str]: ...
+ def _tool_spec_to_native(self, spec: ToolSpec) -> Any:
+ """Convert a :class:`~airflow.providers.common.ai.hooks.base.ToolSpec`
to a pydantic-ai ``Tool``."""
+ return Tool.from_schema(
+ spec.fn,
+ name=spec.name,
+ description=spec.description,
+ sequential=spec.sequential,
+ json_schema=spec.parameters,
+ )
- def create_agent(
- self, output_type: type[Any] = str, *, instructions: str,
**agent_kwargs
- ) -> Agent[None, Any]:
+ def _build_agent(self, request: AgentRunRequest) -> PydanticAgentHandle:
"""
- Create a pydantic-ai Agent configured with this hook's model.
+ Build a pydantic-ai ``Agent`` handle from *request*.
- When ``[common.ai] otel_export_enabled`` is set and the worker has an
- OpenTelemetry exporter configured, the agent is instrumented to emit
- GenAI spans through Airflow's tracing pipeline. See
+ When :attr:`~AgentRunRequest.durable_context` is set, initialises
durable
+ storage and step counter and returns them alongside the native agent
for use
+ by :meth:`run_agent`. When ``[common.ai] otel_export_enabled`` is set
and the
+ worker has an OpenTelemetry exporter configured, the agent is
instrumented to
+ emit GenAI spans through Airflow's tracing pipeline. See
:mod:`airflow.providers.common.ai.observability`.
- :param output_type: The expected output type from the agent (default:
``str``).
- :param instructions: System-level instructions for the agent.
- :param agent_kwargs: Additional keyword arguments passed to the Agent
constructor.
+ :param request: Agent configuration — output type, instructions,
toolsets, extra params.
+
+ Native pydantic-ai ``Tool`` instances supplied in ``request.toolsets``
are passed through
+ unchanged. Airflow tool logging and durable tool-result caching are
applied to
+ framework-neutral callables / ``BaseToolset`` specs and pydantic-ai
``AbstractToolset``
+ instances, but not to native ``Tool`` instances.
"""
- agent = Agent(self.get_conn(), output_type=output_type,
instructions=instructions, **agent_kwargs)
- if "instrument" not in agent_kwargs:
- # Set the public ``agent.instrument`` surface rather than the
+ durable_state = None
+ if request.durable_context is not None:
+ durable_state = self._init_durable(request.durable_context)
+
+ extra_kwargs = dict(request.agent_params or {})
+ if request.toolsets:
+ if "tools" in extra_kwargs:
+ raise ValueError(
+ "agent_params must not include 'tools' when toolsets= is
set on AgentRunRequest."
+ )
+ if "toolsets" in extra_kwargs:
+ raise ValueError(
+ "agent_params must not include 'toolsets' when toolsets=
is set on AgentRunRequest."
+ )
+
+ abstract_items = [ts for ts in request.toolsets if isinstance(ts,
AbstractToolset)]
+ pipeline_items = [ts for ts in request.toolsets if not
isinstance(ts, AbstractToolset)]
Review Comment:
Thanks kaxil Good catch. Bare ToolsetFunc pass through is intentionally not
preserved because AgentRunRequest.toolsets now treats plain callables as
framework-neutral function tools, so a bare callable is ambiguous. The
pydantic-ai-native path still works by wrapping the factory with
DynamicToolset(...), which is an AbstractToolset and already routes through
Agent(toolsets=...). I have add a test/docs note for that explicit path.
--
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]
