From 2408eea7d61258645ba54e7bc32b0ef291529b0b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=BC=E6=AC=A3?= Date: Tue, 14 Nov 2023 16:51:05 +0800 Subject: [PATCH] openai_api.py: compatible with both pydantic v1 and v2 --- openai_api.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/openai_api.py b/openai_api.py index 34714e3..fafd8f0 100644 --- a/openai_api.py +++ b/openai_api.py @@ -416,6 +416,13 @@ async def create_chat_completion(request: ChatCompletionRequest): ) +def _dump_json(data: BaseModel, *args, **kwargs) -> str: + try: + return data.model_dump_json(*args, **kwargs) + except AttributeError: # pydantic<2.0.0 + return data.json(*args, **kwargs) # noqa + + async def predict( query: str, history: List[List[str]], model_id: str, stop_words: List[str], gen_kwargs: Dict, ): @@ -426,7 +433,7 @@ async def predict( chunk = ChatCompletionResponse( model=model_id, choices=[choice_data], object="chat.completion.chunk" ) - yield "{}".format(chunk.model_dump_json(exclude_unset=True)) + yield "{}".format(_dump_json(chunk, exclude_unset=True)) current_length = 0 stop_words_ids = [tokenizer.encode(s) for s in stop_words] if stop_words else None @@ -452,7 +459,7 @@ async def predict( chunk = ChatCompletionResponse( model=model_id, choices=[choice_data], object="chat.completion.chunk" ) - yield "{}".format(chunk.model_dump_json(exclude_unset=True)) + yield "{}".format(_dump_json(chunk, exclude_unset=True)) choice_data = ChatCompletionResponseStreamChoice( index=0, delta=DeltaMessage(), finish_reason="stop" @@ -460,7 +467,7 @@ async def predict( chunk = ChatCompletionResponse( model=model_id, choices=[choice_data], object="chat.completion.chunk" ) - yield "{}".format(chunk.model_dump_json(exclude_unset=True)) + yield "{}".format(_dump_json(chunk, exclude_unset=True)) yield "[DONE]" _gc()