uranusjr commented on code in PR #75369:
URL: https://github.com/apache/airflow/pull/75369#discussion_r4463204726
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -562,6 +586,92 @@ def _provider_factory(pname: str) -> Any:
return infer_model(model_name)
+ def _build_missing_model_error(self) -> ValueError:
+ return ValueError(
+ f"No model specified for connection '{self.llm_conn_id}'. Set
model_id on the "
+ "hook or the Model field on the connection."
+ )
+
+ def _build_readiness_backend(
+ self, model_name: str, backend_factory: Callable[...,
ModelReadinessBackend] | None
+ ) -> ModelReadinessBackend:
+ """
+ Build the readiness backend for the provider-prefixed *model_name*.
+
+ *backend_factory* is called with the unprefixed ``model_name`` and the
connection's
+ ``base_url``; without it only Ollama has a built-in backend.
+ """
+ provider_prefix = model_name.partition(":")[0]
+ if backend_factory is None and provider_prefix != "ollama":
+ raise NotImplementedError(
+ f"ensure_model_ready has no built-in readiness backend for
'{model_name}' on "
+ f"connection '{self.llm_conn_id}'; only Ollama is built in.
Pass backend_factory= a "
+ "callable taking model_name= and base_url= that returns an
object with load(); "
+ "OllamaBackend is available for a model that Ollama serves."
+ )
+ conn, _ = self._get_conn_and_extra()
+ base_url = conn.host or None
+ unprefixed_name = _get_unprefixed_model_name(model_name)
Review Comment:
This is the second time `model_name` is split in this function (first on
line 604). They can be done in one shot.
--
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]
