Update apis/chat_api.py

apis/chat_api.py

@@ -2,11 +2,11 @@ import argparse
 import os
 import sys
 import uvicorn
-import traceback
 
 from fastapi import FastAPI, Depends
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel, Field
+from typing import Union
 from sse_starlette.sse import EventSourceResponse, ServerSentEvent
 from utils.logger import logger
 from networks.message_streamer import MessageStreamer
@@ -25,21 +25,27 @@ class ChatAPIApp:
         self.setup_routes()
 
     def get_available_models(self):
+        # https://platform.openai.com/docs/api-reference/models/list
         # ANCHOR[id=available-models]: Available models
-        self.available_models = [
-            {
-                "id": "mixtral-8x7b",
-                "description": "[mistralai/Mixtral-8x7B-Instruct-v0.1]: https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1",
-            },
-             {
-                "id": "mistral-7b",
-                "description": "[mistralai/Mistral-7B-Instruct-v0.2]: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2"
-             },
-            {
-                "id": "openchat-3.5",
-                "description": "[openchat/openchat-3.5-1210]: https://huggingface.co/openchat/openchat-3.5-1210",
-            },
-        ]
+        self.available_models = {
+            "object": "list",
+            "data": [
+                {
+                    "id": "mixtral-8x7b",
+                    "description": "[mistralai/Mixtral-8x7B-Instruct-v0.1]: https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1",
+                    "object": "model",
+                    "created": 1700000000,
+                    "owned_by": "mistralai",
+                },
+                {
+                    "id": "mistral-7b",
+                    "description": "[mistralai/Mistral-7B-Instruct-v0.2]: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2",
+                    "object": "model",
+                    "created": 1700000000,
+                    "owned_by": "mistralai",
+                },
+            ],
+        }
         return self.available_models
 
     def extract_api_key(
@@ -47,10 +53,20 @@ class ChatAPIApp:
             HTTPBearer(auto_error=False)
         ),
     ):
+        api_key = None
         if credentials:
-            return credentials.credentials
+            api_key = credentials.credentials
         else:
-            return os.getenv("HF_TOKEN") or None
+            api_key = os.getenv("HF_TOKEN")
+
+        if api_key:
+            if api_key.startswith("hf_"):
+                return api_key
+            else:
+                logger.warn(f"Invalid HF Token!")
+        else:
+            logger.warn("Not provide HF Token!")
+        return None
 
     class ChatCompletionsPostItem(BaseModel):
         model: str = Field(
@@ -61,16 +77,16 @@ class ChatAPIApp:
             default=[{"role": "user", "content": "Hello, who are you?"}],
             description="(list) Messages",
         )
-        temperature: float = Field(
-            default=0.01,
+        temperature: Union[float, None] = Field(
+            default=0,
             description="(float) Temperature",
         )
-        max_tokens: int = Field(
-            default=4096,
+        max_tokens: Union[int, None] = Field(
+            default=-1,
             description="(int) Max tokens",
         )
         stream: bool = Field(
-            default=False,
+            default=True,
             description="(bool) Stream",
         )
 
@@ -101,7 +117,7 @@ class ChatAPIApp:
             return data_response
 
     def setup_routes(self):
-        for prefix in ["", "/v1"]:
+        for prefix in ["", "/v1", "/api", "/api/v1"]:
             self.app.get(
                 prefix + "/models",
                 summary="Get available models",
@@ -153,4 +169,4 @@ if __name__ == "__main__":
         uvicorn.run("__main__:app", host=args.server, port=args.port, reload=False)
 
     # python -m apis.chat_api      # [Docker] on product mode
-    # python -m apis.chat_api -d   # [Dev]    on develop mode
+    # python -m apis.chat_api -d   # [Dev]    on develop mode