update file

apis/chat_api.py

@@ -1,88 +1,97 @@
 import argparse
-import markdown2
 import os
 import sys
 import uvicorn
-
+import markdown2
 from pathlib import Path
-from fastapi import FastAPI, Depends
-from fastapi.responses import HTMLResponse
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from fastapi import FastAPI, Depends, HTTPException
+from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials, OAuth2PasswordBearer, OAuth2PasswordRequestForm
 from pydantic import BaseModel, Field
 from typing import Union
+from passlib.context import CryptContext
 from sse_starlette.sse import EventSourceResponse, ServerSentEvent
 from utils.logger import logger
 from networks.message_streamer import MessageStreamer
 from messagers.message_composer import MessageComposer
 from mocks.stream_chat_mocker import stream_chat_mock
 
+class Auth(BaseModel):
+    api_key: str
+    password: str
 
 class ChatAPIApp:
+
+    oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+    # Password hashing context
+    pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+
     def __init__(self):
-        self.app = FastAPI(
-            docs_url="/",
-            title="HuggingFace LLM API",
-            swagger_ui_parameters={"defaultModelsExpandDepth": -1},
-            version="1.0",
-        )
+        self.app = FastAPI(docs_url=None, redoc_url=None)
         self.setup_routes()
+        self.api_key = os.getenv("XCHE_API_KEY")
+        self.password = os.getenv("XCHE_PASSWORD")
+        self.fake_data_db = {
+            self.api_key: {
+                "api_key": self.api_key,
+                "password": self.get_password_hash(self.password)  # Pre-hashed password
+            }
+        }
+
+    def get_password_hash(self, password):
+        return self.pwd_context.hash(password)
+
+    def verify_password(self, plain_password, hashed_password):
+        return self.pwd_context.verify(plain_password, hashed_password)
+
+    def get_api_key(self, db, api_key: str):
+        if api_key in db:
+            api_dict = db[api_key]
+            return Auth(**api_dict)
+
+    def authenticate(self, fake_db, api_key: str, password: str):
+        api_data = self.get_api_key(fake_db, api_key)
+        if not api_data:
+            return False
+        if not self.verify_password(password, api_data.password):
+            return False
+        return api_data
+
+    async def login(self, form_data: OAuth2PasswordRequestForm = Depends()):
+        api_data = self.authenticate(self.fake_data_db, form_data.username, form_data.password)
+        if not api_data:
+            raise HTTPException(
+                status_code=400,
+                detail="Incorrect API KEY or Password",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+        return {"access_token": api_data.api_key, "token_type": "bearer"}
+
+    def check_api_token(self, token: str = Depends(oauth2_scheme)):
+        api_data = self.get_api_key(self.fake_data_db, token)
+        if not api_data:
+            raise HTTPException(status_code=403, detail="Invalid or missing API Key")
+        return api_data
 
     def get_available_models(self):
-        # https://platform.openai.com/docs/api-reference/models/list
-        # ANCHOR[id=available-models]: Available models
         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",
-                },
-                {
-                    "id": "nous-mixtral-8x7b",
-                    "description": "[NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO]: https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
-                    "object": "model",
-                    "created": 1700000000,
-                    "owned_by": "NousResearch",
-                },
-                {
-                    "id": "zephyr-7b-beta",
-                    "description": "[HuggingFaceH4/zephyr-7b-beta]: https://huggingface.co/HuggingFaceH4/zephyr-7b-beta",
-                    "object": "model",
-                    "created": 1700000000,
-                    "owned_by": "TheBloke",
-                },
-                {
-                    "id": "starchat2-15b-v0.1",
-                    "description": "[HuggingFaceH4/starchat2-15b-v0.1]: https://huggingface.co/HuggingFaceH4/starchat2-15b-v0.1",
-                    "object": "model",
-                    "created": 1700000000,
-                    "owned_by": "TheBloke",
-                },
-            ],
+                {"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"},
+                {"id": "nous-mixtral-8x7b", "description": "[NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO]: https://huggingface.co/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO", "object": "model", "created": 1700000000, "owned_by": "NousResearch"},
+                {"id": "zephyr-7b-beta", "description": "[HuggingFaceH4/zephyr-7b-beta]: https://huggingface.co/HuggingFaceH4/zephyr-7b-beta", "object": "model", "created": 1700000000, "owned_by": "TheBloke"},
+                {"id": "starchat2-15b-v0.1", "description": "[HuggingFaceH4/starchat2-15b-v0.1]: https://huggingface.co/HuggingFaceH4/starchat2-15b-v0.1", "object": "model", "created": 1700000000, "owned_by": "TheBloke"},
+            ]
         }
         return self.available_models
 
-    def extract_api_key(
-        credentials: HTTPAuthorizationCredentials = Depends(
-            HTTPBearer(auto_error=False)
-        ),
-    ):
+    def extract_api_key(self, credentials: HTTPAuthorizationCredentials = Depends(HTTPBearer(auto_error=False))):
         api_key = None
         if credentials:
             api_key = credentials.credentials
         else:
             api_key = os.getenv("XCHE_TOKEN")
-
         if api_key:
             if api_key.startswith("hf_"):
                 return api_key
@@ -93,43 +102,19 @@ class ChatAPIApp:
         return None
 
     class ChatCompletionsPostItem(BaseModel):
-        model: str = Field(
-            default="mixtral-8x7b",
-            description="(str) `mixtral-8x7b`",
-        )
-        messages: list = Field(
-            default=[{"role": "user", "content": "Hello, who are you?"}],
-            description="(list) Messages",
-        )
-        temperature: Union[float, None] = Field(
-            default=0.5,
-            description="(float) Temperature",
-        )
-        top_p: Union[float, None] = Field(
-            default=0.95,
-            description="(float) top p",
-        )
-        max_tokens: Union[int, None] = Field(
-            default=-1,
-            description="(int) Max tokens",
-        )
-        use_cache: bool = Field(
-            default=False,
-            description="(bool) Use cache",
-        )
-        stream: bool = Field(
-            default=True,
-            description="(bool) Stream",
-        )
-
-    def chat_completions(
-        self, item: ChatCompletionsPostItem, api_key: str = Depends(extract_api_key)
-    ):
+        model: str = Field(default="mixtral-8x7b", description="(str) `mixtral-8x7b`")
+        messages: list = Field(default=[{"role": "user", "content": "Hello, who are you?"}], description="(list) Messages")
+        temperature: Union[float, None] = Field(default=0.5, description="(float) Temperature")
+        top_p: Union[float, None] = Field(default=0.95, description="(float) top p")
+        max_tokens: Union[int, None] = Field(default=-1, description="(int) Max tokens")
+        use_cache: bool = Field(default=False, description="(bool) Use cache")
+        stream: bool = Field(default=True, description="(bool) Stream")
+
+    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,
@@ -154,9 +139,7 @@ class ChatAPIApp:
         readme_path = Path(__file__).parents[1] / "README.md"
         with open(readme_path, "r", encoding="utf-8") as rf:
             readme_str = rf.read()
-        readme_html = markdown2.markdown(
-            readme_str, extras=["table", "fenced-code-blocks", "highlightjs-lang"]
-        )
+        readme_html = markdown2.markdown(readme_str, extras=["table", "fenced-code-blocks", "highlightjs-lang"])
         return readme_html
 
     def setup_routes(self):
@@ -166,55 +149,20 @@ class ChatAPIApp:
             else:
                 include_in_schema = False
 
-            self.app.get(
-                prefix + "/models",
-                summary="Get available models",
-                include_in_schema=include_in_schema,
-            )(self.get_available_models)
-
-            self.app.post(
-                prefix + "/chat/completions",
-                summary="Chat completions in conversation session",
-                include_in_schema=include_in_schema,
-            )(self.chat_completions)
-        self.app.get(
-            "/readme",
-            summary="README of HF LLM API",
-            response_class=HTMLResponse,
-            include_in_schema=False,
-        )(self.get_readme)
+            self.app.get(prefix + "/models", summary="Get available models", include_in_schema=include_in_schema)(self.get_available_models)
 
+            self.app.post(prefix + "/chat/completions", summary="Chat completions in conversation session", include_in_schema=include_in_schema)(self.chat_completions)
+        self.app.get("/readme", summary="README of HF LLM API", response_class=HTMLResponse, include_in_schema=False)(self.get_readme)
+        self.app.post("/token", include_in_schema=False)(self.login)
 
 class ArgParser(argparse.ArgumentParser):
     def __init__(self, *args, **kwargs):
         super(ArgParser, self).__init__(*args, **kwargs)
-
-        self.add_argument(
-            "-s",
-            "--server",
-            type=str,
-            default="0.0.0.0",
-            help="Server IP for HF LLM Chat API",
-        )
-        self.add_argument(
-            "-p",
-            "--port",
-            type=int,
-            default=7860,
-            help="Server Port for HF LLM Chat API",
-        )
-
-        self.add_argument(
-            "-d",
-            "--dev",
-            default=False,
-            action="store_true",
-            help="Run in dev mode",
-        )
-
+        self.add_argument("-s", "--server", type=str, default="0.0.0.0", help="Server IP for HF LLM Chat API")
+        self.add_argument("-p", "--port", type=int, default=7860, help="Server Port for HF LLM Chat API")
+        self.add_argument("-d", "--dev", default=False, action="store_true", help="Run in dev mode")
         self.args = self.parse_args(sys.argv[1:])
 
-
 app = ChatAPIApp().app
 
 if __name__ == "__main__":
@@ -223,6 +171,3 @@ if __name__ == "__main__":
         uvicorn.run("__main__:app", host=args.server, port=args.port, reload=True)
     else:
         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