nathadfield commented on code in PR #38736:
URL: https://github.com/apache/airflow/pull/38736#discussion_r1555881195
##########
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 could set `n=1` as a default argument which would return
`response.choices[0].message` of type `Message` but if this was overwritten it
could return `response.choices` which would be `list[Message]` but I'm not sure
if that would be a good practice or not.
--
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]
