diff --git a/agents-api/agents_api/clients/litellm.py b/agents-api/agents_api/clients/litellm.py index 463f6fa67..0319f384f 100644 --- a/agents-api/agents_api/clients/litellm.py +++ b/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, ) diff --git a/agents-api/agents_api/routers/sessions/chat.py b/agents-api/agents_api/routers/sessions/chat.py index b7bf96d3a..7708ec7ca 100644 --- a/agents-api/agents_api/routers/sessions/chat.py +++ b/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, ) diff --git a/typespec/chat/endpoints.tsp b/typespec/chat/endpoints.tsp index 64eb6700d..cceaca337 100644 --- a/typespec/chat/endpoints.tsp +++ b/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;