ephraimbuddy commented on code in PR #36014:
URL: https://github.com/apache/airflow/pull/36014#discussion_r1415586415
##########
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:
This means they have to do {"openai_client_kwargs":{"arg1":2,"arg2":
["2","6"], "arg3": '{"a":5}'} which I think is a lot. My preference is to have
them list out the args without openai_client_kwargs. That's more popular in the
hooks here
--
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]
