jeff3071 commented on code in PR #71437:
URL: https://github.com/apache/airflow/pull/71437#discussion_r4062648735
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -220,8 +257,92 @@ def _seed_connection(self, conn: Connection) -> None:
would fetch that same connection a second time the first time it runs,
doubling the
Execution API round trips a fallback chain costs.
"""
- self._conn = conn
- self._conn_extra_dejson = conn.extra_dejson
+ self._connections[conn.conn_id] = conn
+ self._connection_extra_dejson[conn.conn_id] = conn.extra_dejson
+
+ def _get_provider_kwargs_for_model(
+ self, conn: Connection, model_name: str, extra: dict[str, Any]
+ ) -> dict[str, Any]:
+ provider_name, _ = parse_model_id(model_name)
+ provider_config = _PROVIDER_CONNECTION_CONFIGS.get(provider_name)
+ if extra.get("vertexai") is not None:
+ self.log.warning(
+ "The 'vertexai' connection field is ignored; Vertex AI vs.
Generative Language "
+ "API mode is now selected via the model prefix
('google-cloud:' vs. 'google:')."
+ )
+ if provider_config is None:
+ return PydanticAIHook._get_provider_kwargs(conn.password,
conn.host, extra)
+ if provider_config.replacement_fields:
+ ignored_fields = [
+ field for field, value in (("password", conn.password),
("host", conn.host)) if value
+ ]
+ if ignored_fields:
+ self.log.warning(
+ "Connection fields are ignored for provider %r on
connection %r; "
+ "ignored fields: %s; configure these provider-specific
values in extra: %s",
+ provider_name,
+ conn.conn_id,
+ ignored_fields,
+ list(provider_config.replacement_fields),
+ )
+ ignored_extra_fields = [field for field in
provider_config.ignored_extra_fields if extra.get(field)]
+ if ignored_extra_fields:
+ self.log.warning(
+ "Connection extra fields are ignored for provider %r on
connection %r: %s",
+ provider_name,
+ conn.conn_id,
+ ignored_extra_fields,
+ )
+ return provider_config.get_kwargs(conn.password, conn.host, extra)
+
+ def _get_provider_factory_for_model(
+ self, conn: Connection, model_name: str, extra: dict[str, Any]
+ ) -> Callable[[str], Any] | None:
+ provider_name, _ = parse_model_id(model_name)
+ if provider_name == "sentence-transformers":
+ return None
+
+ provider_kwargs = self._get_provider_kwargs_for_model(conn,
model_name, extra)
+ if not provider_kwargs:
+ return None
+
+ self.log.info(
+ "Using explicit connection credentials for model '%s': %s",
+ model_name,
+ list(provider_kwargs),
+ )
+
+ def create_provider(provider: str) -> Any:
+ try:
+ return infer_provider_class(provider)(**provider_kwargs)
+ except TypeError as e:
+ raise TypeError(
+ f"Provider {provider!r} rejected connection
{conn.conn_id!r} fields "
+ f"mapped to kwargs {sorted(provider_kwargs)}: {e}"
+ ) from e
+
+ return create_provider
+
+ def _validate_embedding_connection_provider(self, embed_model_name: str,
extra: dict[str, Any]) -> None:
+ if self.embed_conn_id != self.llm_conn_id:
+ return
+
+ llm_model_name = self.model_id or extra.get("model", "")
+ if not llm_model_name:
+ return
+
+ llm_provider, _ = parse_model_id(llm_model_name)
Review Comment:
I think reusing _qualify_model_name() here is better. A generic pydanticai
connection still produces `llm_provider = None` for a bare model. So it still
pass the guard.
`_qualify_model_name()` already defines the authoritative behavior for
relevant cases:
- recognized `provider:model` names are preserved
- bare native model IDs on vendor connections are qualified with
`model_provider`
- bare model names on a generic connection raise the same actionable
`ValueError` as the normal path
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -220,8 +257,92 @@ def _seed_connection(self, conn: Connection) -> None:
would fetch that same connection a second time the first time it runs,
doubling the
Execution API round trips a fallback chain costs.
"""
- self._conn = conn
- self._conn_extra_dejson = conn.extra_dejson
+ self._connections[conn.conn_id] = conn
+ self._connection_extra_dejson[conn.conn_id] = conn.extra_dejson
+
+ def _get_provider_kwargs_for_model(
+ self, conn: Connection, model_name: str, extra: dict[str, Any]
+ ) -> dict[str, Any]:
+ provider_name, _ = parse_model_id(model_name)
+ provider_config = _PROVIDER_CONNECTION_CONFIGS.get(provider_name)
+ if extra.get("vertexai") is not None:
+ self.log.warning(
+ "The 'vertexai' connection field is ignored; Vertex AI vs.
Generative Language "
+ "API mode is now selected via the model prefix
('google-cloud:' vs. 'google:')."
+ )
+ if provider_config is None:
+ return PydanticAIHook._get_provider_kwargs(conn.password,
conn.host, extra)
+ if provider_config.replacement_fields:
+ ignored_fields = [
+ field for field, value in (("password", conn.password),
("host", conn.host)) if value
+ ]
+ if ignored_fields:
+ self.log.warning(
+ "Connection fields are ignored for provider %r on
connection %r; "
+ "ignored fields: %s; configure these provider-specific
values in extra: %s",
+ provider_name,
+ conn.conn_id,
+ ignored_fields,
+ list(provider_config.replacement_fields),
+ )
+ ignored_extra_fields = [field for field in
provider_config.ignored_extra_fields if extra.get(field)]
+ if ignored_extra_fields:
+ self.log.warning(
+ "Connection extra fields are ignored for provider %r on
connection %r: %s",
+ provider_name,
+ conn.conn_id,
+ ignored_extra_fields,
+ )
+ return provider_config.get_kwargs(conn.password, conn.host, extra)
+
+ def _get_provider_factory_for_model(
+ self, conn: Connection, model_name: str, extra: dict[str, Any]
+ ) -> Callable[[str], Any] | None:
+ provider_name, _ = parse_model_id(model_name)
+ if provider_name == "sentence-transformers":
+ return None
+
+ provider_kwargs = self._get_provider_kwargs_for_model(conn,
model_name, extra)
+ if not provider_kwargs:
+ return None
+
+ self.log.info(
+ "Using explicit connection credentials for model '%s': %s",
+ model_name,
+ list(provider_kwargs),
+ )
+
+ def create_provider(provider: str) -> Any:
+ try:
+ return infer_provider_class(provider)(**provider_kwargs)
+ except TypeError as e:
+ raise TypeError(
+ f"Provider {provider!r} rejected connection
{conn.conn_id!r} fields "
+ f"mapped to kwargs {sorted(provider_kwargs)}: {e}"
+ ) from e
+
+ return create_provider
+
+ def _validate_embedding_connection_provider(self, embed_model_name: str,
extra: dict[str, Any]) -> None:
+ if self.embed_conn_id != self.llm_conn_id:
+ return
+
+ llm_model_name = self.model_id or extra.get("model", "")
+ if not llm_model_name:
+ return
+
+ llm_provider, _ = parse_model_id(llm_model_name)
Review Comment:
I think reusing `_qualify_model_name()` here is better. A generic pydanticai
connection still produces `llm_provider = None` for a bare model. So it still
pass the guard.
`_qualify_model_name()` already defines the authoritative behavior for
relevant cases:
- recognized `provider:model` names are preserved
- bare native model IDs on vendor connections are qualified with
`model_provider`
- bare model names on a generic connection raise the same actionable
`ValueError` as the normal 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]
