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,