From b0dc05c909ae754f637eb644ccc55625d8e1b6e8 Mon Sep 17 00:00:00 2001 From: doria <93405247+dr-frmr@users.noreply.github.com> Date: Wed, 14 Jan 2026 14:59:35 -0500 Subject: [PATCH] feat: message ingestion forced-batching (#327) * feat: implement forced batching for representation tasks and adjust max tokens - Updated `REPRESENTATION_BATCH_MAX_TOKENS` to 1024 in `config.py`. - Enhanced `get_and_claim_work_units` in `QueueManager` to enforce batching based on token thresholds. - Added tests to ensure representation work units are only claimed when token counts meet or exceed the threshold. - Introduced a synthesis prompt for tool execution to improve final response generation. * chore: update queue-status docs to match new behavior * fix: make representation work query efficient * fix: Align alembic models and force dreams on for tests * fix: clean queue between tests --------- Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> --- .../features/advanced/queue-status.mdx | 14 +- src/config.py | 5 +- src/deriver/queue_manager.py | 66 +++- src/models.py | 6 + src/utils/clients.py | 10 + tests/deriver/conftest.py | 21 +- tests/deriver/test_queue_processing.py | 283 ++++++++++++++++++ tests/dreamer/test_dream_scheduler.py | 4 +- 8 files changed, 384 insertions(+), 25 deletions(-) diff --git a/docs/v2.6.0-alpha/documentation/features/advanced/queue-status.mdx b/docs/v2.6.0-alpha/documentation/features/advanced/queue-status.mdx index 5b3d8dfa..8a3e00c5 100644 --- a/docs/v2.6.0-alpha/documentation/features/advanced/queue-status.mdx +++ b/docs/v2.6.0-alpha/documentation/features/advanced/queue-status.mdx @@ -4,11 +4,13 @@ description: Learn how to check the status of Honcho's reasoning icon: "lines-leaning" --- -Whenever messages are stored in Honcho, a background process kicks off to [reason](/v2.6.0-alpha/documentation/core-concepts/reasoning) about the conversation and generate insights. +Whenever messages are stored in Honcho, background processes kick off to [reason](/v2.6.0-alpha/documentation/core-concepts/reasoning) about the conversation and generate insights. -Reasoning is an asynchronous process and, depending on load, may not immediately -generate insights for the latest message you've sent. To help with this, Honcho -provides several utilities to check the status of the queue. +Reasoning is an asynchronous process and will not immediately +generate insights for the latest message you've sent. This is +by design: we want to reason efficiently over batches of messages +rather than assessing each message in a vacuum. Honcho provides +several utilities to check the status of the queue. ```python Python @@ -75,7 +77,7 @@ work_units will be processed in parallel generate an additional work unit for every peer that has `observe_others=True` The `get_queue_status` and `poll_queue_status` methods can take additional -parameters to scope the status to a specific work unit +parameters to scope the status to a specific work unit: ```python Python @@ -104,7 +106,7 @@ export const QueueStatusOptionsSchema = z.object({ Additionally, there are queue status and polling queue status methods available on the session objects in each of the SDKs. -Below are the function signatures for the session level queue status method +Below are the function signatures for the session level queue status method: ```python python diff --git a/src/config.py b/src/config.py index 01741573..f054e5fa 100644 --- a/src/config.py +++ b/src/config.py @@ -267,8 +267,8 @@ class DeriverSettings(BackupLLMSettingsMixin, HonchoSettings): REPRESENTATION_BATCH_MAX_TOKENS: Annotated[ int, - Field(default=4096, ge=128, le=16_384), - ] = 4096 + Field(default=1024, ge=128, le=16_384), + ] = 1024 @model_validator(mode="after") def validate_batch_tokens_vs_context_limit(self): @@ -361,6 +361,7 @@ class DialecticSettings(HonchoSettings): MODEL="gemini-3-flash-preview", THINKING_BUDGET_TOKENS=0, MAX_TOOL_ITERATIONS=5, + TOOL_CHOICE="any", ), "medium": DialecticLevelSettings( PROVIDER="anthropic", diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index b3734edd..5288dfbc 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -10,7 +10,7 @@ import sentry_sdk from dotenv import load_dotenv from nanoid import generate as generate_nanoid from sentry_sdk.integrations.asyncio import AsyncioIntegration -from sqlalchemy import and_, delete, select, update +from sqlalchemy import and_, delete, or_, select, update from sqlalchemy.dialects.postgresql import insert from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql import func @@ -204,26 +204,64 @@ class QueueManager: async def get_and_claim_work_units(self) -> dict[str, str]: """ Get available work units that aren't being processed. + For representation tasks, only returns work units with accumulated tokens + >= REPRESENTATION_BATCH_MAX_TOKENS (forced batching). Returns a dict mapping work_unit_key to aqs_id. """ limit: int = max(0, self.workers - self.get_total_owned_work_units()) if limit == 0: return {} - async with tracked_db( - "get_available_work_units" - ) as db: # Get number of available workers - query = ( - select(models.QueueItem.work_unit_key) - .limit(limit) - .outerjoin( - models.ActiveQueueSession, - models.QueueItem.work_unit_key - == models.ActiveQueueSession.work_unit_key, + + batch_max_tokens = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + + async with tracked_db("get_available_work_units") as db: + representation_prefix = "representation:" + token_stats_subq = ( + select( + models.QueueItem.work_unit_key, + func.sum(models.Message.token_count).label("total_tokens"), + ) + .join( + models.Message, + models.QueueItem.message_id == models.Message.id, ) .where(~models.QueueItem.processed) - .where(models.QueueItem.work_unit_key.isnot(None)) - .where(models.ActiveQueueSession.work_unit_key.is_(None)) - .distinct() + .where(models.QueueItem.work_unit_key.startswith(representation_prefix)) + .group_by(models.QueueItem.work_unit_key) + .subquery() + ) + + work_units_subq = ( + select(models.QueueItem.work_unit_key) + .where(~models.QueueItem.processed) + .group_by(models.QueueItem.work_unit_key) + .subquery() + ) + + query = ( + select(work_units_subq.c.work_unit_key) + .limit(limit) + .outerjoin( + token_stats_subq, + work_units_subq.c.work_unit_key == token_stats_subq.c.work_unit_key, + ) + .where( + ~select(models.ActiveQueueSession.id) + .where( + models.ActiveQueueSession.work_unit_key + == work_units_subq.c.work_unit_key + ) + .exists() + ) + .where( + or_( + ~work_units_subq.c.work_unit_key.startswith( + representation_prefix + ), + func.coalesce(token_stats_subq.c.total_tokens, 0) + >= batch_max_tokens, + ) + ) ) result = await db.execute(query) diff --git a/src/models.py b/src/models.py index d5293db3..8843662c 100644 --- a/src/models.py +++ b/src/models.py @@ -473,6 +473,12 @@ class QueueItem(Base): "processed", "id", ), + Index( + "ux_queue_dream_pending_work_unit_key", + "work_unit_key", + unique=True, + postgresql_where=text("task_type = 'dream' AND processed = false"), + ), ) def __repr__(self) -> str: diff --git a/src/utils/clients.py b/src/utils/clients.py index 9b4be8a4..2edc1a28 100644 --- a/src/utils/clients.py +++ b/src/utils/clients.py @@ -851,6 +851,16 @@ async def _execute_tool_loop( f"Tool execution loop reached max iterations ({max_tool_iterations})" ) + # Add a synthesis prompt to help the model generate a response + # without tool calls - the conversation currently ends with tool results + # and the model may not know to produce text output + synthesis_prompt = ( + "You have reached the maximum number of tool calls. " + "Based on all the information you have gathered, provide your final response now. " + "Do not attempt to call any more tools." + ) + conversation_messages.append({"role": "user", "content": synthesis_prompt}) + # If streaming the final response, use the streaming helper with metadata if stream_final: stream = _stream_final_response( diff --git a/tests/deriver/conftest.py b/tests/deriver/conftest.py index a11d5cc6..8f7b1897 100644 --- a/tests/deriver/conftest.py +++ b/tests/deriver/conftest.py @@ -1,18 +1,35 @@ import asyncio -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence from datetime import datetime, timezone from typing import Any, Literal, TypeAlias, cast from unittest.mock import AsyncMock, MagicMock import pytest from nanoid import generate as generate_nanoid -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.ext.asyncio import AsyncSession from src import crud, models, schemas from src.utils.queue_payload import create_payload from src.utils.work_unit import construct_work_unit_key + +@pytest.fixture(autouse=True) +async def clean_queue_tables(db_session: AsyncSession) -> AsyncGenerator[None, None]: + """Clean up queue-related tables before each test to ensure isolation. + + This prevents webhook queue items and active queue sessions from previous tests + from polluting subsequent tests, which can cause issues when tests have a limit + on how many work units can be claimed (e.g., WORKERS=1). + """ + # Clean up before the test + await db_session.execute(delete(models.ActiveQueueSession)) + await db_session.execute(delete(models.QueueItem)) + await db_session.commit() + + yield + + QueuePayload: TypeAlias = dict[str, Any] QueuePayloadEntry: TypeAlias = QueuePayload | tuple[QueuePayload, int | None] diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index d17681b4..6cb3dc19 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -1061,3 +1061,286 @@ class TestQueueProcessing: processed_batches[1]["payload_count"] == 1 ) # Third message (exceeds limit) assert all(b["task_type"] == "representation" for b in processed_batches) + + @pytest.mark.asyncio + async def test_forced_batching_waits_for_threshold( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + ) -> None: + """Test that representation work units below token threshold are not claimed""" + + session, peers = sample_session_with_peers + peer = peers[0] + + # Create messages with tokens BELOW the threshold + limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + token_counts = [100, 100, 100] # Total 300, way below 4096 + + messages: list[models.Message] = [] + for i, token_count in enumerate(token_counts): + message = models.Message( + session_name=session.name, + workspace_name=session.workspace_name, + peer_name=peer.name, + content=f"Short message {i}", + token_count=token_count, + seq_in_session=i + 1, + ) + db_session.add(message) + messages.append(message) + + await db_session.commit() + for message in messages: + await db_session.refresh(message) + + # Create queue items + queue_items: list[models.QueueItem] = [] + for message in messages: + payload = create_queue_payload( + message=message, + task_type="representation", + observed=peer.name, + observer=peer.name, + ) + work_unit_key = construct_work_unit_key(session.workspace_name, payload) + + queue_item = models.QueueItem( + session_id=session.id, + task_type="representation", + work_unit_key=work_unit_key, + payload=payload, + processed=False, + workspace_name=session.workspace_name, + message_id=message.id, + ) + db_session.add(queue_item) + queue_items.append(queue_item) + + await db_session.commit() + + qm = QueueManager() + + # Try to claim work units - should NOT return the representation work unit + claimed = await qm.get_and_claim_work_units() + + # The representation work unit should NOT be claimed (tokens below threshold) + rep_work_unit_key = queue_items[0].work_unit_key + assert rep_work_unit_key not in claimed + + # Now add more messages to exceed the threshold + more_messages: list[models.Message] = [] + for i in range(10): + message = models.Message( + session_name=session.name, + workspace_name=session.workspace_name, + peer_name=peer.name, + content=f"Longer message {i}", + token_count=limit // 2, # Each message is half the limit + seq_in_session=len(messages) + i + 1, + ) + db_session.add(message) + more_messages.append(message) + + await db_session.commit() + for message in more_messages: + await db_session.refresh(message) + + # Add queue items for new messages + for message in more_messages: + payload = create_queue_payload( + message=message, + task_type="representation", + observed=peer.name, + observer=peer.name, + ) + + queue_item = models.QueueItem( + session_id=session.id, + task_type="representation", + work_unit_key=rep_work_unit_key, # Same work unit + payload=payload, + processed=False, + workspace_name=session.workspace_name, + message_id=message.id, + ) + db_session.add(queue_item) + + await db_session.commit() + + # Now the work unit should be claimable (tokens exceed threshold) + claimed2 = await qm.get_and_claim_work_units() + assert rep_work_unit_key in claimed2 + + @pytest.mark.asyncio + async def test_forced_batching_single_large_message( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + ) -> None: + """Test that a single message >= threshold is immediately claimable""" + + session, peers = sample_session_with_peers + peer = peers[0] + + limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + + # Create a single message that exceeds the threshold + message = models.Message( + session_name=session.name, + workspace_name=session.workspace_name, + peer_name=peer.name, + content="A very long message", + token_count=limit + 1000, # Exceeds threshold + seq_in_session=1, + ) + db_session.add(message) + await db_session.commit() + await db_session.refresh(message) + + # Create queue item + payload = create_queue_payload( + message=message, + task_type="representation", + observed=peer.name, + observer=peer.name, + ) + work_unit_key = construct_work_unit_key(session.workspace_name, payload) + + queue_item = models.QueueItem( + session_id=session.id, + task_type="representation", + work_unit_key=work_unit_key, + payload=payload, + processed=False, + workspace_name=session.workspace_name, + message_id=message.id, + ) + db_session.add(queue_item) + await db_session.commit() + + qm = QueueManager() + + # Single large message should be immediately claimable + claimed = await qm.get_and_claim_work_units() + assert work_unit_key in claimed + + @pytest.mark.asyncio + async def test_forced_batching_bypassed_for_summary( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + ) -> None: + """Test that summary tasks are processed immediately regardless of tokens""" + + session, peers = sample_session_with_peers + peer = peers[0] + + # Create a message with very few tokens + message = models.Message( + session_name=session.name, + workspace_name=session.workspace_name, + peer_name=peer.name, + content="Short", + token_count=10, # Way below threshold + seq_in_session=1, + ) + db_session.add(message) + await db_session.commit() + await db_session.refresh(message) + + # Create a SUMMARY queue item (not representation) + payload = create_queue_payload( + message=message, + task_type="summary", + message_seq_in_session=1, + ) + work_unit_key = construct_work_unit_key(session.workspace_name, payload) + + queue_item = models.QueueItem( + session_id=session.id, + task_type="summary", + work_unit_key=work_unit_key, + payload=payload, + processed=False, + workspace_name=session.workspace_name, + message_id=message.id, + ) + db_session.add(queue_item) + await db_session.commit() + + qm = QueueManager() + + # Summary should be claimable regardless of token count + claimed = await qm.get_and_claim_work_units() + assert work_unit_key in claimed + assert work_unit_key.startswith("summary:") + + @pytest.mark.asyncio + async def test_forced_batching_exact_threshold( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + ) -> None: + """Test that work units exactly at the threshold are claimable""" + + session, peers = sample_session_with_peers + peer = peers[0] + + limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + + # Create messages that sum to exactly the threshold + token_counts = [limit // 2, limit // 2] + + messages: list[models.Message] = [] + for i, token_count in enumerate(token_counts): + message = models.Message( + session_name=session.name, + workspace_name=session.workspace_name, + peer_name=peer.name, + content=f"Message {i}", + token_count=token_count, + seq_in_session=i + 1, + ) + db_session.add(message) + messages.append(message) + + await db_session.commit() + for message in messages: + await db_session.refresh(message) + + # Create queue items + work_unit_key = None + for message in messages: + payload = create_queue_payload( + message=message, + task_type="representation", + observed=peer.name, + observer=peer.name, + ) + if work_unit_key is None: + work_unit_key = construct_work_unit_key(session.workspace_name, payload) + + queue_item = models.QueueItem( + session_id=session.id, + task_type="representation", + work_unit_key=work_unit_key, + payload=payload, + processed=False, + workspace_name=session.workspace_name, + message_id=message.id, + ) + db_session.add(queue_item) + + await db_session.commit() + + qm = QueueManager() + + # Work unit exactly at threshold should be claimable + claimed = await qm.get_and_claim_work_units() + assert work_unit_key is not None + assert work_unit_key in claimed diff --git a/tests/dreamer/test_dream_scheduler.py b/tests/dreamer/test_dream_scheduler.py index b713e85b..06f904c1 100644 --- a/tests/dreamer/test_dream_scheduler.py +++ b/tests/dreamer/test_dream_scheduler.py @@ -17,7 +17,9 @@ def dream_scheduler(): DreamScheduler.reset_singleton() scheduler = DreamScheduler() set_dream_scheduler(scheduler) - yield scheduler + # Patch DREAM.ENABLED to True so tests work regardless of local config + with patch("src.dreamer.dream_scheduler.settings.DREAM.ENABLED", True): + yield scheduler # Cleanup DreamScheduler.reset_singleton()