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:
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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."""
|
||||||
|
|||||||
Reference in New Issue
Block a user