pankajkoti commented on code in PR #36014:
URL: https://github.com/apache/airflow/pull/36014#discussion_r1415608376
##########
airflow/providers/openai/hooks/openai.py:
##########
@@ -41,37 +42,40 @@ class OpenAIHook(BaseHook):
def __init__(self, conn_id: str = default_conn_name, *args: Any, **kwargs:
Any) -> None:
super().__init__(*args, **kwargs)
self.conn_id = conn_id
- openai.api_key = self._get_api_key()
- api_base = self._get_api_base()
- if api_base:
- openai.api_base = api_base
- @staticmethod
- def get_ui_field_behaviour() -> dict[str, Any]:
+ @classmethod
+ def get_ui_field_behaviour(cls) -> dict[str, Any]:
"""Return custom field behaviour."""
return {
- "hidden_fields": ["schema", "port", "login", "extra"],
+ "hidden_fields": ["schema", "port", "login"],
"relabeling": {"password": "API Key"},
"placeholders": {},
}
def test_connection(self) -> tuple[bool, str]:
try:
- openai.Model.list()
+ self.conn.models.list()
return True, "Connection established!"
except Exception as e:
return False, str(e)
- def _get_api_key(self) -> str:
- """Get the OpenAI API key from the connection."""
- conn = self.get_connection(self.conn_id)
- if not conn.password:
- raise ValueError("OpenAI API key not found in connection")
- return str(conn.password)
+ @cached_property
+ def conn(self) -> OpenAI:
+ """Return an OpenAI connection object."""
+ return self.get_conn()
- def _get_api_base(self) -> None | str:
+ def get_conn(self) -> OpenAI:
+ """Return an OpenAI connection object."""
conn = self.get_connection(self.conn_id)
- return conn.host
+ extras = conn.extra_dejson
+ openai_client_kwargs = extras.get("openai_client_kwargs", {})
Review Comment:
I did consider the alternative of listing out the arguments directly. The
reason for using the flexible nested dict `openai_client_kwargs` is mainly
driven by the rapidly evolving nature of the OpenAI library. With each update,
the client may introduce changes such as renaming existing arguments, adding
new ones, or dropping some altogether. By encapsulating the arguments within a
nested dict, our hook becomes less prone to immediate disruptions caused by
such changes and avoids higher maintenance overhead as it offers adaptability
to changes in the OpenAI library without requiring immediate updates on our end.
Since fields under `extra` already go in a dict and are not direct fields in
the UI form, I don't think having another nested key `opeai_client_kwargs`
would be a lot.
It's just that instead of
extra=`{"arg1":2,"arg2": ["2","6"], "arg3": '{"a":5}'}`,
it goes one level deep
extra=`{"openai_client_kwargs":{"arg1":2,"arg2": ["2","6"], "arg3":
'{"a":5}'}}`
and offers low maintenance.
--
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]
