scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

fixing PR comments

Browse source
This commit is contained in:
devarakondasrikanth 2026-03-13 20:28:52 -07:00
parent 6e89167fd4
commit e12ab768e9
2 changed files with 17 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/proxy_server.py
View file

@ -5565,7 +5565,8 @@ async def async_data_generator(
request_data=request_data
)
model_mismatch_logged = False
# Chunks for post-guardrail log: use exclude_none=True only so stream_chunk_builder gets required keys
# Chunks for post-guardrail log: use exclude_none=True only so stream_chunk_builder gets required keys.
# This list holds the full stream in memory until the background post-guardrail log task completes.
_streaming_chunks_for_log: List[Dict[str, Any]] = []
# Use a running string instead of list + join to avoid O(n^2) overhead.
# Previously "".join(str_so_far_parts) was called every chunk, re-joining

26
tests/test_litellm/proxy/hooks/test_async_post_guardrail_log_success_event.py
View file

@ -85,18 +85,22 @@ async def test_post_guardrail_log_base_custom_logger_not_invoked():
"""Base CustomLogger (no override) is not invoked; only overriders are called."""
base_logger = CustomLogger()
overriding_logger = PostGuardrailLogger()
# Base first, then override: only overriding should be called
with patch("litellm.callbacks", [base_logger, overriding_logger]):
from litellm.proxy.utils import ProxyLogging
from litellm.caching.caching import DualCache
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
await proxy_logging.async_post_guardrail_log_success_event(
data={"model": "gpt-4"},
response=ModelResponse(id="r1", choices=[], model="gpt-4"),
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
)
with patch.object(
base_logger,
"async_post_guardrail_log_success_event",
new_callable=AsyncMock,
) as mock_base_method:
with patch("litellm.callbacks", [base_logger, overriding_logger]):
from litellm.proxy.utils import ProxyLogging
from litellm.caching.caching import DualCache
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
await proxy_logging.async_post_guardrail_log_success_event(
data={"model": "gpt-4"},
response=ModelResponse(id="r1", choices=[], model="gpt-4"),
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
)
mock_base_method.assert_not_called()
assert overriding_logger.called is True