diff options
Diffstat (limited to '.venv/lib/python3.12/site-packages/litellm/llms/azure_ai/embed')
3 files changed, 392 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/litellm/llms/azure_ai/embed/__init__.py b/.venv/lib/python3.12/site-packages/litellm/llms/azure_ai/embed/__init__.py new file mode 100644 index 00000000..e0d67acb --- /dev/null +++ b/.venv/lib/python3.12/site-packages/litellm/llms/azure_ai/embed/__init__.py @@ -0,0 +1 @@ +from .handler import AzureAIEmbedding diff --git a/.venv/lib/python3.12/site-packages/litellm/llms/azure_ai/embed/cohere_transformation.py b/.venv/lib/python3.12/site-packages/litellm/llms/azure_ai/embed/cohere_transformation.py new file mode 100644 index 00000000..38b0dbbe --- /dev/null +++ b/.venv/lib/python3.12/site-packages/litellm/llms/azure_ai/embed/cohere_transformation.py @@ -0,0 +1,99 @@ +""" +Transformation logic from OpenAI /v1/embeddings format to Azure AI Cohere's /v1/embed. + +Why separate file? Make it easy to see how transformation works + +Convers +- Cohere request format + +Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-embed-text.html +""" + +from typing import List, Optional, Tuple + +from litellm.types.llms.azure_ai import ImageEmbeddingInput, ImageEmbeddingRequest +from litellm.types.llms.openai import EmbeddingCreateParams +from litellm.types.utils import EmbeddingResponse, Usage +from litellm.utils import is_base64_encoded + + +class AzureAICohereConfig: + def __init__(self) -> None: + pass + + def _map_azure_model_group(self, model: str) -> str: + + if model == "offer-cohere-embed-multili-paygo": + return "Cohere-embed-v3-multilingual" + elif model == "offer-cohere-embed-english-paygo": + return "Cohere-embed-v3-english" + + return model + + def _transform_request_image_embeddings( + self, input: List[str], optional_params: dict + ) -> ImageEmbeddingRequest: + """ + Assume all str in list is base64 encoded string + """ + image_input: List[ImageEmbeddingInput] = [] + for i in input: + embedding_input = ImageEmbeddingInput(image=i) + image_input.append(embedding_input) + return ImageEmbeddingRequest(input=image_input, **optional_params) + + def _transform_request( + self, input: List[str], optional_params: dict, model: str + ) -> Tuple[ImageEmbeddingRequest, EmbeddingCreateParams, List[int]]: + """ + Return the list of input to `/image/embeddings`, `/v1/embeddings`, list of image_embedding_idx for recombination + """ + image_embeddings: List[str] = [] + image_embedding_idx: List[int] = [] + for idx, i in enumerate(input): + """ + - is base64 -> route to image embeddings + - is ImageEmbeddingInput -> route to image embeddings + - else -> route to `/v1/embeddings` + """ + if is_base64_encoded(i): + image_embeddings.append(i) + image_embedding_idx.append(idx) + + ## REMOVE IMAGE EMBEDDINGS FROM input list + filtered_input = [ + item for idx, item in enumerate(input) if idx not in image_embedding_idx + ] + + v1_embeddings_request = EmbeddingCreateParams( + input=filtered_input, model=model, **optional_params + ) + image_embeddings_request = self._transform_request_image_embeddings( + input=image_embeddings, optional_params=optional_params + ) + + return image_embeddings_request, v1_embeddings_request, image_embedding_idx + + def _transform_response(self, response: EmbeddingResponse) -> EmbeddingResponse: + additional_headers: Optional[dict] = response._hidden_params.get( + "additional_headers" + ) + if additional_headers: + # CALCULATE USAGE + input_tokens: Optional[str] = additional_headers.get( + "llm_provider-num_tokens" + ) + if input_tokens: + if response.usage: + response.usage.prompt_tokens = int(input_tokens) + else: + response.usage = Usage(prompt_tokens=int(input_tokens)) + + # SET MODEL + base_model: Optional[str] = additional_headers.get( + "llm_provider-azureml-model-group" + ) + if base_model: + response.model = self._map_azure_model_group(base_model) + + return response diff --git a/.venv/lib/python3.12/site-packages/litellm/llms/azure_ai/embed/handler.py b/.venv/lib/python3.12/site-packages/litellm/llms/azure_ai/embed/handler.py new file mode 100644 index 00000000..f33c979c --- /dev/null +++ b/.venv/lib/python3.12/site-packages/litellm/llms/azure_ai/embed/handler.py @@ -0,0 +1,292 @@ +from typing import List, Optional, Union + +from openai import OpenAI + +import litellm +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + get_async_httpx_client, +) +from litellm.llms.openai.openai import OpenAIChatCompletion +from litellm.types.llms.azure_ai import ImageEmbeddingRequest +from litellm.types.utils import EmbeddingResponse +from litellm.utils import convert_to_model_response_object + +from .cohere_transformation import AzureAICohereConfig + + +class AzureAIEmbedding(OpenAIChatCompletion): + + def _process_response( + self, + image_embedding_responses: Optional[List], + text_embedding_responses: Optional[List], + image_embeddings_idx: List[int], + model_response: EmbeddingResponse, + input: List, + ): + combined_responses = [] + if ( + image_embedding_responses is not None + and text_embedding_responses is not None + ): + # Combine and order the results + text_idx = 0 + image_idx = 0 + + for idx in range(len(input)): + if idx in image_embeddings_idx: + combined_responses.append(image_embedding_responses[image_idx]) + image_idx += 1 + else: + combined_responses.append(text_embedding_responses[text_idx]) + text_idx += 1 + + model_response.data = combined_responses + elif image_embedding_responses is not None: + model_response.data = image_embedding_responses + elif text_embedding_responses is not None: + model_response.data = text_embedding_responses + + response = AzureAICohereConfig()._transform_response(response=model_response) # type: ignore + + return response + + async def async_image_embedding( + self, + model: str, + data: ImageEmbeddingRequest, + timeout: float, + logging_obj, + model_response: litellm.EmbeddingResponse, + optional_params: dict, + api_key: Optional[str], + api_base: Optional[str], + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + ) -> EmbeddingResponse: + if client is None or not isinstance(client, AsyncHTTPHandler): + client = get_async_httpx_client( + llm_provider=litellm.LlmProviders.AZURE_AI, + params={"timeout": timeout}, + ) + + url = "{}/images/embeddings".format(api_base) + + response = await client.post( + url=url, + json=data, # type: ignore + headers={"Authorization": "Bearer {}".format(api_key)}, + ) + + embedding_response = response.json() + embedding_headers = dict(response.headers) + returned_response: EmbeddingResponse = convert_to_model_response_object( # type: ignore + response_object=embedding_response, + model_response_object=model_response, + response_type="embedding", + stream=False, + _response_headers=embedding_headers, + ) + return returned_response + + def image_embedding( + self, + model: str, + data: ImageEmbeddingRequest, + timeout: float, + logging_obj, + model_response: EmbeddingResponse, + optional_params: dict, + api_key: Optional[str], + api_base: Optional[str], + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + ): + if api_base is None: + raise ValueError( + "api_base is None. Please set AZURE_AI_API_BASE or dynamically via `api_base` param, to make the request." + ) + if api_key is None: + raise ValueError( + "api_key is None. Please set AZURE_AI_API_KEY or dynamically via `api_key` param, to make the request." + ) + + if client is None or not isinstance(client, HTTPHandler): + client = HTTPHandler(timeout=timeout, concurrent_limit=1) + + url = "{}/images/embeddings".format(api_base) + + response = client.post( + url=url, + json=data, # type: ignore + headers={"Authorization": "Bearer {}".format(api_key)}, + ) + + embedding_response = response.json() + embedding_headers = dict(response.headers) + returned_response: EmbeddingResponse = convert_to_model_response_object( # type: ignore + response_object=embedding_response, + model_response_object=model_response, + response_type="embedding", + stream=False, + _response_headers=embedding_headers, + ) + return returned_response + + async def async_embedding( + self, + model: str, + input: List, + timeout: float, + logging_obj, + model_response: litellm.EmbeddingResponse, + optional_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + client=None, + ) -> EmbeddingResponse: + + ( + image_embeddings_request, + v1_embeddings_request, + image_embeddings_idx, + ) = AzureAICohereConfig()._transform_request( + input=input, optional_params=optional_params, model=model + ) + + image_embedding_responses: Optional[List] = None + text_embedding_responses: Optional[List] = None + + if image_embeddings_request["input"]: + image_response = await self.async_image_embedding( + model=model, + data=image_embeddings_request, + timeout=timeout, + logging_obj=logging_obj, + model_response=model_response, + optional_params=optional_params, + api_key=api_key, + api_base=api_base, + client=client, + ) + + image_embedding_responses = image_response.data + if image_embedding_responses is None: + raise Exception("/image/embeddings route returned None Embeddings.") + + if v1_embeddings_request["input"]: + response: EmbeddingResponse = await super().embedding( # type: ignore + model=model, + input=input, + timeout=timeout, + logging_obj=logging_obj, + model_response=model_response, + optional_params=optional_params, + api_key=api_key, + api_base=api_base, + client=client, + aembedding=True, + ) + text_embedding_responses = response.data + if text_embedding_responses is None: + raise Exception("/v1/embeddings route returned None Embeddings.") + + return self._process_response( + image_embedding_responses=image_embedding_responses, + text_embedding_responses=text_embedding_responses, + image_embeddings_idx=image_embeddings_idx, + model_response=model_response, + input=input, + ) + + def embedding( + self, + model: str, + input: List, + timeout: float, + logging_obj, + model_response: EmbeddingResponse, + optional_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + client=None, + aembedding=None, + max_retries: Optional[int] = None, + ) -> EmbeddingResponse: + """ + - Separate image url from text + -> route image url call to `/image/embeddings` + -> route text call to `/v1/embeddings` (OpenAI route) + + assemble result in-order, and return + """ + if aembedding is True: + return self.async_embedding( # type: ignore + model, + input, + timeout, + logging_obj, + model_response, + optional_params, + api_key, + api_base, + client, + ) + + ( + image_embeddings_request, + v1_embeddings_request, + image_embeddings_idx, + ) = AzureAICohereConfig()._transform_request( + input=input, optional_params=optional_params, model=model + ) + + image_embedding_responses: Optional[List] = None + text_embedding_responses: Optional[List] = None + + if image_embeddings_request["input"]: + image_response = self.image_embedding( + model=model, + data=image_embeddings_request, + timeout=timeout, + logging_obj=logging_obj, + model_response=model_response, + optional_params=optional_params, + api_key=api_key, + api_base=api_base, + client=client, + ) + + image_embedding_responses = image_response.data + if image_embedding_responses is None: + raise Exception("/image/embeddings route returned None Embeddings.") + + if v1_embeddings_request["input"]: + response: EmbeddingResponse = super().embedding( # type: ignore + model, + input, + timeout, + logging_obj, + model_response, + optional_params, + api_key, + api_base, + client=( + client + if client is not None and isinstance(client, OpenAI) + else None + ), + aembedding=aembedding, + ) + + text_embedding_responses = response.data + if text_embedding_responses is None: + raise Exception("/v1/embeddings route returned None Embeddings.") + + return self._process_response( + image_embedding_responses=image_embedding_responses, + text_embedding_responses=text_embedding_responses, + image_embeddings_idx=image_embeddings_idx, + model_response=model_response, + input=input, + ) |