From 078e2d341bd99e075e532bc2b688f044907a2f57 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Mar 2025 15:12:01 -0700 Subject: [PATCH 1/7] feat(cost_calculator.py): support reading litellm response cost header in client sdk allows consistent cost tracking when sdk is calling proxy --- litellm/cost_calculator.py | 27 +++++++++++++++++++++- tests/litellm/test_cost_calculator.py | 32 +++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 tests/litellm/test_cost_calculator.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 58600ea14f..e17a94c87e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -808,6 +808,23 @@ def completion_cost( # noqa: PLR0915 raise e +def get_response_cost_from_hidden_params( + hidden_params: Union[dict, BaseModel] +) -> Optional[float]: + if isinstance(hidden_params, BaseModel): + _hidden_params_dict = hidden_params.model_dump() + else: + _hidden_params_dict = hidden_params + + additional_headers = _hidden_params_dict.get("additional_headers", {}) + if additional_headers and "x-litellm-response-cost" in additional_headers: + response_cost = additional_headers["x-litellm-response-cost"] + if response_cost is None: + return None + return float(additional_headers["x-litellm-response-cost"]) + return None + + def response_cost_calculator( response_object: Union[ ModelResponse, @@ -844,7 +861,7 @@ def response_cost_calculator( base_model: Optional[str] = None, custom_pricing: Optional[bool] = None, prompt: str = "", -) -> Optional[float]: +) -> float: """ Returns - float or None: cost of response @@ -856,6 +873,14 @@ def response_cost_calculator( else: if isinstance(response_object, BaseModel): response_object._hidden_params["optional_params"] = optional_params + + if hasattr(response_object, "_hidden_params"): + provider_response_cost = get_response_cost_from_hidden_params( + response_object._hidden_params + ) + if provider_response_cost is not None: + return provider_response_cost + response_cost = completion_cost( completion_response=response_object, model=model, diff --git a/tests/litellm/test_cost_calculator.py b/tests/litellm/test_cost_calculator.py new file mode 100644 index 0000000000..9c9f6d9043 --- /dev/null +++ b/tests/litellm/test_cost_calculator.py @@ -0,0 +1,32 @@ +import json +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from unittest.mock import MagicMock, patch + +from pydantic import BaseModel + +from litellm.cost_calculator import response_cost_calculator + + +def test_cost_calculator(): + class MockResponse(BaseModel): + _hidden_params = {"additional_headers": {"x-litellm-response-cost": 1000}} + + result = response_cost_calculator( + response_object=MockResponse(), + model="", + custom_llm_provider=None, + call_type="", + optional_params={}, + cache_hit=None, + base_model=None, + ) + + assert result == 1000 From 8e27b2026ab351e1a0c43e751f8ec2b08547f9d5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Mar 2025 15:48:31 -0700 Subject: [PATCH 2/7] fix(http_handler.py): support reading ssl security level from env var Allows user to specify lower security settings --- litellm/llms/custom_httpx/http_handler.py | 26 +++++++++++++++-- .../llms/custom_httpx/test_http_handler.py | 29 +++++++++++++++++++ 2 files changed, 53 insertions(+), 2 deletions(-) create mode 100644 tests/litellm/llms/custom_httpx/test_http_handler.py diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 736b85dc53..cafa530130 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1,10 +1,12 @@ import asyncio import os +import ssl import time from typing import TYPE_CHECKING, Any, Callable, List, Mapping, Optional, Union import httpx from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport +from httpx._types import VerifyTypes import litellm from litellm.litellm_core_utils.logging_utils import track_llm_api_timing @@ -94,7 +96,7 @@ class AsyncHTTPHandler: event_hooks: Optional[Mapping[str, List[Callable[..., Any]]]] = None, concurrent_limit=1000, client_alias: Optional[str] = None, # name for client in logs - ssl_verify: Optional[Union[bool, str]] = None, + ssl_verify: Optional[VerifyTypes] = None, ): self.timeout = timeout self.event_hooks = event_hooks @@ -111,13 +113,33 @@ class AsyncHTTPHandler: timeout: Optional[Union[float, httpx.Timeout]], concurrent_limit: int, event_hooks: Optional[Mapping[str, List[Callable[..., Any]]]], - ssl_verify: Optional[Union[bool, str]] = None, + ssl_verify: Optional[VerifyTypes] = None, ) -> httpx.AsyncClient: # SSL certificates (a.k.a CA bundle) used to verify the identity of requested hosts. # /path/to/certificate.pem if ssl_verify is None: ssl_verify = os.getenv("SSL_VERIFY", litellm.ssl_verify) + + ssl_security_level = os.getenv("SSL_SECURITY_LEVEL") + + # If ssl_verify is not False and we need a lower security level + if ( + not ssl_verify + and ssl_security_level + and isinstance(ssl_security_level, str) + ): + # Create a custom SSL context with reduced security level + custom_ssl_context = ssl.create_default_context() + custom_ssl_context.set_ciphers(ssl_security_level) + + # If ssl_verify is a path to a CA bundle, load it into our custom context + if isinstance(ssl_verify, str) and os.path.exists(ssl_verify): + custom_ssl_context.load_verify_locations(cafile=ssl_verify) + + # Use our custom SSL context instead of the original ssl_verify value + ssl_verify = custom_ssl_context + # An SSL certificate used by the requested host to authenticate the client. # /path/to/client.pem cert = os.getenv("SSL_CERTIFICATE", litellm.ssl_certificate) diff --git a/tests/litellm/llms/custom_httpx/test_http_handler.py b/tests/litellm/llms/custom_httpx/test_http_handler.py new file mode 100644 index 0000000000..9ab8e24f85 --- /dev/null +++ b/tests/litellm/llms/custom_httpx/test_http_handler.py @@ -0,0 +1,29 @@ +import io +import os +import pathlib +import ssl +import sys +from unittest.mock import MagicMock + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + + +@pytest.mark.asyncio +async def test_ssl_security_level(monkeypatch): + # Set environment variable for SSL security level + monkeypatch.setenv("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1") + + # Create async client with SSL verification disabled to isolate SSL context testing + client = AsyncHTTPHandler(ssl_verify=False) + + # Get the SSL context from the client + ssl_context = client.client._transport._pool._ssl_context + + # Verify that the SSL context exists and has the correct cipher string + assert isinstance(ssl_context, ssl.SSLContext) From 057c774c140b6badf64658ca0914900b49b710dd Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Mar 2025 16:42:32 -0700 Subject: [PATCH 3/7] fix(http_handler.py): fix typing error --- litellm/llms/custom_httpx/http_handler.py | 1 - litellm/types/llms/custom_http.py | 7 +++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index cafa530130..34d70434d5 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any, Callable, List, Mapping, Optional, Union import httpx from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport -from httpx._types import VerifyTypes import litellm from litellm.litellm_core_utils.logging_utils import track_llm_api_timing diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index 33b7392850..5eec187dd4 100644 --- a/litellm/types/llms/custom_http.py +++ b/litellm/types/llms/custom_http.py @@ -1,6 +1,6 @@ +import ssl from enum import Enum - -import litellm +from typing import Union class httpxSpecialProvider(str, Enum): @@ -19,3 +19,6 @@ class httpxSpecialProvider(str, Enum): SecretManager = "secret_manager" PassThroughEndpoint = "pass_through_endpoint" PromptFactory = "prompt_factory" + + +VerifyTypes = Union[str, bool, ssl.SSLContext] From 594d2ad433272b1ac02749985b56bc539e9079cc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Mar 2025 16:55:32 -0700 Subject: [PATCH 4/7] docs(config_settings.md): document new env var --- docs/my-website/docs/proxy/config_settings.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 9e24437449..cbd0706970 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -499,6 +499,7 @@ router_settings: | SMTP_USERNAME | Username for SMTP authentication (do not set if SMTP does not require auth) | SPEND_LOGS_URL | URL for retrieving spend logs | SSL_CERTIFICATE | Path to the SSL certificate file +| SSL_SECURITY_LEVEL | [BETA] Security level for SSL/TLS connections. E.g. `DEFAULT@SECLEVEL=1` | SSL_VERIFY | Flag to enable or disable SSL certificate verification | SUPABASE_KEY | API key for Supabase service | SUPABASE_URL | Base URL for Supabase instance From dd9e79adbd45f583d764dede7c3a3cdf0707eb82 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Mar 2025 17:26:21 -0700 Subject: [PATCH 5/7] fix(streaming_handler.py): emit deep copy of completed chunk --- litellm/litellm_core_utils/logging_utils.py | 3 +++ .../litellm_core_utils/streaming_handler.py | 23 +++++++++++++------ tests/local_testing/test_caching.py | 4 ++-- 3 files changed, 21 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 3c934a4276..c7512ea146 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -78,6 +78,9 @@ def _assemble_complete_response_from_streaming_chunks( Union[ModelResponse, TextCompletionResponse] ] = None + if isinstance(result, ModelResponse): + return result + if result.choices[0].finish_reason is not None: # if it's the last chunk streaming_chunks.append(result) try: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 15d94b31a9..8fc63db5eb 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1472,6 +1472,15 @@ class CustomStreamWrapper: """ self.logging_loop = loop + def cache_streaming_response(self, processed_chunk, cache_hit: bool): + """ + Caches the streaming response + """ + if not cache_hit and self.logging_obj._llm_caching_handler is not None: + self.logging_obj._llm_caching_handler._sync_add_streaming_response_to_cache( + processed_chunk + ) + def run_success_logging_and_cache_storage(self, processed_chunk, cache_hit: bool): """ Runs success logging in a thread and adds the response to the cache @@ -1503,12 +1512,6 @@ class CustomStreamWrapper: ## SYNC LOGGING self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) - ## Sync store in cache - if self.logging_obj._llm_caching_handler is not None: - self.logging_obj._llm_caching_handler._sync_add_streaming_response_to_cache( - processed_chunk - ) - def finish_reason_handler(self): model_response = self.model_response_creator() _finish_reason = self.received_finish_reason or self.intermittent_finish_reason @@ -1603,9 +1606,15 @@ class CustomStreamWrapper: "usage", getattr(complete_streaming_response, "usage"), ) + self.cache_streaming_response( + processed_chunk=complete_streaming_response.model_copy( + deep=True + ), + cache_hit=cache_hit, + ) executor.submit( self.logging_obj.success_handler, - complete_streaming_response, + complete_streaming_response.model_copy(deep=True), None, None, cache_hit, diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py index df2afdc167..7c6e400c05 100644 --- a/tests/local_testing/test_caching.py +++ b/tests/local_testing/test_caching.py @@ -2151,7 +2151,7 @@ def test_logging_turn_off_message_logging_streaming(): mock_obj = Cache(type="local") litellm.cache = mock_obj - with patch.object(mock_obj, "add_cache", new=MagicMock()) as mock_client: + with patch.object(mock_obj, "add_cache") as mock_client: print(f"mock_obj.add_cache: {mock_obj.add_cache}") resp = litellm.completion( @@ -2167,7 +2167,7 @@ def test_logging_turn_off_message_logging_streaming(): time.sleep(1) mock_client.assert_called_once() - + print(f"mock_client.call_args: {mock_client.call_args}") assert mock_client.call_args.args[0].choices[0].message.content == "hello" From c4b2e0ae3d9a1c479bde78ac5f1440a4a0b83637 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Mar 2025 18:10:39 -0700 Subject: [PATCH 6/7] fix(streaming_handler.py): support logging complete streaming response on cache hit --- litellm/caching/caching_handler.py | 2 +- .../litellm_core_utils/streaming_handler.py | 24 ++++--- tests/local_testing/test_caching.py | 62 ++++++++++++++----- 3 files changed, 66 insertions(+), 22 deletions(-) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 2a958c9eee..09fabf1c12 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -790,6 +790,7 @@ class LLMCachingHandler: - Else append the chunk to self.async_streaming_chunks """ + complete_streaming_response: Optional[ Union[ModelResponse, TextCompletionResponse] ] = _assemble_complete_response_from_streaming_chunks( @@ -800,7 +801,6 @@ class LLMCachingHandler: streaming_chunks=self.async_streaming_chunks, is_async=True, ) - # if a complete_streaming_response is assembled, add it to the cache if complete_streaming_response is not None: await self.async_set_cache( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 8fc63db5eb..56e64d1859 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1481,6 +1481,15 @@ class CustomStreamWrapper: processed_chunk ) + async def async_cache_streaming_response(self, processed_chunk, cache_hit: bool): + """ + Caches the streaming response + """ + if not cache_hit and self.logging_obj._llm_caching_handler is not None: + await self.logging_obj._llm_caching_handler._add_streaming_response_to_cache( + processed_chunk + ) + def run_success_logging_and_cache_storage(self, processed_chunk, cache_hit: bool): """ Runs success logging in a thread and adds the response to the cache @@ -1711,13 +1720,6 @@ class CustomStreamWrapper: if processed_chunk is None: continue - if self.logging_obj._llm_caching_handler is not None: - asyncio.create_task( - self.logging_obj._llm_caching_handler._add_streaming_response_to_cache( - processed_chunk=cast(ModelResponse, processed_chunk), - ) - ) - choice = processed_chunk.choices[0] if isinstance(choice, StreamingChoices): self.response_uptil_now += choice.delta.get("content", "") or "" @@ -1788,6 +1790,14 @@ class CustomStreamWrapper: "usage", getattr(complete_streaming_response, "usage"), ) + asyncio.create_task( + self.async_cache_streaming_response( + processed_chunk=complete_streaming_response.model_copy( + deep=True + ), + cache_hit=cache_hit, + ) + ) if self.sent_stream_usage is False and self.send_stream_usage is True: self.sent_stream_usage = True return response diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py index 7c6e400c05..ac04d06c12 100644 --- a/tests/local_testing/test_caching.py +++ b/tests/local_testing/test_caching.py @@ -8,6 +8,7 @@ from dotenv import load_dotenv load_dotenv() import os +import json sys.path.insert( 0, os.path.abspath("../..") @@ -2146,29 +2147,62 @@ async def test_redis_proxy_batch_redis_get_cache(): assert "cache_key" in response._hidden_params -def test_logging_turn_off_message_logging_streaming(): +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_logging_turn_off_message_logging_streaming(sync_mode): litellm.turn_off_message_logging = True mock_obj = Cache(type="local") litellm.cache = mock_obj - with patch.object(mock_obj, "add_cache") as mock_client: + with patch.object(mock_obj, "add_cache") as mock_client, patch.object( + mock_obj, "async_add_cache" + ) as mock_async_client: print(f"mock_obj.add_cache: {mock_obj.add_cache}") - resp = litellm.completion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hi"}], - mock_response="hello", - stream=True, - ) + if sync_mode is True: + resp = litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello", + stream=True, + ) - for chunk in resp: - continue + for chunk in resp: + continue - time.sleep(1) + time.sleep(1) + mock_client.assert_called_once() + print(f"mock_client.call_args: {mock_client.call_args}") + assert mock_client.call_args.args[0].choices[0].message.content == "hello" + else: + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello", + stream=True, + ) - mock_client.assert_called_once() - print(f"mock_client.call_args: {mock_client.call_args}") - assert mock_client.call_args.args[0].choices[0].message.content == "hello" + async for chunk in resp: + continue + + await asyncio.sleep(1) + + mock_async_client.assert_called_once() + print(f"mock_async_client.call_args: {mock_async_client.call_args.args[0]}") + print( + f"mock_async_client.call_args: {json.loads(mock_async_client.call_args.args[0])}" + ) + json_mock = json.loads(mock_async_client.call_args.args[0]) + try: + assert json_mock["choices"][0]["message"]["content"] == "hello" + except Exception as e: + print( + f"mock_async_client.call_args.args[0]: {mock_async_client.call_args.args[0]}" + ) + print( + f"mock_async_client.call_args.args[0]['choices']: {mock_async_client.call_args.args[0]['choices']}" + ) + raise e def test_basic_caching_import(): From 359b8298f82b710a77803e42eb2e7c60f4fb1522 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Mar 2025 18:17:01 -0700 Subject: [PATCH 7/7] test(test_assemble_streaming_responses.py): update test to use correct type --- .../test_assemble_streaming_responses.py | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/tests/logging_callback_tests/test_assemble_streaming_responses.py b/tests/logging_callback_tests/test_assemble_streaming_responses.py index 1101350fa2..20e46db229 100644 --- a/tests/logging_callback_tests/test_assemble_streaming_responses.py +++ b/tests/logging_callback_tests/test_assemble_streaming_responses.py @@ -24,7 +24,14 @@ import pytest from respx import MockRouter import litellm -from litellm import Choices, Message, ModelResponse, TextCompletionResponse, TextChoices +from litellm import ( + Choices, + Message, + ModelResponse, + ModelResponseStream, + TextCompletionResponse, + TextChoices, +) from litellm.litellm_core_utils.logging_utils import ( _assemble_complete_response_from_streaming_chunks, @@ -63,7 +70,7 @@ def test_assemble_complete_response_from_streaming_chunks_1(is_async): "system_fingerprint": None, "usage": None, } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = ModelResponseStream(**chunk) complete_streaming_response = _assemble_complete_response_from_streaming_chunks( result=chunk, start_time=datetime.now(), @@ -103,7 +110,7 @@ def test_assemble_complete_response_from_streaming_chunks_1(is_async): "system_fingerprint": None, "usage": None, } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = ModelResponseStream(**chunk) complete_streaming_response = _assemble_complete_response_from_streaming_chunks( result=chunk, start_time=datetime.now(), @@ -164,7 +171,7 @@ def test_assemble_complete_response_from_streaming_chunks_2(is_async): "system_fingerprint": None, "usage": None, } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = ModelResponseStream(**chunk) chunk = _text_completion_stream_wrapper.convert_to_text_completion_object(chunk) complete_streaming_response = _assemble_complete_response_from_streaming_chunks( @@ -206,7 +213,7 @@ def test_assemble_complete_response_from_streaming_chunks_2(is_async): "system_fingerprint": None, "usage": None, } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = ModelResponseStream(**chunk) chunk = _text_completion_stream_wrapper.convert_to_text_completion_object(chunk) complete_streaming_response = _assemble_complete_response_from_streaming_chunks( result=chunk, @@ -261,7 +268,7 @@ def test_assemble_complete_response_from_streaming_chunks_3(is_async): "system_fingerprint": None, "usage": None, } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = ModelResponseStream(**chunk) complete_streaming_response = _assemble_complete_response_from_streaming_chunks( result=chunk, start_time=datetime.now(), @@ -338,7 +345,7 @@ def test_assemble_complete_response_from_streaming_chunks_4(is_async): "system_fingerprint": None, "usage": None, } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = ModelResponseStream(**chunk) # remove attribute id from chunk del chunk.object