nathadfield commented on code in PR #38736:
URL: https://github.com/apache/airflow/pull/38736#discussion_r1557326945
##########
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
Review Comment:
I suppose not. My initial thought was that I just wanted to return the
message object but, if there can be more than one, then returning the list and
then handling this outside probably seems like the correct thing to do.
--
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]
