From 5a7822539c24613b709a6882cc42fdf608df438f Mon Sep 17 00:00:00 2001 From: ChengZi Date: Thu, 16 Nov 2023 11:56:11 +0800 Subject: [PATCH] update openai version Signed-off-by: ChengZi --- openai_chat.py | 31 +++++++++++++++++++++++-------- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/openai_chat.py b/openai_chat.py index d4fd3d9..3d1e23b 100644 --- a/openai_chat.py +++ b/openai_chat.py @@ -14,7 +14,7 @@ import os from typing import List - +from packaging.version import Version import openai from towhee.operator.base import PyOperator @@ -33,13 +33,28 @@ class OpenaiChat(PyOperator): def __call__(self, messages: List[dict]): messages = self.parse_inputs(messages) - response = openai.ChatCompletion.create( - model=self._model, - messages=messages, - n=1, - stream=self.stream, - **self.kwargs - ) + if Version(openai.__version__) >= Version('1.0.0'): + from openai import OpenAI + client = OpenAI() + response = client.chat.completions.create( + model=self._model, + messages=messages, + n=1, + stream=self.stream, + **self.kwargs + ) + if self.stream: + response = (res.dict() for res in response) + else: + response = response.dict() + else: + response = openai.ChatCompletion.create( + model=self._model, + messages=messages, + n=1, + stream=self.stream, + **self.kwargs + ) if self.stream: return self.stream_output(response) else: