Lee-W commented on code in PR #38736:
URL: https://github.com/apache/airflow/pull/38736#discussion_r1556009587
##########
airflow/providers/openai/hooks/openai.py:
##########
@@ -77,6 +89,165 @@ def get_conn(self) -> OpenAI:
**openai_client_kwargs,
)
+ def create_chat_completion(
+ self,
+ messages: list[
+ ChatCompletionSystemMessageParam
+ | ChatCompletionUserMessageParam
+ | ChatCompletionAssistantMessageParam
+ | ChatCompletionToolMessageParam
+ | ChatCompletionFunctionMessageParam
+ ],
+ model: str = "gpt-3.5-turbo",
+ **kwargs: Any,
+ ) -> ChatCompletionMessage:
+ """
+ Create a model response for the given chat conversation and returns
the message.
+
+ :param messages: A list of messages comprising the conversation so far
+ :param model: ID of the model to use
+ """
+ response = self.conn.chat.completions.create(model=model,
messages=messages, **kwargs)
+ return response.choices[0].message
+
+ def create_assistant(self, model: str = "gpt-3.5-turbo", **kwargs: Any) ->
Assistant:
+ """Create an OpenAI assistant using the given model.
+
+ :param model: The OpenAI model for the assistant to use.
+ """
+ assistant = self.conn.beta.assistants.create(model=model, **kwargs)
+ return assistant
+
+ def get_assistant(self, assistant_id: str) -> Assistant:
+ """
+ Get an OpenAI assistant.
+
+ :param assistant_id: The ID of the assistant to retrieve.
+ """
+ assistant =
self.conn.beta.assistants.retrieve(assistant_id=assistant_id)
+ return assistant
+
+ def get_assistants(self, **kwargs) -> list[Assistant]:
+ """Get a list of Assistant objects."""
+ assistants = self.conn.beta.assistants.list(**kwargs)
+ return assistants.data
+
+ def get_assistant_with_name(self, assistant_name: str) -> Assistant | None:
Review Comment:
```suggestion
def get_assistant_by_name(self, assistant_name: str) -> Assistant | None:
```
Not sure whether it sounds more reasonable 🤔
##########
airflow/providers/openai/hooks/openai.py:
##########
@@ -77,6 +89,165 @@ def get_conn(self) -> OpenAI:
**openai_client_kwargs,
)
+ def create_chat_completion(
+ self,
+ messages: list[
+ ChatCompletionSystemMessageParam
+ | ChatCompletionUserMessageParam
+ | ChatCompletionAssistantMessageParam
+ | ChatCompletionToolMessageParam
+ | ChatCompletionFunctionMessageParam
+ ],
+ model: str = "gpt-3.5-turbo",
+ **kwargs: Any,
+ ) -> ChatCompletionMessage:
+ """
+ Create a model response for the given chat conversation and returns
the message.
+
+ :param messages: A list of messages comprising the conversation so far
+ :param model: ID of the model to use
+ """
+ response = self.conn.chat.completions.create(model=model,
messages=messages, **kwargs)
+ return response.choices[0].message
+
+ def create_assistant(self, model: str = "gpt-3.5-turbo", **kwargs: Any) ->
Assistant:
+ """Create an OpenAI assistant using the given model.
+
+ :param model: The OpenAI model for the assistant to use.
+ """
+ assistant = self.conn.beta.assistants.create(model=model, **kwargs)
+ return assistant
+
+ def get_assistant(self, assistant_id: str) -> Assistant:
+ """
+ Get an OpenAI assistant.
+
+ :param assistant_id: The ID of the assistant to retrieve.
+ """
+ assistant =
self.conn.beta.assistants.retrieve(assistant_id=assistant_id)
+ return assistant
+
+ def get_assistants(self, **kwargs) -> list[Assistant]:
Review Comment:
```suggestion
def get_assistants(self, **kwargs: Any) -> list[Assistant]:
```
--
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: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
