💎 [Feature] Enable gpt-3.5 in chat_api

apis/chat_api.py

@@ -12,21 +12,24 @@ from fastapi.responses import HTMLResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel, Field
 from sse_starlette.sse import EventSourceResponse, ServerSentEvent
+from tclogger import logger
+
+from constants.models import AVAILABLE_MODELS_DICTS
+from constants.envs import CONFIG
 
 from messagers.message_composer import MessageComposer
 from mocks.stream_chat_mocker import stream_chat_mock
-from networks.message_streamer import MessageStreamer
-from utils.logger import logger
-from constants.models import AVAILABLE_MODELS_DICTS
+from networks.huggingface_streamer import HuggingfaceStreamer
+from networks.openai_streamer import OpenaiStreamer
 
 
 class ChatAPIApp:
     def __init__(self):
         self.app = FastAPI(
             docs_url="/",
-            title="HuggingFace LLM API",
+            title=CONFIG["app_name"],
             swagger_ui_parameters={"defaultModelsExpandDepth": -1},
-            version="1.0",
+            version=CONFIG["version"],
         )
         self.setup_routes()
 
@@ -86,19 +89,22 @@ class ChatAPIApp:
     def chat_completions(
         self, item: ChatCompletionsPostItem, api_key: str = Depends(extract_api_key)
     ):
-        streamer = MessageStreamer(model=item.model)
-        composer = MessageComposer(model=item.model)
-        composer.merge(messages=item.messages)
-        # streamer.chat = stream_chat_mock
-
-        stream_response = streamer.chat_response(
-            prompt=composer.merged_str,
-            temperature=item.temperature,
-            top_p=item.top_p,
-            max_new_tokens=item.max_tokens,
-            api_key=api_key,
-            use_cache=item.use_cache,
-        )
+        if item.model == "gpt-3.5":
+            streamer = OpenaiStreamer()
+            stream_response = streamer.chat_response(messages=item.messages)
+        else:
+            streamer = HuggingfaceStreamer(model=item.model)
+            composer = MessageComposer(model=item.model)
+            composer.merge(messages=item.messages)
+            stream_response = streamer.chat_response(
+                prompt=composer.merged_str,
+                temperature=item.temperature,
+                top_p=item.top_p,
+                max_new_tokens=item.max_tokens,
+                api_key=api_key,
+                use_cache=item.use_cache,
+            )
+
         if item.stream:
             event_source_response = EventSourceResponse(
                 streamer.chat_return_generator(stream_response),
@@ -152,17 +158,17 @@ class ArgParser(argparse.ArgumentParser):
 
         self.add_argument(
             "-s",
-            "--server",
+            "--host",
             type=str,
-            default="0.0.0.0",
-            help="Server IP for HF LLM Chat API",
+            default=CONFIG["host"],
+            help=f"Host for {CONFIG['app_name']}",
         )
         self.add_argument(
             "-p",
             "--port",
             type=int,
-            default=23333,
-            help="Server Port for HF LLM Chat API",
+            default=CONFIG["port"],
+            help=f"Port for {CONFIG['app_name']}",
         )
 
         self.add_argument(
@@ -181,9 +187,9 @@ app = ChatAPIApp().app
 if __name__ == "__main__":
     args = ArgParser().args
     if args.dev:
-        uvicorn.run("__main__:app", host=args.server, port=args.port, reload=True)
+        uvicorn.run("__main__:app", host=args.host, port=args.port, reload=True)
     else:
-        uvicorn.run("__main__:app", host=args.server, port=args.port, reload=False)
+        uvicorn.run("__main__:app", host=args.host, port=args.port, reload=False)
 
     # python -m apis.chat_api      # [Docker] on product mode
     # python -m apis.chat_api -d   # [Dev]    on develop mode