Add custom api key support to chat endpoint

agents-api/agents_api/clients/litellm.py

@@ -14,9 +14,8 @@
 @wraps(_acompletion)
 @beartype
 async def acompletion(
-    *, model: str, messages: list[dict], **kwargs
+    *, model: str, messages: list[dict], custom_api_key: str = None, **kwargs
 ) -> ModelResponse | CustomStreamWrapper:
-    model = f"openai/{model}"  # This is here because litellm proxy expects this format
 
     supported_params = get_supported_openai_params(model)
     settings = {k: v for k, v in kwargs.items() if k in supported_params}
@@ -25,6 +24,6 @@ async def acompletion(
         model=model,
         messages=messages,
         **settings,
-        base_url=litellm_url,
-        api_key=litellm_master_key,
+        base_url=None if custom_api_key else litellm_url,
+        api_key=custom_api_key or litellm_master_key,
     )

agents-api/agents_api/routers/sessions/chat.py

@@ -1,7 +1,7 @@
-from typing import Annotated
+from typing import Annotated, Optional
 from uuid import UUID, uuid4
 
-from fastapi import BackgroundTasks, Depends
+from fastapi import BackgroundTasks, Depends, Header
 from starlette.status import HTTP_201_CREATED
 
 from ...autogen.openapi_model import (
@@ -33,6 +33,7 @@ async def chat(
     session_id: UUID,
     chat_input: ChatInput,
     background_tasks: BackgroundTasks,
+    x_custom_api_key: Optional[str] = Header(None, alias="X-Custom-Api-Key"),
 ) -> ChatResponse:
     if chat_input.stream:
         raise NotImplementedError("Streaming is not yet implemented")
@@ -107,6 +108,7 @@ async def chat(
         tools=tools or None,
         user=str(developer.id),  # For tracking usage
         tags=developer.tags,  # For filtering models in litellm
+        custom_api_key=x_custom_api_key,
         **settings,
     )
 

typespec/chat/endpoints.tsp

@@ -23,6 +23,10 @@ interface Endpoints {
         @doc("The session ID")
         id: uuid;
 
+        @header
+        @doc("Custom API key")
+        "X-Custom-Api-Key"?: string;
+
         @bodyRoot
         @doc("Request to generate a response from the model")
         body: ChatInput;