mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix: also wrap plan.response_override and plan.terminate paths, add unit tests
Address review feedback: - Wrap plan.response_override and plan.terminate early-return paths with _maybe_wrap_in_fake_stream() to ensure all exit points handle the streaming conversion correctly - Add unit tests for _maybe_wrap_in_fake_stream() covering: - Dict response with flag=True -> wrapped in FakeAnthropicMessagesStreamIterator - Flag=False -> response returned unchanged - Non-dict response -> returned unchanged - logging_obj=None -> response returned unchanged
This commit is contained in:
parent
6ee175dd24
commit
023d437685
2 changed files with 77 additions and 2 deletions
|
|
@ -4885,14 +4885,18 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
if plan.response_override is not None:
|
||||
return plan.response_override
|
||||
return self._maybe_wrap_in_fake_stream(
|
||||
plan.response_override, logging_obj
|
||||
)
|
||||
if plan.terminate:
|
||||
verbose_logger.debug(
|
||||
"Agentic loop terminated by callback=%s reason=%s",
|
||||
callback.__class__.__name__,
|
||||
plan.stop_reason,
|
||||
)
|
||||
return response
|
||||
return self._maybe_wrap_in_fake_stream(
|
||||
response, logging_obj
|
||||
)
|
||||
if not plan.run_agentic_loop:
|
||||
continue
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,71 @@
|
|||
"""
|
||||
Unit tests for _maybe_wrap_in_fake_stream in BaseLLMHTTPHandler.
|
||||
|
||||
Tests that agentic loop responses are correctly wrapped in
|
||||
FakeAnthropicMessagesStreamIterator when the original request was streaming.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
)
|
||||
|
||||
|
||||
class TestMaybeWrapInFakeStream:
|
||||
def setup_method(self):
|
||||
self.handler = BaseLLMHTTPHandler()
|
||||
|
||||
def test_wraps_dict_when_converted_stream_flag_is_true(self):
|
||||
"""When websearch_interception_converted_stream is True and response is dict, wrap it."""
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {
|
||||
"websearch_interception_converted_stream": True
|
||||
}
|
||||
response = {
|
||||
"id": "msg_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "hello"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
|
||||
result = self.handler._maybe_wrap_in_fake_stream(response, logging_obj)
|
||||
|
||||
assert isinstance(result, FakeAnthropicMessagesStreamIterator)
|
||||
|
||||
def test_returns_response_unchanged_when_flag_is_false(self):
|
||||
"""When flag is False, return response as-is."""
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {
|
||||
"websearch_interception_converted_stream": False
|
||||
}
|
||||
response = {"id": "msg_123", "content": []}
|
||||
|
||||
result = self.handler._maybe_wrap_in_fake_stream(response, logging_obj)
|
||||
|
||||
assert result is response
|
||||
|
||||
def test_returns_response_unchanged_when_not_dict(self):
|
||||
"""When response is not a dict (e.g., already a stream), return as-is."""
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {
|
||||
"websearch_interception_converted_stream": True
|
||||
}
|
||||
response = MagicMock() # Not a dict
|
||||
|
||||
result = self.handler._maybe_wrap_in_fake_stream(response, logging_obj)
|
||||
|
||||
assert result is response
|
||||
|
||||
def test_returns_response_when_logging_obj_is_none(self):
|
||||
"""When logging_obj is None, return response as-is."""
|
||||
response = {"id": "msg_123", "content": []}
|
||||
|
||||
result = self.handler._maybe_wrap_in_fake_stream(response, None)
|
||||
|
||||
assert result is response
|
||||
Loading…
Add table
Reference in a new issue