From 8c8a8c41032b0ee968f2cbd0f52ff29c4312a90b Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Wed, 21 Jan 2026 13:22:53 -0500 Subject: [PATCH] fix: use separate db tx for create and save summary (#334) --- src/utils/summarizer.py | 284 +++++++++++++----------- tests/integration/test_token_metrics.py | 4 - 2 files changed, 151 insertions(+), 137 deletions(-) diff --git a/src/utils/summarizer.py b/src/utils/summarizer.py index c8bca593..29cc02cb 100644 --- a/src/utils/summarizer.py +++ b/src/utils/summarizer.py @@ -85,7 +85,7 @@ class SummaryType(Enum): def short_summary_prompt( - messages: list[models.Message], + formatted_messages: str, output_words: int, previous_summary_text: str, ) -> str: @@ -109,7 +109,7 @@ Return only the summary without any explanation or meta-commentary. -{_format_messages(messages)} +{formatted_messages} Produce as thorough a summary as possible in {output_words} words or less. @@ -117,7 +117,7 @@ Produce as thorough a summary as possible in {output_words} words or less. def long_summary_prompt( - messages: list[models.Message], + formatted_messages: str, output_words: int, previous_summary_text: str, ) -> str: @@ -143,7 +143,7 @@ Return only the summary without any explanation or meta-commentary. -{_format_messages(messages)} +{formatted_messages} Produce as thorough a summary as possible in {output_words} words or less. @@ -156,7 +156,7 @@ def estimate_short_summary_prompt_tokens() -> int: try: return estimate_tokens( short_summary_prompt( - messages=[], + formatted_messages="", output_words=0, previous_summary_text="", ) @@ -172,7 +172,7 @@ def estimate_long_summary_prompt_tokens() -> int: try: return estimate_tokens( long_summary_prompt( - messages=[], + formatted_messages="", output_words=0, previous_summary_text="", ) @@ -184,7 +184,7 @@ def estimate_long_summary_prompt_tokens() -> int: @conditional_observe(name="Create Short Summary") async def create_short_summary( - messages: list[models.Message], + formatted_messages: str, input_tokens: int, previous_summary: str | None = None, ) -> HonchoLLMCallResponse[str]: @@ -200,7 +200,9 @@ async def create_short_summary( else: previous_summary_text = "There is no previous summary -- the messages are the beginning of the conversation." - prompt = short_summary_prompt(messages, output_words, previous_summary_text) + prompt = short_summary_prompt( + formatted_messages, output_words, previous_summary_text + ) return await honcho_llm_call( llm_settings=settings.SUMMARY, @@ -211,7 +213,7 @@ async def create_short_summary( @conditional_observe(name="Create Long Summary") async def create_long_summary( - messages: list[models.Message], + formatted_messages: str, previous_summary: str | None = None, ) -> HonchoLLMCallResponse[str]: # the word/token ratio is roughly 4:3 so we multiply by 0.75. @@ -223,7 +225,9 @@ async def create_long_summary( else: previous_summary_text = "There is no previous summary -- the messages are the beginning of the conversation." - prompt = long_summary_prompt(messages, output_words, previous_summary_text) + prompt = long_summary_prompt( + formatted_messages, output_words, previous_summary_text + ) return await honcho_llm_call( llm_settings=settings.SUMMARY, @@ -267,93 +271,85 @@ async def summarize_if_needed( if should_create_long is False and should_create_short is False: return - # If both summaries need to be created, run them in parallel with separate database sessions + # If both summaries need to be created, run them in parallel if should_create_long and should_create_short: - async def create_long_summary(): - async with tracked_db("create_long_summary") as db_session: - await _create_and_save_summary( - db_session, - workspace_name, - session_name, - message_id=message_id, - message_seq_in_session=message_seq_in_session, - message_public_id=message_public_id, - summary_type=SummaryType.LONG, - configuration=configuration, - ) - accumulate_metric( - f"summary_{workspace_name}_{message_id}", - "long_summary_up_to_message", - message_seq_in_session, - "count", - ) + async def create_long_summary_task(): + await _create_and_save_summary( + workspace_name, + session_name, + message_id=message_id, + message_seq_in_session=message_seq_in_session, + message_public_id=message_public_id, + summary_type=SummaryType.LONG, + configuration=configuration, + ) + accumulate_metric( + f"summary_{workspace_name}_{message_id}", + "long_summary_up_to_message", + message_seq_in_session, + "count", + ) - async def create_short_summary(): - async with tracked_db("create_short_summary") as db_session: - await _create_and_save_summary( - db_session, - workspace_name, - session_name, - message_id=message_id, - message_seq_in_session=message_seq_in_session, - message_public_id=message_public_id, - summary_type=SummaryType.SHORT, - configuration=configuration, - ) - accumulate_metric( - f"summary_{workspace_name}_{message_id}", - "short_summary_up_to_message", - message_seq_in_session, - "count", - ) + async def create_short_summary_task(): + await _create_and_save_summary( + workspace_name, + session_name, + message_id=message_id, + message_seq_in_session=message_seq_in_session, + message_public_id=message_public_id, + summary_type=SummaryType.SHORT, + configuration=configuration, + ) + accumulate_metric( + f"summary_{workspace_name}_{message_id}", + "short_summary_up_to_message", + message_seq_in_session, + "count", + ) await asyncio.gather( - create_long_summary(), - create_short_summary(), + create_long_summary_task(), + create_short_summary_task(), return_exceptions=True, ) else: - async with tracked_db("create_summary") as db: - # If only one summary needs to be created, run them individually - if should_create_long: - await _create_and_save_summary( - db, - workspace_name, - session_name, - message_id=message_id, - message_seq_in_session=message_seq_in_session, - message_public_id=message_public_id, - summary_type=SummaryType.LONG, - configuration=configuration, - ) - accumulate_metric( - f"summary_{workspace_name}_{message_id}", - "long_summary_up_to_message", - message_seq_in_session, - "count", - ) - elif should_create_short: - await _create_and_save_summary( - db, - workspace_name, - session_name, - message_id=message_id, - message_seq_in_session=message_seq_in_session, - message_public_id=message_public_id, - summary_type=SummaryType.SHORT, - configuration=configuration, - ) - accumulate_metric( - f"summary_{workspace_name}_{message_id}", - "short_summary_up_to_message", - message_seq_in_session, - "count", - ) + # If only one summary needs to be created, run individually + if should_create_long: + await _create_and_save_summary( + workspace_name, + session_name, + message_id=message_id, + message_seq_in_session=message_seq_in_session, + message_public_id=message_public_id, + summary_type=SummaryType.LONG, + configuration=configuration, + ) + accumulate_metric( + f"summary_{workspace_name}_{message_id}", + "long_summary_up_to_message", + message_seq_in_session, + "count", + ) + elif should_create_short: + await _create_and_save_summary( + workspace_name, + session_name, + message_id=message_id, + message_seq_in_session=message_seq_in_session, + message_public_id=message_public_id, + summary_type=SummaryType.SHORT, + configuration=configuration, + ) + accumulate_metric( + f"summary_{workspace_name}_{message_id}", + "short_summary_up_to_message", + message_seq_in_session, + "count", + ) async def _create_and_save_summary( - db: AsyncSession, workspace_name: str, session_name: str, *, @@ -372,51 +368,62 @@ async def _create_and_save_summary( """ logger.debug("Creating new %s summary", summary_type.name) - # Time summarization step summary_start = time.perf_counter() - latest_summary = await get_summary(db, workspace_name, session_name, summary_type) - if latest_summary: - latest_summary_message_id = latest_summary["message_id"] - # Skip if latest summary already covers message. - if latest_summary_message_id >= message_id: + async with tracked_db("summary.fetch_data") as db: + latest_summary = await get_summary( + db, workspace_name, session_name, summary_type + ) + if latest_summary: + latest_summary_message_id = latest_summary["message_id"] + # Skip if latest summary already covers message. + if latest_summary_message_id >= message_id: + return + + previous_summary_text = latest_summary["content"] if latest_summary else None + + # Calculate the sequence range for messages to summarize + # We want to get the last N messages where N is the configured summary interval + messages_per_summary = ( + configuration.summary.messages_per_long_summary + if summary_type == SummaryType.LONG + else configuration.summary.messages_per_short_summary + ) + start_seq = max(message_seq_in_session - messages_per_summary + 1, 1) + + messages: list[Message] = await crud.get_messages_by_seq_range( + db, + workspace_name, + session_name, + start_seq=start_seq, + end_seq=message_seq_in_session, + ) + if not messages: + logger.warning("No messages to summarize for message %s", message_id) return - previous_summary_text = latest_summary["content"] if latest_summary else None + # Extract values before closing session + formatted_messages = _format_messages(messages) + last_message_id = messages[-1].id + last_message_content_preview = messages[-1].content[:30] + message_count = len(messages) - # Calculate the sequence range for messages to summarize - # We want to get the last N messages where N is the configured summary interval - messages_per_summary = ( - configuration.summary.messages_per_long_summary - if summary_type == SummaryType.LONG - else configuration.summary.messages_per_short_summary - ) - start_seq = max(message_seq_in_session - messages_per_summary + 1, 1) - - messages: list[Message] = await crud.get_messages_by_seq_range( - db, - workspace_name, - session_name, - start_seq=start_seq, - end_seq=message_seq_in_session, - ) - if not messages: - logger.warning("No messages to summarize for message %s", message_id) - return - - messages_tokens = sum([message.token_count for message in messages]) - previous_summary_tokens = latest_summary["token_count"] if latest_summary else 0 - input_tokens = messages_tokens + previous_summary_tokens + messages_tokens = sum([message.token_count for message in messages]) + previous_summary_tokens = latest_summary["token_count"] if latest_summary else 0 + input_tokens = messages_tokens + previous_summary_tokens new_summary, is_fallback = await _create_summary( - messages=messages, + formatted_messages=formatted_messages, previous_summary_text=previous_summary_text, summary_type=summary_type, input_tokens=input_tokens, message_public_id=message_public_id, + last_message_id=last_message_id, + last_message_content_preview=last_message_content_preview, + message_count=message_count, ) - # Only track tokens and save summary if this was a real LLM call + # Step 3: Save to database with new transaction if not is_fallback: # Get base prompt tokens based on summary type if summary_type == SummaryType.SHORT: @@ -440,13 +447,14 @@ async def _create_and_save_summary( component=prometheus.DeriverComponents.OUTPUT_TOTAL.value, ).inc(new_summary["token_count"]) - # Save summary to database - await _save_summary( - db, - new_summary, - workspace_name, - session_name, - ) + # Save summary to database with new transaction + async with tracked_db("summary.save") as db: + await _save_summary( + db, + new_summary, + workspace_name, + session_name, + ) accumulate_metric( f"summary_{workspace_name}_{message_id}", @@ -471,19 +479,27 @@ async def _create_and_save_summary( async def _create_summary( - messages: list[models.Message], + formatted_messages: str, previous_summary_text: str | None, summary_type: SummaryType, input_tokens: int, message_public_id: str, + last_message_id: int, + last_message_content_preview: str, + message_count: int, ) -> tuple[Summary, bool]: """ Generate a summary of the provided messages using an LLM. Args: - messages_since_last: List of messages to summarize - last_summary_text: Optional previous summary to provide context + formatted_messages: Pre-formatted message string + previous_summary_text: Optional previous summary to provide context summary_type: Type of summary to create ("short" or "long") + input_tokens: Token count for input + message_public_id: Public ID of the last message + last_message_id: ID of the last message + last_message_content_preview: Preview of last message content for fallback + message_count: Number of messages for fallback Returns: A tuple of (Summary, is_fallback) where is_fallback indicates if @@ -495,10 +511,12 @@ async def _create_summary( try: if summary_type == SummaryType.SHORT: response = await create_short_summary( - messages, input_tokens, previous_summary_text + formatted_messages, input_tokens, previous_summary_text ) else: - response = await create_long_summary(messages, previous_summary_text) + response = await create_long_summary( + formatted_messages, previous_summary_text + ) summary_text = response.content summary_tokens = response.output_tokens @@ -512,8 +530,8 @@ async def _create_summary( logger.exception("Error generating summary!") # Fallback to a basic summary in case of error summary_text = ( - f"Conversation with {len(messages)} messages about {messages[-1].content[:30]}..." - if messages + f"Conversation with {message_count} messages about {last_message_content_preview}..." + if message_count > 0 else "" ) summary_tokens = 50 @@ -522,7 +540,7 @@ async def _create_summary( return ( Summary( content=summary_text, - message_id=messages[-1].id if messages else 0, + message_id=last_message_id, summary_type=summary_type.value, created_at=utc_now_iso(), token_count=summary_tokens, diff --git a/tests/integration/test_token_metrics.py b/tests/integration/test_token_metrics.py index a1908c65..d4a55338 100644 --- a/tests/integration/test_token_metrics.py +++ b/tests/integration/test_token_metrics.py @@ -409,7 +409,6 @@ class TestDeriverSummaryMetrics: ), ): await _create_and_save_summary( - db=db_session, workspace_name=workspace.name, session_name=session.name, message_id=last_message.id, @@ -471,7 +470,6 @@ class TestDeriverSummaryMetrics: ), ): await _create_and_save_summary( - db=db_session, workspace_name=workspace.name, session_name=session.name, message_id=last_message.id, @@ -541,7 +539,6 @@ class TestDeriverSummaryMetrics: ), ): await _create_and_save_summary( - db=db_session, workspace_name=workspace.name, session_name=session.name, message_id=last_message.id, @@ -605,7 +602,6 @@ class TestDeriverSummaryMetrics: new=AsyncMock(return_value=(mock_summary, True)), # is_fallback=True ): await _create_and_save_summary( - db=db_session, workspace_name=workspace.name, session_name=session.name, message_id=last_message.id,