diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 69ffac5d6f9..0284a17baf8 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -792,6 +792,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_key=api_key, additional_args={ "complete_input_dict": data, + "api_base": api_base, "headers": {"api_key": api_key, "azure_ad_token": azure_ad_token}, }, ) diff --git a/litellm/llms/cohere/embed/handler.py b/litellm/llms/cohere/embed/handler.py index 34fde444880..6b372c775e4 100644 --- a/litellm/llms/cohere/embed/handler.py +++ b/litellm/llms/cohere/embed/handler.py @@ -159,7 +159,7 @@ def embedding( logging_obj.pre_call( input=input, api_key=api_key, - additional_args={"complete_input_dict": data}, + additional_args={"complete_input_dict": data, "api_base": embed_url}, ) ## COMPLETION CALL diff --git a/tests/unit/llms/azure/test_azure_embedding.py b/tests/unit/llms/azure/test_azure_embedding.py index 96f8290e97f..50f686c29da 100644 --- a/tests/unit/llms/azure/test_azure_embedding.py +++ b/tests/unit/llms/azure/test_azure_embedding.py @@ -1,6 +1,8 @@ +import asyncio import json import os import sys +import threading from typing import Final from unittest.mock import MagicMock, patch @@ -8,12 +10,14 @@ import httpx import pytest import respx from httpx import Response +from openai import AsyncAzureOpenAI, AzureOpenAI sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) ) import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.azure.azure import AzureChatCompletion from litellm.types.utils import EmbeddingResponse, Usage @@ -205,3 +209,82 @@ def test_azure_embedding_max_retries_zero_sends_one_request(respx_mock: respx.Mo ) assert route.call_count == 1 + + +_EMBEDDING_API_BASE: Final = "https://my-resource.openai.azure.com" + + +def _embeddings_endpoint(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "text-embedding-3-large", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + }, + ) + + +class _LoggedApiBase(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.api_bases: list[str] = [] + self.logged = threading.Event() + + def log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + self.api_bases.append(kwargs["standard_logging_object"]["api_base"]) + self.logged.set() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + self.log_success_event(kwargs, response_obj, start_time, end_time) + + +@pytest.fixture +def logged_api_base(monkeypatch: pytest.MonkeyPatch) -> _LoggedApiBase: + logger: Final = _LoggedApiBase() + monkeypatch.setattr(litellm, "callbacks", [logger]) + return logger + + +def test_azure_embedding_logs_the_deployment_api_base(logged_api_base: _LoggedApiBase) -> None: + client: Final = AzureOpenAI( + api_key="azure-test-key", + api_version="2024-02-01", + azure_endpoint=_EMBEDDING_API_BASE, + http_client=httpx.Client(transport=httpx.MockTransport(_embeddings_endpoint)), + ) + + litellm.embedding( + model="azure/text-embedding-3-large", + input=["hello"], + api_base=_EMBEDDING_API_BASE, + api_key="azure-test-key", + api_version="2024-02-01", + client=client, + ) + + assert logged_api_base.logged.wait(timeout=5) + assert logged_api_base.api_bases == [_EMBEDDING_API_BASE] + + +@pytest.mark.asyncio +async def test_azure_aembedding_logs_the_deployment_api_base(logged_api_base: _LoggedApiBase) -> None: + client: Final = AsyncAzureOpenAI( + api_key="azure-test-key", + api_version="2024-02-01", + azure_endpoint=_EMBEDDING_API_BASE, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(_embeddings_endpoint)), + ) + + await litellm.aembedding( + model="azure/text-embedding-3-large", + input=["hello"], + api_base=_EMBEDDING_API_BASE, + api_key="azure-test-key", + api_version="2024-02-01", + client=client, + ) + + assert await asyncio.to_thread(logged_api_base.logged.wait, 5) + assert logged_api_base.api_bases == [_EMBEDDING_API_BASE] diff --git a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py index 904ca3bb8ca..dfa35a12f61 100644 --- a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py @@ -1,6 +1,7 @@ import json import asyncio import os +import threading from typing import Final from unittest.mock import Mock, patch from unittest.mock import AsyncMock, MagicMock @@ -11,6 +12,7 @@ import respx import litellm from litellm.caching.caching import Cache +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.bedrock.embed.embedding import BedrockEmbedding @@ -190,6 +192,33 @@ def test_bedrock_cohere_embedding_returns_the_scripted_vector(respx_mock: respx. assert response.data == [{"object": "embedding", "index": 0, "embedding": [0.3, 0.2, 0.1]}] +def test_bedrock_cohere_embedding_logs_the_runtime_endpoint_as_api_base( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + url: Final = "https://bedrock-runtime.us-east-1.amazonaws.com/model/cohere.embed-multilingual-v3/invoke" + respx_mock.post(url).respond(json={"embeddings": [[0.3, 0.2, 0.1]], "inputTextTokenCount": 4}) + logged: Final = threading.Event() + api_bases: Final[list[str]] = [] + + class LoggedApiBase(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + api_bases.append(kwargs["standard_logging_object"]["api_base"]) + logged.set() + + monkeypatch.setattr(litellm, "callbacks", [LoggedApiBase()]) + + litellm.embedding( + model="bedrock/cohere.embed-multilingual-v3", + input=["hello"], + api_key="test-bearer-token", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + + assert logged.wait(timeout=5) + assert api_bases == [url] + + GOVCLOUD_WEST_RUNTIME: Final = "https://bedrock-runtime.us-gov-west-1.amazonaws.com"