From 375af3bc1a695bfe8f5a77a9c5ea655f90c77293 Mon Sep 17 00:00:00 2001 From: josc146 Date: Mon, 22 May 2023 11:24:57 +0800 Subject: [PATCH] improve compatible API --- backend-python/routes/completion.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/backend-python/routes/completion.py b/backend-python/routes/completion.py index 3b5af55..7d7e132 100644 --- a/backend-python/routes/completion.py +++ b/backend-python/routes/completion.py @@ -21,6 +21,7 @@ class ChatCompletionBody(ModelConfigBody): messages: List[Message] model: str = "rwkv" stream: bool = False + stop: str = None completion_lock = Lock() @@ -70,7 +71,9 @@ async def chat_completions(body: ChatCompletionBody, request: Request): set_rwkv_config(model, body) if body.stream: for response, delta in rwkv_generate( - model, completion_text, stop="\n\nBob" + model, + completion_text, + stop="\n\nBob" if body.stop is None else body.stop, ): if await request.is_disconnected(): break @@ -107,7 +110,9 @@ async def chat_completions(body: ChatCompletionBody, request: Request): else: response = None for response, delta in rwkv_generate( - model, completion_text, stop="\n\nBob" + model, + completion_text, + stop="\n\nBob" if body.stop is None else body.stop, ): if await request.is_disconnected(): break