cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2895825191
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -19,42 +19,41 @@
from typing import TYPE_CHECKING, Any, TypeVar, overload
from pydantic_ai import Agent
-from pydantic_ai.models import Model, infer_model
-from pydantic_ai.providers import Provider, infer_provider,
infer_provider_class
+from pydantic_ai.models import infer_model
+from pydantic_ai.providers import infer_provider, infer_provider_class
from airflow.providers.common.compat.sdk import BaseHook
OutputT = TypeVar("OutputT")
if TYPE_CHECKING:
- from pydantic_ai.models import KnownModelName
+ from pydantic_ai.models import KnownModelName, Model
class PydanticAIHook(BaseHook):
"""
Hook for LLM access via pydantic-ai.
- Manages connection credentials and model creation. Uses pydantic-ai's
- model inference to support any provider (OpenAI, Anthropic, Google,
- Bedrock, Ollama, vLLM, etc.).
+ Covers providers that use a standard ``api_key`` + optional ``base_url``
+ (OpenAI, Anthropic, Groq, Mistral, DeepSeek, Ollama, vLLM, …).
- Connection fields:
- - **Model** (conn-field): Model in ``provider:model`` format (e.g.
``"anthropic:claude-sonnet-4-20250514"``)
- - **password**: API key (OpenAI, Anthropic, Groq, Mistral, etc.)
- - **host**: Base URL (optional — for custom endpoints like Ollama,
vLLM, Azure)
+ For cloud providers with non-standard auth use the dedicated subclasses:
+ :class:`PydanticAIAzureHook`, :class:`PydanticAIBedrockHook`,
+ :class:`PydanticAIVertexHook`.
- Cloud providers (Bedrock, Vertex) that use native auth chains should leave
- password empty and configure environment-based auth (``AWS_PROFILE``,
- ``GOOGLE_APPLICATION_CREDENTIALS``).
+ Connection fields:
+ - **password**: API key
+ - **host**: Base URL (optional, e.g. ``https://api.openai.com/v1``)
+ - **extra** JSON: ``{"model": "openai:gpt-5.3"}``
:param llm_conn_id: Airflow connection ID for the LLM provider.
:param model_id: Model identifier in ``provider:model`` format (e.g.
``"openai:gpt-5.3"``).
Overrides the model stored in the connection's extra field.
"""
conn_name_attr = "llm_conn_id"
- default_conn_name = "pydanticai_default"
- conn_type = "pydanticai"
+ default_conn_name = "pydantic_ai_default"
+ conn_type = "pydantic_ai"
Review Comment:
It mixed my previous commits which before your #62817 pr(I think I didn't
merge correctly). I changed other hooks but I forgot this :) I'm rolling back
right now ;)
--
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]
