diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 9f293905d72..347301e7b37 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/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="", diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index ab037b0f1e9..eac0e42ca4a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/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 diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 0038451cca8..bd2303b6f6f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/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 diff --git a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py index 37bae5880b7..624ad0c3e18 100644 --- a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py +++ b/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 = {