cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2901765798
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -101,8 +101,15 @@ def __init__(
@cached_property
def llm_hook(self) -> PydanticAIHook:
- """Return PydanticAIHook for the configured LLM connection."""
- return PydanticAIHook(llm_conn_id=self.llm_conn_id,
model_id=self.model_id)
+ """
+ Return the correct PydanticAIHook subclass for the configured
connection.
+
+ Delegates to :meth:`~PydanticAIHook.get_hook` which looks up
+ the connection's ``conn_type`` and instantiates the matching subclass
+ (e.g.
:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook`
+ for ``pydanticaiazure`` connections).
+ """
+ return PydanticAIHook.get_hook(self.llm_conn_id)
def execute(self, context: Context) -> Any:
agent: Agent[None, Any] = self.llm_hook.create_agent(
Review Comment:
I thought it would be better to pass this information to the class using
`hook_params`, and I didn't want to add it later because any exceptions that
might arise could be reflected within the class itself.
```
hook_params = {
"model_id": self.model_id,
}
return PydanticAIHook.get_hook(self.llm_conn_id, hook_params=hook_params)
```
--
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]
