support non stream for chat

smallcloudai · May 24, 2024 · 9fe3a72 · 9fe3a72
1 parent 90035f7
commit 9fe3a72
Showing 1 changed file with 27 additions and 1 deletion.
diff --git a/refact_webgui/webgui/selfhost_fastapi_completions.py b/refact_webgui/webgui/selfhost_fastapi_completions.py
@@ -79,6 +79,8 @@ class ChatContext(NlpSamplingParams):
     messages: List[ChatMessage]
     tools: Optional[List[Dict[str, Any]]] = None
     tool_choice: Optional[str] = None
+    stream: bool = True
+    stop: Optional[Any]
     n: int = 1
 
 
@@ -516,7 +518,31 @@ async def litellm_streamer(post: ChatContext):
                     log(err_msg)
                     yield prefix + json.dumps({"error": err_msg}) + postfix
 
-            response_streamer = litellm_streamer(post)
+            async def litellm_non_streamer(post: ChatContext):
+                try:
+                    self._integrations_env_setup()
+                    model_response = await litellm.acompletion(
+                        model=model_name, messages=[m.dict() for m in post.messages], stream=False,
+                        temperature=post.temperature, top_p=post.top_p,
+                        max_tokens=min(model_dict.get('T_out', post.max_tokens), post.max_tokens),
+                        tools=post.tools,
+                        tool_choice=post.tool_choice,
+                        stop=post.stop
+                    )
+                    finish_reason = None
+                    try:
+                        data = model_response.dict()
+                        print(data)
+                        finish_reason = data["choices"][0]["finish_reason"]
+                    except json.JSONDecodeError:
+                        data = {"choices": [{"finish_reason": finish_reason}]}
+                    yield json.dumps(data)
+                except BaseException as e:
+                    err_msg = f"litellm error: {e}"
+                    log(err_msg)
+                    yield prefix + json.dumps({"error": err_msg}) + postfix
+
+            response_streamer = litellm_streamer(post) if post.stream else litellm_non_streamer(post)
         else:
             async def chat_completion_streamer(post: ChatContext):
                 post_url = "http://127.0.0.1:8001/v1/chat"