mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-02 14:22:16 +00:00
fix(lint): resolve PLR0915 too-many-statements in 4 files
Extract helpers to reduce statement count below the 50-statement limit: - a2a_protocol/main.py: extract _execute_a2a_send_with_retry() (56 → 43) - fine_tuning/main.py: extract _resolve_fine_tuning_timeout() (53 → 48) - generic_guardrail_api.py: extract _build_request_headers() (51 → 49) - mcp_streaming_iterator.py: extract _handle_initial_response_phase() (73 → 31) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Julio Quinteros
co-authored by
Claude Sonnet 4.6
Claude Sonnet 4.6
parent
b7f43d411a
commit
7512f7dfc3
@@ -162,6 +162,46 @@ async def _send_message_via_completion_bridge(
|
||||
return LiteLLMSendMessageResponse.from_dict(response_dict)
|
||||
|
||||
|
||||
async def _execute_a2a_send_with_retry(
|
||||
a2a_client: Any,
|
||||
request: Any,
|
||||
agent_card: Any,
|
||||
card_url: Optional[str],
|
||||
api_base: Optional[str],
|
||||
agent_name: Optional[str],
|
||||
) -> Any:
|
||||
"""Send an A2A message with retry logic for localhost URL errors."""
|
||||
a2a_response = None
|
||||
for _ in range(2): # max 2 attempts: original + 1 retry
|
||||
try:
|
||||
a2a_response = await a2a_client.send_message(request)
|
||||
break # success, exit retry loop
|
||||
except A2ALocalhostURLError as e:
|
||||
a2a_client = handle_a2a_localhost_retry(
|
||||
error=e,
|
||||
agent_card=agent_card,
|
||||
a2a_client=a2a_client,
|
||||
is_streaming=False,
|
||||
)
|
||||
card_url = agent_card.url if agent_card else None
|
||||
except Exception as e:
|
||||
try:
|
||||
map_a2a_exception(e, card_url, api_base, model=agent_name)
|
||||
except A2ALocalhostURLError as localhost_err:
|
||||
a2a_client = handle_a2a_localhost_retry(
|
||||
error=localhost_err,
|
||||
agent_card=agent_card,
|
||||
a2a_client=a2a_client,
|
||||
is_streaming=False,
|
||||
)
|
||||
card_url = agent_card.url if agent_card else None
|
||||
continue
|
||||
except Exception:
|
||||
raise
|
||||
assert a2a_response is not None
|
||||
return a2a_response
|
||||
|
||||
|
||||
@client
|
||||
async def asend_message(
|
||||
a2a_client: Optional["A2AClientType"] = None,
|
||||
@@ -279,44 +319,17 @@ async def asend_message(
|
||||
if getattr(message, "context_id", None) is None:
|
||||
message.context_id = context_id
|
||||
|
||||
# Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL
|
||||
a2a_response = None
|
||||
for _ in range(2): # max 2 attempts: original + 1 retry
|
||||
try:
|
||||
a2a_response = await a2a_client.send_message(request)
|
||||
break # success, exit retry loop
|
||||
except A2ALocalhostURLError as e:
|
||||
# Localhost URL error - fix and retry
|
||||
a2a_client = handle_a2a_localhost_retry(
|
||||
error=e,
|
||||
agent_card=agent_card,
|
||||
a2a_client=a2a_client,
|
||||
is_streaming=False,
|
||||
)
|
||||
card_url = agent_card.url if agent_card else None
|
||||
except Exception as e:
|
||||
# Map exception - will raise A2ALocalhostURLError if applicable
|
||||
try:
|
||||
map_a2a_exception(e, card_url, api_base, model=agent_name)
|
||||
except A2ALocalhostURLError as localhost_err:
|
||||
# Localhost URL error - fix and retry
|
||||
a2a_client = handle_a2a_localhost_retry(
|
||||
error=localhost_err,
|
||||
agent_card=agent_card,
|
||||
a2a_client=a2a_client,
|
||||
is_streaming=False,
|
||||
)
|
||||
card_url = agent_card.url if agent_card else None
|
||||
continue
|
||||
except Exception:
|
||||
# Re-raise the mapped exception
|
||||
raise
|
||||
a2a_response = await _execute_a2a_send_with_retry(
|
||||
a2a_client=a2a_client,
|
||||
request=request,
|
||||
agent_card=agent_card,
|
||||
card_url=card_url,
|
||||
api_base=api_base,
|
||||
agent_name=agent_name,
|
||||
)
|
||||
|
||||
verbose_logger.info(f"A2A send_message completed, request_id={request.id}")
|
||||
|
||||
# a2a_response is guaranteed to be set if we reach here (loop breaks on success or raises)
|
||||
assert a2a_response is not None
|
||||
|
||||
# Wrap in LiteLLM response type for _hidden_params support
|
||||
response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response)
|
||||
|
||||
|
||||
+19
-15
@@ -126,6 +126,21 @@ async def acreate_fine_tuning_job(
|
||||
raise e
|
||||
|
||||
|
||||
def _resolve_fine_tuning_timeout(
|
||||
timeout: Any,
|
||||
custom_llm_provider: str,
|
||||
) -> float:
|
||||
"""Normalise a raw timeout value to a float (seconds) for fine-tuning calls."""
|
||||
timeout = timeout or 600
|
||||
if isinstance(timeout, httpx.Timeout):
|
||||
if not supports_httpx_timeout(custom_llm_provider):
|
||||
return float(timeout.read or 600)
|
||||
return timeout # type: ignore[return-value]
|
||||
if timeout is None:
|
||||
return 600.0
|
||||
return float(timeout)
|
||||
|
||||
|
||||
@client
|
||||
def create_fine_tuning_job(
|
||||
model: str,
|
||||
@@ -164,21 +179,10 @@ def create_fine_tuning_job(
|
||||
_oai_hyperparameters: Hyperparameters = Hyperparameters(
|
||||
**hyperparameters
|
||||
) # Typed Hyperparameters for OpenAI Spec
|
||||
### TIMEOUT LOGIC ###
|
||||
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
|
||||
# set timeout for 10 minutes by default
|
||||
|
||||
if (
|
||||
timeout is not None
|
||||
and isinstance(timeout, httpx.Timeout)
|
||||
and supports_httpx_timeout(custom_llm_provider) is False
|
||||
):
|
||||
read_timeout = timeout.read or 600
|
||||
timeout = read_timeout # default 10 min timeout
|
||||
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
|
||||
timeout = float(timeout) # type: ignore
|
||||
elif timeout is None:
|
||||
timeout = 600.0
|
||||
timeout = _resolve_fine_tuning_timeout(
|
||||
optional_params.timeout or kwargs.get("request_timeout", 600),
|
||||
custom_llm_provider,
|
||||
)
|
||||
|
||||
# OpenAI
|
||||
if custom_llm_provider == "openai":
|
||||
|
||||
+8
-4
@@ -312,6 +312,13 @@ class GenericGuardrailAPI(CustomGuardrail):
|
||||
return_inputs.update(inputs)
|
||||
return return_inputs
|
||||
|
||||
def _build_request_headers(self) -> dict:
|
||||
"""Build HTTP headers for the guardrail API request."""
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if self.headers:
|
||||
headers.update(self.headers)
|
||||
return headers
|
||||
|
||||
def _build_guardrail_return_inputs(
|
||||
self,
|
||||
*,
|
||||
@@ -416,10 +423,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
||||
model=model,
|
||||
)
|
||||
|
||||
# Prepare headers
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if self.headers:
|
||||
headers.update(self.headers)
|
||||
headers = self._build_request_headers()
|
||||
|
||||
try:
|
||||
# Make the API request
|
||||
|
||||
@@ -404,74 +404,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
||||
|
||||
# Phase 1: Initial Response Stream (emit standard OpenAI events first)
|
||||
if self.phase == "initial_response":
|
||||
# Create the initial response iterator if not already created
|
||||
if self.base_iterator is None:
|
||||
await self._create_initial_response_iterator()
|
||||
|
||||
if self.base_iterator is None:
|
||||
# LLM call failed — still emit MCP discovery events before finishing
|
||||
if self.mcp_discovery_events:
|
||||
self.phase = "mcp_discovery"
|
||||
else:
|
||||
self.phase = "finished"
|
||||
raise StopAsyncIteration
|
||||
|
||||
if self.base_iterator:
|
||||
# Check if base_iterator is actually iterable
|
||||
if hasattr(self.base_iterator, "__anext__"):
|
||||
try:
|
||||
chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined]
|
||||
|
||||
# Capture the response ID from the first event to ensure consistency
|
||||
if self._cached_response_id is None and hasattr(chunk, 'response'):
|
||||
response_obj = getattr(chunk, 'response', None)
|
||||
if response_obj and hasattr(response_obj, 'id'):
|
||||
self._cached_response_id = response_obj.id
|
||||
verbose_logger.debug(f"Cached response ID: {self._cached_response_id}")
|
||||
|
||||
# After emitting response.output_item.added, transition to MCP discovery
|
||||
# Check if this is the output_item.added event
|
||||
if not self.initial_events_emitted and hasattr(chunk, 'type'):
|
||||
chunk_type = getattr(chunk, 'type', None)
|
||||
if chunk_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED:
|
||||
self.initial_events_emitted = True
|
||||
# Transition to MCP discovery phase after returning this chunk
|
||||
self.phase = "mcp_discovery"
|
||||
return chunk
|
||||
|
||||
# If auto-execution is enabled, check for completed responses
|
||||
if self.should_auto_execute and self._is_response_completed(
|
||||
chunk
|
||||
):
|
||||
# Collect the response for tool execution
|
||||
response_obj = getattr(chunk, "response", None)
|
||||
if isinstance(response_obj, ResponsesAPIResponse):
|
||||
self.collected_response = response_obj
|
||||
# Move to tool execution phase after emitting this chunk
|
||||
self.phase = "tool_execution"
|
||||
await self._generate_tool_execution_events()
|
||||
|
||||
return chunk
|
||||
except StopAsyncIteration:
|
||||
# Initial response ended, move to next phase
|
||||
if self.should_auto_execute and self.collected_response:
|
||||
self.phase = "tool_execution"
|
||||
await self._generate_tool_execution_events()
|
||||
else:
|
||||
self.phase = "finished"
|
||||
raise
|
||||
else:
|
||||
# base_iterator is not async iterable (likely a ResponsesAPIResponse)
|
||||
# Collect it for tool execution if needed
|
||||
if self.should_auto_execute and isinstance(
|
||||
self.base_iterator, ResponsesAPIResponse
|
||||
):
|
||||
self.collected_response = self.base_iterator
|
||||
self.phase = "tool_execution"
|
||||
await self._generate_tool_execution_events()
|
||||
else:
|
||||
self.phase = "finished"
|
||||
raise StopAsyncIteration
|
||||
result = await self._handle_initial_response_phase()
|
||||
if result is not None:
|
||||
return result
|
||||
|
||||
# Phase 2: MCP Discovery Events (after response.output_item.added)
|
||||
if self.phase == "mcp_discovery":
|
||||
@@ -523,6 +458,76 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
||||
# Should not reach here
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def _handle_initial_response_phase(
|
||||
self,
|
||||
) -> Optional[ResponsesAPIStreamingResponse]:
|
||||
"""
|
||||
Handle Phase 1: Initial Response Stream.
|
||||
|
||||
Returns a chunk to emit, or None to fall through to the next phase.
|
||||
Raises StopAsyncIteration when the stream is exhausted with no auto-execution.
|
||||
"""
|
||||
if self.base_iterator is None:
|
||||
await self._create_initial_response_iterator()
|
||||
|
||||
if self.base_iterator is None:
|
||||
# LLM call failed — still emit MCP discovery events before finishing
|
||||
if self.mcp_discovery_events:
|
||||
self.phase = "mcp_discovery"
|
||||
else:
|
||||
self.phase = "finished"
|
||||
raise StopAsyncIteration
|
||||
return None
|
||||
|
||||
if self.base_iterator:
|
||||
if hasattr(self.base_iterator, "__anext__"):
|
||||
try:
|
||||
chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined]
|
||||
|
||||
# Capture the response ID from the first event to ensure consistency
|
||||
if self._cached_response_id is None and hasattr(chunk, "response"):
|
||||
response_obj = getattr(chunk, "response", None)
|
||||
if response_obj and hasattr(response_obj, "id"):
|
||||
self._cached_response_id = response_obj.id
|
||||
verbose_logger.debug(f"Cached response ID: {self._cached_response_id}")
|
||||
|
||||
# After emitting response.output_item.added, transition to MCP discovery
|
||||
if not self.initial_events_emitted and hasattr(chunk, "type"):
|
||||
chunk_type = getattr(chunk, "type", None)
|
||||
if chunk_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED:
|
||||
self.initial_events_emitted = True
|
||||
self.phase = "mcp_discovery"
|
||||
return chunk
|
||||
|
||||
# If auto-execution is enabled, check for completed responses
|
||||
if self.should_auto_execute and self._is_response_completed(chunk):
|
||||
response_obj = getattr(chunk, "response", None)
|
||||
if isinstance(response_obj, ResponsesAPIResponse):
|
||||
self.collected_response = response_obj
|
||||
self.phase = "tool_execution"
|
||||
await self._generate_tool_execution_events()
|
||||
|
||||
return chunk
|
||||
except StopAsyncIteration:
|
||||
if self.should_auto_execute and self.collected_response:
|
||||
self.phase = "tool_execution"
|
||||
await self._generate_tool_execution_events()
|
||||
else:
|
||||
self.phase = "finished"
|
||||
raise
|
||||
else:
|
||||
# base_iterator is not async iterable (likely a ResponsesAPIResponse)
|
||||
if self.should_auto_execute and isinstance(
|
||||
self.base_iterator, ResponsesAPIResponse
|
||||
):
|
||||
self.collected_response = self.base_iterator
|
||||
self.phase = "tool_execution"
|
||||
await self._generate_tool_execution_events()
|
||||
else:
|
||||
self.phase = "finished"
|
||||
raise StopAsyncIteration
|
||||
return None
|
||||
|
||||
def _is_response_completed(self, chunk: ResponsesAPIStreamingResponse) -> bool:
|
||||
"""Check if this chunk indicates the response is completed"""
|
||||
from litellm.types.llms.openai import ResponsesAPIStreamEvents
|
||||
|
||||
Reference in New Issue
Block a user