uranusjr commented on code in PR #75369:
URL: https://github.com/apache/airflow/pull/75369#discussion_r4463185700
##########
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)
+ if backend_factory is not None:
+ return backend_factory(model_name=unprefixed_name,
base_url=base_url)
+ # pydantic-ai's OllamaProvider falls back to OLLAMA_BASE_URL when no
base_url is passed.
+ if not (ollama_url := base_url or os.getenv("OLLAMA_BASE_URL")):
+ raise ValueError(
+ f"Connection '{self.llm_conn_id}' needs a host, or the
OLLAMA_BASE_URL environment "
+ "variable, to load the model."
+ )
+ return OllamaBackend(ollama_url, unprefixed_name)
+
+ def ensure_model_ready(
+ self,
+ *,
+ timeout: float = 300.0,
+ backend_factory: Callable[..., ModelReadinessBackend] | None = None,
+ ) -> None:
+ """
+ Block until this connection's local model is loaded, asking the server
to load it if needed.
+
+ Sends a load request and waits for it to return; a model that is
already loaded answers
+ right away. Any error from the server is raised as is.
+
+ Ollama is built in; for any other server pass *backend_factory*. Only
this hook's own
+ connection is loaded -- connections in ``fallback_conn_ids`` are not,
so a fallback model
+ is never loaded ahead of need.
+
+ The built-in Ollama backend does not send the connection password, so
a server behind
+ authentication is not supported.
+
+ Blocks and holds the worker slot for the whole wait.
Review Comment:
Or the other way around, the documentation can reference the docstring.
--
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]
