diff --git a/proxy/routes/proxy.py b/proxy/routes/proxy.py index 6196a1e..bdd376a 100644 --- a/proxy/routes/proxy.py +++ b/proxy/routes/proxy.py @@ -1,49 +1,36 @@ """Proxy service to forward requests to the appropriate language model provider""" -from enum import Enum - from fastapi import APIRouter, Depends, Header, HTTPException, Request proxy = APIRouter() -class LLMProvider(str, Enum): - """Supported language model providers""" - - OPEN_AI = "openai" - ANTHROPIC = "anthropic" - - @classmethod - def is_valid(cls, provider: str) -> bool: - """Check if a provider is a valid LLM provider""" - return provider in {provider.value for provider in cls} - - def validate_headers(invariant_authorization: str = Header(None)): - """Require the Invariant-Authorization header to be present""" + """Require the invariant-authorization header to be present""" if invariant_authorization is None: raise HTTPException( - status_code=400, detail="Missing Invariant-Authorization header" + status_code=400, detail="Missing invariant-authorization header" ) return invariant_authorization +allowed_openai_endpoints = {"chat/completions", "moderations"} + + @proxy.post( - "/{username}/{dataset_name}/{llm_provider}", + "/{username}/{dataset_name}/openai/{endpoint:path}", dependencies=[Depends(validate_headers)], ) -async def chat_completion( +async def openai_proxy( request: Request, username: str, dataset_name: str, - llm_provider: str, + endpoint: str, ): """Proxy call to a language model provider""" - - if not LLMProvider.is_valid(llm_provider): - raise HTTPException( - status_code=400, - detail=f"Unsupported LLM provider '{llm_provider}'.", - ) - - return {"message": f"Upload {dataset_name} for {username} to {llm_provider}"} + print("headers: ", dict(request.headers)) + # print("request: ", await request.json()) + print(f"Proxying to OpenAI endpoint: {endpoint}") + if endpoint not in allowed_openai_endpoints: + raise HTTPException(status_code=404, detail="Not supported OpenAI endpoint") + return {"message": f"Upload {dataset_name} for {username} to openai"}