Implement openai api compatible authentication

docmeth02 · docmeth02 · commit 673e4fa93207 · 2023-12-14T20:26:53.000+01:00
diff --git a/llama_cpp/server/app.py b/llama_cpp/server/app.py
@@ -14,11 +14,12 @@
 import anyio
 from anyio.streams.memory import MemoryObjectSendStream
 from starlette.concurrency import run_in_threadpool, iterate_in_threadpool
-from fastapi import Depends, FastAPI, APIRouter, Request, Response
+from fastapi import Depends, FastAPI, APIRouter, Request, Response, HTTPException, status
 from fastapi.middleware import Middleware
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi.routing import APIRoute
+from fastapi.security import HTTPBearer
 from pydantic import BaseModel, Field
 from pydantic_settings import BaseSettings
 from sse_starlette.sse import EventSourceResponse
@@ -161,6 +162,10 @@ class Settings(BaseSettings):
         default=True,
         description="Whether to interrupt requests when a new request is received.",
     )
+    api_key: Optional[str] = Field(
+        default=None,
+        description="API key for authentication. If set all requests need to be authenticated."
+    )
 
 
 class ErrorResponse(TypedDict):
@@ -312,6 +317,9 @@ async def custom_route_handler(request: Request) -> Response:
                 elapsed_time_ms = int((time.perf_counter() - start_sec) * 1000)
                 response.headers["openai-processing-ms"] = f"{elapsed_time_ms}"
                 return response
+            except HTTPException as unauthorized:
+                # api key check failed
+                raise unauthorized
             except Exception as exc:
                 json_body = await request.json()
                 try:
@@ -656,6 +664,27 @@ def _logit_bias_tokens_to_input_ids(
     return to_bias
 
 
+# Setup Bearer authentication scheme
+bearer_scheme = HTTPBearer(auto_error=False)
+
+
+async def authenticate(settings: Settings = Depends(get_settings), authorization: Optional[str] = Depends(bearer_scheme)):
+    # Skip API key check if it's not set in settings
+    if settings.api_key is None:
+        return True
+
+    # check bearer credentials against the api_key
+    if authorization and authorization.credentials == settings.api_key:
+        # api key is valid
+        return authorization.credentials
+
+    # raise http error 401
+    raise HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Invalid API key",
+    )
+
+
 @router.post(
     "/v1/completions",
     summary="Completion"
@@ -665,6 +694,7 @@ async def create_completion(
     request: Request,
     body: CreateCompletionRequest,
     llama: llama_cpp.Llama = Depends(get_llama),
+    authenticated: str = Depends(authenticate),
 ) -> llama_cpp.Completion:
     if isinstance(body.prompt, list):
         assert len(body.prompt) <= 1
@@ -738,7 +768,9 @@ class CreateEmbeddingRequest(BaseModel):
     summary="Embedding"
 )
 async def create_embedding(
-    request: CreateEmbeddingRequest, llama: llama_cpp.Llama = Depends(get_llama)
+    request: CreateEmbeddingRequest,
+    llama: llama_cpp.Llama = Depends(get_llama),
+    authenticated: str = Depends(authenticate),
 ):
     return await run_in_threadpool(
         llama.create_embedding, **request.model_dump(exclude={"user"})
@@ -832,6 +864,7 @@ async def create_chat_completion(
     body: CreateChatCompletionRequest,
     llama: llama_cpp.Llama = Depends(get_llama),
     settings: Settings = Depends(get_settings),
+    authenticated: str = Depends(authenticate),
 ) -> llama_cpp.ChatCompletion:
     exclude = {
         "n",
@@ -893,6 +926,7 @@ class ModelList(TypedDict):
 @router.get("/v1/models", summary="Models")
 async def get_models(
     settings: Settings = Depends(get_settings),
+    authenticated: str = Depends(authenticate),
 ) -> ModelList:
     assert llama is not None
     return {
