From fdecafd3475f6a6eb1a83c4ade465ffdaa8c5c6e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 2 Sep 2024 15:51:21 -0700 Subject: [PATCH 1/7] new streaming handler fn --- .../streaming_handler.py | 94 +++++++++++++++++++ 1 file changed, 94 insertions(+) create mode 100644 litellm/proxy/pass_through_endpoints/streaming_handler.py diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py new file mode 100644 index 00000000000..ba0359317d7 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -0,0 +1,94 @@ +import asyncio +import json +from datetime import datetime +from enum import Enum +from typing import AsyncIterable, Dict, List, Optional, Union + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator as VertexAIIterator, +) +from litellm.types.utils import GenericStreamingChunk + + +class ModelIteratorType(Enum): + VERTEX_AI = "vertexAI" + # Add more iterator types here as needed + + +MODEL_ITERATORS: Dict[ModelIteratorType, type] = { + ModelIteratorType.VERTEX_AI: VertexAIIterator, + # Add more mappings here as needed +} + + +def get_litellm_chunk( + model_iterator: VertexAIIterator, + custom_stream_wrapper: litellm.utils.CustomStreamWrapper, + chunk_dict: Dict, +) -> Optional[Dict]: + generic_chunk: GenericStreamingChunk = model_iterator.chunk_parser(chunk_dict) + if generic_chunk: + return custom_stream_wrapper.chunk_creator(chunk=generic_chunk) + return None + + +async def chunk_processor( + aiter_bytes: AsyncIterable[bytes], + litellm_logging_obj: LiteLLMLoggingObj, + iterator_type: ModelIteratorType, + start_time: datetime, +) -> AsyncIterable[bytes]: + + IteratorClass = MODEL_ITERATORS[iterator_type] + model_iterator = IteratorClass(sync_stream=False, streaming_response=aiter_bytes) + custom_stream_wrapper = litellm.utils.CustomStreamWrapper( + completion_stream=aiter_bytes, model=None, logging_obj=litellm_logging_obj + ) + buffer = b"" + all_chunks = [] + async for chunk in aiter_bytes: + buffer += chunk + try: + _decoded_chunk = chunk.decode("utf-8") + _chunk_dict = json.loads(_decoded_chunk) + litellm_chunk = get_litellm_chunk( + model_iterator, custom_stream_wrapper, _chunk_dict + ) + if litellm_chunk: + all_chunks.append(litellm_chunk) + except json.JSONDecodeError: + pass + finally: + yield chunk # Yield the original bytes + + # Process any remaining data in the buffer + if buffer: + try: + _chunk_dict = json.loads(buffer.decode("utf-8")) + + if isinstance(_chunk_dict, list): + for _chunk in _chunk_dict: + litellm_chunk = get_litellm_chunk( + model_iterator, custom_stream_wrapper, _chunk + ) + if litellm_chunk: + all_chunks.append(litellm_chunk) + elif isinstance(_chunk_dict, dict): + litellm_chunk = get_litellm_chunk( + model_iterator, custom_stream_wrapper, _chunk_dict + ) + if litellm_chunk: + all_chunks.append(litellm_chunk) + except json.JSONDecodeError: + pass + + complete_streaming_response = litellm.stream_chunk_builder(chunks=all_chunks) + + end_time = datetime.now() + await litellm_logging_obj.async_success_handler( + result=complete_streaming_response, + start_time=start_time, + end_time=end_time, + ) From 73d0a7844432884853b48526defee541b2978da0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 2 Sep 2024 15:51:52 -0700 Subject: [PATCH 2/7] use chunk_processort --- .../pass_through_endpoints.py | 21 ++++++++++++++----- .../tests/test_vertex_sdk_forward_headers.py | 9 ++++++-- 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index f34efdcf390..e138df0096a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -22,6 +22,9 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, +) from litellm.proxy._types import ( ConfigFieldInfo, ConfigFieldUpdate, @@ -32,6 +35,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from .streaming_handler import ModelIteratorType, chunk_processor from .success_handler import PassThroughEndpointLogging router = APIRouter() @@ -416,9 +420,13 @@ async def pass_through_request( status_code=e.response.status_code, detail=await e.response.aread() ) - # Create an async generator to yield the response content async def stream_response() -> AsyncIterable[bytes]: - async for chunk in response.aiter_bytes(): + async for chunk in chunk_processor( + response.aiter_bytes(), + litellm_logging_obj=logging_obj, + iterator_type=ModelIteratorType.VERTEX_AI, + start_time=start_time, + ): yield chunk return StreamingResponse( @@ -454,10 +462,13 @@ async def pass_through_request( status_code=e.response.status_code, detail=await e.response.aread() ) - # streaming response - # Create an async generator to yield the response content async def stream_response() -> AsyncIterable[bytes]: - async for chunk in response.aiter_bytes(): + async for chunk in chunk_processor( + response.aiter_bytes(), + litellm_logging_obj=logging_obj, + iterator_type=ModelIteratorType.VERTEX_AI, + start_time=start_time, + ): yield chunk return StreamingResponse( diff --git a/litellm/proxy/tests/test_vertex_sdk_forward_headers.py b/litellm/proxy/tests/test_vertex_sdk_forward_headers.py index 0799ef8eb8b..7aa87905ab6 100644 --- a/litellm/proxy/tests/test_vertex_sdk_forward_headers.py +++ b/litellm/proxy/tests/test_vertex_sdk_forward_headers.py @@ -10,7 +10,12 @@ vertexai.init( api_transport="rest", ) -model = GenerativeModel(model_name="gemini-1.0-pro") -response = model.generate_content("hi") +model = GenerativeModel(model_name="gemini-1.5-flash-001") +response = model.generate_content( + "hi tell me a joke and a very long story", stream=True +) print("response", response) + +for chunk in response: + print(chunk) From a6d4a27207af843bfc947adc6f92e7941585de16 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 2 Sep 2024 16:11:20 -0700 Subject: [PATCH 3/7] pass through track usage for streaming endpoints --- .../pass_through_endpoints.py | 4 +++ .../streaming_handler.py | 29 ++++++++++++++----- 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index e138df0096a..99c6faad0f5 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -426,6 +426,8 @@ async def pass_through_request( litellm_logging_obj=logging_obj, iterator_type=ModelIteratorType.VERTEX_AI, start_time=start_time, + passthrough_success_handler_obj=pass_through_endpoint_logging, + url_route=str(url), ): yield chunk @@ -468,6 +470,8 @@ async def pass_through_request( litellm_logging_obj=logging_obj, iterator_type=ModelIteratorType.VERTEX_AI, start_time=start_time, + passthrough_success_handler_obj=pass_through_endpoint_logging, + url_route=str(url), ): yield chunk diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index ba0359317d7..8513e2702b6 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -11,6 +11,8 @@ from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_stu ) from litellm.types.utils import GenericStreamingChunk +from .success_handler import PassThroughEndpointLogging + class ModelIteratorType(Enum): VERTEX_AI = "vertexAI" @@ -28,6 +30,7 @@ def get_litellm_chunk( custom_stream_wrapper: litellm.utils.CustomStreamWrapper, chunk_dict: Dict, ) -> Optional[Dict]: + generic_chunk: GenericStreamingChunk = model_iterator.chunk_parser(chunk_dict) if generic_chunk: return custom_stream_wrapper.chunk_creator(chunk=generic_chunk) @@ -39,6 +42,8 @@ async def chunk_processor( litellm_logging_obj: LiteLLMLoggingObj, iterator_type: ModelIteratorType, start_time: datetime, + passthrough_success_handler_obj: PassThroughEndpointLogging, + url_route: str, ) -> AsyncIterable[bytes]: IteratorClass = MODEL_ITERATORS[iterator_type] @@ -84,11 +89,21 @@ async def chunk_processor( except json.JSONDecodeError: pass - complete_streaming_response = litellm.stream_chunk_builder(chunks=all_chunks) - - end_time = datetime.now() - await litellm_logging_obj.async_success_handler( - result=complete_streaming_response, - start_time=start_time, - end_time=end_time, + complete_streaming_response: litellm.ModelResponse = litellm.stream_chunk_builder( + chunks=all_chunks + ) + end_time = datetime.now() + + if passthrough_success_handler_obj.is_vertex_route(url_route): + _model = passthrough_success_handler_obj.extract_model_from_url(url_route) + complete_streaming_response.model = _model + litellm_logging_obj.model = _model + litellm_logging_obj.model_call_details["model"] = _model + + asyncio.create_task( + litellm_logging_obj.async_success_handler( + result=complete_streaming_response, + start_time=start_time, + end_time=end_time, + ) ) From e9427205ef1a0ab7736c84491d757c8b5995139e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 2 Sep 2024 16:17:49 -0700 Subject: [PATCH 4/7] add test for pass through streaming usage tracking --- tests/pass_through_tests/test_vertex_ai.py | 34 ++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/pass_through_tests/test_vertex_ai.py b/tests/pass_through_tests/test_vertex_ai.py index 25fd60c72e3..40998dc2f17 100644 --- a/tests/pass_through_tests/test_vertex_ai.py +++ b/tests/pass_through_tests/test_vertex_ai.py @@ -117,3 +117,37 @@ async def test_basic_vertex_ai_pass_through_with_spendlog(): ) pass + + +@pytest.mark.asyncio() +async def test_basic_vertex_ai_pass_through_streaming_with_spendlog(): + + spend_before = await call_spend_logs_endpoint() or 0.0 + print("spend_before", spend_before) + load_vertex_ai_credentials() + + vertexai.init( + project="adroit-crow-413218", + location="us-central1", + api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex-ai", + api_transport="rest", + ) + + model = GenerativeModel(model_name="gemini-1.0-pro") + response = model.generate_content("hi", stream=True) + + for chunk in response: + print("chunk", chunk) + + print("response", response) + + await asyncio.sleep(20) + spend_after = await call_spend_logs_endpoint() + print("spend_after", spend_after) + assert ( + spend_after > spend_before + ), "Spend should be greater than before. spend_before: {}, spend_after: {}".format( + spend_before, spend_after + ) + + pass From 42b95c5979cc35369ce75c3ec7b53d257a45524e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 2 Sep 2024 16:36:19 -0700 Subject: [PATCH 5/7] code cleanup --- .../pass_through_endpoints.py | 15 +- .../streaming_handler.py | 134 +++++++++--------- litellm/proxy/pass_through_endpoints/types.py | 6 + 3 files changed, 88 insertions(+), 67 deletions(-) create mode 100644 litellm/proxy/pass_through_endpoints/types.py diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 99c6faad0f5..1dc9784350a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -35,8 +35,9 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from .streaming_handler import ModelIteratorType, chunk_processor +from .streaming_handler import chunk_processor from .success_handler import PassThroughEndpointLogging +from .types import EndpointType router = APIRouter() @@ -288,6 +289,12 @@ def get_response_headers(headers: httpx.Headers) -> dict: return return_headers +def get_endpoint_type(url: str) -> EndpointType: + if ("generateContent") in url or ("streamGenerateContent") in url: + return EndpointType.VERTEX_AI + return EndpointType.GENERIC + + async def pass_through_request( request: Request, target: str, @@ -311,6 +318,8 @@ async def pass_through_request( request=request, headers=headers, forward_headers=forward_headers ) + endpoint_type: EndpointType = get_endpoint_type(str(url)) + _parsed_body = None if custom_body: _parsed_body = custom_body @@ -424,7 +433,7 @@ async def pass_through_request( async for chunk in chunk_processor( response.aiter_bytes(), litellm_logging_obj=logging_obj, - iterator_type=ModelIteratorType.VERTEX_AI, + endpoint_type=endpoint_type, start_time=start_time, passthrough_success_handler_obj=pass_through_endpoint_logging, url_route=str(url), @@ -468,7 +477,7 @@ async def pass_through_request( async for chunk in chunk_processor( response.aiter_bytes(), litellm_logging_obj=logging_obj, - iterator_type=ModelIteratorType.VERTEX_AI, + endpoint_type=endpoint_type, start_time=start_time, passthrough_success_handler_obj=pass_through_endpoint_logging, url_route=str(url), diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 8513e2702b6..ab1d5d813ff 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -12,17 +12,7 @@ from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_stu from litellm.types.utils import GenericStreamingChunk from .success_handler import PassThroughEndpointLogging - - -class ModelIteratorType(Enum): - VERTEX_AI = "vertexAI" - # Add more iterator types here as needed - - -MODEL_ITERATORS: Dict[ModelIteratorType, type] = { - ModelIteratorType.VERTEX_AI: VertexAIIterator, - # Add more mappings here as needed -} +from .types import EndpointType def get_litellm_chunk( @@ -37,73 +27,89 @@ def get_litellm_chunk( return None +def get_iterator_class_from_endpoint_type( + endpoint_type: EndpointType, +) -> Optional[type]: + if endpoint_type == EndpointType.VERTEX_AI: + return VertexAIIterator + return None + + async def chunk_processor( aiter_bytes: AsyncIterable[bytes], litellm_logging_obj: LiteLLMLoggingObj, - iterator_type: ModelIteratorType, + endpoint_type: EndpointType, start_time: datetime, passthrough_success_handler_obj: PassThroughEndpointLogging, url_route: str, ) -> AsyncIterable[bytes]: - IteratorClass = MODEL_ITERATORS[iterator_type] - model_iterator = IteratorClass(sync_stream=False, streaming_response=aiter_bytes) - custom_stream_wrapper = litellm.utils.CustomStreamWrapper( - completion_stream=aiter_bytes, model=None, logging_obj=litellm_logging_obj - ) - buffer = b"" - all_chunks = [] - async for chunk in aiter_bytes: - buffer += chunk - try: - _decoded_chunk = chunk.decode("utf-8") - _chunk_dict = json.loads(_decoded_chunk) - litellm_chunk = get_litellm_chunk( - model_iterator, custom_stream_wrapper, _chunk_dict - ) - if litellm_chunk: - all_chunks.append(litellm_chunk) - except json.JSONDecodeError: - pass - finally: - yield chunk # Yield the original bytes - - # Process any remaining data in the buffer - if buffer: - try: - _chunk_dict = json.loads(buffer.decode("utf-8")) - - if isinstance(_chunk_dict, list): - for _chunk in _chunk_dict: - litellm_chunk = get_litellm_chunk( - model_iterator, custom_stream_wrapper, _chunk - ) - if litellm_chunk: - all_chunks.append(litellm_chunk) - elif isinstance(_chunk_dict, dict): + iteratorClass = get_iterator_class_from_endpoint_type(endpoint_type) + if iteratorClass is None: + # Generic endpoint - litellm does not do any tracking / logging for this + async for chunk in aiter_bytes: + yield chunk + else: + # known streaming endpoint - litellm will do tracking / logging for this + model_iterator = iteratorClass( + sync_stream=False, streaming_response=aiter_bytes + ) + custom_stream_wrapper = litellm.utils.CustomStreamWrapper( + completion_stream=aiter_bytes, model=None, logging_obj=litellm_logging_obj + ) + buffer = b"" + all_chunks = [] + async for chunk in aiter_bytes: + buffer += chunk + try: + _decoded_chunk = chunk.decode("utf-8") + _chunk_dict = json.loads(_decoded_chunk) litellm_chunk = get_litellm_chunk( model_iterator, custom_stream_wrapper, _chunk_dict ) if litellm_chunk: all_chunks.append(litellm_chunk) - except json.JSONDecodeError: - pass + except json.JSONDecodeError: + pass + finally: + yield chunk # Yield the original bytes - complete_streaming_response: litellm.ModelResponse = litellm.stream_chunk_builder( - chunks=all_chunks - ) - end_time = datetime.now() + # Process any remaining data in the buffer + if buffer: + try: + _chunk_dict = json.loads(buffer.decode("utf-8")) - if passthrough_success_handler_obj.is_vertex_route(url_route): - _model = passthrough_success_handler_obj.extract_model_from_url(url_route) - complete_streaming_response.model = _model - litellm_logging_obj.model = _model - litellm_logging_obj.model_call_details["model"] = _model + if isinstance(_chunk_dict, list): + for _chunk in _chunk_dict: + litellm_chunk = get_litellm_chunk( + model_iterator, custom_stream_wrapper, _chunk + ) + if litellm_chunk: + all_chunks.append(litellm_chunk) + elif isinstance(_chunk_dict, dict): + litellm_chunk = get_litellm_chunk( + model_iterator, custom_stream_wrapper, _chunk_dict + ) + if litellm_chunk: + all_chunks.append(litellm_chunk) + except json.JSONDecodeError: + pass - asyncio.create_task( - litellm_logging_obj.async_success_handler( - result=complete_streaming_response, - start_time=start_time, - end_time=end_time, + complete_streaming_response: litellm.ModelResponse = ( + litellm.stream_chunk_builder(chunks=all_chunks) + ) + end_time = datetime.now() + + if passthrough_success_handler_obj.is_vertex_route(url_route): + _model = passthrough_success_handler_obj.extract_model_from_url(url_route) + complete_streaming_response.model = _model + litellm_logging_obj.model = _model + litellm_logging_obj.model_call_details["model"] = _model + + asyncio.create_task( + litellm_logging_obj.async_success_handler( + result=complete_streaming_response, + start_time=start_time, + end_time=end_time, + ) ) - ) diff --git a/litellm/proxy/pass_through_endpoints/types.py b/litellm/proxy/pass_through_endpoints/types.py new file mode 100644 index 00000000000..662788af087 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/types.py @@ -0,0 +1,6 @@ +from enum import Enum + + +class EndpointType(str, Enum): + VERTEX_AI = "vertex-ai" + GENERIC = "generic" From f89487a496220d9e745b7675c3819a8bf7811918 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 2 Sep 2024 17:08:03 -0700 Subject: [PATCH 6/7] fix linting error --- litellm/proxy/pass_through_endpoints/streaming_handler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index ab1d5d813ff..4420bd1d70a 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -95,9 +95,9 @@ async def chunk_processor( except json.JSONDecodeError: pass - complete_streaming_response: litellm.ModelResponse = ( - litellm.stream_chunk_builder(chunks=all_chunks) - ) + complete_streaming_response: Optional[ + Union[litellm.ModelResponse, litellm.TextCompletionResponse] + ] = litellm.stream_chunk_builder(chunks=all_chunks) end_time = datetime.now() if passthrough_success_handler_obj.is_vertex_route(url_route): From fa6b09f1474aff1c41085a5874cfa46958a761b2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 2 Sep 2024 18:13:32 -0700 Subject: [PATCH 7/7] fix linting error --- litellm/proxy/pass_through_endpoints/streaming_handler.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 4420bd1d70a..b7faa21e46a 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -98,6 +98,8 @@ async def chunk_processor( complete_streaming_response: Optional[ Union[litellm.ModelResponse, litellm.TextCompletionResponse] ] = litellm.stream_chunk_builder(chunks=all_chunks) + if complete_streaming_response is None: + complete_streaming_response = litellm.ModelResponse() end_time = datetime.now() if passthrough_success_handler_obj.is_vertex_route(url_route):