phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

add support for using google ai sdk with litellm

Browse source
This commit is contained in:
Ishaan Jaff 2024-11-22 10:56:42 -08:00
parent d047a9b2c9
commit a332cc1861
3 changed files with 16 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/_types.py
View file

@ -2075,6 +2075,7 @@ class SpecialHeaders(enum.Enum):
openai_authorization = "Authorization" openai_authorization = "Authorization"
azure_authorization = "API-Key" azure_authorization = "API-Key"
anthropic_authorization = "x-api-key" anthropic_authorization = "x-api-key"
google_ai_studio_authorization = "x-goog-api-key"
class LitellmDataForBackendLLMCall(TypedDict, total=False): class LitellmDataForBackendLLMCall(TypedDict, total=False):

10
litellm/proxy/auth/user_api_key_auth.py
View file

@ -95,6 +95,11 @@ anthropic_api_key_header = APIKeyHeader(
auto_error=False, auto_error=False,
description="If anthropic client used.", description="If anthropic client used.",
) )
google_ai_studio_api_key_header = APIKeyHeader(
name=SpecialHeaders.google_ai_studio_authorization.value,
auto_error=False,
description="If google ai studio client used.",
)
def _get_bearer_token( def _get_bearer_token(
@ -197,6 +202,9 @@ async def user_api_key_auth( # noqa: PLR0915
anthropic_api_key_header: Optional[str] = fastapi.Security( anthropic_api_key_header: Optional[str] = fastapi.Security(
anthropic_api_key_header anthropic_api_key_header
), ),
google_ai_studio_api_key_header: Optional[str] = fastapi.Security(
google_ai_studio_api_key_header
),
) -> UserAPIKeyAuth: ) -> UserAPIKeyAuth:
from litellm.proxy.proxy_server import ( from litellm.proxy.proxy_server import (
general_settings, general_settings,
@ -233,6 +241,8 @@ async def user_api_key_auth( # noqa: PLR0915
api_key = azure_api_key_header api_key = azure_api_key_header
elif isinstance(anthropic_api_key_header, str): elif isinstance(anthropic_api_key_header, str):
api_key = anthropic_api_key_header api_key = anthropic_api_key_header
elif isinstance(google_ai_studio_api_key_header, str):
api_key = google_ai_studio_api_key_header
elif pass_through_endpoints is not None: elif pass_through_endpoints is not None:
for endpoint in pass_through_endpoints: for endpoint in pass_through_endpoints:
if endpoint.get("path", "") == route: if endpoint.get("path", "") == route:

7
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -61,10 +61,13 @@ async def gemini_proxy_route(
fastapi_response: Response, fastapi_response: Response,
): ):
## CHECK FOR LITELLM API KEY IN THE QUERY PARAMS - ?..key=LITELLM_API_KEY ## CHECK FOR LITELLM API KEY IN THE QUERY PARAMS - ?..key=LITELLM_API_KEY
api_key = request.query_params.get("key") google_ai_studio_api_key = request.query_params.get("key") or request.headers.get(
"x-goog-api-key"
)
user_api_key_dict = await user_api_key_auth( user_api_key_dict = await user_api_key_auth(
request=request, api_key="Bearer {}".format(api_key) request=request,
google_ai_studio_api_key_header=google_ai_studio_api_key,
) )
base_target_url = "https://generativelanguage.googleapis.com" base_target_url = "https://generativelanguage.googleapis.com"