dabla commented on code in PR #74150:
URL: https://github.com/apache/airflow/pull/74150#discussion_r4293020378
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -669,6 +705,57 @@ def create_agent(
agent.instrument = settings
return agent
+ @overload
+ async def acreate_agent(
+ self, output_type: type[OutputT], *, instructions: str, **agent_kwargs
+ ) -> Agent[object, OutputT]: ...
+
+ @overload
+ async def acreate_agent(self, *, instructions: str, **agent_kwargs) ->
Agent[object, str]: ...
+
+ @overload
+ async def acreate_agent(
+ self,
+ output_type: type[OutputT],
+ *,
+ spec_file: str | Path,
+ instructions: str | None = ...,
+ **agent_kwargs,
+ ) -> Agent[object, OutputT]: ...
+
+ @overload
+ async def acreate_agent(
+ self,
+ *,
+ spec_file: str | Path,
+ instructions: str | None = ...,
+ **agent_kwargs,
+ ) -> Agent[object, str]: ...
+
+ async def acreate_agent(
Review Comment:
Added a section on using the hook from an async task.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -318,6 +326,31 @@ def get_conn(self) -> Model:
)
return self._model
+ async def aget_conn(self) -> Model:
Review Comment:
Added `aget_hook` with `get_hook`'s type check. `get_hook` now also caches
the connection on the hook, so it's fetched once either way.
--
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]
