VertexAI Anthropic - streaming passthrough cost tracking (#11734)

* feat(vertex_passthrough_logging_handler.py): initial anthropic passthrough streaming cost tracking support

* fix: fix linting errors

* test: update test

litellm/llms/base_llm/base_model_iterator.py

@@ -37,10 +37,8 @@ class BaseModelResponseIterator:
     def __iter__(self):
         return self
 
-    def _handle_string_chunk(
-        self, str_line: str
-    ) -> Union[GenericStreamingChunk, ModelResponseStream]:
-        # chunk is a str at this point
+    @staticmethod
+    def _string_to_dict_parser(str_line: str) -> Optional[dict]:
         stripped_json_chunk: Optional[dict] = None
         stripped_chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(
             str_line
@@ -52,7 +50,15 @@ class BaseModelResponseIterator:
                 stripped_json_chunk = None
         except json.JSONDecodeError:
             stripped_json_chunk = None
+        return stripped_json_chunk
 
+    def _handle_string_chunk(
+        self, str_line: str
+    ) -> Union[GenericStreamingChunk, ModelResponseStream]:
+        # chunk is a str at this point
+        stripped_json_chunk = BaseModelResponseIterator._string_to_dict_parser(
+            str_line=str_line
+        )
         if "[DONE]" in str_line:
             return GenericStreamingChunk(
                 text="",

litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py

@@ -135,7 +135,7 @@ class VertexPassthroughLoggingHandler:
             _json_response = httpx_response.json()
 
             litellm_prediction_response = ModelResponse()
-            
+
             if vertex_publisher_or_api_spec is not None:
                 vertex_ai_partner_model_config = get_vertex_ai_partner_model_config(
                     model=model,
@@ -206,6 +206,7 @@ class VertexPassthroughLoggingHandler:
                 all_chunks=all_chunks,
                 litellm_logging_obj=litellm_logging_obj,
                 model=model,
+                url_route=url_route,
             )
         )
 
@@ -240,12 +241,37 @@ class VertexPassthroughLoggingHandler:
         all_chunks: List[str],
         litellm_logging_obj: LiteLLMLoggingObj,
         model: str,
+        url_route: str,
     ) -> Optional[Union[ModelResponse, TextCompletionResponse]]:
-        vertex_iterator = VertexModelResponseIterator(
-            streaming_response=None,
-            sync_stream=False,
-            logging_obj=litellm_logging_obj,
-        )
+        parsed_chunks = []
+        if "generateContent" in url_route or "streamGenerateContent" in url_route:
+            vertex_iterator: Any = VertexModelResponseIterator(
+                streaming_response=None,
+                sync_stream=False,
+                logging_obj=litellm_logging_obj,
+            )
+            chunk_parsing_logic: Any = vertex_iterator._common_chunk_parsing_logic
+            parsed_chunks = [chunk_parsing_logic(chunk) for chunk in all_chunks]
+        elif "rawPredict" in url_route or "streamRawPredict" in url_route:
+            from litellm.llms.anthropic.chat.handler import ModelResponseIterator
+            from litellm.llms.base_llm.base_model_iterator import (
+                BaseModelResponseIterator,
+            )
+
+            vertex_iterator = ModelResponseIterator(
+                streaming_response=None,
+                sync_stream=False,
+            )
+            chunk_parsing_logic = vertex_iterator.chunk_parser
+            for chunk in all_chunks:
+                dict_chunk = BaseModelResponseIterator._string_to_dict_parser(chunk)
+                if dict_chunk is None:
+                    continue
+                parsed_chunks.append(chunk_parsing_logic(dict_chunk))
+        else:
+            return None
+        if len(parsed_chunks) == 0:
+            return None
         litellm_custom_stream_wrapper = litellm.CustomStreamWrapper(
             completion_stream=vertex_iterator,
             model=model,
@@ -253,11 +279,17 @@ class VertexPassthroughLoggingHandler:
             custom_llm_provider="vertex_ai",
         )
         all_openai_chunks = []
-        for chunk in all_chunks:
-            generic_chunk = vertex_iterator._common_chunk_parsing_logic(chunk)
-            litellm_chunk = litellm_custom_stream_wrapper.chunk_creator(
-                chunk=generic_chunk
-            )
+        for parsed_chunk in parsed_chunks:
+            try:
+                litellm_chunk = litellm_custom_stream_wrapper.chunk_creator(
+                    chunk=parsed_chunk
+                )
+            except Exception as e:
+                verbose_proxy_logger.error(
+                    "Error creating litellm chunk from vertex passthrough endpoint: %s",
+                    str(e),
+                )
+                continue
             if litellm_chunk is not None:
                 all_openai_chunks.append(litellm_chunk)
 
@@ -315,6 +347,7 @@ class VertexPassthroughLoggingHandler:
         response_cost = litellm.completion_cost(
             completion_response=litellm_model_response,
             model=model,
+            custom_llm_provider="vertex_ai",
         )
 
         kwargs["response_cost"] = response_cost

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -303,7 +303,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
     @staticmethod
     def get_endpoint_type(url: str) -> EndpointType:
         parsed_url = urlparse(url)
-        if ("generateContent") in url or ("streamGenerateContent") in url:
+        if (
+            ("generateContent") in url
+            or ("streamGenerateContent") in url
+            or ("rawPredict") in url
+            or ("streamRawPredict") in url
+        ):
             return EndpointType.VERTEX_AI
         elif parsed_url.hostname == "api.anthropic.com":
             return EndpointType.ANTHROPIC

tests/logging_callback_tests/test_unit_tests_init_callbacks.py

@@ -35,6 +35,7 @@ from litellm.integrations.mlflow import MlflowLogger
 from litellm.integrations.argilla import ArgillaLogger
 from litellm.integrations.deepeval.deepeval import DeepEvalLogger
 from litellm.integrations.s3_v2 import S3Logger
+from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger
 from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
 from litellm.integrations.vector_stores.bedrock_vector_store import BedrockVectorStore
 from litellm.integrations.langfuse.langfuse_prompt_management import (
@@ -44,10 +45,18 @@ from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLog
 from litellm.integrations.agentops import AgentOps
 from litellm.integrations.humanloop import HumanloopLogger
 from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler
-from litellm_enterprise.enterprise_callbacks.generic_api_callback import GenericAPILogger
-from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ResendEmailLogger
-from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import SMTPEmailLogger
-from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import PagerDutyAlerting
+from litellm_enterprise.enterprise_callbacks.generic_api_callback import (
+    GenericAPILogger,
+)
+from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import (
+    ResendEmailLogger,
+)
+from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import (
+    SMTPEmailLogger,
+)
+from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import (
+    PagerDutyAlerting,
+)
 from unittest.mock import patch
 
 # clear prometheus collectors / registry
@@ -91,6 +100,7 @@ callback_class_str_to_classType = {
     "smtp_email": SMTPEmailLogger,
     "deepeval": DeepEvalLogger,
     "s3_v2": S3Logger,
+    "langfuse_otel": LangfuseOtelLogger,
 }
 
 expected_env_vars = {