Add summarization timeout to prevent hung LLM calls

Adds a configurable summarization_timeout (default 120s) that cancels
  summary generation if the LLM hangs. On timeout, an error result is
  returned so _summarization_in_progress resets and future
  summarizations are unblocked.
This commit is contained in:
Mark Backman
2026-02-26 21:31:04 -05:00
parent a489bfaf00
commit 50710e9c3f
5 changed files with 40 additions and 6 deletions

View File

@@ -2019,6 +2019,8 @@ class LLMContextSummaryRequestFrame(ControlFrame):
the summary text. the summary text.
summarization_prompt: System prompt instructing the LLM how to generate summarization_prompt: System prompt instructing the LLM how to generate
the summary. the summary.
summarization_timeout: Maximum time in seconds for the LLM to generate a
summary. None means no timeout.
""" """
request_id: str request_id: str
@@ -2026,6 +2028,7 @@ class LLMContextSummaryRequestFrame(ControlFrame):
min_messages_to_keep: int min_messages_to_keep: int
target_context_tokens: int target_context_tokens: int
summarization_prompt: str summarization_prompt: str
summarization_timeout: Optional[float] = None
@dataclass @dataclass

View File

@@ -218,6 +218,7 @@ class LLMContextSummarizer(BaseObject):
min_messages_to_keep=min_keep, min_messages_to_keep=min_keep,
target_context_tokens=self._config.target_context_tokens, target_context_tokens=self._config.target_context_tokens,
summarization_prompt=self._config.summary_prompt, summarization_prompt=self._config.summary_prompt,
summarization_timeout=self._config.summarization_timeout,
) )
# Emit event for aggregator to broadcast # Emit event for aggregator to broadcast

View File

@@ -1284,20 +1284,35 @@ class LLMAssistantAggregator(LLMContextAggregator):
frame: The summarization request frame. frame: The summarization request frame.
""" """
try: try:
summary, last_index = await llm._generate_summary(frame) if frame.summarization_timeout:
summary, last_index = await asyncio.wait_for(
llm._generate_summary(frame),
timeout=frame.summarization_timeout,
)
else:
summary, last_index = await llm._generate_summary(frame)
result_frame = LLMContextSummaryResultFrame( result_frame = LLMContextSummaryResultFrame(
request_id=frame.request_id, request_id=frame.request_id,
summary=summary, summary=summary,
last_summarized_index=last_index, last_summarized_index=last_index,
) )
except Exception as e: except asyncio.TimeoutError:
error = f"Error generating context summary: {e}" error = f"Context summarization timed out after {frame.summarization_timeout}s"
await self.push_error(error, exception=e) logger.error(f"{self}: {error}")
result_frame = LLMContextSummaryResultFrame( result_frame = LLMContextSummaryResultFrame(
request_id=frame.request_id, request_id=frame.request_id,
summary="", summary="",
last_summarized_index=-1, last_summarized_index=-1,
error=f"Error generating context summary: {e}", error=error,
)
except Exception as e:
error = f"Error generating context summary: {e}"
await self.push_error(error_msg=error, exception=e)
result_frame = LLMContextSummaryResultFrame(
request_id=frame.request_id,
summary="",
last_summarized_index=-1,
error=error,
) )
if self._summarizer: if self._summarizer:

View File

@@ -437,7 +437,17 @@ class LLMService(UserTurnCompletionLLMServiceMixin, AIService):
error = None error = None
try: try:
summary, last_index = await self._generate_summary(frame) if frame.summarization_timeout:
summary, last_index = await asyncio.wait_for(
self._generate_summary(frame),
timeout=frame.summarization_timeout,
)
else:
summary, last_index = await self._generate_summary(frame)
except asyncio.TimeoutError:
await self.push_error(
error_msg=f"Context summarization timed out after {frame.summarization_timeout}s"
)
except Exception as e: except Exception as e:
error = f"Error generating context summary: {e}" error = f"Error generating context summary: {e}"
await self.push_error(error, exception=e) await self.push_error(error, exception=e)

View File

@@ -86,6 +86,10 @@ class LLMContextSummarizationConfig:
pipeline's primary LLM. Useful for routing summarization to a pipeline's primary LLM. Useful for routing summarization to a
cheaper/faster model (e.g., Gemini Flash) while keeping an cheaper/faster model (e.g., Gemini Flash) while keeping an
expensive model for conversation. If None, uses the pipeline LLM. expensive model for conversation. If None, uses the pipeline LLM.
summarization_timeout: Maximum time in seconds to wait for the LLM to
generate a summary. If the call exceeds this timeout, summarization
is aborted with an error and future summarizations are unblocked.
Set to None to disable the timeout.
""" """
max_context_tokens: int = 8000 max_context_tokens: int = 8000
@@ -95,6 +99,7 @@ class LLMContextSummarizationConfig:
summarization_prompt: Optional[str] = None summarization_prompt: Optional[str] = None
summary_message_template: str = "Conversation summary: {summary}" summary_message_template: str = "Conversation summary: {summary}"
llm: Optional["LLMService"] = None llm: Optional["LLMService"] = None
summarization_timeout: Optional[float] = 120.0
def __post_init__(self): def __post_init__(self):
"""Validate configuration parameters.""" """Validate configuration parameters."""