gopidesupavan commented on code in PR #67788:
URL: https://github.com/apache/airflow/pull/67788#discussion_r3329423587
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -179,16 +180,54 @@ def create_agent(
@overload
def create_agent(self, *, instructions: str, **agent_kwargs) ->
Agent[None, str]: ...
+ @overload
+ def create_agent(
+ self,
+ output_type: type[OutputT],
+ *,
+ spec_file: str | Path,
+ instructions: str | None = ...,
+ **agent_kwargs,
+ ) -> Agent[None, OutputT]: ...
+
+ @overload
def create_agent(
- self, output_type: type[Any] = str, *, instructions: str,
**agent_kwargs
+ self,
+ *,
+ spec_file: str | Path,
+ instructions: str | None = ...,
+ **agent_kwargs,
+ ) -> Agent[None, str]: ...
+
+ def create_agent(
+ self,
+ output_type: type[Any] = str,
+ *,
+ instructions: str | None = None,
+ spec_file: str | Path | None = None,
+ **agent_kwargs,
) -> Agent[None, Any]:
"""
Create a pydantic-ai Agent configured with this hook's model.
:param output_type: The expected output type from the agent (default:
``str``).
:param instructions: System-level instructions for the agent.
+ Required when *spec_file* is not given. When *spec_file* is
given, this
+ value overrides the instructions in the file; omit to use the file
value.
+ :param spec_file: Path to a YAML or JSON ``AgentSpec`` file. When
supplied,
+ delegates to ``Agent.from_file(spec_file, model=..., ...)``.
:param agent_kwargs: Additional keyword arguments passed to the Agent
constructor.
"""
+ if spec_file is not None:
+ return Agent.from_file(
+ spec_file,
+ model=self.get_conn(),
+ output_type=output_type,
+ instructions=instructions,
+ **agent_kwargs,
+ )
+ if instructions is None:
+ raise ValueError("instructions is required when spec_file is not
provided.")
Review Comment:
Yeah updating..
--
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]
