Fix tests

This commit is contained in:
Sameer Kankute 2026-04-17 08:45:48 +05:30
parent 1926a8b778
commit 171a3a8ff1
No known key found for this signature in database
3 changed files with 59 additions and 116 deletions

View file

@ -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"

View file

@ -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:

View file

@ -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": {