From accdc79fdbe101fab68891dc78c5f3e9437fefb7 Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Tue, 7 Oct 2025 12:37:22 -0400 Subject: [PATCH] feat: Modify the Summary BaseModel to use public_id of message rather than internal ID (#224) * feat: Modify the Summary BaseModel to use public_id of message rather than internal ID * fix: fallback invocation * fix: strict validation * fix: rm ID from Message schema * fix: make SDK changes --- sdks/python/src/honcho/session_context.py | 2 +- sdks/typescript/src/session_context.ts | 28 +++++++++++------------ src/deriver/consumer.py | 28 +++++++++++++++++++++++ src/deriver/queue_payload.py | 10 ++++++++ src/routers/messages.py | 2 ++ src/schemas.py | 7 +++++- src/utils/summarizer.py | 13 +++++++++++ tests/deriver/conftest.py | 1 + tests/deriver/test_queue_processing.py | 3 +++ 9 files changed, 78 insertions(+), 16 deletions(-) diff --git a/sdks/python/src/honcho/session_context.py b/sdks/python/src/honcho/session_context.py index 1cd9e7f5..45017075 100644 --- a/sdks/python/src/honcho/session_context.py +++ b/sdks/python/src/honcho/session_context.py @@ -13,7 +13,7 @@ class Summary(BaseModel): """Represents a summary of a session's conversation.""" content: str = Field(..., description="The summary text") - message_id: int = Field( + message_id: str = Field( ..., description="The ID of the message that this summary covers up to" ) summary_type: str = Field(..., description="The type of summary (short or long)") diff --git a/sdks/typescript/src/session_context.ts b/sdks/typescript/src/session_context.ts index ecd9b8c1..3dd4c136 100644 --- a/sdks/typescript/src/session_context.ts +++ b/sdks/typescript/src/session_context.ts @@ -3,7 +3,7 @@ import type { Peer } from './peer' export interface SummaryData { content: string - message_id: number + message_id: string summary_type: string created_at: string token_count: number @@ -21,7 +21,7 @@ export class Summary { /** * The ID of the message that this summary covers up to. */ - readonly messageId: number + readonly messageId: string /** * The type of summary (short or long). @@ -138,9 +138,9 @@ export class SessionContext { const assistantId = typeof assistant === 'string' ? assistant : assistant.id const summaryMessage = this.summary ? { - role: 'system', - content: `${this.summary.content}`, - } + role: 'system', + content: `${this.summary.content}`, + } : null const messages = this.messages.map((message) => ({ role: message.peer_id === assistantId ? 'assistant' : 'user', @@ -172,20 +172,20 @@ export class SessionContext { const assistantId = typeof assistant === 'string' ? assistant : assistant.id const summaryMessage = this.summary ? { - role: 'user', - content: `${this.summary.content}`, - } + role: 'user', + content: `${this.summary.content}`, + } : null const messages = this.messages.map((message) => message.peer_id === assistantId ? { - role: 'assistant', - content: message.content, - } + role: 'assistant', + content: message.content, + } : { - role: 'user', - content: `${message.peer_id}: ${message.content}`, - } + role: 'user', + content: `${message.peer_id}: ${message.content}`, + } ) return summaryMessage ? [summaryMessage, ...messages] : messages } diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py index ae2e5442..95e363d7 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -5,7 +5,9 @@ import sentry_sdk from langfuse import get_client from pydantic import ValidationError from rich.console import Console +from sqlalchemy import select +from src import models from src.config import settings from src.dependencies import tracked_db from src.deriver.deriver import process_representation_tasks_batch @@ -52,6 +54,30 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None: queue_payload, ) raise ValueError(f"Invalid payload structure: {str(e)}") from e + + message_public_id = validated.message_public_id + if not message_public_id: + logger.info( + "Fetching message public ID for message %s", validated.message_id + ) + async with tracked_db(operation_name="summary_fallback") as db: + stmt = ( + select(models.Message) + .where(models.Message.workspace_name == validated.workspace_name) + .where(models.Message.session_name == validated.session_name) + .where(models.Message.id == validated.message_id) + ) + result = await db.execute(stmt) + + message = result.scalar_one_or_none() + if message is None: + logger.error( + "Failed to fetch message with ID %s for process_summary_task", + validated.message_id, + ) + return + message_public_id = message.public_id + with sentry_sdk.start_transaction(name="process_summary_task", op="deriver"): if settings.LANGFUSE_PUBLIC_KEY: with lf.start_as_current_span( @@ -70,6 +96,7 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None: validated.session_name, validated.message_id, validated.message_seq_in_session, + message_public_id, ) log_performance_metrics( f"summary_{validated.workspace_name}_{validated.message_id}" @@ -80,6 +107,7 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None: validated.session_name, validated.message_id, validated.message_seq_in_session, + message_public_id, ) log_performance_metrics( f"summary_{validated.workspace_name}_{validated.message_id}" diff --git a/src/deriver/queue_payload.py b/src/deriver/queue_payload.py index b787457d..c3c111e0 100644 --- a/src/deriver/queue_payload.py +++ b/src/deriver/queue_payload.py @@ -37,6 +37,8 @@ class SummaryPayload(BasePayload): session_name: str message_id: int message_seq_in_session: int + # Optional for backward compatibility with older queue items + message_public_id: str | None = None class WebhookPayload(BasePayload): @@ -124,12 +126,20 @@ def create_payload( elif task_type == "summary": if message_seq_in_session is None: raise ValueError("message_seq_in_session is required for summary tasks") + message_public_id = message.get("message_public_id") + if message_public_id is not None and ( + not isinstance(message_public_id, str) or not message_public_id.strip() + ): + raise ValueError( + "message_public_id must be a non-empty string if provided" + ) validated_payload = SummaryPayload( workspace_name=workspace_name, session_name=session_name, message_id=message_id, message_seq_in_session=message_seq_in_session, + message_public_id=message_public_id, ) # Convert back to dict for compatibility with JSON serialization diff --git a/src/routers/messages.py b/src/routers/messages.py index 3c59dbb3..2ae174c1 100644 --- a/src/routers/messages.py +++ b/src/routers/messages.py @@ -66,6 +66,7 @@ async def create_messages_for_session( "content": message.content, "peer_name": message.peer_name, "created_at": message.created_at, + "message_public_id": message.public_id, } for message in created_messages ] @@ -127,6 +128,7 @@ async def create_messages_with_file( "content": message.content, "peer_name": message.peer_name, "created_at": message.created_at, + "message_public_id": message.public_id, } for message in created_messages ] diff --git a/src/schemas.py b/src/schemas.py index a9c2f8aa..cf4f382c 100644 --- a/src/schemas.py +++ b/src/schemas.py @@ -235,7 +235,12 @@ class Session(SessionBase): class Summary(BaseModel): content: str = Field(description="The summary text") message_id: int = Field( - description="The ID of the message that this summary covers up to" + description="The internal ID of the message that this summary covers up to", + exclude=True, + ) + message_public_id: str = Field( + description="The public ID of the message that this summary covers up to", + serialization_alias="message_id", ) summary_type: str = Field(description="The type of summary (short or long)") created_at: str = Field( diff --git a/src/utils/summarizer.py b/src/utils/summarizer.py index a40eea1d..6b7b2cf3 100644 --- a/src/utils/summarizer.py +++ b/src/utils/summarizer.py @@ -39,6 +39,7 @@ class Summary(TypedDict): summary_type: str created_at: str token_count: int + message_public_id: str def to_schema_summary(s: Summary) -> schemas.Summary: @@ -48,6 +49,7 @@ def to_schema_summary(s: Summary) -> schemas.Summary: summary_type=s["summary_type"], created_at=s["created_at"], token_count=s["token_count"], + message_public_id=s["message_public_id"], ) @@ -180,6 +182,7 @@ async def summarize_if_needed( session_name: str, message_id: int, message_seq_in_session: int, + message_public_id: str, ) -> None: """ Create short/long summaries if thresholds met. @@ -209,6 +212,7 @@ async def summarize_if_needed( session_name, message_id, SummaryType.LONG, + message_public_id, ) logger.info( "Saved long summary for session %s covering up to message %s (%s in session)", @@ -225,6 +229,7 @@ async def summarize_if_needed( session_name, message_id, SummaryType.SHORT, + message_public_id, ) logger.info( "Saved short summary for session %s covering up to message %s (%s in session)", @@ -248,6 +253,7 @@ async def summarize_if_needed( session_name, message_id, SummaryType.LONG, + message_public_id, ) logger.info( "Saved long summary for session %s covering up to message %s (%s in session)", @@ -262,6 +268,7 @@ async def summarize_if_needed( session_name, message_id, SummaryType.SHORT, + message_public_id, ) logger.info( "Saved short summary for session %s covering up to message %s (%s in session)", @@ -277,6 +284,7 @@ async def _create_and_save_summary( session_name: str, message_id: int, summary_type: SummaryType, + message_public_id: str, ) -> None: """ Create a new summary and save it to the database. @@ -311,6 +319,7 @@ async def _create_and_save_summary( previous_summary_text=previous_summary_text, summary_type=summary_type, input_tokens=input_tokens, + message_public_id=message_public_id, ) await _save_summary( @@ -334,6 +343,7 @@ async def _create_summary( previous_summary_text: str | None, summary_type: SummaryType, input_tokens: int, + message_public_id: str, ) -> Summary: """ Generate a summary of the provided messages using an LLM. @@ -390,6 +400,7 @@ async def _create_summary( summary_type=summary_type.value, created_at=utc_now_iso(), token_count=summary_tokens, + message_public_id=message_public_id, ) @@ -610,6 +621,7 @@ async def get_session_context( summary_type=latest_long_summary["summary_type"], created_at=latest_long_summary["created_at"], token_count=latest_long_summary["token_count"], + message_public_id=latest_long_summary["message_public_id"], ) messages_tokens = token_limit - latest_long_summary["token_count"] messages_start_id = latest_long_summary["message_id"] @@ -622,6 +634,7 @@ async def get_session_context( summary_type=latest_short_summary["summary_type"], created_at=latest_short_summary["created_at"], token_count=latest_short_summary["token_count"], + message_public_id=latest_short_summary["message_public_id"], ) messages_tokens = token_limit - latest_short_summary["token_count"] messages_start_id = latest_short_summary["message_id"] diff --git a/tests/deriver/conftest.py b/tests/deriver/conftest.py index 3025a217..5578152d 100644 --- a/tests/deriver/conftest.py +++ b/tests/deriver/conftest.py @@ -135,6 +135,7 @@ def create_queue_payload() -> Callable[..., Any]: "message_id": message.id, "content": message.content, "created_at": message.created_at or datetime.now(timezone.utc), + "message_public_id": message.public_id, } return create_payload( diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index 9e4ae473..8309f04c 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -3,6 +3,7 @@ from typing import Any from unittest.mock import patch import pytest +from nanoid import generate as generate_nanoid from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -706,6 +707,7 @@ class TestQueueProcessing: workspace_name=session.workspace_name, peer_name=peer.name, content="First summary message", + public_id=generate_nanoid(), ), models.Message( id=1000, @@ -713,6 +715,7 @@ class TestQueueProcessing: workspace_name=session.workspace_name, peer_name=peer.name, content="Second summary message", + public_id=generate_nanoid(), ), ]