diff --git a/proxy/routes/proxy.py b/proxy/routes/proxy.py index b460a06..6196a1e 100644 --- a/proxy/routes/proxy.py +++ b/proxy/routes/proxy.py @@ -2,7 +2,7 @@ from enum import Enum -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, Depends, Header, HTTPException, Request proxy = APIRouter() @@ -19,8 +19,25 @@ class LLMProvider(str, Enum): return provider in {provider.value for provider in cls} -@proxy.post("/{username}/{dataset_name}/{llm_provider}") -async def chat_completion(username: str, dataset_name: str, llm_provider: str): +def validate_headers(invariant_authorization: str = Header(None)): + """Require the Invariant-Authorization header to be present""" + if invariant_authorization is None: + raise HTTPException( + status_code=400, detail="Missing Invariant-Authorization header" + ) + return invariant_authorization + + +@proxy.post( + "/{username}/{dataset_name}/{llm_provider}", + dependencies=[Depends(validate_headers)], +) +async def chat_completion( + request: Request, + username: str, + dataset_name: str, + llm_provider: str, +): """Proxy call to a language model provider""" if not LLMProvider.is_valid(llm_provider):