:gem: [Feature] Support call hf api with api_key via HTTP Bearer

apis/chat_api.py

@@ -2,7 +2,8 @@ import argparse
 import uvicorn
 import sys
 
-from fastapi import FastAPI
+from fastapi import FastAPI, Depends
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel, Field
 from sse_starlette.sse import EventSourceResponse, ServerSentEvent
 from utils.logger import logger
@@ -38,6 +39,16 @@ class ChatAPIApp:
         ]
         return self.available_models
 
+    def extract_api_key(
+        credentials: HTTPAuthorizationCredentials = Depends(
+            HTTPBearer(auto_error=False)
+        ),
+    ):
+        if credentials:
+            return credentials.credentials
+        else:
+            return None
+
     class ChatCompletionsPostItem(BaseModel):
         model: str = Field(
             default="mixtral-8x7b",
@@ -60,7 +71,9 @@ class ChatAPIApp:
             description="(bool) Stream",
         )
 
-    def chat_completions(self, item: ChatCompletionsPostItem):
+    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)
@@ -70,6 +83,7 @@ class ChatAPIApp:
             prompt=composer.merged_str,
             temperature=item.temperature,
             max_new_tokens=item.max_tokens,
+            api_key=api_key,
         )
         if item.stream:
             event_source_response = EventSourceResponse(

networks/message_streamer.py

@@ -36,6 +36,7 @@ class MessageStreamer:
         prompt: str = None,
         temperature: float = 0.01,
         max_new_tokens: int = 8192,
+        api_key: str = None,
     ):
         # 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>'
@@ -45,6 +46,13 @@ class MessageStreamer:
         self.request_headers = {
             "Content-Type": "application/json",
         }
+
+        if api_key:
+            logger.note(
+                f"Using API Key: {api_key[:3]}{(len(api_key)-7)*'*'}{api_key[-4:]}"
+            )
+            self.request_headers["Authorization"] = f"Bearer {api_key}"
+
         # References:
         #   huggingface_hub/inference/_client.py:
         #     class InferenceClient > def text_generation()