Fix black test
Some checks failed
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Has been cancelled
Unit Tests: Security / security (push) Has been cancelled
Unit Tests: Proxy DB Operations / budgets (push) Has been cancelled
Unit Tests: Proxy DB Operations / custom-logging (push) Has been cancelled
Unit Tests: Proxy DB Operations / logging-misc (push) Has been cancelled
Unit Tests: Proxy DB Operations / auth-checks (push) Has been cancelled
Unit Tests: Proxy DB Operations / db-and-spend (push) Has been cancelled
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Has been cancelled
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Has been cancelled
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Has been cancelled
Unit Tests: Proxy DB Operations / key-generation (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-runtime (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-server-core (push) Has been cancelled
Unit Tests: Proxy DB Operations / schema-migration (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-utils (push) Has been cancelled

This commit is contained in:
Sameer Kankute 2026-04-24 20:20:45 +05:30
parent 0dba636f98
commit bf7b9fb35c
No known key found for this signature in database
5 changed files with 158 additions and 100 deletions

View file

@ -185,9 +185,9 @@ class AdvisorInterceptionLogger(CustomLogger):
return self._wrap_as_streaming_if_needed(response)
return None
custom_llm_provider = request_data.get("custom_llm_provider", "") or request_data.get(
"litellm_params", {}
).get("custom_llm_provider", "")
custom_llm_provider = request_data.get(
"custom_llm_provider", ""
) or request_data.get("litellm_params", {}).get("custom_llm_provider", "")
if not custom_llm_provider:
try:
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model)
@ -197,14 +197,16 @@ class AdvisorInterceptionLogger(CustomLogger):
tools = request_data.get("tools")
stream = bool(request_data.get("stream", False))
should_run, tools_dict = await self.async_should_run_chat_completion_agentic_loop(
response=response,
model=model,
messages=messages,
tools=tools if isinstance(tools, list) else None,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=request_data,
should_run, tools_dict = (
await self.async_should_run_chat_completion_agentic_loop(
response=response,
model=model,
messages=messages,
tools=tools if isinstance(tools, list) else None,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=request_data,
)
)
if not should_run:
if isinstance(call_id, str):
@ -265,8 +267,14 @@ class AdvisorInterceptionLogger(CustomLogger):
# to the orchestration loop below.
if custom_llm_provider in ADVISOR_NATIVE_PROVIDERS:
call_id_check = kwargs.get("litellm_call_id")
advisor_cfg = self._advisor_config_by_call_id.get(call_id_check, {}) if isinstance(call_id_check, str) else {}
advisor_model_check = advisor_cfg.get("advisor_model") or self.default_advisor_model or ""
advisor_cfg = (
self._advisor_config_by_call_id.get(call_id_check, {})
if isinstance(call_id_check, str)
else {}
)
advisor_model_check = (
advisor_cfg.get("advisor_model") or self.default_advisor_model or ""
)
if self._is_native_anthropic_advisor_model(advisor_model_check):
return False, {}
@ -323,7 +331,9 @@ class AdvisorInterceptionLogger(CustomLogger):
"""
advisor_config = tools.get("advisor_config", {}) or {}
max_uses = int(advisor_config.get("max_uses", ADVISOR_MAX_USES))
advisor_model = advisor_config.get("advisor_model") or self.default_advisor_model
advisor_model = (
advisor_config.get("advisor_model") or self.default_advisor_model
)
if not advisor_model:
raise ValueError(
"No advisor model configured. Either:\n"
@ -357,7 +367,9 @@ class AdvisorInterceptionLogger(CustomLogger):
if not advisor_calls:
final_executor_cost = self._safe_get_response_cost(current_response)
advisor_first_call_cost = max(
total_response_cost - advisor_subcall_cost - final_executor_cost,
total_response_cost
- advisor_subcall_cost
- final_executor_cost,
0.0,
)
self._set_response_cost_if_possible(
@ -427,10 +439,12 @@ class AdvisorInterceptionLogger(CustomLogger):
)
)
advisor_text = self._extract_text_content(advisor_response)
advisor_interactions.append({
"tool_use_id": advisor_call["id"],
"advisor_text": advisor_text,
})
advisor_interactions.append(
{
"tool_use_id": advisor_call["id"],
"advisor_text": advisor_text,
}
)
tool_messages.append(
{
"role": "tool",
@ -439,7 +453,9 @@ class AdvisorInterceptionLogger(CustomLogger):
}
)
current_messages = current_messages + [assistant_message] + tool_messages
current_messages = (
current_messages + [assistant_message] + tool_messages
)
optional_params_clean = {
k: v
@ -664,10 +680,10 @@ class AdvisorInterceptionLogger(CustomLogger):
call_id = kwargs.get("litellm_call_id")
if isinstance(call_id, str):
self._advisor_config_by_call_id[call_id] = {
"advisor_model": advisor_model,
"max_uses": int(max_uses),
"api_key": api_key,
"api_base": api_base,
"advisor_model": advisor_model,
"max_uses": int(max_uses),
"api_key": api_key,
"api_base": api_base,
}
if kwargs.get("stream"):
kwargs["stream"] = False
@ -702,7 +718,9 @@ class AdvisorInterceptionLogger(CustomLogger):
"api_base": api_base,
}
def _extract_advisor_tool_calls(self, response: Any) -> Tuple[List[Dict], List[Dict]]:
def _extract_advisor_tool_calls(
self, response: Any
) -> Tuple[List[Dict], List[Dict]]:
message = self._extract_first_choice_message(response)
if not message:
return [], []
@ -793,19 +811,23 @@ class AdvisorInterceptionLogger(CustomLogger):
for interaction in advisor_interactions:
tool_use_id = interaction["tool_use_id"]
advisor_text = interaction["advisor_text"]
advisor_results.append({
"type": "server_tool_use",
"id": tool_use_id,
"name": "advisor",
})
advisor_results.append({
"type": "advisor_tool_result",
"tool_use_id": tool_use_id,
"content": {
"type": "advisor_result",
"text": advisor_text,
},
})
advisor_results.append(
{
"type": "server_tool_use",
"id": tool_use_id,
"name": "advisor",
}
)
advisor_results.append(
{
"type": "advisor_tool_result",
"tool_use_id": tool_use_id,
"content": {
"type": "advisor_result",
"text": advisor_text,
},
}
)
message = AdvisorInterceptionLogger._extract_first_choice_message_obj(response)
if message is None:
@ -868,10 +890,14 @@ class AdvisorInterceptionLogger(CustomLogger):
"id": getattr(tc, "id", None),
"type": getattr(tc, "type", None),
"function": {
"name": getattr(function, "name", None) if function else None,
"arguments": getattr(function, "arguments", None)
if function
else None,
"name": (
getattr(function, "name", None) if function else None
),
"arguments": (
getattr(function, "arguments", None)
if function
else None
),
},
}
)
@ -956,8 +982,7 @@ class AdvisorInterceptionLogger(CustomLogger):
if usage is not None:
input_tokens = (
AdvisorInterceptionLogger._get_usage_value(usage, "prompt_tokens")
or 0
AdvisorInterceptionLogger._get_usage_value(usage, "prompt_tokens") or 0
)
output_tokens = (
AdvisorInterceptionLogger._get_usage_value(usage, "completion_tokens")

View file

@ -93,10 +93,14 @@ def is_advisor_tool(tool: Any) -> bool:
"""
Check whether a tool is an advisor tool in any supported format.
"""
tool_type = tool.get("type") if isinstance(tool, dict) else getattr(tool, "type", None)
tool_type = (
tool.get("type") if isinstance(tool, dict) else getattr(tool, "type", None)
)
if tool_type == ANTHROPIC_ADVISOR_TOOL_TYPE:
return True
if is_advisor_tool_chat_completion(tool):
return True
tool_name = tool.get("name") if isinstance(tool, dict) else getattr(tool, "name", None)
tool_name = (
tool.get("name") if isinstance(tool, dict) else getattr(tool, "name", None)
)
return tool_name in _ADVISOR_TOOL_NAMES

View file

@ -8,7 +8,18 @@
import asyncio
import contextvars
from functools import partial
from typing import Any, AsyncIterator, Coroutine, Dict, List, Optional, Tuple, Type, Union, cast
from typing import (
Any,
AsyncIterator,
Coroutine,
Dict,
List,
Optional,
Tuple,
Type,
Union,
cast,
)
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj

View file

@ -173,19 +173,21 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
try:
while True:
# --- Executor call (always non-streaming) ---
executor_response: AnthropicMessagesResponse = await _call_messages_handler(
model=model,
messages=current_messages,
tools=executor_tools,
stream=False,
max_tokens=max_tokens,
custom_llm_provider=custom_llm_provider,
metadata={
**metadata_base,
"advisor_sub_call": False,
"parent_request_id": parent_request_id,
},
**kwargs,
executor_response: AnthropicMessagesResponse = (
await _call_messages_handler(
model=model,
messages=current_messages,
tools=executor_tools,
stream=False,
max_tokens=max_tokens,
custom_llm_provider=custom_llm_provider,
metadata={
**metadata_base,
"advisor_sub_call": False,
"parent_request_id": parent_request_id,
},
**kwargs,
)
)
executor_cost = _get_response_cost(executor_response, model=model)
@ -256,17 +258,19 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
# Use the resolved model so router routing / cost lookup hit the
# real underlying deployment; the alias is kept only for the
# client-visible iteration entry below.
advisor_response: AnthropicMessagesResponse = await _call_advisor_with_router(
model=resolved_advisor_model,
messages=advisor_messages,
max_tokens=max_tokens,
metadata={
**metadata_base,
"advisor_sub_call": True,
"parent_request_id": parent_request_id,
},
api_key=advisor_api_key,
api_base=advisor_api_base,
advisor_response: AnthropicMessagesResponse = (
await _call_advisor_with_router(
model=resolved_advisor_model,
messages=advisor_messages,
max_tokens=max_tokens,
metadata={
**metadata_base,
"advisor_sub_call": True,
"parent_request_id": parent_request_id,
},
api_key=advisor_api_key,
api_base=advisor_api_base,
)
)
advisor_call_cost = _get_response_cost(
@ -284,12 +288,14 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
advisor_text = _extract_response_text(advisor_response)
# Record the interaction for later injection into the final response.
advisor_interactions.append({
"tool_use_id": advisor_use_block.get(
"id", f"srvtoolu_{uuid.uuid4().hex[:24]}"
),
"advisor_text": advisor_text,
})
advisor_interactions.append(
{
"tool_use_id": advisor_use_block.get(
"id", f"srvtoolu_{uuid.uuid4().hex[:24]}"
),
"advisor_text": advisor_text,
}
)
# --- Inject advisor result and continue loop ---
current_messages = _inject_advisor_turn(
@ -498,9 +504,9 @@ def _finalize_orchestrated_response(
from litellm.types.utils import CallTypes
litellm_logging_obj.call_type = CallTypes.anthropic_messages.value
litellm_logging_obj.model_call_details[
"call_type"
] = CallTypes.anthropic_messages.value
litellm_logging_obj.model_call_details["call_type"] = (
CallTypes.anthropic_messages.value
)
except Exception:
pass
@ -549,7 +555,9 @@ def _fire_async_success_logging(
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
start_time = getattr(litellm_logging_obj, "start_time", None) or _dt.datetime.now()
start_time = (
getattr(litellm_logging_obj, "start_time", None) or _dt.datetime.now()
)
end_time = _dt.datetime.now()
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=litellm_logging_obj.async_success_handler(
@ -599,7 +607,11 @@ def _openai_response_to_anthropic_dict(response: Any) -> Dict:
choices = response.choices if hasattr(response, "choices") else []
content_blocks: List[Dict] = []
for choice in choices:
msg = choice.message if hasattr(choice, "message") else choice.get("message", {})
msg = (
choice.message
if hasattr(choice, "message")
else choice.get("message", {})
)
text = msg.content if hasattr(msg, "content") else msg.get("content", "")
if text:
content_blocks.append({"type": "text", "text": text})
@ -614,7 +626,9 @@ def _openai_response_to_anthropic_dict(response: Any) -> Dict:
hidden_params = getattr(response, "_hidden_params", None)
if hidden_params is not None:
anthropic_dict["_hidden_params"] = (
dict(hidden_params) if isinstance(hidden_params, dict) else hidden_params
dict(hidden_params)
if isinstance(hidden_params, dict)
else hidden_params
)
return anthropic_dict
except Exception:
@ -693,20 +707,24 @@ def _inject_advisor_blocks_into_response(
# the Anthropic response to OpenAI format). We did not observe a real
# tool-call payload here — the advisor was invoked via a synthetic
# tool — so an empty object is the correct stub.
content.append({
"type": "server_tool_use",
"id": tool_use_id,
"name": "advisor",
"input": {},
})
content.append({
"type": "advisor_tool_result",
"tool_use_id": tool_use_id,
"content": {
"type": "advisor_result",
"text": advisor_text,
},
})
content.append(
{
"type": "server_tool_use",
"id": tool_use_id,
"name": "advisor",
"input": {},
}
)
content.append(
{
"type": "advisor_tool_result",
"tool_use_id": tool_use_id,
"content": {
"type": "advisor_result",
"text": advisor_text,
},
}
)
def _extract_response_text(response: Any) -> str:

View file

@ -454,9 +454,9 @@ class BaseLLMHTTPHandler:
# Check if stream was converted for Advisor interception
if litellm_params.get("_advisor_interception_converted_stream", False):
logging_obj.model_call_details[
"_advisor_interception_converted_stream"
] = True
logging_obj.model_call_details["_advisor_interception_converted_stream"] = (
True
)
if acompletion is True:
if stream is True: