From bc0023a40905d2e5dc38e9dbefd9cc0f800ec434 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 17 Aug 2024 10:46:59 -0700 Subject: [PATCH 1/3] feat(google_ai_studio_endpoints.py): support pass-through endpoint for all google ai studio requests New Feature --- litellm/cost_calculator.py | 9 ++- litellm/litellm_core_utils/litellm_logging.py | 3 +- .../pass_through_endpoints.py | 62 +++++++++++---- litellm/proxy/proxy_server.py | 4 + .../google_ai_studio_endpoints.py | 79 +++++++++++++++++++ litellm/tests/test_proxy_server.py | 49 ++++++++++++ 6 files changed, 186 insertions(+), 20 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 4b5ac51db9..d6f9adb008 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -412,7 +412,7 @@ def get_replicate_completion_pricing(completion_response=None, total_time=0.0): def _select_model_name_for_cost_calc( model: Optional[str], - completion_response: Union[BaseModel, dict], + completion_response: Union[BaseModel, dict, str], base_model: Optional[str] = None, custom_pricing: Optional[bool] = None, ) -> Optional[str]: @@ -428,7 +428,12 @@ def _select_model_name_for_cost_calc( if base_model is not None: return base_model - return_model = model or completion_response.get("model", "") # type: ignore + return_model = model + if isinstance(completion_response, str): + return return_model + + elif return_model is None: + return_model = completion_response.get("model", "") # type: ignore if hasattr(completion_response, "_hidden_params"): if ( completion_response._hidden_params.get("model", None) is not None diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 74c6d0db01..77837a7898 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -274,6 +274,7 @@ class Logging: headers = {} data = additional_args.get("complete_input_dict", {}) api_base = str(additional_args.get("api_base", "")) + query_params = additional_args.get("query_params", {}) if "key=" in api_base: # Find the position of "key=" in the string key_index = api_base.find("key=") + 4 @@ -2362,7 +2363,7 @@ def get_standard_logging_object_payload( return payload except Exception as e: - verbose_logger.error( + verbose_logger.exception( "Error creating standard logging object - {}".format(str(e)) ) return None diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fd18c707f3..5ee6cc6d6b 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -273,6 +273,7 @@ async def pass_through_request( custom_headers: dict, user_api_key_dict: UserAPIKeyAuth, forward_headers: Optional[bool] = False, + query_params: Optional[dict] = None, ): try: import time @@ -308,23 +309,9 @@ async def pass_through_request( ) async_client = httpx.AsyncClient() - response = await async_client.request( - method=request.method, - url=url, - headers=headers, - params=request.query_params, - json=_parsed_body, - ) - if response.status_code >= 300: - raise HTTPException(status_code=response.status_code, detail=response.text) - - content = await response.aread() - - ## LOG SUCCESS - start_time = time.time() - end_time = time.time() # create logging object + start_time = time.time() logging_obj = Logging( model="unknown", messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}], @@ -334,6 +321,7 @@ async def pass_through_request( litellm_call_id=str(uuid.uuid4()), function_id="1245", ) + # done for supporting 'parallel_request_limiter.py' with pass-through endpoints kwargs = { "litellm_params": { @@ -354,6 +342,44 @@ async def pass_through_request( call_type="pass_through_endpoint", ) + # combine url with query params for logging + + requested_query_params = query_params or request.query_params.__dict__ + requested_query_params_str = "&".join( + f"{k}={v}" for k, v in requested_query_params.items() + ) + + if "?" in str(url): + logging_url = str(url) + "&" + requested_query_params_str + else: + logging_url = str(url) + "?" + requested_query_params_str + + logging_obj.pre_call( + input=[{"role": "user", "content": "no-message-pass-through-endpoint"}], + api_key="", + additional_args={ + "complete_input_dict": _parsed_body, + "api_base": logging_url, + "headers": headers, + }, + ) + + response = await async_client.request( + method=request.method, + url=url, + headers=headers, + params=requested_query_params, + json=_parsed_body, + ) + + if response.status_code >= 300: + raise HTTPException(status_code=response.status_code, detail=response.text) + + content = await response.aread() + + ## LOG SUCCESS + end_time = time.time() + await logging_obj.async_success_handler( result="", start_time=start_time, @@ -431,17 +457,19 @@ def create_pass_through_route( except Exception: verbose_proxy_logger.debug("Defaulting to target being a url.") - async def endpoint_func( + async def endpoint_func( # type: ignore request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + query_params: Optional[dict] = None, ): - return await pass_through_request( + return await pass_through_request( # type: ignore request=request, target=target, custom_headers=custom_headers or {}, user_api_key_dict=user_api_key_dict, forward_headers=_forward_headers, + query_params=query_params, ) return endpoint_func diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f09ae7d350..60df87e729 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -227,6 +227,9 @@ from litellm.proxy.utils import ( send_email, update_spend, ) +from litellm.proxy.vertex_ai_endpoints.google_ai_studio_endpoints import ( + router as gemini_router, +) from litellm.proxy.vertex_ai_endpoints.vertex_endpoints import router as vertex_router from litellm.proxy.vertex_ai_endpoints.vertex_endpoints import set_default_vertex_config from litellm.router import ( @@ -9704,6 +9707,7 @@ def cleanup_router_config_variables(): app.include_router(router) app.include_router(fine_tuning_router) app.include_router(vertex_router) +app.include_router(gemini_router) app.include_router(pass_through_router) app.include_router(health_router) app.include_router(key_management_router) diff --git a/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py b/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py index 1bb966964d..3b6105b447 100644 --- a/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py +++ b/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py @@ -3,3 +3,82 @@ What is this? Google AI Studio Pass-Through Endpoints """ + +""" +1. Create pass-through endpoints for any LITELLM_BASE_URL/gemini/ map to https://generativelanguage.googleapis.com/ +""" + +import ast +import asyncio +import traceback +from datetime import datetime, timedelta, timezone +from typing import List, Optional +from urllib.parse import urlencode + +import fastapi +import httpx +from fastapi import ( + APIRouter, + Depends, + File, + Form, + Header, + HTTPException, + Request, + Response, + UploadFile, + status, +) +from starlette.datastructures import QueryParams + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.batches.main import FileObject +from litellm.fine_tuning.main import vertex_fine_tuning_apis_instance +from litellm.proxy._types import * +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + create_pass_through_route, +) + +router = APIRouter() +default_vertex_config = None + + +@router.api_route("/gemini/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE"]) +async def gemini_proxy_route( + endpoint: str, + request: Request, + fastapi_response: Response, +): + ## CHECK FOR LITELLM API KEY IN THE QUERY PARAMS - ?..key=LITELLM_API_KEY + api_key = request.query_params.get("key") + + user_api_key_dict = await user_api_key_auth( + request=request, api_key="Bearer {}".format(api_key) + ) + + base_target_url = "https://generativelanguage.googleapis.com" + encoded_endpoint = httpx.URL(endpoint).path + + # Ensure endpoint starts with '/' for proper URL construction + if not encoded_endpoint.startswith("/"): + encoded_endpoint = "/" + encoded_endpoint + + # Construct the full target URL using httpx + base_url = httpx.URL(base_target_url) + updated_url = base_url.copy_with(path=encoded_endpoint) + + # Add or update query parameters + gemini_api_key = litellm.utils.get_secret(secret_name="GEMINI_API_KEY") + # Merge query parameters, giving precedence to those in updated_url + merged_params = dict(request.query_params) + merged_params.update({"key": gemini_api_key}) + + endpoint_func = create_pass_through_route( + endpoint=endpoint, + target=str(updated_url), + ) # dynamically construct pass-through endpoint based on incoming path + return await endpoint_func( + request, fastapi_response, user_api_key_dict, query_params=merged_params + ) diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index 9a1c091267..28f3aad632 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -1166,3 +1166,52 @@ async def test_add_callback_via_key_litellm_pre_call_utils(prisma_client): assert new_data["success_callback"] == ["langfuse"] assert "langfuse_public_key" in new_data assert "langfuse_secret_key" in new_data + + +@pytest.mark.asyncio +async def test_gemini_pass_through_endpoint(): + from starlette.datastructures import URL + + from litellm.proxy.vertex_ai_endpoints.google_ai_studio_endpoints import ( + Request, + Response, + gemini_proxy_route, + ) + + body = b""" + { + "contents": [{ + "parts":[{ + "text": "The quick brown fox jumps over the lazy dog." + }] + }] + } + """ + + # Construct the scope dictionary + scope = { + "type": "http", + "method": "POST", + "path": "/gemini/v1beta/models/gemini-1.5-flash:countTokens", + "query_string": b"key=sk-1234", + "headers": [ + (b"content-type", b"application/json"), + ], + } + + # Create a new Request object + async def async_receive(): + return {"type": "http.request", "body": body, "more_body": False} + + request = Request( + scope=scope, + receive=async_receive, + ) + + resp = await gemini_proxy_route( + endpoint="v1beta/models/gemini-1.5-flash:countTokens?key=sk-1234", + request=request, + fastapi_response=Response(), + ) + + print(resp.body) From fd44cf8d261ddf50e04130ab25edc2af0f3e65cd Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 17 Aug 2024 12:46:57 -0700 Subject: [PATCH 2/3] feat(pass_through_endpoints.py): support streaming requests --- .../pass_through_endpoints.py | 66 ++++++++++++++++--- .../google_ai_studio_endpoints.py | 16 ++++- 2 files changed, 70 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 5ee6cc6d6b..511ce1ee9f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3,7 +3,7 @@ import asyncio import json import traceback from base64 import b64encode -from typing import List, Optional +from typing import AsyncIterable, List, Optional import httpx from fastapi import ( @@ -267,6 +267,17 @@ def forward_headers_from_request( return headers +def get_response_headers(headers: httpx.Headers) -> dict: + excluded_headers = {"transfer-encoding", "content-encoding"} + return_headers = { + key: value + for key, value in headers.items() + if key.lower() not in excluded_headers + } + + return return_headers + + async def pass_through_request( request: Request, target: str, @@ -274,6 +285,7 @@ async def pass_through_request( user_api_key_dict: UserAPIKeyAuth, forward_headers: Optional[bool] = False, query_params: Optional[dict] = None, + stream: Optional[bool] = None, ): try: import time @@ -292,7 +304,7 @@ async def pass_through_request( body_str = request_body.decode() try: _parsed_body = ast.literal_eval(body_str) - except: + except Exception: _parsed_body = json.loads(body_str) verbose_proxy_logger.debug( @@ -364,6 +376,27 @@ async def pass_through_request( }, ) + if stream: + req = async_client.build_request( + "POST", + url, + json=_parsed_body, + params=requested_query_params, + headers=headers, + ) + response = await async_client.send(req, stream=stream) + + # Create an async generator to yield the response content + async def stream_response() -> AsyncIterable[bytes]: + async for chunk in response.aiter_bytes(): + yield chunk + + return StreamingResponse( + stream_response(), + headers=get_response_headers(response.headers), + status_code=response.status_code, + ) + response = await async_client.request( method=request.method, url=url, @@ -372,6 +405,22 @@ async def pass_through_request( json=_parsed_body, ) + if ( + response.headers.get("content-type") is not None + and response.headers["content-type"] == "text/event-stream" + ): + # streaming response + # Create an async generator to yield the response content + async def stream_response() -> AsyncIterable[bytes]: + async for chunk in response.aiter_bytes(): + yield chunk + + return StreamingResponse( + stream_response(), + headers=get_response_headers(response.headers), + status_code=response.status_code, + ) + if response.status_code >= 300: raise HTTPException(status_code=response.status_code, detail=response.text) @@ -387,17 +436,10 @@ async def pass_through_request( cache_hit=False, ) - excluded_headers = {"transfer-encoding", "content-encoding"} - headers = { - key: value - for key, value in response.headers.items() - if key.lower() not in excluded_headers - } - return Response( content=content, status_code=response.status_code, - headers=headers, + headers=get_response_headers(response.headers), ) except Exception as e: verbose_proxy_logger.exception( @@ -462,6 +504,9 @@ def create_pass_through_route( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), query_params: Optional[dict] = None, + stream: Optional[ + bool + ] = None, # if pass-through endpoint is a streaming request ): return await pass_through_request( # type: ignore request=request, @@ -470,6 +515,7 @@ def create_pass_through_route( user_api_key_dict=user_api_key_dict, forward_headers=_forward_headers, query_params=query_params, + stream=stream, ) return endpoint_func diff --git a/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py b/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py index 3b6105b447..9eb0f9bd1f 100644 --- a/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py +++ b/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py @@ -75,10 +75,22 @@ async def gemini_proxy_route( merged_params = dict(request.query_params) merged_params.update({"key": gemini_api_key}) + ## check for streaming + is_streaming_request = False + if "stream" in str(updated_url): + is_streaming_request = True + + ## CREATE PASS-THROUGH endpoint_func = create_pass_through_route( endpoint=endpoint, target=str(updated_url), ) # dynamically construct pass-through endpoint based on incoming path - return await endpoint_func( - request, fastapi_response, user_api_key_dict, query_params=merged_params + received_value = await endpoint_func( + request, + fastapi_response, + user_api_key_dict, + query_params=merged_params, + stream=is_streaming_request, ) + + return received_value From 0df41653f3f8ad8699b2c6c0649cbdb010afa2a3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 17 Aug 2024 13:47:05 -0700 Subject: [PATCH 3/3] docs(google_ai_studio.md): add docs on google ai studio pass through endpoints --- .../docs/pass_through/google_ai_studio.md | 223 ++++++++++++++++++ .../docs/{ => pass_through}/vertex_ai.md | 0 docs/my-website/sidebars.js | 3 +- 3 files changed, 225 insertions(+), 1 deletion(-) create mode 100644 docs/my-website/docs/pass_through/google_ai_studio.md rename docs/my-website/docs/{ => pass_through}/vertex_ai.md (100%) diff --git a/docs/my-website/docs/pass_through/google_ai_studio.md b/docs/my-website/docs/pass_through/google_ai_studio.md new file mode 100644 index 0000000000..e37fa1218d --- /dev/null +++ b/docs/my-website/docs/pass_through/google_ai_studio.md @@ -0,0 +1,223 @@ +# Google AI Studio (Pass-Through) + +Pass-through endpoints for Google AI Studio - call provider-specific endpoint, in native format (no translation). + +Just replace `https://generativelanguage.googleapis.com` with `LITELLM_PROXY_BASE_URL/gemini` 🚀 + +#### **Example Usage** +```bash +http://0.0.0.0:4000/gemini/v1beta/models/gemini-1.5-flash:countTokens?key=sk-anything' \ +-H 'Content-Type: application/json' \ +-d '{ + "contents": [{ + "parts":[{ + "text": "The quick brown fox jumps over the lazy dog." + }] + }] +}' +``` + +Supports **ALL** Google AI Studio Endpoints (including streaming). + +[**See All Google AI Studio Endpoints**](https://ai.google.dev/api) + +## Quick Start + +Let's call the Gemini [`/countTokens` endpoint](https://ai.google.dev/api/tokens#method:-models.counttokens) + +1. Add Gemini API Key to your environment + +```bash +export GEMINI_API_KEY="" +``` + +2. Start LiteLLM Proxy + +```bash +litellm + +# RUNNING on http://0.0.0.0:4000 +``` + +3. Test it! + +Let's call the Google AI Studio token counting endpoint + +```bash +http://0.0.0.0:4000/gemini/v1beta/models/gemini-1.5-flash:countTokens?key=anything' \ +-H 'Content-Type: application/json' \ +-d '{ + "contents": [{ + "parts":[{ + "text": "The quick brown fox jumps over the lazy dog." + }] + }] +}' +``` + + +## Examples + +Anything after `http://0.0.0.0:4000/gemini` is treated as a provider-specific route, and handled accordingly. + +Key Changes: + +| **Original Endpoint** | **Replace With** | +|------------------------------------------------------|-----------------------------------| +| `https://generativelanguage.googleapis.com` | `http://0.0.0.0:4000/gemini` (LITELLM_PROXY_BASE_URL="http://0.0.0.0:4000") | +| `key=$GOOGLE_API_KEY` | `key=anything` (use `key=LITELLM_VIRTUAL_KEY` if Virtual Keys are setup on proxy) | + + +### **Example 1: Counting tokens** + +#### LiteLLM Proxy Call + +```bash +curl http://0.0.0.0:4000/gemini/v1beta/models/gemini-1.5-flash:countTokens?key=anything \ + -H 'Content-Type: application/json' \ + -X POST \ + -d '{ + "contents": [{ + "parts":[{ + "text": "The quick brown fox jumps over the lazy dog." + }], + }], + }' +``` + +#### Direct Google AI Studio Call + +```bash +curl https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:countTokens?key=$GOOGLE_API_KEY \ + -H 'Content-Type: application/json' \ + -X POST \ + -d '{ + "contents": [{ + "parts":[{ + "text": "The quick brown fox jumps over the lazy dog." + }], + }], + }' +``` + +### **Example 2: Generate content** + +#### LiteLLM Proxy Call + +```bash +curl "http://0.0.0.0:4000/gemini/v1beta/models/gemini-1.5-flash:generateContent?key=anything" \ + -H 'Content-Type: application/json' \ + -X POST \ + -d '{ + "contents": [{ + "parts":[{"text": "Write a story about a magic backpack."}] + }] + }' 2> /dev/null +``` + +#### Direct Google AI Studio Call + +```bash +curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key=$GOOGLE_API_KEY" \ + -H 'Content-Type: application/json' \ + -X POST \ + -d '{ + "contents": [{ + "parts":[{"text": "Write a story about a magic backpack."}] + }] + }' 2> /dev/null +``` + +### **Example 3: Caching** + + +```bash +curl -X POST "http://0.0.0.0:4000/gemini/v1beta/models/gemini-1.5-flash-001:generateContent?key=anything" \ +-H 'Content-Type: application/json' \ +-d '{ + "contents": [ + { + "parts":[{ + "text": "Please summarize this transcript" + }], + "role": "user" + }, + ], + "cachedContent": "'$CACHE_NAME'" + }' +``` + +#### Direct Google AI Studio Call + +```bash +curl -X POST "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-001:generateContent?key=$GOOGLE_API_KEY" \ +-H 'Content-Type: application/json' \ +-d '{ + "contents": [ + { + "parts":[{ + "text": "Please summarize this transcript" + }], + "role": "user" + }, + ], + "cachedContent": "'$CACHE_NAME'" + }' +``` + + +## Advanced - Use with Virtual Keys + +Pre-requisites +- [Setup proxy with DB](../proxy/virtual_keys.md#setup) + +Use this, to avoid giving developers the raw Google AI Studio key, but still letting them use Google AI Studio endpoints. + +### Usage + +1. Setup environment + +```bash +export DATABASE_URL="" +export LITELLM_MASTER_KEY="" +export GEMINI_API_KEY="" +``` + +```bash +litellm + +# RUNNING on http://0.0.0.0:4000 +``` + +2. Generate virtual key + +```bash +curl -X POST 'http://0.0.0.0:4000/key/generate' \ +-H 'Authorization: Bearer sk-1234' \ +-H 'Content-Type: application/json' \ +-d '{}' +``` + +Expected Response + +```bash +{ + ... + "key": "sk-1234ewknldferwedojwojw" +} +``` + +3. Test it! + + +```bash +http://0.0.0.0:4000/gemini/v1beta/models/gemini-1.5-flash:countTokens?key=sk-1234ewknldferwedojwojw' \ +-H 'Content-Type: application/json' \ +-d '{ + "contents": [{ + "parts":[{ + "text": "The quick brown fox jumps over the lazy dog." + }] + }] +}' +``` \ No newline at end of file diff --git a/docs/my-website/docs/vertex_ai.md b/docs/my-website/docs/pass_through/vertex_ai.md similarity index 100% rename from docs/my-website/docs/vertex_ai.md rename to docs/my-website/docs/pass_through/vertex_ai.md diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index bfa8953d4c..a4c97e29bb 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -192,7 +192,8 @@ const sidebars = { "batches", "fine_tuning", "anthropic_completion", - "vertex_ai" + "pass_through/vertex_ai", + "pass_through/google_ai_studio" ], }, "scheduler",