jeff3071 commented on code in PR #71437:
URL: https://github.com/apache/airflow/pull/71437#discussion_r3965649544
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -160,30 +181,57 @@ def get_conn(self) -> Model:
provider_kwargs = self._get_provider_kwargs(api_key, base_url, extra)
if provider_kwargs:
- _kwargs = provider_kwargs # capture for closure
self.log.info(
"Using explicit credentials for provider with model '%s': %s",
model_name,
list(provider_kwargs),
)
-
- def _provider_factory(pname: str) -> Any:
- try:
- return infer_provider_class(pname)(**_kwargs)
- except TypeError:
- self.log.warning(
- "Provider '%s' rejected kwargs %s; falling back to
env-var auth",
- pname,
- list(_kwargs),
- )
- return infer_provider(pname)
-
- self._model = infer_model(model_name,
provider_factory=_provider_factory)
+ self._model = infer_model(
+ model_name,
+
provider_factory=self._create_provider_factory(provider_kwargs),
+ )
return self._model
self._model = infer_model(model_name)
return self._model
+ def get_embedder(self) -> Embedder:
Review Comment:
Thanks for the feedback.
`test_connection()` now validates both the model and embedding model. I also
added two unit tests covering these cases.
I’d be happy to address the same issue in `LangChainHook` and
`LlamaIndexHook` in a separate PR.
--
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]
