fix(spend_tracking): populate api_base in SpendLogs for embedding calls

Azure embeddings (sync and async) and the Cohere embed handler used by Bedrock Cohere models called logging pre_call without api_base, so the logged api_base and the SpendLogs api_base column stayed empty for these requests. pre_call is the only place that reads api_base from additional_args, so the fix is there. Fixes #23768.

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},
                 },
             )

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

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]

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
@@ -189,6 +191,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"