nathadfield commented on code in PR #38736:
URL: https://github.com/apache/airflow/pull/38736#discussion_r1555871738
##########
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:
```
n
integer or null
Optional
Defaults to 1
How many chat completion choices to generate for each input message. Note
that you will be charged based on the number of generated tokens across all of
the choices. Keep n as 1 to minimize costs.
```
--
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]