mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
Fix tests
This commit is contained in:
parent
1926a8b778
commit
171a3a8ff1
3 changed files with 59 additions and 116 deletions
|
|
@ -9,10 +9,9 @@ interceptor detection, loop logic, and message assembly all run for real.
|
|||
"""
|
||||
|
||||
from typing import Dict
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
import litellm
|
||||
|
||||
ADVISOR_TOOL = {
|
||||
"type": "advisor_20260301",
|
||||
|
|
@ -35,9 +34,7 @@ def _text_resp(text: str, model: str = "gpt-4o-mini") -> Dict:
|
|||
}
|
||||
|
||||
|
||||
def _advisor_call_resp(
|
||||
question: str = "How do I approach this?", tool_id: str = "tid_01"
|
||||
) -> Dict:
|
||||
def _advisor_call_resp(question: str = "How do I approach this?", tool_id: str = "tid_01") -> Dict:
|
||||
return {
|
||||
"id": "msg_int_test",
|
||||
"type": "message",
|
||||
|
|
@ -72,20 +69,22 @@ async def test_full_dispatch_interceptor_fires_and_loop_completes():
|
|||
anthropic_messages,
|
||||
)
|
||||
|
||||
call_count = 0
|
||||
executor_call_count = 0
|
||||
|
||||
async def mock_handler(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
return _advisor_call_resp() # executor: calls advisor
|
||||
if call_count == 2:
|
||||
return _text_resp("Use trial division.", model="claude-opus-4-6") # advisor
|
||||
return _text_resp("def is_prime(n): ...") # executor: final
|
||||
async def mock_messages(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
nonlocal executor_call_count
|
||||
executor_call_count += 1
|
||||
if executor_call_count == 1:
|
||||
return _advisor_call_resp()
|
||||
return _text_resp("def is_prime(n): ...")
|
||||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_messages_handler",
|
||||
side_effect=mock_handler,
|
||||
side_effect=mock_messages,
|
||||
), patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_advisor_with_router",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_text_resp("Use trial division.", model="claude-opus-4-6"),
|
||||
):
|
||||
result = await anthropic_messages(
|
||||
model="openai/gpt-4o-mini",
|
||||
|
|
@ -96,14 +95,14 @@ async def test_full_dispatch_interceptor_fires_and_loop_completes():
|
|||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
# 2 calls to _call_messages_handler: executor -> executor-final.
|
||||
# Advisor subcall runs via _call_advisor_with_router (acompletion path).
|
||||
assert call_count == 2
|
||||
assert executor_call_count == 2
|
||||
|
||||
assert isinstance(result, dict)
|
||||
content = result.get("content", [])
|
||||
text_blocks = [b for b in content if b.get("type") == "text"]
|
||||
advisor_uses = [b for b in content if b.get("type") == "tool_use" and b.get("name") == "advisor"]
|
||||
advisor_uses = [
|
||||
b for b in content if b.get("type") == "tool_use" and b.get("name") == "consult_advisor"
|
||||
]
|
||||
|
||||
assert len(text_blocks) >= 1, "Final response must have text"
|
||||
assert len(advisor_uses) == 0, "No advisor tool_use blocks must appear in final output"
|
||||
|
|
@ -129,15 +128,16 @@ async def test_max_uses_enforced_through_full_handler():
|
|||
|
||||
advisor_tool_capped = {**ADVISOR_TOOL, "max_uses": 1}
|
||||
|
||||
async def mock_handler(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
# Advisor always returns text; executor always calls advisor
|
||||
if tools is None:
|
||||
return _text_resp("Some advice.", model="claude-opus-4-6")
|
||||
async def mock_messages(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
return _advisor_call_resp()
|
||||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_messages_handler",
|
||||
side_effect=mock_handler,
|
||||
side_effect=mock_messages,
|
||||
), patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_advisor_with_router",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_text_resp("Some advice.", model="claude-opus-4-6"),
|
||||
):
|
||||
with pytest.raises(AdvisorMaxIterationsError):
|
||||
await anthropic_messages(
|
||||
|
|
@ -187,66 +187,3 @@ async def test_anthropic_provider_bypasses_interceptor():
|
|||
content = result.get("content", []) if isinstance(result, dict) else []
|
||||
text_blocks = [b for b in content if b.get("type") == "text"]
|
||||
assert any("Native anthropic" in b.get("text", "") for b in text_blocks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_messages_interceptor_path_still_runs_non_advisor_pre_request_hooks():
|
||||
"""
|
||||
Pre-request hooks should still run for intercepted /messages requests, while
|
||||
advisor tool conversion must be skipped so interceptors can still see the
|
||||
native advisor_20260301 tool.
|
||||
"""
|
||||
from litellm.integrations.advisor_interception.handler import (
|
||||
AdvisorInterceptionLogger,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
|
||||
anthropic_messages,
|
||||
)
|
||||
|
||||
class _NonAdvisorHookLogger(CustomLogger):
|
||||
async def async_pre_request_hook(self, model, messages, kwargs):
|
||||
kwargs["metadata"] = {"from_non_advisor_hook": True}
|
||||
return kwargs
|
||||
|
||||
class _CaptureInterceptor:
|
||||
def __init__(self):
|
||||
self.seen_kwargs = {}
|
||||
self.seen_tools = None
|
||||
|
||||
def can_handle(self, tools, custom_llm_provider):
|
||||
return True
|
||||
|
||||
async def handle(self, **kwargs):
|
||||
self.seen_kwargs = kwargs
|
||||
self.seen_tools = kwargs.get("tools")
|
||||
return _text_resp("interceptor handled request")
|
||||
|
||||
capture_interceptor = _CaptureInterceptor()
|
||||
original_callbacks = litellm.callbacks
|
||||
try:
|
||||
litellm.callbacks = [
|
||||
AdvisorInterceptionLogger(default_advisor_model="claude-opus-4-6"),
|
||||
_NonAdvisorHookLogger(),
|
||||
]
|
||||
with patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.handler.get_messages_interceptors",
|
||||
return_value=[capture_interceptor],
|
||||
):
|
||||
result = await anthropic_messages(
|
||||
model="openai/gpt-4o-mini",
|
||||
messages=MESSAGES,
|
||||
tools=[ADVISOR_TOOL],
|
||||
stream=False,
|
||||
max_tokens=256,
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
finally:
|
||||
litellm.callbacks = original_callbacks
|
||||
|
||||
assert isinstance(result, dict)
|
||||
assert capture_interceptor.seen_kwargs.get("metadata") == {
|
||||
"from_non_advisor_hook": True
|
||||
}
|
||||
assert isinstance(capture_interceptor.seen_tools, list)
|
||||
assert capture_interceptor.seen_tools[0].get("type") == "advisor_20260301"
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ def _make_advisor_tool_use_response(
|
|||
tool_id: str = "toolu_advisor_01",
|
||||
model: str = "openai/gpt-4o-mini",
|
||||
) -> Dict:
|
||||
# Must match ``_SYNTHETIC_ADVISOR_TOOL_NAME`` (consult_advisor) in the handler.
|
||||
return {
|
||||
"id": "msg_test",
|
||||
"type": "message",
|
||||
|
|
@ -53,7 +54,7 @@ def _make_advisor_tool_use_response(
|
|||
{
|
||||
"type": "tool_use",
|
||||
"id": tool_id,
|
||||
"name": "advisor",
|
||||
"name": "consult_advisor",
|
||||
"input": {"question": question},
|
||||
}
|
||||
],
|
||||
|
|
@ -172,21 +173,23 @@ async def test_loop_one_advisor_call():
|
|||
"def is_prime(n):\n import math\n if n < 2: return False\n for i in range(2, int(math.sqrt(n))+1):\n if n % i == 0: return False\n return True"
|
||||
)
|
||||
|
||||
call_count = 0
|
||||
executor_call_count = 0
|
||||
|
||||
async def mock_call(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
async def mock_messages(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
nonlocal executor_call_count
|
||||
executor_call_count += 1
|
||||
if executor_call_count == 1:
|
||||
return advisor_tool_use_resp # executor: calls advisor
|
||||
if call_count == 2:
|
||||
return advisor_advice_resp # advisor: returns advice
|
||||
return final_resp # executor: final answer
|
||||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_messages_handler",
|
||||
side_effect=mock_call,
|
||||
):
|
||||
side_effect=mock_messages,
|
||||
), patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_advisor_with_router",
|
||||
new_callable=AsyncMock,
|
||||
return_value=advisor_advice_resp,
|
||||
) as mock_advisor:
|
||||
h = AdvisorOrchestrationHandler()
|
||||
result = await h.handle(
|
||||
model="openai/gpt-4o-mini",
|
||||
|
|
@ -197,15 +200,18 @@ async def test_loop_one_advisor_call():
|
|||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
assert call_count == 3
|
||||
assert executor_call_count == 2
|
||||
assert mock_advisor.await_count == 1
|
||||
content = result.get("content", [])
|
||||
texts = [b for b in content if b.get("type") == "text"]
|
||||
assert len(texts) == 1
|
||||
assert "is_prime" in texts[0]["text"]
|
||||
|
||||
# No advisor tool_use blocks in final response
|
||||
# No synthetic advisor tool_use blocks in final response
|
||||
advisor_uses = [
|
||||
b for b in content if b.get("type") == "tool_use" and b.get("name") == "advisor"
|
||||
b
|
||||
for b in content
|
||||
if b.get("type") == "tool_use" and b.get("name") == "consult_advisor"
|
||||
]
|
||||
assert len(advisor_uses) == 0
|
||||
|
||||
|
|
@ -228,19 +234,16 @@ async def test_loop_max_uses_raises():
|
|||
advisor_tool_use_resp = _make_advisor_tool_use_response()
|
||||
advisor_advice_resp = _make_text_response("Here is my advice.")
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def mock_call(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
# Executor calls always return advisor tool_use; advisor always returns text
|
||||
if tools is None:
|
||||
return advisor_advice_resp
|
||||
async def mock_messages(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
return advisor_tool_use_resp
|
||||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_messages_handler",
|
||||
side_effect=mock_call,
|
||||
side_effect=mock_messages,
|
||||
), patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_advisor_with_router",
|
||||
new_callable=AsyncMock,
|
||||
return_value=advisor_advice_resp,
|
||||
):
|
||||
h = AdvisorOrchestrationHandler()
|
||||
with pytest.raises(AdvisorMaxIterationsError):
|
||||
|
|
@ -389,14 +392,14 @@ async def test_advisor_tool_translated_for_executor():
|
|||
|
||||
captured_tools = []
|
||||
|
||||
async def mock_call(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
async def mock_messages(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
if tools:
|
||||
captured_tools.extend(tools)
|
||||
return _make_text_response("Done.")
|
||||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_messages_handler",
|
||||
side_effect=mock_call,
|
||||
side_effect=mock_messages,
|
||||
):
|
||||
h = AdvisorOrchestrationHandler()
|
||||
await h.handle(
|
||||
|
|
@ -409,7 +412,7 @@ async def test_advisor_tool_translated_for_executor():
|
|||
)
|
||||
|
||||
assert len(captured_tools) > 0
|
||||
advisor_tool = next(t for t in captured_tools if t.get("name") == "advisor")
|
||||
advisor_tool = next(t for t in captured_tools if t.get("name") == "consult_advisor")
|
||||
# Must NOT have the advisor_20260301 type (provider won't understand it)
|
||||
assert advisor_tool.get("type") != "advisor_20260301"
|
||||
# Must have a description and input_schema
|
||||
|
|
@ -495,14 +498,16 @@ async def test_max_uses_none_falls_back_to_default():
|
|||
advisor_tool_use_resp = _make_advisor_tool_use_response()
|
||||
advisor_advice_resp = _make_text_response("Here is advice.")
|
||||
|
||||
async def mock_call(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
if tools is None:
|
||||
return advisor_advice_resp
|
||||
async def mock_messages(model, messages, tools, stream, max_tokens, **kwargs):
|
||||
return advisor_tool_use_resp
|
||||
|
||||
with patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_messages_handler",
|
||||
side_effect=mock_call,
|
||||
side_effect=mock_messages,
|
||||
), patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_advisor_with_router",
|
||||
new_callable=AsyncMock,
|
||||
return_value=advisor_advice_resp,
|
||||
):
|
||||
h = AdvisorOrchestrationHandler()
|
||||
with pytest.raises(AdvisorMaxIterationsError) as exc_info:
|
||||
|
|
|
|||
|
|
@ -833,6 +833,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
|
|||
},
|
||||
"supports_native_streaming": {"type": "boolean"},
|
||||
"supports_native_structured_output": {"type": "boolean"},
|
||||
"supports_native_advisor_tool": {"type": "boolean"},
|
||||
"tiered_pricing": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue