From d29a7087f1c7b4dd3a549b09757f6d44fc2d0c1c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 27 Aug 2024 16:53:11 -0700 Subject: [PATCH 01/16] feat(vertex_ai_and_google_ai_studio): Support Google AI Studio Embeddings endpoint Closes https://github.com/BerriAI/litellm/issues/5385 --- .../common_utils.py | 61 ++++++++++++++- .../transformation.py} | 0 .../vertex_and_google_ai_studio_gemini.py | 74 +++++++++++-------- litellm/main.py | 7 +- .../tests/test_amazing_vertex_completion.py | 8 +- 5 files changed, 110 insertions(+), 40 deletions(-) rename litellm/llms/vertex_ai_and_google_ai_studio/{gemini_transformation.py => gemini/transformation.py} (100%) rename litellm/llms/vertex_ai_and_google_ai_studio/{ => gemini}/vertex_and_google_ai_studio_gemini.py (97%) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py b/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py index 8faf7a3afa..7e2f9b29d0 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Literal, Tuple import httpx @@ -37,3 +37,62 @@ def get_supports_system_message( supports_system_message = False return supports_system_message + + +from typing import Literal, Optional + +all_gemini_url_modes = Literal["chat", "embedding"] + + +def _get_vertex_url( + mode: all_gemini_url_modes, + model: str, + stream: Optional[bool], + vertex_project: Optional[str], + vertex_location: Optional[str], + vertex_api_version: Literal["v1", "v1beta1"], +) -> Tuple[str, str]: + if mode == "chat": + ### SET RUNTIME ENDPOINT ### + endpoint = "generateContent" + if stream is True: + endpoint = "streamGenerateContent" + url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}?alt=sse" + else: + url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" + + # if model is only numeric chars then it's a fine tuned gemini model + # model = 4965075652664360960 + # send to this url: url = f"https://{vertex_location}-aiplatform.googleapis.com/{version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}" + if model.isdigit(): + # It's a fine-tuned Gemini model + url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}" + if stream is True: + url += "?alt=sse" + + return url, endpoint + + +def _get_gemini_url( + mode: all_gemini_url_modes, + model: str, + stream: Optional[bool], + gemini_api_key: Optional[str], +) -> Tuple[str, str]: + if mode == "chat": + _gemini_model_name = "models/{}".format(model) + endpoint = "generateContent" + if stream is True: + endpoint = "streamGenerateContent" + url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}&alt=sse".format( + _gemini_model_name, endpoint, gemini_api_key + ) + else: + url = ( + "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format( + _gemini_model_name, endpoint, gemini_api_key + ) + ) + elif mode == "embedding": + pass + return url, endpoint diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini_transformation.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/transformation.py similarity index 100% rename from litellm/llms/vertex_ai_and_google_ai_studio/gemini_transformation.py rename to litellm/llms/vertex_ai_and_google_ai_studio/gemini/transformation.py diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py similarity index 97% rename from litellm/llms/vertex_ai_and_google_ai_studio/vertex_and_google_ai_studio_gemini.py rename to litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py index 5392f253f8..d897f5bfbd 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py @@ -54,10 +54,16 @@ from litellm.types.llms.vertex_ai import ( from litellm.types.utils import GenericStreamingChunk from litellm.utils import CustomStreamWrapper, ModelResponse, Usage -from ..base import BaseLLM -from .common_utils import VertexAIError, get_supports_system_message -from .context_caching.vertex_ai_context_caching import ContextCachingEndpoints -from .gemini_transformation import transform_system_message +from ...base import BaseLLM +from ..common_utils import ( + VertexAIError, + _get_gemini_url, + _get_vertex_url, + all_gemini_url_modes, + get_supports_system_message, +) +from ..context_caching.vertex_ai_context_caching import ContextCachingEndpoints +from .transformation import transform_system_message context_caching_endpoints = ContextCachingEndpoints() @@ -309,6 +315,7 @@ class GoogleAIStudioGeminiConfig: # key diff from VertexAI - 'frequency_penalty "n", "stop", ] + def _map_function(self, value: List[dict]) -> List[Tools]: gtool_func_declarations = [] googleSearchRetrieval: Optional[dict] = None @@ -1164,6 +1171,7 @@ class VertexLLM(BaseLLM): custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], api_base: Optional[str], should_use_v1beta1_features: Optional[bool] = False, + mode: all_gemini_url_modes = "chat", ) -> Tuple[Optional[str], str]: """ Internal function. Returns the token and url for the call. @@ -1174,18 +1182,13 @@ class VertexLLM(BaseLLM): token, url """ if custom_llm_provider == "gemini": - _gemini_model_name = "models/{}".format(model) auth_header = None - endpoint = "generateContent" - if stream is True: - endpoint = "streamGenerateContent" - url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}&alt=sse".format( - _gemini_model_name, endpoint, gemini_api_key - ) - else: - url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format( - _gemini_model_name, endpoint, gemini_api_key - ) + url, endpoint = _get_gemini_url( + mode=mode, + model=model, + stream=stream, + gemini_api_key=gemini_api_key, + ) else: auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project @@ -1193,23 +1196,17 @@ class VertexLLM(BaseLLM): vertex_location = self.get_vertex_region(vertex_region=vertex_location) ### SET RUNTIME ENDPOINT ### - version = "v1beta1" if should_use_v1beta1_features is True else "v1" - endpoint = "generateContent" - litellm.utils.print_verbose("vertex_project - {}".format(vertex_project)) - if stream is True: - endpoint = "streamGenerateContent" - url = f"https://{vertex_location}-aiplatform.googleapis.com/{version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}?alt=sse" - else: - url = f"https://{vertex_location}-aiplatform.googleapis.com/{version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" - - # if model is only numeric chars then it's a fine tuned gemini model - # model = 4965075652664360960 - # send to this url: url = f"https://{vertex_location}-aiplatform.googleapis.com/{version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}" - if model.isdigit(): - # It's a fine-tuned Gemini model - url = f"https://{vertex_location}-aiplatform.googleapis.com/{version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}" - if stream is True: - url += "?alt=sse" + version: Literal["v1beta1", "v1"] = ( + "v1beta1" if should_use_v1beta1_features is True else "v1" + ) + url, endpoint = _get_vertex_url( + mode=mode, + model=model, + stream=stream, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_api_version=version, + ) if ( api_base is not None @@ -1793,8 +1790,10 @@ class VertexLLM(BaseLLM): input: Union[list, str], print_verbose, model_response: litellm.EmbeddingResponse, + custom_llm_provider: Literal["gemini", "vertex_ai"], optional_params: dict, api_key: Optional[str] = None, + api_base: Optional[str] = None, logging_obj=None, encoding=None, vertex_project=None, @@ -1804,6 +1803,17 @@ class VertexLLM(BaseLLM): timeout=300, client=None, ): + auth_header, url = self._get_token_and_url( + model=model, + gemini_api_key=api_key, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_credentials=vertex_credentials, + stream=None, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + should_use_v1beta1_features=False, + ) if client is None: _params = {} diff --git a/litellm/main.py b/litellm/main.py index a77a03522a..b83a583f4a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -126,12 +126,12 @@ from .llms.vertex_ai_and_google_ai_studio import ( vertex_ai_anthropic, vertex_ai_non_gemini, ) +from .llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( + VertexLLM, +) from .llms.vertex_ai_and_google_ai_studio.vertex_ai_partner_models.main import ( VertexAIPartnerModels, ) -from .llms.vertex_ai_and_google_ai_studio.vertex_and_google_ai_studio_gemini import ( - VertexLLM, -) from .llms.watsonx import IBMWatsonXAI from .types.llms.openai import HttpxBinaryResponseContent from .types.utils import ( @@ -3568,6 +3568,7 @@ def embedding( vertex_credentials=vertex_credentials, aembedding=aembedding, print_verbose=print_verbose, + custom_llm_provider="vertex_ai", ) else: response = vertex_ai_non_gemini.embedding( diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 53005fac09..f542d18d16 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -28,7 +28,7 @@ from litellm import ( completion_cost, embedding, ) -from litellm.llms.vertex_ai_and_google_ai_studio.vertex_and_google_ai_studio_gemini import ( +from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( _gemini_convert_messages_with_history, ) from litellm.tests.test_streaming import streaming_format_tests @@ -2065,7 +2065,7 @@ def test_prompt_factory_nested(): def test_get_token_url(): - from litellm.llms.vertex_ai_and_google_ai_studio.vertex_and_google_ai_studio_gemini import ( + from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( VertexLLM, ) @@ -2087,7 +2087,7 @@ def test_get_token_url(): vertex_credentials=vertex_credentials, gemini_api_key="", custom_llm_provider="vertex_ai_beta", - should_use_v1beta1_features=should_use_v1beta1_features, + should_use_vertex_v1beta1_features=should_use_v1beta1_features, api_base=None, model="", stream=False, @@ -2107,7 +2107,7 @@ def test_get_token_url(): vertex_credentials=vertex_credentials, gemini_api_key="", custom_llm_provider="vertex_ai_beta", - should_use_v1beta1_features=should_use_v1beta1_features, + should_use_vertex_v1beta1_features=should_use_v1beta1_features, api_base=None, model="", stream=False, From 77e6da78a1cd613df8f07e8d0630dcf34ed52f4c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 27 Aug 2024 17:35:56 -0700 Subject: [PATCH 02/16] fix: initial commit --- litellm/__init__.py | 5 +- litellm/llms/fine_tuning_apis/vertex_ai.py | 2 +- litellm/llms/text_to_speech/vertex_ai.py | 2 +- .../common_utils.py | 10 +- .../context_caching/transformation.py | 6 +- .../gemini/embeddings_handler.py | 121 ++++++++++++++++++ .../gemini/embeddings_transformation.py | 5 + .../vertex_and_google_ai_studio_gemini.py | 36 ++---- .../vertex_ai_anthropic.py | 4 +- litellm/main.py | 21 +++ litellm/tests/test_embedding.py | 16 +++ 11 files changed, 192 insertions(+), 36 deletions(-) create mode 100644 litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py create mode 100644 litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index a627061cfe..591d5873cf 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -848,7 +848,7 @@ from .llms.gemini import GeminiConfig from .llms.nlp_cloud import NLPCloudConfig from .llms.aleph_alpha import AlephAlphaConfig from .llms.petals import PetalsConfig -from .llms.vertex_ai_and_google_ai_studio.vertex_and_google_ai_studio_gemini import ( +from .llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, GoogleAIStudioGeminiConfig, VertexAIConfig, @@ -862,9 +862,6 @@ from .llms.vertex_ai_and_google_ai_studio.vertex_ai_anthropic import ( from .llms.vertex_ai_and_google_ai_studio.vertex_ai_partner_models.llama3.transformation import ( VertexAILlama3Config, ) -from .llms.vertex_ai_and_google_ai_studio.vertex_ai_partner_models.ai21.transformation import ( - VertexAIAi21Config, -) from .llms.sagemaker.sagemaker import SagemakerConfig from .llms.ollama import OllamaConfig from .llms.ollama_chat import OllamaChatConfig diff --git a/litellm/llms/fine_tuning_apis/vertex_ai.py b/litellm/llms/fine_tuning_apis/vertex_ai.py index e87a9bf3c4..618cf510af 100644 --- a/litellm/llms/fine_tuning_apis/vertex_ai.py +++ b/litellm/llms/fine_tuning_apis/vertex_ai.py @@ -8,7 +8,7 @@ from openai.types.fine_tuning.fine_tuning_job import FineTuningJob, Hyperparamet from litellm._logging import verbose_logger from litellm.llms.base import BaseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.llms.vertex_ai_and_google_ai_studio.vertex_and_google_ai_studio_gemini import ( +from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( VertexLLM, ) from litellm.types.llms.openai import FineTuningJobCreate diff --git a/litellm/llms/text_to_speech/vertex_ai.py b/litellm/llms/text_to_speech/vertex_ai.py index b9fca53250..0aac32eb50 100644 --- a/litellm/llms/text_to_speech/vertex_ai.py +++ b/litellm/llms/text_to_speech/vertex_ai.py @@ -13,7 +13,7 @@ from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, ) from litellm.llms.openai import HttpxBinaryResponseContent -from litellm.llms.vertex_ai_and_google_ai_studio.vertex_and_google_ai_studio_gemini import ( +from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( VertexLLM, ) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py b/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py index 7e2f9b29d0..d8607b4a8a 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py @@ -69,6 +69,9 @@ def _get_vertex_url( url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}" if stream is True: url += "?alt=sse" + elif mode == "embedding": + endpoint = "predict" + url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" return url, endpoint @@ -79,8 +82,8 @@ def _get_gemini_url( stream: Optional[bool], gemini_api_key: Optional[str], ) -> Tuple[str, str]: + _gemini_model_name = "models/{}".format(model) if mode == "chat": - _gemini_model_name = "models/{}".format(model) endpoint = "generateContent" if stream is True: endpoint = "streamGenerateContent" @@ -94,5 +97,8 @@ def _get_gemini_url( ) ) elif mode == "embedding": - pass + endpoint = "embedContent" + url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format( + _gemini_model_name, endpoint, gemini_api_key + ) return url, endpoint diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/context_caching/transformation.py b/litellm/llms/vertex_ai_and_google_ai_studio/context_caching/transformation.py index 944ae00bc9..d394cafd3a 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/context_caching/transformation.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/context_caching/transformation.py @@ -11,8 +11,10 @@ from litellm.types.llms.vertex_ai import CachedContentRequestBody, SystemInstruc from litellm.utils import is_cached_message from ..common_utils import VertexAIError, get_supports_system_message -from ..gemini_transformation import transform_system_message -from ..vertex_and_google_ai_studio_gemini import _gemini_convert_messages_with_history +from ..gemini.transformation import transform_system_message +from ..gemini.vertex_and_google_ai_studio_gemini import ( + _gemini_convert_messages_with_history, +) def separate_cached_messages( diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py new file mode 100644 index 0000000000..bc0d4ac16f --- /dev/null +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py @@ -0,0 +1,121 @@ +""" +Google AI Studio Embeddings Endpoint +""" + +import json +from typing import Literal, Optional, Union + +import httpx + +import litellm +from litellm import EmbeddingResponse +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +from .vertex_and_google_ai_studio_gemini import VertexLLM + + +class GoogleEmbeddings(VertexLLM): + def text_embeddings( + self, + model: str, + input: Union[list, str], + print_verbose, + model_response: EmbeddingResponse, + custom_llm_provider: Literal["gemini", "vertex_ai"], + optional_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + logging_obj=None, + encoding=None, + vertex_project=None, + vertex_location=None, + vertex_credentials=None, + aembedding=False, + timeout=300, + client=None, + ) -> EmbeddingResponse: + return model_response + auth_header, url = self._get_token_and_url( + model=model, + gemini_api_key=api_key, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_credentials=vertex_credentials, + stream=None, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + should_use_v1beta1_features=False, + mode="embedding", + ) + + if client is None: + _params = {} + if timeout is not None: + if isinstance(timeout, float) or isinstance(timeout, int): + _httpx_timeout = httpx.Timeout(timeout) + _params["timeout"] = _httpx_timeout + else: + _params["timeout"] = httpx.Timeout(timeout=600.0, connect=5.0) + + sync_handler: HTTPHandler = HTTPHandler(**_params) # type: ignore + else: + sync_handler = client # type: ignore + + optional_params = optional_params or {} + + # request_data = VertexMultimodalEmbeddingRequest() + + # if "instances" in optional_params: + # request_data["instances"] = optional_params["instances"] + # elif isinstance(input, list): + # request_data["instances"] = input + # else: + # # construct instances + # vertex_request_instance = Instance(**optional_params) + + # if isinstance(input, str): + # vertex_request_instance["text"] = input + + # request_data["instances"] = [vertex_request_instance] + + # headers = { + # "Content-Type": "application/json; charset=utf-8", + # "Authorization": f"Bearer {auth_header}", + # } + + # ## LOGGING + # logging_obj.pre_call( + # input=input, + # api_key="", + # additional_args={ + # "complete_input_dict": request_data, + # "api_base": url, + # "headers": headers, + # }, + # ) + + # if aembedding is True: + # pass + + # response = sync_handler.post( + # url=url, + # headers=headers, + # data=json.dumps(request_data), + # ) + + # if response.status_code != 200: + # raise Exception(f"Error: {response.status_code} {response.text}") + + # _json_response = response.json() + # if "predictions" not in _json_response: + # raise litellm.InternalServerError( + # message=f"embedding response does not contain 'predictions', got {_json_response}", + # llm_provider="vertex_ai", + # model=model, + # ) + # _predictions = _json_response["predictions"] + + # model_response.data = _predictions + # model_response.model = model + + # return model_response diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py new file mode 100644 index 0000000000..2e3d156f51 --- /dev/null +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py @@ -0,0 +1,5 @@ +""" +Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /embedContent format. + +Why separate file? Make it easy to see how transformation works +""" diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py index d897f5bfbd..819a94cb0c 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py @@ -1813,6 +1813,7 @@ class VertexLLM(BaseLLM): custom_llm_provider=custom_llm_provider, api_base=api_base, should_use_v1beta1_features=False, + mode="embedding", ) if client is None: @@ -1828,11 +1829,6 @@ class VertexLLM(BaseLLM): else: sync_handler = client # type: ignore - url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:predict" - - auth_header, _ = self._ensure_access_token( - credentials=vertex_credentials, project_id=vertex_project - ) optional_params = optional_params or {} request_data = VertexMultimodalEmbeddingRequest() @@ -1850,30 +1846,22 @@ class VertexLLM(BaseLLM): request_data["instances"] = [vertex_request_instance] - request_str = f"\n curl -X POST \\\n -H \"Authorization: Bearer {auth_header[:10] + 'XXXXXXXXXX'}\" \\\n -H \"Content-Type: application/json; charset=utf-8\" \\\n -d {request_data} \\\n \"{url}\"" - logging_obj.pre_call( - input=[], - api_key=None, - additional_args={ - "complete_input_dict": optional_params, - "request_str": request_str, - }, - ) - - logging_obj.pre_call( - input=[], - api_key=None, - additional_args={ - "complete_input_dict": optional_params, - "request_str": request_str, - }, - ) - headers = { "Content-Type": "application/json; charset=utf-8", "Authorization": f"Bearer {auth_header}", } + ## LOGGING + logging_obj.pre_call( + input=input, + api_key="", + additional_args={ + "complete_input_dict": request_data, + "api_base": url, + "headers": headers, + }, + ) + if aembedding is True: return self.async_multimodal_embedding( model=model, diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_anthropic.py b/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_anthropic.py index b13b87bc67..e85160a43c 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_anthropic.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_anthropic.py @@ -205,7 +205,7 @@ def get_vertex_client( vertex_credentials: Optional[str], ) -> Tuple[Any, Optional[str]]: args = locals() - from litellm.llms.vertex_ai_and_google_ai_studio.vertex_and_google_ai_studio_gemini import ( + from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( VertexLLM, ) @@ -270,7 +270,7 @@ def completion( from anthropic import AnthropicVertex from litellm.llms.anthropic import AnthropicChatCompletion - from litellm.llms.vertex_ai_and_google_ai_studio.vertex_and_google_ai_studio_gemini import ( + from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( VertexLLM, ) except: diff --git a/litellm/main.py b/litellm/main.py index b83a583f4a..8896f1faf1 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3134,6 +3134,7 @@ async def aembedding(*args, **kwargs) -> EmbeddingResponse: or custom_llm_provider == "fireworks_ai" or custom_llm_provider == "ollama" or custom_llm_provider == "vertex_ai" + or custom_llm_provider == "gemini" or custom_llm_provider == "databricks" or custom_llm_provider == "watsonx" or custom_llm_provider == "cohere" @@ -3528,6 +3529,26 @@ def embedding( client=client, aembedding=aembedding, ) + elif custom_llm_provider == "gemini": + + gemini_api_key = api_key or get_secret("GEMINI_API_KEY") or litellm.api_key + + response = vertex_chat_completion.multimodal_embedding( # type: ignore + model=model, + input=input, + encoding=encoding, + logging_obj=logging, + optional_params=optional_params, + model_response=EmbeddingResponse(), + vertex_project=None, + vertex_location=None, + vertex_credentials=None, + aembedding=aembedding, + print_verbose=print_verbose, + custom_llm_provider="gemini", + api_key=gemini_api_key, + ) + elif custom_llm_provider == "vertex_ai": vertex_ai_project = ( optional_params.pop("vertex_project", None) diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 31268395f1..2fbc70f024 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -686,6 +686,22 @@ async def test_triton_embeddings(): pytest.fail(f"Error occurred: {e}") +@pytest.mark.asyncio +async def test_gemini_embeddings(): + try: + litellm.set_verbose = True + response = await litellm.aembedding( + model="gemini/text-embedding-004", + input=["good morning from litellm"], + ) + print(f"response: {response}") + + # stubbed endpoint is setup to return this + assert response.data[0]["embedding"] == [0.1, 0.2] + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_databricks_embeddings(sync_mode): From 5b29ddd2a677d41b4cca43ac769b9e530931eb2d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 27 Aug 2024 18:14:56 -0700 Subject: [PATCH 03/16] fix(embeddings_handler.py): initial working commit for google ai studio text embeddings /embedContent endpoint --- .../gemini/embeddings_handler.py | 98 ++++++++++--------- .../gemini/embeddings_transformation.py | 22 +++++ litellm/main.py | 6 +- litellm/tests/test_embedding.py | 3 +- litellm/types/llms/openai.py | 4 + litellm/types/llms/vertex_ai.py | 26 +++++ 6 files changed, 111 insertions(+), 48 deletions(-) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py index bc0d4ac16f..98cebfc313 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py @@ -10,7 +10,14 @@ import httpx import litellm from litellm import EmbeddingResponse from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.types.llms.vertex_ai import ( + VertexAITextEmbeddingsRequestBody, + VertexAITextEmbeddingsResponseObject, +) +from litellm.types.utils import Embedding +from litellm.utils import get_formatted_prompt +from .embeddings_transformation import transform_openai_input_gemini_content from .vertex_and_google_ai_studio_gemini import VertexLLM @@ -34,7 +41,7 @@ class GoogleEmbeddings(VertexLLM): timeout=300, client=None, ) -> EmbeddingResponse: - return model_response + auth_header, url = self._get_token_and_url( model=model, gemini_api_key=api_key, @@ -63,59 +70,58 @@ class GoogleEmbeddings(VertexLLM): optional_params = optional_params or {} - # request_data = VertexMultimodalEmbeddingRequest() + ### TRANSFORMATION ### + content = transform_openai_input_gemini_content(input=input) - # if "instances" in optional_params: - # request_data["instances"] = optional_params["instances"] - # elif isinstance(input, list): - # request_data["instances"] = input - # else: - # # construct instances - # vertex_request_instance = Instance(**optional_params) + request_data: VertexAITextEmbeddingsRequestBody = { + "content": content, + **optional_params, + } - # if isinstance(input, str): - # vertex_request_instance["text"] = input + headers = { + "Content-Type": "application/json; charset=utf-8", + } - # request_data["instances"] = [vertex_request_instance] + ## LOGGING + logging_obj.pre_call( + input=input, + api_key="", + additional_args={ + "complete_input_dict": request_data, + "api_base": url, + "headers": headers, + }, + ) - # headers = { - # "Content-Type": "application/json; charset=utf-8", - # "Authorization": f"Bearer {auth_header}", - # } + if aembedding is True: + pass - # ## LOGGING - # logging_obj.pre_call( - # input=input, - # api_key="", - # additional_args={ - # "complete_input_dict": request_data, - # "api_base": url, - # "headers": headers, - # }, - # ) + response = sync_handler.post( + url=url, + headers=headers, + data=json.dumps(request_data), + ) - # if aembedding is True: - # pass + if response.status_code != 200: + raise Exception(f"Error: {response.status_code} {response.text}") - # response = sync_handler.post( - # url=url, - # headers=headers, - # data=json.dumps(request_data), - # ) + _json_response = response.json() + _predictions = VertexAITextEmbeddingsResponseObject(**_json_response) # type: ignore - # if response.status_code != 200: - # raise Exception(f"Error: {response.status_code} {response.text}") + model_response.data = [ + Embedding( + embedding=_predictions["embedding"]["values"], + index=0, + object="embedding", + ) + ] - # _json_response = response.json() - # if "predictions" not in _json_response: - # raise litellm.InternalServerError( - # message=f"embedding response does not contain 'predictions', got {_json_response}", - # llm_provider="vertex_ai", - # model=model, - # ) - # _predictions = _json_response["predictions"] + model_response.model = model - # model_response.data = _predictions - # model_response.model = model + input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") + prompt_tokens = litellm.token_counter(model=model, text=input_text) + model_response.usage = litellm.Usage( + prompt_tokens=prompt_tokens, total_tokens=prompt_tokens + ) - # return model_response + return model_response diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py index 2e3d156f51..dd5abfa380 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py @@ -3,3 +3,25 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /embe Why separate file? Make it easy to see how transformation works """ + +from typing import List + +from litellm.types.llms.openai import EmbeddingInput +from litellm.types.llms.vertex_ai import ContentType, PartType + +from ..common_utils import VertexAIError + + +def transform_openai_input_gemini_content(input: EmbeddingInput) -> ContentType: + """ + The content to embed. Only the parts.text fields will be counted. + """ + if isinstance(input, str): + return ContentType(parts=[PartType(text=input)]) + elif isinstance(input, list) and len(input) == 1: + return ContentType(parts=[PartType(text=input[0])]) + else: + raise VertexAIError( + status_code=422, + message="/embedContent only generates a single text embedding vector. File an issue, to add support for /batchEmbedContent - https://github.com/BerriAI/litellm/issues", + ) diff --git a/litellm/main.py b/litellm/main.py index 8896f1faf1..e9a3d2898b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -126,6 +126,9 @@ from .llms.vertex_ai_and_google_ai_studio import ( vertex_ai_anthropic, vertex_ai_non_gemini, ) +from .llms.vertex_ai_and_google_ai_studio.gemini.embeddings_handler import ( + GoogleEmbeddings, +) from .llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( VertexLLM, ) @@ -172,6 +175,7 @@ triton_chat_completions = TritonChatCompletion() bedrock_chat_completion = BedrockLLM() bedrock_converse_chat_completion = BedrockConverseLLM() vertex_chat_completion = VertexLLM() +google_embeddings = GoogleEmbeddings() vertex_partner_models_chat_completion = VertexAIPartnerModels() vertex_text_to_speech = VertexTextToSpeechAPI() watsonxai = IBMWatsonXAI() @@ -3533,7 +3537,7 @@ def embedding( gemini_api_key = api_key or get_secret("GEMINI_API_KEY") or litellm.api_key - response = vertex_chat_completion.multimodal_embedding( # type: ignore + response = google_embeddings.text_embeddings( # type: ignore model=model, input=input, encoding=encoding, diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 2fbc70f024..c318264d45 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -697,7 +697,8 @@ async def test_gemini_embeddings(): print(f"response: {response}") # stubbed endpoint is setup to return this - assert response.data[0]["embedding"] == [0.1, 0.2] + assert isinstance(response.data[0]["embedding"], list) + assert response.usage.prompt_tokens > 0 except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 470f72c5b6..138441a7eb 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -30,6 +30,7 @@ from openai.types.beta.threads.message import Message as OpenAIMessage from openai.types.beta.threads.message_content import MessageContent from openai.types.beta.threads.run import Run from openai.types.chat import ChatCompletionChunk +from openai.types.embedding import Embedding as OpenAIEmbedding from pydantic import BaseModel, Field from typing_extensions import Dict, Required, TypedDict, override @@ -47,6 +48,9 @@ FileTypes = Union[ ] +EmbeddingInput = Union[str, List[str]] + + class NotGiven: """ A sentinel singleton class used to distinguish omitted keyword arguments diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 90730d75fe..bacb4d2252 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -336,3 +336,29 @@ class VertexMultimodalEmbeddingRequest(TypedDict, total=False): class VertexAICachedContentResponseObject(TypedDict): name: str model: str + + +class TaskTypeEnum(Enum): + TASK_TYPE_UNSPECIFIED = "TASK_TYPE_UNSPECIFIED" + RETRIEVAL_QUERY = "RETRIEVAL_QUERY" + RETRIEVAL_DOCUMENT = "RETRIEVAL_DOCUMENT" + SEMANTIC_SIMILARITY = "SEMANTIC_SIMILARITY" + CLASSIFICATION = "CLASSIFICATION" + CLUSTERING = "CLUSTERING" + QUESTION_ANSWERING = "QUESTION_ANSWERING" + FACT_VERIFICATION = "FACT_VERIFICATION" + + +class VertexAITextEmbeddingsRequestBody(TypedDict, total=False): + content: Required[ContentType] + taskType: TaskTypeEnum + title: str + outputDimensionality: int + + +class ContentEmbeddings(TypedDict): + values: List[int] + + +class VertexAITextEmbeddingsResponseObject(TypedDict): + embedding: ContentEmbeddings From bb42146ffe3995e689e453f323ec76a83d61b9d4 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 27 Aug 2024 18:31:57 -0700 Subject: [PATCH 04/16] feat(embeddings_handler.py): support async gemini embeddings --- .../gemini/embeddings_handler.py | 82 +++++++++++++++---- .../gemini/embeddings_transformation.py | 34 +++++++- litellm/tests/test_embedding.py | 17 ++-- 3 files changed, 109 insertions(+), 24 deletions(-) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py index 98cebfc313..2b26d6c04d 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py @@ -9,7 +9,8 @@ import httpx import litellm from litellm import EmbeddingResponse -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.llms.openai import EmbeddingInput from litellm.types.llms.vertex_ai import ( VertexAITextEmbeddingsRequestBody, VertexAITextEmbeddingsResponseObject, @@ -17,7 +18,10 @@ from litellm.types.llms.vertex_ai import ( from litellm.types.utils import Embedding from litellm.utils import get_formatted_prompt -from .embeddings_transformation import transform_openai_input_gemini_content +from .embeddings_transformation import ( + process_response, + transform_openai_input_gemini_content, +) from .vertex_and_google_ai_studio_gemini import VertexLLM @@ -94,7 +98,16 @@ class GoogleEmbeddings(VertexLLM): ) if aembedding is True: - pass + return self.async_text_embeddings( # type: ignore + model=model, + api_base=api_base, + url=url, + data=request_data, + model_response=model_response, + timeout=timeout, + headers=headers, + input=input, + ) response = sync_handler.post( url=url, @@ -108,20 +121,53 @@ class GoogleEmbeddings(VertexLLM): _json_response = response.json() _predictions = VertexAITextEmbeddingsResponseObject(**_json_response) # type: ignore - model_response.data = [ - Embedding( - embedding=_predictions["embedding"]["values"], - index=0, - object="embedding", - ) - ] - - model_response.model = model - - input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") - prompt_tokens = litellm.token_counter(model=model, text=input_text) - model_response.usage = litellm.Usage( - prompt_tokens=prompt_tokens, total_tokens=prompt_tokens + return process_response( + model=model, + model_response=model_response, + _predictions=_predictions, + input=input, ) - return model_response + async def async_text_embeddings( + self, + model: str, + api_base: Optional[str], + url: str, + data: VertexAITextEmbeddingsRequestBody, + model_response: EmbeddingResponse, + input: EmbeddingInput, + timeout: Optional[Union[float, httpx.Timeout]], + headers={}, + client: Optional[AsyncHTTPHandler] = None, + ) -> EmbeddingResponse: + if client is None: + _params = {} + if timeout is not None: + if isinstance(timeout, float) or isinstance(timeout, int): + _httpx_timeout = httpx.Timeout(timeout) + _params["timeout"] = _httpx_timeout + else: + _params["timeout"] = httpx.Timeout(timeout=600.0, connect=5.0) + + async_handler: AsyncHTTPHandler = AsyncHTTPHandler(**_params) # type: ignore + else: + async_handler = client # type: ignore + + response = await async_handler.post( + url=url, + headers=headers, + data=json.dumps(data), + ) + + if response.status_code != 200: + raise Exception(f"Error: {response.status_code} {response.text}") + + _json_response = response.json() + _predictions = VertexAITextEmbeddingsResponseObject(**_json_response) # type: ignore + + return process_response( + model=model, + model_response=model_response, + _predictions=_predictions, + input=input, + ) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py index dd5abfa380..198811578b 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py @@ -6,8 +6,15 @@ Why separate file? Make it easy to see how transformation works from typing import List +from litellm import EmbeddingResponse from litellm.types.llms.openai import EmbeddingInput -from litellm.types.llms.vertex_ai import ContentType, PartType +from litellm.types.llms.vertex_ai import ( + ContentType, + PartType, + VertexAITextEmbeddingsResponseObject, +) +from litellm.types.utils import Embedding, Usage +from litellm.utils import get_formatted_prompt, token_counter from ..common_utils import VertexAIError @@ -25,3 +32,28 @@ def transform_openai_input_gemini_content(input: EmbeddingInput) -> ContentType: status_code=422, message="/embedContent only generates a single text embedding vector. File an issue, to add support for /batchEmbedContent - https://github.com/BerriAI/litellm/issues", ) + + +def process_response( + input: EmbeddingInput, + model_response: EmbeddingResponse, + model: str, + _predictions: VertexAITextEmbeddingsResponseObject, +) -> EmbeddingResponse: + model_response.data = [ + Embedding( + embedding=_predictions["embedding"]["values"], + index=0, + object="embedding", + ) + ] + + model_response.model = model + + input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") + prompt_tokens = token_counter(model=model, text=input_text) + model_response.usage = Usage( + prompt_tokens=prompt_tokens, total_tokens=prompt_tokens + ) + + return model_response diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index c318264d45..a17b22f489 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -686,14 +686,21 @@ async def test_triton_embeddings(): pytest.fail(f"Error occurred: {e}") +@pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_gemini_embeddings(): +async def test_gemini_embeddings(sync_mode): try: litellm.set_verbose = True - response = await litellm.aembedding( - model="gemini/text-embedding-004", - input=["good morning from litellm"], - ) + if sync_mode: + response = litellm.embedding( + model="gemini/text-embedding-004", + input=["good morning from litellm"], + ) + else: + response = await litellm.aembedding( + model="gemini/text-embedding-004", + input=["good morning from litellm"], + ) print(f"response: {response}") # stubbed endpoint is setup to return this From a6ce27ca290fc33905496af878259a9af6f3c977 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 27 Aug 2024 19:23:50 -0700 Subject: [PATCH 05/16] feat(batch_embed_content_transformation.py): support google ai studio /batchEmbedContent endpoint Allows for multiple strings to be given for embedding --- .../common_utils.py | 8 +- .../embeddings/batch_embed_content_handler.py | 167 ++++++++++++++++++ .../batch_embed_content_transformation.py | 68 +++++++ .../embed_content_handler.py} | 11 +- .../embed_content_transformation.py} | 14 +- litellm/main.py | 53 ++++-- litellm/tests/test_embedding.py | 9 +- litellm/types/llms/vertex_ai.py | 12 ++ 8 files changed, 303 insertions(+), 39 deletions(-) create mode 100644 litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_handler.py create mode 100644 litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_transformation.py rename litellm/llms/vertex_ai_and_google_ai_studio/{gemini/embeddings_handler.py => embeddings/embed_content_handler.py} (94%) rename litellm/llms/vertex_ai_and_google_ai_studio/{gemini/embeddings_transformation.py => embeddings/embed_content_transformation.py} (69%) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py b/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py index d8607b4a8a..2fef2233c0 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py @@ -41,7 +41,7 @@ def get_supports_system_message( from typing import Literal, Optional -all_gemini_url_modes = Literal["chat", "embedding"] +all_gemini_url_modes = Literal["chat", "embedding", "batch_embedding"] def _get_vertex_url( @@ -101,4 +101,10 @@ def _get_gemini_url( url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format( _gemini_model_name, endpoint, gemini_api_key ) + elif mode == "batch_embedding": + endpoint = "batchEmbedContents" + url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format( + _gemini_model_name, endpoint, gemini_api_key + ) + return url, endpoint diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_handler.py new file mode 100644 index 0000000000..9535c5594b --- /dev/null +++ b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_handler.py @@ -0,0 +1,167 @@ +""" +Google AI Studio /batchEmbedContents Embeddings Endpoint +""" + +import json +from typing import List, Literal, Optional, Union + +import httpx + +from litellm import EmbeddingResponse +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.llms.openai import EmbeddingInput +from litellm.types.llms.vertex_ai import ( + VertexAIBatchEmbeddingsRequestBody, + VertexAIBatchEmbeddingsResponseObject, +) + +from ..gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from .batch_embed_content_transformation import ( + process_response, + transform_openai_input_gemini_content, +) + + +class GoogleBatchEmbeddings(VertexLLM): + def batch_embeddings( + self, + model: str, + input: List[str], + print_verbose, + model_response: EmbeddingResponse, + custom_llm_provider: Literal["gemini", "vertex_ai"], + optional_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + logging_obj=None, + encoding=None, + vertex_project=None, + vertex_location=None, + vertex_credentials=None, + aembedding=False, + timeout=300, + client=None, + ) -> EmbeddingResponse: + + auth_header, url = self._get_token_and_url( + model=model, + gemini_api_key=api_key, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_credentials=vertex_credentials, + stream=None, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + should_use_v1beta1_features=False, + mode="batch_embedding", + ) + + if client is None: + _params = {} + if timeout is not None: + if isinstance(timeout, float) or isinstance(timeout, int): + _httpx_timeout = httpx.Timeout(timeout) + _params["timeout"] = _httpx_timeout + else: + _params["timeout"] = httpx.Timeout(timeout=600.0, connect=5.0) + + sync_handler: HTTPHandler = HTTPHandler(**_params) # type: ignore + else: + sync_handler = client # type: ignore + + optional_params = optional_params or {} + + ### TRANSFORMATION ### + request_data = transform_openai_input_gemini_content( + input=input, model=model, optional_params=optional_params + ) + + headers = { + "Content-Type": "application/json; charset=utf-8", + } + + ## LOGGING + logging_obj.pre_call( + input=input, + api_key="", + additional_args={ + "complete_input_dict": request_data, + "api_base": url, + "headers": headers, + }, + ) + + if aembedding is True: + return self.async_batch_embeddings( # type: ignore + model=model, + api_base=api_base, + url=url, + data=request_data, + model_response=model_response, + timeout=timeout, + headers=headers, + input=input, + ) + + response = sync_handler.post( + url=url, + headers=headers, + data=json.dumps(request_data), + ) + + if response.status_code != 200: + raise Exception(f"Error: {response.status_code} {response.text}") + + _json_response = response.json() + _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore + + return process_response( + model=model, + model_response=model_response, + _predictions=_predictions, + input=input, + ) + + async def async_batch_embeddings( + self, + model: str, + api_base: Optional[str], + url: str, + data: VertexAIBatchEmbeddingsRequestBody, + model_response: EmbeddingResponse, + input: EmbeddingInput, + timeout: Optional[Union[float, httpx.Timeout]], + headers={}, + client: Optional[AsyncHTTPHandler] = None, + ) -> EmbeddingResponse: + if client is None: + _params = {} + if timeout is not None: + if isinstance(timeout, float) or isinstance(timeout, int): + _httpx_timeout = httpx.Timeout(timeout) + _params["timeout"] = _httpx_timeout + else: + _params["timeout"] = httpx.Timeout(timeout=600.0, connect=5.0) + + async_handler: AsyncHTTPHandler = AsyncHTTPHandler(**_params) # type: ignore + else: + async_handler = client # type: ignore + + response = await async_handler.post( + url=url, + headers=headers, + data=json.dumps(data), + ) + + if response.status_code != 200: + raise Exception(f"Error: {response.status_code} {response.text}") + + _json_response = response.json() + _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore + + return process_response( + model=model, + model_response=model_response, + _predictions=_predictions, + input=input, + ) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_transformation.py new file mode 100644 index 0000000000..e17c79991e --- /dev/null +++ b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_transformation.py @@ -0,0 +1,68 @@ +""" +Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batchEmbedContents format. + +Why separate file? Make it easy to see how transformation works +""" + +from typing import List + +from litellm import EmbeddingResponse +from litellm.types.llms.openai import EmbeddingInput +from litellm.types.llms.vertex_ai import ( + ContentType, + EmbedContentRequest, + PartType, + VertexAIBatchEmbeddingsRequestBody, + VertexAIBatchEmbeddingsResponseObject, +) +from litellm.types.utils import Embedding, Usage +from litellm.utils import get_formatted_prompt, token_counter + +from ..common_utils import VertexAIError + + +def transform_openai_input_gemini_content( + input: List[str], model: str, optional_params: dict +) -> VertexAIBatchEmbeddingsRequestBody: + """ + The content to embed. Only the parts.text fields will be counted. + """ + gemini_model_name = "models/{}".format(model) + requests: List[EmbedContentRequest] = [] + for i in input: + request = EmbedContentRequest( + model=gemini_model_name, + content=ContentType(parts=[PartType(text=i)]), + **optional_params + ) + requests.append(request) + + return VertexAIBatchEmbeddingsRequestBody(requests=requests) + + +def process_response( + input: EmbeddingInput, + model_response: EmbeddingResponse, + model: str, + _predictions: VertexAIBatchEmbeddingsResponseObject, +) -> EmbeddingResponse: + + openai_embeddings: List[Embedding] = [] + for embedding in _predictions["embeddings"]: + openai_embedding = Embedding( + embedding=embedding["values"], + index=0, + object="embedding", + ) + openai_embeddings.append(openai_embedding) + + model_response.data = openai_embeddings + model_response.model = model + + input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") + prompt_tokens = token_counter(model=model, text=input_text) + model_response.usage = Usage( + prompt_tokens=prompt_tokens, total_tokens=prompt_tokens + ) + + return model_response diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_handler.py similarity index 94% rename from litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py rename to litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_handler.py index 2b26d6c04d..7c1d474352 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_handler.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_handler.py @@ -1,5 +1,5 @@ """ -Google AI Studio Embeddings Endpoint +Google AI Studio /embedContent Embeddings Endpoint """ import json @@ -7,7 +7,6 @@ from typing import Literal, Optional, Union import httpx -import litellm from litellm import EmbeddingResponse from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.openai import EmbeddingInput @@ -15,21 +14,19 @@ from litellm.types.llms.vertex_ai import ( VertexAITextEmbeddingsRequestBody, VertexAITextEmbeddingsResponseObject, ) -from litellm.types.utils import Embedding -from litellm.utils import get_formatted_prompt -from .embeddings_transformation import ( +from ..gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from .embed_content_transformation import ( process_response, transform_openai_input_gemini_content, ) -from .vertex_and_google_ai_studio_gemini import VertexLLM class GoogleEmbeddings(VertexLLM): def text_embeddings( self, model: str, - input: Union[list, str], + input: str, print_verbose, model_response: EmbeddingResponse, custom_llm_provider: Literal["gemini", "vertex_ai"], diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_transformation.py similarity index 69% rename from litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py rename to litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_transformation.py index 198811578b..bbda553175 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/embeddings_transformation.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_transformation.py @@ -4,8 +4,6 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /embe Why separate file? Make it easy to see how transformation works """ -from typing import List - from litellm import EmbeddingResponse from litellm.types.llms.openai import EmbeddingInput from litellm.types.llms.vertex_ai import ( @@ -19,19 +17,11 @@ from litellm.utils import get_formatted_prompt, token_counter from ..common_utils import VertexAIError -def transform_openai_input_gemini_content(input: EmbeddingInput) -> ContentType: +def transform_openai_input_gemini_content(input: str) -> ContentType: """ The content to embed. Only the parts.text fields will be counted. """ - if isinstance(input, str): - return ContentType(parts=[PartType(text=input)]) - elif isinstance(input, list) and len(input) == 1: - return ContentType(parts=[PartType(text=input[0])]) - else: - raise VertexAIError( - status_code=422, - message="/embedContent only generates a single text embedding vector. File an issue, to add support for /batchEmbedContent - https://github.com/BerriAI/litellm/issues", - ) + return ContentType(parts=[PartType(text=input)]) def process_response( diff --git a/litellm/main.py b/litellm/main.py index e9a3d2898b..bf1b0ede8c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -126,7 +126,10 @@ from .llms.vertex_ai_and_google_ai_studio import ( vertex_ai_anthropic, vertex_ai_non_gemini, ) -from .llms.vertex_ai_and_google_ai_studio.gemini.embeddings_handler import ( +from .llms.vertex_ai_and_google_ai_studio.embeddings.batch_embed_content_handler import ( + GoogleBatchEmbeddings, +) +from .llms.vertex_ai_and_google_ai_studio.embeddings.embed_content_handler import ( GoogleEmbeddings, ) from .llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( @@ -176,6 +179,7 @@ bedrock_chat_completion = BedrockLLM() bedrock_converse_chat_completion = BedrockConverseLLM() vertex_chat_completion = VertexLLM() google_embeddings = GoogleEmbeddings() +google_batch_embeddings = GoogleBatchEmbeddings() vertex_partner_models_chat_completion = VertexAIPartnerModels() vertex_text_to_speech = VertexTextToSpeechAPI() watsonxai = IBMWatsonXAI() @@ -3537,21 +3541,38 @@ def embedding( gemini_api_key = api_key or get_secret("GEMINI_API_KEY") or litellm.api_key - response = google_embeddings.text_embeddings( # type: ignore - model=model, - input=input, - encoding=encoding, - logging_obj=logging, - optional_params=optional_params, - model_response=EmbeddingResponse(), - vertex_project=None, - vertex_location=None, - vertex_credentials=None, - aembedding=aembedding, - print_verbose=print_verbose, - custom_llm_provider="gemini", - api_key=gemini_api_key, - ) + if isinstance(input, str): + response = google_embeddings.text_embeddings( # type: ignore + model=model, + input=input, + encoding=encoding, + logging_obj=logging, + optional_params=optional_params, + model_response=EmbeddingResponse(), + vertex_project=None, + vertex_location=None, + vertex_credentials=None, + aembedding=aembedding, + print_verbose=print_verbose, + custom_llm_provider="gemini", + api_key=gemini_api_key, + ) + else: + response = google_batch_embeddings.batch_embeddings( # type: ignore + model=model, + input=input, + encoding=encoding, + logging_obj=logging, + optional_params=optional_params, + model_response=EmbeddingResponse(), + vertex_project=None, + vertex_location=None, + vertex_credentials=None, + aembedding=aembedding, + print_verbose=print_verbose, + custom_llm_provider="gemini", + api_key=gemini_api_key, + ) elif custom_llm_provider == "vertex_ai": vertex_ai_project = ( diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index a17b22f489..667674b752 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -687,19 +687,22 @@ async def test_triton_embeddings(): @pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + "input", ["good morning from litellm", ["good morning from litellm"]] # +) @pytest.mark.asyncio -async def test_gemini_embeddings(sync_mode): +async def test_gemini_embeddings(sync_mode, input): try: litellm.set_verbose = True if sync_mode: response = litellm.embedding( model="gemini/text-embedding-004", - input=["good morning from litellm"], + input=input, ) else: response = await litellm.aembedding( model="gemini/text-embedding-004", - input=["good morning from litellm"], + input=input, ) print(f"response: {response}") diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index bacb4d2252..aeda867979 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -362,3 +362,15 @@ class ContentEmbeddings(TypedDict): class VertexAITextEmbeddingsResponseObject(TypedDict): embedding: ContentEmbeddings + + +class EmbedContentRequest(VertexAITextEmbeddingsRequestBody): + model: Required[str] + + +class VertexAIBatchEmbeddingsRequestBody(TypedDict, total=False): + requests: List[EmbedContentRequest] + + +class VertexAIBatchEmbeddingsResponseObject(TypedDict): + embeddings: List[ContentEmbeddings] From bd4f63eebf16dd2429dd13d1e41bbd332ffe0137 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 27 Aug 2024 19:35:03 -0700 Subject: [PATCH 06/16] fix(__init__.py): fix import --- litellm/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm/__init__.py b/litellm/__init__.py index 591d5873cf..a3727eb2e8 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -862,6 +862,10 @@ from .llms.vertex_ai_and_google_ai_studio.vertex_ai_anthropic import ( from .llms.vertex_ai_and_google_ai_studio.vertex_ai_partner_models.llama3.transformation import ( VertexAILlama3Config, ) +from .llms.vertex_ai_and_google_ai_studio.vertex_ai_partner_models.ai21.transformation import ( + VertexAIAi21Config, +) + from .llms.sagemaker.sagemaker import SagemakerConfig from .llms.ollama import OllamaConfig from .llms.ollama_chat import OllamaChatConfig From e1db58b8e5156cb79e4a03ad479e642cf8ad03ce Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 27 Aug 2024 21:46:05 -0700 Subject: [PATCH 07/16] fix(main.py): simplify to just use `/batchEmbedContent` --- .../embeddings/batch_embed_content_handler.py | 2 +- .../batch_embed_content_transformation.py | 14 +- .../embeddings/embed_content_handler.py | 170 ------------------ .../embed_content_transformation.py | 49 ----- litellm/main.py | 51 ++---- litellm/proxy/auth/rds_iam_token.py | 2 +- 6 files changed, 28 insertions(+), 260 deletions(-) delete mode 100644 litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_handler.py delete mode 100644 litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_transformation.py diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_handler.py index 9535c5594b..d05688deea 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_handler.py @@ -26,7 +26,7 @@ class GoogleBatchEmbeddings(VertexLLM): def batch_embeddings( self, model: str, - input: List[str], + input: EmbeddingInput, print_verbose, model_response: EmbeddingResponse, custom_llm_provider: Literal["gemini", "vertex_ai"], diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_transformation.py index e17c79991e..f1785e58f1 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/batch_embed_content_transformation.py @@ -22,20 +22,28 @@ from ..common_utils import VertexAIError def transform_openai_input_gemini_content( - input: List[str], model: str, optional_params: dict + input: EmbeddingInput, model: str, optional_params: dict ) -> VertexAIBatchEmbeddingsRequestBody: """ The content to embed. Only the parts.text fields will be counted. """ gemini_model_name = "models/{}".format(model) requests: List[EmbedContentRequest] = [] - for i in input: + if isinstance(input, str): request = EmbedContentRequest( model=gemini_model_name, - content=ContentType(parts=[PartType(text=i)]), + content=ContentType(parts=[PartType(text=input)]), **optional_params ) requests.append(request) + else: + for i in input: + request = EmbedContentRequest( + model=gemini_model_name, + content=ContentType(parts=[PartType(text=i)]), + **optional_params + ) + requests.append(request) return VertexAIBatchEmbeddingsRequestBody(requests=requests) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_handler.py deleted file mode 100644 index 7c1d474352..0000000000 --- a/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_handler.py +++ /dev/null @@ -1,170 +0,0 @@ -""" -Google AI Studio /embedContent Embeddings Endpoint -""" - -import json -from typing import Literal, Optional, Union - -import httpx - -from litellm import EmbeddingResponse -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.types.llms.openai import EmbeddingInput -from litellm.types.llms.vertex_ai import ( - VertexAITextEmbeddingsRequestBody, - VertexAITextEmbeddingsResponseObject, -) - -from ..gemini.vertex_and_google_ai_studio_gemini import VertexLLM -from .embed_content_transformation import ( - process_response, - transform_openai_input_gemini_content, -) - - -class GoogleEmbeddings(VertexLLM): - def text_embeddings( - self, - model: str, - input: str, - print_verbose, - model_response: EmbeddingResponse, - custom_llm_provider: Literal["gemini", "vertex_ai"], - optional_params: dict, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - logging_obj=None, - encoding=None, - vertex_project=None, - vertex_location=None, - vertex_credentials=None, - aembedding=False, - timeout=300, - client=None, - ) -> EmbeddingResponse: - - auth_header, url = self._get_token_and_url( - model=model, - gemini_api_key=api_key, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_credentials=vertex_credentials, - stream=None, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - should_use_v1beta1_features=False, - mode="embedding", - ) - - if client is None: - _params = {} - if timeout is not None: - if isinstance(timeout, float) or isinstance(timeout, int): - _httpx_timeout = httpx.Timeout(timeout) - _params["timeout"] = _httpx_timeout - else: - _params["timeout"] = httpx.Timeout(timeout=600.0, connect=5.0) - - sync_handler: HTTPHandler = HTTPHandler(**_params) # type: ignore - else: - sync_handler = client # type: ignore - - optional_params = optional_params or {} - - ### TRANSFORMATION ### - content = transform_openai_input_gemini_content(input=input) - - request_data: VertexAITextEmbeddingsRequestBody = { - "content": content, - **optional_params, - } - - headers = { - "Content-Type": "application/json; charset=utf-8", - } - - ## LOGGING - logging_obj.pre_call( - input=input, - api_key="", - additional_args={ - "complete_input_dict": request_data, - "api_base": url, - "headers": headers, - }, - ) - - if aembedding is True: - return self.async_text_embeddings( # type: ignore - model=model, - api_base=api_base, - url=url, - data=request_data, - model_response=model_response, - timeout=timeout, - headers=headers, - input=input, - ) - - response = sync_handler.post( - url=url, - headers=headers, - data=json.dumps(request_data), - ) - - if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") - - _json_response = response.json() - _predictions = VertexAITextEmbeddingsResponseObject(**_json_response) # type: ignore - - return process_response( - model=model, - model_response=model_response, - _predictions=_predictions, - input=input, - ) - - async def async_text_embeddings( - self, - model: str, - api_base: Optional[str], - url: str, - data: VertexAITextEmbeddingsRequestBody, - model_response: EmbeddingResponse, - input: EmbeddingInput, - timeout: Optional[Union[float, httpx.Timeout]], - headers={}, - client: Optional[AsyncHTTPHandler] = None, - ) -> EmbeddingResponse: - if client is None: - _params = {} - if timeout is not None: - if isinstance(timeout, float) or isinstance(timeout, int): - _httpx_timeout = httpx.Timeout(timeout) - _params["timeout"] = _httpx_timeout - else: - _params["timeout"] = httpx.Timeout(timeout=600.0, connect=5.0) - - async_handler: AsyncHTTPHandler = AsyncHTTPHandler(**_params) # type: ignore - else: - async_handler = client # type: ignore - - response = await async_handler.post( - url=url, - headers=headers, - data=json.dumps(data), - ) - - if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") - - _json_response = response.json() - _predictions = VertexAITextEmbeddingsResponseObject(**_json_response) # type: ignore - - return process_response( - model=model, - model_response=model_response, - _predictions=_predictions, - input=input, - ) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_transformation.py b/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_transformation.py deleted file mode 100644 index bbda553175..0000000000 --- a/litellm/llms/vertex_ai_and_google_ai_studio/embeddings/embed_content_transformation.py +++ /dev/null @@ -1,49 +0,0 @@ -""" -Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /embedContent format. - -Why separate file? Make it easy to see how transformation works -""" - -from litellm import EmbeddingResponse -from litellm.types.llms.openai import EmbeddingInput -from litellm.types.llms.vertex_ai import ( - ContentType, - PartType, - VertexAITextEmbeddingsResponseObject, -) -from litellm.types.utils import Embedding, Usage -from litellm.utils import get_formatted_prompt, token_counter - -from ..common_utils import VertexAIError - - -def transform_openai_input_gemini_content(input: str) -> ContentType: - """ - The content to embed. Only the parts.text fields will be counted. - """ - return ContentType(parts=[PartType(text=input)]) - - -def process_response( - input: EmbeddingInput, - model_response: EmbeddingResponse, - model: str, - _predictions: VertexAITextEmbeddingsResponseObject, -) -> EmbeddingResponse: - model_response.data = [ - Embedding( - embedding=_predictions["embedding"]["values"], - index=0, - object="embedding", - ) - ] - - model_response.model = model - - input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") - prompt_tokens = token_counter(model=model, text=input_text) - model_response.usage = Usage( - prompt_tokens=prompt_tokens, total_tokens=prompt_tokens - ) - - return model_response diff --git a/litellm/main.py b/litellm/main.py index bf1b0ede8c..c86bce8e26 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -129,9 +129,6 @@ from .llms.vertex_ai_and_google_ai_studio import ( from .llms.vertex_ai_and_google_ai_studio.embeddings.batch_embed_content_handler import ( GoogleBatchEmbeddings, ) -from .llms.vertex_ai_and_google_ai_studio.embeddings.embed_content_handler import ( - GoogleEmbeddings, -) from .llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( VertexLLM, ) @@ -178,7 +175,6 @@ triton_chat_completions = TritonChatCompletion() bedrock_chat_completion = BedrockLLM() bedrock_converse_chat_completion = BedrockConverseLLM() vertex_chat_completion = VertexLLM() -google_embeddings = GoogleEmbeddings() google_batch_embeddings = GoogleBatchEmbeddings() vertex_partner_models_chat_completion = VertexAIPartnerModels() vertex_text_to_speech = VertexTextToSpeechAPI() @@ -3541,38 +3537,21 @@ def embedding( gemini_api_key = api_key or get_secret("GEMINI_API_KEY") or litellm.api_key - if isinstance(input, str): - response = google_embeddings.text_embeddings( # type: ignore - model=model, - input=input, - encoding=encoding, - logging_obj=logging, - optional_params=optional_params, - model_response=EmbeddingResponse(), - vertex_project=None, - vertex_location=None, - vertex_credentials=None, - aembedding=aembedding, - print_verbose=print_verbose, - custom_llm_provider="gemini", - api_key=gemini_api_key, - ) - else: - response = google_batch_embeddings.batch_embeddings( # type: ignore - model=model, - input=input, - encoding=encoding, - logging_obj=logging, - optional_params=optional_params, - model_response=EmbeddingResponse(), - vertex_project=None, - vertex_location=None, - vertex_credentials=None, - aembedding=aembedding, - print_verbose=print_verbose, - custom_llm_provider="gemini", - api_key=gemini_api_key, - ) + response = google_batch_embeddings.batch_embeddings( # type: ignore + model=model, + input=input, + encoding=encoding, + logging_obj=logging, + optional_params=optional_params, + model_response=EmbeddingResponse(), + vertex_project=None, + vertex_location=None, + vertex_credentials=None, + aembedding=aembedding, + print_verbose=print_verbose, + custom_llm_provider="gemini", + api_key=gemini_api_key, + ) elif custom_llm_provider == "vertex_ai": vertex_ai_project = ( diff --git a/litellm/proxy/auth/rds_iam_token.py b/litellm/proxy/auth/rds_iam_token.py index f65fc4a99d..19dd6f7a3f 100644 --- a/litellm/proxy/auth/rds_iam_token.py +++ b/litellm/proxy/auth/rds_iam_token.py @@ -149,7 +149,7 @@ def init_rds_client( # boto3 automatically reads env variables client = boto3.client( - service_name="bedrock-runtime", + service_name="rds", region_name=region_name, config=config, ) From 83c5b48842ac325344309ec8d0bc64dabfeafde9 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 28 Aug 2024 12:15:57 -0700 Subject: [PATCH 08/16] fix(rds_iam_token.py): fix boto3 client init for rds --- litellm/proxy/auth/rds_iam_token.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/auth/rds_iam_token.py b/litellm/proxy/auth/rds_iam_token.py index f65fc4a99d..19dd6f7a3f 100644 --- a/litellm/proxy/auth/rds_iam_token.py +++ b/litellm/proxy/auth/rds_iam_token.py @@ -149,7 +149,7 @@ def init_rds_client( # boto3 automatically reads env variables client = boto3.client( - service_name="bedrock-runtime", + service_name="rds", region_name=region_name, config=config, ) From 0861180bfe86b4119be640364354a81b3a7b5681 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 28 Aug 2024 12:18:45 -0700 Subject: [PATCH 09/16] fix(rds_iam_token.py): support common aws env var's - AWS_ROLE_ARN, AWS_WEB_IDENTITY_TOKEN_FILE --- litellm/proxy/auth/rds_iam_token.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/rds_iam_token.py b/litellm/proxy/auth/rds_iam_token.py index 19dd6f7a3f..ec3a424b9f 100644 --- a/litellm/proxy/auth/rds_iam_token.py +++ b/litellm/proxy/auth/rds_iam_token.py @@ -168,8 +168,10 @@ def generate_iam_auth_token(db_host, db_port, db_user) -> str: aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"), aws_session_name=os.getenv("AWS_SESSION_NAME"), aws_profile_name=os.getenv("AWS_PROFILE_NAME"), - aws_role_name=os.getenv("AWS_ROLE_NAME"), - aws_web_identity_token=os.getenv("AWS_WEB_IDENTITY_TOKEN"), + aws_role_name=os.getenv("AWS_ROLE_NAME", os.getenv("AWS_ROLE_ARN")), + aws_web_identity_token=os.getenv( + "AWS_WEB_IDENTITY_TOKEN", os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") + ), ) token = boto_client.generate_db_auth_token( From 325538f8a4975c2251094d6a94ba41b36e3a62ee Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 28 Aug 2024 12:27:32 -0700 Subject: [PATCH 10/16] fix(key_management_endpoints.py): expose 'key' param, for setting your own key value --- litellm/proxy/_types.py | 1 + .../management_endpoints/key_management_endpoints.py | 9 ++++++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f2bef739c1..dd038d80bd 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -587,6 +587,7 @@ class GenerateRequestBase(LiteLLMBase): class GenerateKeyRequest(GenerateRequestBase): key_alias: Optional[str] = None + key: Optional[str] = None duration: Optional[str] = None aliases: Optional[dict] = {} config: Optional[dict] = {} diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 9bb07cfee3..00e17400c3 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -55,6 +55,7 @@ async def generate_key_fn( Parameters: - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - key_alias: Optional[str] - User defined key alias + - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you. - team_id: Optional[str] - The team id of the key - user_id: Optional[str] - The user id of the key - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models) @@ -728,6 +729,9 @@ async def generate_key_helper_fn( max_budget: Optional[float] = None, # max_budget is used to Budget Per user budget_duration: Optional[str] = None, # max_budget is used to Budget Per user token: Optional[str] = None, + key: Optional[ + str + ] = None, # dev-friendly alt param for 'token'. Exposed on `/key/generate` for setting key value yourself. user_id: Optional[str] = None, team_id: Optional[str] = None, user_email: Optional[str] = None, @@ -763,7 +767,10 @@ async def generate_key_helper_fn( ) if token is None: - token = f"sk-{secrets.token_urlsafe(16)}" + if key is not None: + token = key + else: + token = f"sk-{secrets.token_urlsafe(16)}" if duration is None: # allow tokens that never expire expires = None From 023d0f76643a2a8b2d6e207d0f3f3daa3306881d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 28 Aug 2024 12:54:24 -0700 Subject: [PATCH 11/16] build(model_prices_and_context_window.json): fix bedrock/llama3-1 pricing --- litellm/model_prices_and_context_window_backup.json | 8 ++++---- model_prices_and_context_window.json | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6f40ee1ee3..d0001819ec 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4030,8 +4030,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 2048, - "input_cost_per_token": 0.0000004, - "output_cost_per_token": 0.0000006, + "input_cost_per_token": 0.00000022, + "output_cost_per_token": 0.00000022, "litellm_provider": "bedrock", "mode": "chat", "supports_function_calling": true, @@ -4041,8 +4041,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 2048, - "input_cost_per_token": 0.00000265, - "output_cost_per_token": 0.0000035, + "input_cost_per_token": 0.00000099, + "output_cost_per_token": 0.00000099, "litellm_provider": "bedrock", "mode": "chat", "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6f40ee1ee3..d0001819ec 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4030,8 +4030,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 2048, - "input_cost_per_token": 0.0000004, - "output_cost_per_token": 0.0000006, + "input_cost_per_token": 0.00000022, + "output_cost_per_token": 0.00000022, "litellm_provider": "bedrock", "mode": "chat", "supports_function_calling": true, @@ -4041,8 +4041,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 2048, - "input_cost_per_token": 0.00000265, - "output_cost_per_token": 0.0000035, + "input_cost_per_token": 0.00000099, + "output_cost_per_token": 0.00000099, "litellm_provider": "bedrock", "mode": "chat", "supports_function_calling": true, From 17646b50ec0c8bf2aa1f8d603a6fec20ab1fb1e0 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 28 Aug 2024 13:07:56 -0700 Subject: [PATCH 12/16] build(model_prices_and_context_window.json): bedrock/llama3 models - region-based pricing --- ...odel_prices_and_context_window_backup.json | 128 +++++++++++++++++- model_prices_and_context_window.json | 128 +++++++++++++++++- 2 files changed, 254 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d0001819ec..df996dac03 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4012,11 +4012,74 @@ "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.0000004, + "input_cost_per_token": 0.0000003, "output_cost_per_token": 0.0000006, "litellm_provider": "bedrock", "mode": "chat" }, + "bedrock/us-east-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.0000003, + "output_cost_per_token": 0.0000006, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/us-west-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.0000003, + "output_cost_per_token": 0.0000006, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/ap-south-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000036, + "output_cost_per_token": 0.00000072, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/ca-central-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000035, + "output_cost_per_token": 0.00000069, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/eu-west-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000032, + "output_cost_per_token": 0.00000065, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/eu-west-2/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000039, + "output_cost_per_token": 0.00000078, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/sa-east-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.0000005, + "output_cost_per_token": 0.00000101, + "litellm_provider": "bedrock", + "mode": "chat" + }, "meta.llama3-70b-instruct-v1:0": { "max_tokens": 8192, "max_input_tokens": 8192, @@ -4026,6 +4089,69 @@ "litellm_provider": "bedrock", "mode": "chat" }, + "bedrock/us-east-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000265, + "output_cost_per_token": 0.0000035, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000265, + "output_cost_per_token": 0.0000035, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/ap-south-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000318, + "output_cost_per_token": 0.0000042, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/ca-central-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000305, + "output_cost_per_token": 0.00000403, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/eu-west-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000286, + "output_cost_per_token": 0.00000378, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/eu-west-2/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000345, + "output_cost_per_token": 0.00000455, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/sa-east-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000445, + "output_cost_per_token": 0.00000588, + "litellm_provider": "bedrock", + "mode": "chat" + }, "meta.llama3-1-8b-instruct-v1:0": { "max_tokens": 128000, "max_input_tokens": 128000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d0001819ec..df996dac03 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4012,11 +4012,74 @@ "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.0000004, + "input_cost_per_token": 0.0000003, "output_cost_per_token": 0.0000006, "litellm_provider": "bedrock", "mode": "chat" }, + "bedrock/us-east-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.0000003, + "output_cost_per_token": 0.0000006, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/us-west-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.0000003, + "output_cost_per_token": 0.0000006, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/ap-south-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000036, + "output_cost_per_token": 0.00000072, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/ca-central-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000035, + "output_cost_per_token": 0.00000069, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/eu-west-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000032, + "output_cost_per_token": 0.00000065, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/eu-west-2/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000039, + "output_cost_per_token": 0.00000078, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/sa-east-1/meta.llama3-8b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.0000005, + "output_cost_per_token": 0.00000101, + "litellm_provider": "bedrock", + "mode": "chat" + }, "meta.llama3-70b-instruct-v1:0": { "max_tokens": 8192, "max_input_tokens": 8192, @@ -4026,6 +4089,69 @@ "litellm_provider": "bedrock", "mode": "chat" }, + "bedrock/us-east-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000265, + "output_cost_per_token": 0.0000035, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000265, + "output_cost_per_token": 0.0000035, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/ap-south-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000318, + "output_cost_per_token": 0.0000042, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/ca-central-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000305, + "output_cost_per_token": 0.00000403, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/eu-west-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000286, + "output_cost_per_token": 0.00000378, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/eu-west-2/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000345, + "output_cost_per_token": 0.00000455, + "litellm_provider": "bedrock", + "mode": "chat" + }, + "bedrock/sa-east-1/meta.llama3-70b-instruct-v1:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "input_cost_per_token": 0.00000445, + "output_cost_per_token": 0.00000588, + "litellm_provider": "bedrock", + "mode": "chat" + }, "meta.llama3-1-8b-instruct-v1:0": { "max_tokens": 128000, "max_input_tokens": 128000, From cc411f1e9729807e1c0fe315c12e07c159a5e9a2 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 28 Aug 2024 13:33:13 -0700 Subject: [PATCH 13/16] docs(reliability.md): cleanup docs --- docs/my-website/docs/proxy/reliability.md | 132 ++++++++++++---------- 1 file changed, 72 insertions(+), 60 deletions(-) diff --git a/docs/my-website/docs/proxy/reliability.md b/docs/my-website/docs/proxy/reliability.md index cb6550a478..7a2c65a90c 100644 --- a/docs/my-website/docs/proxy/reliability.md +++ b/docs/my-website/docs/proxy/reliability.md @@ -274,6 +274,17 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ ## Advanced ### Fallbacks + Retries + Timeouts + Cooldowns +To set fallbacks, just do: + +``` +litellm_settings: + fallbacks: [{"zephyr-beta": ["gpt-3.5-turbo"]}] +``` + +**Covers all errors (429, 500, etc.)** + +[**See Code**]() + **Set via config** ```yaml model_list: @@ -302,10 +313,70 @@ litellm_settings: num_retries: 3 # retry call 3 times on each model_name (e.g. zephyr-beta) request_timeout: 10 # raise Timeout error if call takes longer than 10s. Sets litellm.request_timeout fallbacks: [{"zephyr-beta": ["gpt-3.5-turbo"]}] # fallback to gpt-3.5-turbo if call fails num_retries - context_window_fallbacks: [{"zephyr-beta": ["gpt-3.5-turbo-16k"]}, {"gpt-3.5-turbo": ["gpt-3.5-turbo-16k"]}] # fallback to gpt-3.5-turbo-16k if context window error allowed_fails: 3 # cooldown model if it fails > 1 call in a minute. cooldown_time: 30 # how long to cooldown model if fails/min > allowed_fails ``` + +### Test Fallbacks! + +Check if your fallbacks are working as expected. + +#### **Regular Fallbacks** +```bash +curl -X POST 'http://0.0.0.0:4000/chat/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-1234' \ +-D '{ + "model": "my-bad-model", + "messages": [ + { + "role": "user", + "content": "ping" + } + ], + "mock_testing_fallbacks": true # 👈 KEY CHANGE +} +' +``` + +#### **Content Policy Fallbacks** +```bash +curl -X POST 'http://0.0.0.0:4000/chat/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-1234' \ +-D '{ + "model": "my-bad-model", + "messages": [ + { + "role": "user", + "content": "ping" + } + ], + "mock_testing_content_policy_fallbacks": true # 👈 KEY CHANGE +} +' +``` + +#### **Context Window Fallbacks** + +```bash +curl -X POST 'http://0.0.0.0:4000/chat/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-1234' \ +-D '{ + "model": "my-bad-model", + "messages": [ + { + "role": "user", + "content": "ping" + } + ], + "mock_testing_context_window_fallbacks": true # 👈 KEY CHANGE +} +' +``` + + ### Context Window Fallbacks (Pre-Call Checks + Fallbacks) **Before call is made** check if a call is within model context window with **`enable_pre_call_checks: true`**. @@ -493,65 +564,6 @@ This will default to claude-opus in case any model fails. A model-specific fallbacks (e.g. {"gpt-3.5-turbo-small": ["claude-opus"]}) overrides default fallback. -### Test Fallbacks! - -Check if your fallbacks are working as expected. - -#### **Regular Fallbacks** -```bash -curl -X POST 'http://0.0.0.0:4000/chat/completions' \ --H 'Content-Type: application/json' \ --H 'Authorization: Bearer sk-1234' \ --D '{ - "model": "my-bad-model", - "messages": [ - { - "role": "user", - "content": "ping" - } - ], - "mock_testing_fallbacks": true # 👈 KEY CHANGE -} -' -``` - -#### **Content Policy Fallbacks** -```bash -curl -X POST 'http://0.0.0.0:4000/chat/completions' \ --H 'Content-Type: application/json' \ --H 'Authorization: Bearer sk-1234' \ --D '{ - "model": "my-bad-model", - "messages": [ - { - "role": "user", - "content": "ping" - } - ], - "mock_testing_content_policy_fallbacks": true # 👈 KEY CHANGE -} -' -``` - -#### **Context Window Fallbacks** - -```bash -curl -X POST 'http://0.0.0.0:4000/chat/completions' \ --H 'Content-Type: application/json' \ --H 'Authorization: Bearer sk-1234' \ --D '{ - "model": "my-bad-model", - "messages": [ - { - "role": "user", - "content": "ping" - } - ], - "mock_testing_context_window_fallbacks": true # 👈 KEY CHANGE -} -' -``` - ### EU-Region Filtering (Pre-Call Checks) **Before call is made** check if a call is within model context window with **`enable_pre_call_checks: true`**. From 4ce59f1a9611beee508944f3fc7d34835cbe2bb2 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 28 Aug 2024 14:27:07 -0700 Subject: [PATCH 14/16] test(test_amazing_vertex_completion.py): update test to not pick experimental gemini models --- litellm/tests/test_amazing_vertex_completion.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index fa33bab3b6..de1f1007dd 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -50,6 +50,8 @@ VERTEX_MODELS_TO_NOT_TEST = [ "text-bison@001", "gemini-1.5-pro", "gemini-1.5-pro-preview-0215", + "gemini-pro-flash", + "gemini-pro-experimental", ] @@ -444,7 +446,9 @@ async def test_async_vertexai_response(): test_models = random.sample(test_models, 1) test_models += litellm.vertex_language_models # always test gemini-pro for model in test_models: - print(f"model being tested in async call: {model}") + print( + f"model being tested in async call: {model}, litellm.vertex_language_models: {litellm.vertex_language_models}" + ) if model in VERTEX_MODELS_TO_NOT_TEST or ( "gecko" in model or "32k" in model or "ultra" in model or "002" in model ): From dd9c5d10bdc089b91ebff8e32cfb4d9d5f06f1c9 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 28 Aug 2024 18:07:37 -0700 Subject: [PATCH 15/16] fix(vertex_ai_partner_models.py): fix vertex import --- .../vertex_ai_partner_models/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_partner_models/main.py index 983a3c4ad3..60c1fa607d 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_partner_models/main.py @@ -83,7 +83,7 @@ class VertexAIPartnerModels(BaseLLM): from litellm.llms.databricks import DatabricksChatCompletion from litellm.llms.openai import OpenAIChatCompletion from litellm.llms.text_completion_codestral import CodestralTextCompletion - from litellm.llms.vertex_ai_and_google_ai_studio.vertex_and_google_ai_studio_gemini import ( + from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( VertexLLM, ) except Exception: From fd1cca207f3c242fac69e8abc7adcd1e6fdb883f Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 28 Aug 2024 21:12:39 -0700 Subject: [PATCH 16/16] test(test_amazing_vertex_completion.py): fix test --- litellm/tests/test_amazing_vertex_completion.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 4c5c3dc8e0..02683de158 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -2112,7 +2112,7 @@ def test_get_token_url(): vertex_credentials=vertex_credentials, gemini_api_key="", custom_llm_provider="vertex_ai_beta", - should_use_vertex_v1beta1_features=should_use_v1beta1_features, + should_use_v1beta1_features=should_use_v1beta1_features, api_base=None, model="", stream=False, @@ -2132,7 +2132,7 @@ def test_get_token_url(): vertex_credentials=vertex_credentials, gemini_api_key="", custom_llm_provider="vertex_ai_beta", - should_use_vertex_v1beta1_features=should_use_v1beta1_features, + should_use_v1beta1_features=should_use_v1beta1_features, api_base=None, model="", stream=False,