:gem: [Feature] Support no-stream mode with dict response

apis/chat_api.py

@@ -64,19 +64,24 @@ class ChatAPIApp:
         streamer = MessageStreamer(model=item.model)
         composer = MessageComposer(model=item.model)
         composer.merge(messages=item.messages)
-        event_source_response = EventSourceResponse(
-            streamer.chat(
-                prompt=composer.merged_str,
-                temperature=item.temperature,
-                max_new_tokens=item.max_tokens,
-                stream=item.stream,
-                yield_output=True,
-            ),
-            media_type="text/event-stream",
-            ping=2000,
-            ping_message_factory=lambda: ServerSentEvent(**{"comment": ""}),
+        # streamer.chat = stream_chat_mock
+
+        stream_response = streamer.chat_response(
+            prompt=composer.merged_str,
+            temperature=item.temperature,
+            max_new_tokens=item.max_tokens,
         )
-        return event_source_response
+        if item.stream:
+            event_source_response = EventSourceResponse(
+                streamer.chat_return_generator(stream_response),
+                media_type="text/event-stream",
+                ping=2000,
+                ping_message_factory=lambda: ServerSentEvent(**{"comment": ""}),
+            )
+            return event_source_response
+        else:
+            data_response = streamer.chat_return_dict(stream_response)
+            return data_response
 
     def setup_routes(self):
         for prefix in ["", "/v1"]:

messagers/message_outputer.py

@@ -7,20 +7,22 @@ class OpenaiStreamOutputer:
     * https://platform.openai.com/docs/api-reference/chat/create
     """
 
-    def data_to_string(self, data={}, content_type=""):
-        data_str = f"{json.dumps(data)}"
-
-        return data_str
-
-    def output(self, content=None, content_type="Completions") -> str:
-        data = {
+    def __init__(self):
+        self.default_data = {
             "created": 1700000000,
             "id": "chatcmpl-hugginface",
             "object": "chat.completion.chunk",
-            # "content_type": content_type,
+            # "content_type": "Completions",
             "model": "hugginface",
             "choices": [],
         }
+
+    def data_to_string(self, data={}, content_type=""):
+        data_str = f"{json.dumps(data)}"
+        return data_str
+
+    def output(self, content=None, content_type="Completions") -> str:
+        data = self.default_data.copy()
         if content_type == "Role":
             data["choices"] = [
                 {

mocks/stream_chat_mocker.py

@@ -2,10 +2,11 @@ import time
 from utils.logger import logger
 
 
-def stream_chat_mock():
-    for i in range(8):
+def stream_chat_mock(*args, **kwargs):
+    logger.note(msg=str(args) + str(kwargs))
+    for i in range(10):
         content = f"W{i+1} "
-        time.sleep(1.5)
+        time.sleep(0.1)
         logger.mesg(content, end="")
         yield content
     logger.mesg("")

networks/message_streamer.py

@@ -31,13 +31,11 @@ class MessageStreamer:
         content = data["token"]["text"]
         return content
 
-    def chat(
+    def chat_response(
         self,
         prompt: str = None,
         temperature: float = 0.01,
         max_new_tokens: int = 8192,
-        stream: bool = True,
-        yield_output: bool = False,
     ):
         # https://huggingface.co/docs/api-inference/detailed_parameters?code=curl
         # curl --proxy http://<server>:<port> https://api-inference.huggingface.co/models/<org>/<model_name> -X POST -d '{"inputs":"who are you?","parameters":{"max_new_token":64}}' -H 'Content-Type: application/json' -H 'Authorization: Bearer <HF_TOKEN>'
@@ -60,24 +58,57 @@ class MessageStreamer:
                 "max_new_tokens": max_new_tokens,
                 "return_full_text": False,
             },
-            "stream": stream,
+            "stream": True,
         }
         logger.back(self.request_url)
         enver.set_envs(proxies=True)
-        stream = requests.post(
+        stream_response = requests.post(
             self.request_url,
             headers=self.request_headers,
             json=self.request_body,
             proxies=enver.requests_proxies,
-            stream=stream,
+            stream=True,
         )
-        status_code = stream.status_code
+        status_code = stream_response.status_code
         if status_code == 200:
             logger.success(status_code)
         else:
             logger.err(status_code)
 
-        for line in stream.iter_lines():
+        return stream_response
+
+    def chat_return_dict(self, stream_response):
+        # https://platform.openai.com/docs/guides/text-generation/chat-completions-response-format
+        final_output = self.message_outputer.default_data.copy()
+        final_output["choices"] = [
+            {
+                "index": 0,
+                "finish_reason": "stop",
+                "message": {
+                    "role": "assistant",
+                    "content": "",
+                },
+            }
+        ]
+        logger.back(final_output)
+
+        for line in stream_response.iter_lines():
+            if not line:
+                continue
+            content = self.parse_line(line)
+
+            if content.strip() == "</s>":
+                logger.success("\n[Finished]")
+                break
+            else:
+                logger.back(content, end="")
+                final_output["choices"][0]["message"]["content"] += content
+
+        return final_output
+
+    def chat_return_generator(self, stream_response):
+        is_finished = False
+        for line in stream_response.iter_lines():
             if not line:
                 continue
 
@@ -86,12 +117,15 @@ class MessageStreamer:
             if content.strip() == "</s>":
                 content_type = "Finished"
                 logger.success("\n[Finished]")
+                is_finished = True
             else:
                 content_type = "Completions"
                 logger.back(content, end="")
 
-            if yield_output:
-                output = self.message_outputer.output(
-                    content=content, content_type=content_type
-                )
-                yield output
+            output = self.message_outputer.output(
+                content=content, content_type=content_type
+            )
+            yield output
+
+        if not is_finished:
+            yield self.message_outputer.output(content="", content_type="Finished")