(fix) Pass through spend tracking - ensure custom_llm_provider is tracked for Vertex, Google AI Studio, Anthropic (#8882)

* fix track custom llm provider on pass through routes

* fix use correct provider for google ai studio

* fix tracking custom llm provider on pass through route

* ui fix get provider logo

* update tests to track custom llm provider

* test_anthropic_streaming_with_headers

* Potential fix for code scanning alert no. 2263: Incomplete URL substring sanitization

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

---------

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

Author: ishaan-jaff

litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py

@@ -123,6 +123,9 @@ def _create_anthropic_response_logging_payload(
             litellm_model_response.id = logging_obj.litellm_call_id
             litellm_model_response.model = model
             logging_obj.model_call_details["model"] = model
+            logging_obj.model_call_details["custom_llm_provider"] = (
+                litellm.LlmProviders.ANTHROPIC.value
+            )
             return kwargs
         except Exception as e:
             verbose_proxy_logger.exception(

litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py

@@ -2,7 +2,7 @@
 import re
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
-
+from urllib.parse import urlparse
 import httpx
 
 import litellm
@@ -66,6 +66,9 @@ def vertex_passthrough_handler(
                 start_time=start_time,
                 end_time=end_time,
                 logging_obj=logging_obj,
+                custom_llm_provider=VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(
+                    url_route
+                ),
             )
 
             return {
@@ -167,6 +170,9 @@ def _handle_logging_vertex_collected_chunks(
             start_time=start_time,
             end_time=end_time,
             logging_obj=litellm_logging_obj,
+            custom_llm_provider=VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(
+                url_route
+            ),
         )
 
         return {
@@ -213,6 +219,13 @@ def extract_model_from_url(url: str) -> str:
             return match.group(1)
         return "unknown"
 
+    @staticmethod
+    def _get_custom_llm_provider_from_url(url: str) -> str:
+        parsed_url = urlparse(url)
+        if parsed_url.hostname and parsed_url.hostname.endswith("generativelanguage.googleapis.com"):
+            return litellm.LlmProviders.GEMINI.value
+        return litellm.LlmProviders.VERTEX_AI.value
+
     @staticmethod
     def _create_vertex_response_logging_payload_for_generate_content(
         litellm_model_response: Union[ModelResponse, TextCompletionResponse],
@@ -221,6 +234,7 @@ def _create_vertex_response_logging_payload_for_generate_content(
         start_time: datetime,
         end_time: datetime,
         logging_obj: LiteLLMLoggingObj,
+        custom_llm_provider: str,
     ):
         """
         Create the standard logging object for Vertex passthrough generateContent (streaming and non-streaming)
@@ -240,4 +254,5 @@ def _create_vertex_response_logging_payload_for_generate_content(
         litellm_model_response.id = logging_obj.litellm_call_id
         logging_obj.model = litellm_model_response.model or model
         logging_obj.model_call_details["model"] = logging_obj.model
+        logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
         return kwargs

tests/pass_through_tests/test_anthropic_passthrough.py

@@ -162,6 +162,7 @@ async def test_anthropic_basic_completion_with_headers():
                 ), "Should have user API key in metadata"
 
                 assert "claude" in log_entry["model"]
+                assert log_entry["custom_llm_provider"] == "anthropic"
 
 
 @pytest.mark.asyncio
@@ -230,7 +231,7 @@ async def test_anthropic_streaming_with_headers():
             print("anthropic_api_output_tokens", anthropic_api_output_tokens)
 
             # Wait for spend to be logged
-            await asyncio.sleep(10)
+            await asyncio.sleep(20)
 
             # Check spend logs for this specific request
             async with session.get(
@@ -301,3 +302,4 @@ async def test_anthropic_streaming_with_headers():
                 assert "claude" in log_entry["model"]
 
                 assert log_entry["end_user"] == "test-user-1"
+                assert log_entry["custom_llm_provider"] == "anthropic"

tests/pass_through_tests/test_gemini_with_spend.test.js

@@ -62,6 +62,7 @@ describe('Gemini AI Tests', () => {
         expect(spendData[0].request_tags).toEqual(['gemini-js-sdk', 'pass-through-endpoint']);
         expect(spendData[0].metadata).toHaveProperty('user_api_key');
         expect(spendData[0].model).toContain('gemini');
+        expect(spendData[0].custom_llm_provider).toBe('gemini');
         expect(spendData[0].spend).toBeGreaterThan(0);
     }, 25000);
 
@@ -119,5 +120,6 @@ describe('Gemini AI Tests', () => {
         expect(spendData[0].metadata).toHaveProperty('user_api_key');
         expect(spendData[0].model).toContain('gemini');
         expect(spendData[0].spend).toBeGreaterThan(0);
+        expect(spendData[0].custom_llm_provider).toBe('gemini');
     }, 25000);
 });

tests/pass_through_tests/test_vertex_with_spend.test.js

@@ -121,6 +121,7 @@ describe('Vertex AI Tests', () => {
         expect(spendData[0].metadata).toHaveProperty('user_api_key');
         expect(spendData[0].model).toContain('gemini');
         expect(spendData[0].spend).toBeGreaterThan(0);
+        expect(spendData[0].custom_llm_provider).toBe('vertex_ai');
     }, 25000);
 
     test('should successfully generate streaming content with tags', async () => {
@@ -190,5 +191,6 @@ describe('Vertex AI Tests', () => {
         expect(spendData[0].metadata).toHaveProperty('user_api_key');
         expect(spendData[0].model).toContain('gemini');
         expect(spendData[0].spend).toBeGreaterThan(0);
+        expect(spendData[0].custom_llm_provider).toBe('vertex_ai');
     }, 25000);
 });

ui/litellm-dashboard/src/components/provider_info_helpers.tsx

@@ -62,6 +62,13 @@ export const getProviderLogoAndName = (providerValue: string): { logo: string, d
         return { logo: "", displayName: "-" };
     }
 
+    // Handle special case for "gemini" provider value
+    if (providerValue.toLowerCase() === "gemini") {
+        const displayName = Providers.Google_AI_Studio;
+        const logo = providerLogoMap[displayName];
+        return { logo, displayName };
+    }
+
     // Find the enum key by matching provider_map values
     const enumKey = Object.keys(provider_map).find(
         key => provider_map[key].toLowerCase() === providerValue.toLowerCase()

ui/litellm-dashboard/src/components/view_logs/index.tsx

@@ -598,6 +598,10 @@ function RequestViewer({ row }: { row: Row<LogEntry> }) {
             <span className="font-medium w-1/3">Model:</span>
             <span>{row.original.model}</span>
           </div>
+          <div className="flex">
+            <span className="font-medium w-1/3">Custom LLM Provider:</span>
+            <span>{row.original.custom_llm_provider}</span>
+          </div>
           <div className="flex">
             <span className="font-medium w-1/3">Api Base:</span>
             <span>{row.original.api_base}</span>