fix: defensively get summary, better summary logging, gemini client token count (#236)

* fix: remove default max_distance from get_working_rep, get token count from gemini client properly, log summaries better

* fix: defensively get message public id (new field)

* fix: add defensive check to to_schema_summary

---------

Co-authored-by: Rajat Ahuja <rahuja445@gmail.com>

src/crud/representation.py

@@ -300,9 +300,7 @@ class RepresentationManager:
                 db,
                 query=include_semantic_query,
                 top_k=semantic_observations,
-                max_distance=semantic_search_max_distance
-                if semantic_search_max_distance is not None
-                else 0.3,
+                max_distance=semantic_search_max_distance,
             )
             representation.merge_representation(
                 Representation.from_documents(semantic_docs)
@@ -331,7 +329,7 @@ class RepresentationManager:
         db: AsyncSession,
         query: str,
         top_k: int,
-        max_distance: float,
+        max_distance: float | None = None,
         level: str | None = None,
         conversation_context: str = "",
     ) -> list[models.Document]:
@@ -343,8 +341,8 @@ class RepresentationManager:
                     query,
                     level,
                     conversation_context,
-                    max_distance,
                     top_k,
+                    max_distance,
                 )
             else:
                 documents = await crud.query_documents(
@@ -428,8 +426,8 @@ class RepresentationManager:
         query: str,
         level: str,
         conversation_context: str,
-        max_distance: float,
         count: int,
+        max_distance: float | None = None,
     ) -> list[models.Document]:
         """Query documents for a specific level."""
         documents = await crud.query_documents(

src/utils/clients.py

@@ -519,8 +519,8 @@ async def honcho_llm_call_inner(
                 # Safely extract response data
                 text_content = gemini_response.text if gemini_response.text else ""
                 token_count = (
-                    gemini_response.candidates[0].token_count or 0
-                    if gemini_response.candidates
+                    gemini_response.usage_metadata.candidates_token_count or 0
+                    if gemini_response.usage_metadata
                     else 0
                 )
                 finish_reason = (
@@ -547,8 +547,8 @@ async def honcho_llm_call_inner(
                 )
 
                 token_count = (
-                    gemini_response.candidates[0].token_count or 0
-                    if gemini_response.candidates
+                    gemini_response.usage_metadata.candidates_token_count or 0
+                    if gemini_response.usage_metadata
                     else 0
                 )
                 finish_reason = (

src/utils/summarizer.py

@@ -49,7 +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"],
+        message_public_id=s.get("message_public_id", ""),
     )
 
 
@@ -195,9 +195,6 @@ async def summarize_if_needed(
         session_name: The session name
         message_id: The message ID
     """
-
-    logger.debug("Checking if summaries should be created for session %s", session_name)
-
     should_create_long: bool = message_seq_in_session % MESSAGES_PER_LONG_SUMMARY == 0
     should_create_short: bool = message_seq_in_session % MESSAGES_PER_SHORT_SUMMARY == 0
 
@@ -214,11 +211,11 @@ async def summarize_if_needed(
                     SummaryType.LONG,
                     message_public_id,
                 )
-                logger.info(
-                    "Saved long summary for session %s covering up to message %s (%s in session)",
-                    session_name,
-                    message_id,
+                accumulate_metric(
+                    f"summary_{workspace_name}_{message_id}",
+                    "long_summary_up_to_message",
                     message_seq_in_session,
+                    "count",
                 )
 
         async def create_short_summary():
@@ -231,11 +228,11 @@ async def summarize_if_needed(
                     SummaryType.SHORT,
                     message_public_id,
                 )
-                logger.info(
-                    "Saved short summary for session %s covering up to message %s (%s in session)",
-                    session_name,
-                    message_id,
+                accumulate_metric(
+                    f"summary_{workspace_name}_{message_id}",
+                    "short_summary_up_to_message",
                     message_seq_in_session,
+                    "count",
                 )
 
         await asyncio.gather(
@@ -255,11 +252,11 @@ async def summarize_if_needed(
                     SummaryType.LONG,
                     message_public_id,
                 )
-                logger.info(
-                    "Saved long summary for session %s covering up to message %s (%s in session)",
-                    session_name,
-                    message_id,
+                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(
@@ -270,11 +267,11 @@ async def summarize_if_needed(
                     SummaryType.SHORT,
                     message_public_id,
                 )
-                logger.info(
-                    "Saved short summary for session %s covering up to message %s (%s in session)",
-                    session_name,
-                    message_id,
+                accumulate_metric(
+                    f"summary_{workspace_name}_{message_id}",
+                    "short_summary_up_to_message",
                     message_seq_in_session,
+                    "count",
                 )
 
 
@@ -294,7 +291,7 @@ async def _create_and_save_summary(
     4. Save the new summary to the database
     """
 
-    logger.info("Creating new %s summary", summary_type.name)
+    logger.debug("Creating new %s summary", summary_type.name)
     # Time summarization step
     summary_start = time.perf_counter()
 
@@ -329,6 +326,19 @@ async def _create_and_save_summary(
         session_name,
     )
 
+    accumulate_metric(
+        f"summary_{workspace_name}_{message_id}",
+        f"{summary_type.name}_summary_text",
+        new_summary["content"],
+        "blob",
+    )
+    accumulate_metric(
+        f"summary_{workspace_name}_{message_id}",
+        f"{summary_type.name}_summary_size",
+        new_summary["token_count"],
+        "tokens",
+    )
+
     summary_duration = (time.perf_counter() - summary_start) * 1000
     accumulate_metric(
         f"summary_{workspace_name}_{message_id}",
@@ -374,9 +384,6 @@ async def _create_summary(
             logger.error(
                 "Generated summary is empty! This may indicate a token limit issue."
             )
-
-        logger.info("Summary text: %s", summary_text)
-        logger.info("Summary size: %s tokens", summary_tokens)
     except Exception:
         logger.exception("Error generating summary!")
         # Fallback to a basic summary in case of error
@@ -387,13 +394,6 @@ async def _create_summary(
         )
         summary_tokens = 50
 
-    accumulate_metric(
-        f"summary_{messages[-1].workspace_name}_{messages[-1].id}",
-        f"{summary_type.name}_summary_size",
-        response.output_tokens if response else f"{summary_tokens} (est.)",
-        "tokens",
-    )
-
     return Summary(
         content=summary_text,
         message_id=messages[-1].id if messages else 0,
@@ -621,7 +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"],
+                message_public_id=latest_long_summary.get("message_public_id", ""),
             )
             messages_tokens = token_limit - latest_long_summary["token_count"]
             messages_start_id = latest_long_summary["message_id"]
@@ -634,7 +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"],
+                message_public_id=latest_short_summary.get("message_public_id", ""),
             )
             messages_tokens = token_limit - latest_short_summary["token_count"]
             messages_start_id = latest_short_summary["message_id"]

tests/utils/test_clients.py

@@ -542,9 +542,11 @@ class TestGoogleClient:
         mock_response.text = "Hello from Gemini"
         mock_finish_reason = Mock()
         mock_finish_reason.name = "STOP"
-        mock_response.candidates = [
-            Mock(token_count=5, finish_reason=mock_finish_reason)
-        ]
+        mock_response.candidates = [Mock(finish_reason=mock_finish_reason)]
+        # Mock the usage_metadata with candidates_token_count
+        mock_usage_metadata = Mock()
+        mock_usage_metadata.candidates_token_count = 5
+        mock_response.usage_metadata = mock_usage_metadata
         # Mock the async aio interface
         mock_aio = Mock()
         mock_aio.models.generate_content = AsyncMock(return_value=mock_response)
@@ -572,9 +574,11 @@ class TestGoogleClient:
         mock_response.text = '{"result": "success"}'
         mock_finish_reason = Mock()
         mock_finish_reason.name = "STOP"
-        mock_response.candidates = [
-            Mock(token_count=10, finish_reason=mock_finish_reason)
-        ]
+        mock_response.candidates = [Mock(finish_reason=mock_finish_reason)]
+        # Mock the usage_metadata with candidates_token_count
+        mock_usage_metadata = Mock()
+        mock_usage_metadata.candidates_token_count = 10
+        mock_response.usage_metadata = mock_usage_metadata
         # Mock the async aio interface
         mock_aio = Mock()
         mock_aio.models.generate_content = AsyncMock(return_value=mock_response)
@@ -606,9 +610,11 @@ class TestGoogleClient:
         mock_response.parsed = mock_parsed
         mock_finish_reason = Mock()
         mock_finish_reason.name = "STOP"
-        mock_response.candidates = [
-            Mock(token_count=15, finish_reason=mock_finish_reason)
-        ]
+        mock_response.candidates = [Mock(finish_reason=mock_finish_reason)]
+        # Mock the usage_metadata with candidates_token_count
+        mock_usage_metadata = Mock()
+        mock_usage_metadata.candidates_token_count = 15
+        mock_response.usage_metadata = mock_usage_metadata
         # Mock the async aio interface
         mock_aio = Mock()
         mock_aio.models.generate_content = AsyncMock(return_value=mock_response)
@@ -691,6 +697,8 @@ class TestGoogleClient:
         mock_response = Mock()
         mock_response.text = "Response text"
         mock_response.candidates = []  # Empty candidates
+        # Mock usage_metadata as None to test fallback
+        mock_response.usage_metadata = None
         # Mock the async aio interface
         mock_aio = Mock()
         mock_aio.models.generate_content = AsyncMock(return_value=mock_response)