diff --git a/src/crud/representation.py b/src/crud/representation.py index 598d067f..09eaeae7 100644 --- a/src/crud/representation.py +++ b/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( diff --git a/src/utils/clients.py b/src/utils/clients.py index dd30f02e..920af15a 100644 --- a/src/utils/clients.py +++ b/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 = ( diff --git a/src/utils/summarizer.py b/src/utils/summarizer.py index 51c71629..be21009d 100644 --- a/src/utils/summarizer.py +++ b/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"] diff --git a/tests/utils/test_clients.py b/tests/utils/test_clients.py index abc85f5d..c6552926 100644 --- a/tests/utils/test_clients.py +++ b/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)