diff --git a/migrations/versions/394e7c39362d_add_token_count_to_queueitem_table.py b/migrations/versions/394e7c39362d_add_token_count_to_queueitem_table.py deleted file mode 100644 index 306c772b..00000000 --- a/migrations/versions/394e7c39362d_add_token_count_to_queueitem_table.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Add token_count to QueueItem table - -Revision ID: 394e7c39362d -Revises: 88b0fb10906f -Create Date: 2025-08-27 10:49:26.591473 - -""" - -import json -from collections.abc import Sequence - -import sqlalchemy as sa -from alembic import op - -from migrations.utils import column_exists -from src.config import settings - -# revision identifiers, used by Alembic. -revision: str = "394e7c39362d" -down_revision: str | None = "88b0fb10906f" -branch_labels: str | Sequence[str] | None = None -depends_on: str | Sequence[str] | None = None - -schema = settings.DB.SCHEMA - - -def upgrade() -> None: - op.add_column( - "queue", - sa.Column("token_count", sa.Integer(), nullable=False, server_default="0"), - schema=schema, - ) - - bind = op.get_bind() - - BATCH_SIZE = 500 # Process 500 items at a time - last_seen_id: int = 0 - - while True: - items_stmt = sa.text( - f""" - SELECT id, payload FROM {schema}.queue - WHERE processed = false - AND token_count = 0 - AND task_type IN ('representation', 'summary') - AND (payload->>'message_id') IS NOT NULL - AND id > :after_id - ORDER BY id - LIMIT :batch_size - """ - ).columns(id=sa.Integer, payload=sa.JSON) - - items_to_backfill = ( - bind.execute( - items_stmt, - {"batch_size": BATCH_SIZE, "after_id": last_seen_id}, - ) - .mappings() - .all() - ) - - if not items_to_backfill: - break # No more items to process - last_seen_id = items_to_backfill[-1]["id"] - - def _extract_message_id(payload: object) -> int | None: - if isinstance(payload, dict): - return payload.get("message_id") # pyright: ignore - if payload is None: - return None - try: - parsed = json.loads(str(payload)) - except (TypeError, json.JSONDecodeError): - return None - if isinstance(parsed, dict): - return parsed.get("message_id") # pyright: ignore - return None - - message_id_map = { - row["id"]: _extract_message_id(row["payload"]) for row in items_to_backfill - } - # Drop rows where message_id couldn't be extracted - message_id_map = { - queue_id: message_id - for queue_id, message_id in message_id_map.items() - if message_id is not None - } - - # Get token counts for the message IDs - token_ids = list(message_id_map.values()) - if not token_ids: - continue - - placeholders = ", ".join(f":id_{idx}" for idx in range(len(token_ids))) - token_counts_stmt = sa.text( - f""" - SELECT id, token_count FROM {schema}.messages - WHERE id IN ({placeholders}) - """ - ) - bind_params = {f"id_{idx}": token_id for idx, token_id in enumerate(token_ids)} - token_counts_result = bind.execute(token_counts_stmt, bind_params) - token_map = {msg_id: count for msg_id, count in token_counts_result.fetchall()} - - # Prepare parameters for the bulk update, skipping any queue items whose - # message_id was not found in the messages table. - update_params = [ - (qid, token_map.get(mid)) - for qid, mid in message_id_map.items() - if token_map.get(mid) is not None - ] - - if update_params: - update_stmt = sa.text( - f"UPDATE {schema}.queue SET token_count = :token_count WHERE id = :queue_id" - ) - - bind.execute( - update_stmt, - [ - {"queue_id": queue_id, "token_count": token_count} - for queue_id, token_count in update_params - ], - ) - - # If we fetched fewer items than the batch size, we are on the last batch - if len(items_to_backfill) < BATCH_SIZE: - break - - -def downgrade() -> None: - inspector = sa.inspect(op.get_bind()) - if column_exists("queue", "token_count", inspector): - op.drop_column("queue", "token_count", schema=schema) diff --git a/src/deriver/enqueue.py b/src/deriver/enqueue.py index 5974fc9e..c097e99c 100644 --- a/src/deriver/enqueue.py +++ b/src/deriver/enqueue.py @@ -165,7 +165,6 @@ def create_representation_record( "payload": processed_payload, "session_id": session_id, "task_type": "representation", - "token_count": message.get("token_count"), } @@ -198,7 +197,6 @@ def create_summary_record( "payload": processed_payload, "session_id": session_id, "task_type": "summary", - "token_count": message.get("token_count"), } diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 660165e3..f396efe3 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -8,7 +8,7 @@ from logging import getLogger import sentry_sdk from dotenv import load_dotenv from sentry_sdk.integrations.asyncio import AsyncioIntegration -from sqlalchemy import delete, select, update +from sqlalchemy import Integer, delete, select, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql import func @@ -234,50 +234,24 @@ class QueueManager: async with self.semaphore: message_count = 0 try: + from src.deriver.utils import parse_work_unit_key + + parsed_key = parse_work_unit_key(work_unit_key) + task_type = parsed_key["task_type"] + while not self.shutdown_event.is_set(): - candidate_messages = await self.get_message_batch( + messages_to_process: list[QueueItem] = await self.get_message_batch( work_unit_key, - limit=10, # hard limit of 10 messages per batch + task_type, ) - if not candidate_messages: - logger.debug(f"No more messages for work unit {work_unit_key}") - break - - next_message = candidate_messages[0] - task_type = next_message.task_type - - messages_to_process: list[QueueItem] = [] - - if task_type != "representation": - messages_to_process.append(next_message) - else: - # It's a representation task, build a batch - token_count = 0 - max_tokens = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS - - for msg in candidate_messages: - msg_tokens = msg.token_count or 0 - # Always process at least one message, even if over limit - if ( - not messages_to_process - or token_count + msg_tokens <= max_tokens - ): - messages_to_process.append(msg) - token_count += msg_tokens - else: - break - if not messages_to_process: - logger.warning( - "No messages to process, breaking loop for work unit %s to prevent infinite loop.", - work_unit_key, - ) + logger.debug(f"No more messages for work unit {work_unit_key}") break # Process the batch/single item try: - raw_payloads = [msg.payload for msg in messages_to_process] - await process_items(task_type, raw_payloads) + payloads = [msg.payload for msg in messages_to_process] + await process_items(task_type, payloads) except Exception as e: logger.error( f"Error processing tasks for work unit {work_unit_key}: {e}", @@ -286,7 +260,6 @@ class QueueManager: if settings.SENTRY.ENABLED: sentry_sdk.capture_exception(e) - # Mark messages as processed (only for non-LLM errors) await self.mark_messages_as_processed( messages_to_process, work_unit_key ) @@ -343,19 +316,79 @@ class QueueManager: @sentry_sdk.trace async def get_message_batch( - self, work_unit_key: str, limit: int + self, work_unit_key: str, task_type: str ) -> list[QueueItem]: - """Get a batch of unprocessed messages for a specific work unit ordered by id.""" + """ + Get a batch of unprocessed messages for a specific work unit ordered by id. + For representation tasks, this will be a batch of messages up to REPRESENTATION_BATCH_MAX_TOKENS. + For other tasks, it will be a single message. + """ async with tracked_db("get_message_batch") as db: - query = ( - select(models.QueueItem) - .where(models.QueueItem.work_unit_key == work_unit_key) - .where(~models.QueueItem.processed) - .order_by(models.QueueItem.id) - .limit(limit) - ) - result = await db.execute(query) - messages = result.scalars().all() + if task_type != "representation": + # For non-representation tasks, just get the next single message. + query = ( + select(models.QueueItem) + .where(models.QueueItem.work_unit_key == work_unit_key) + .where(~models.QueueItem.processed) + .order_by(models.QueueItem.id) + .limit(1) + ) + result = await db.execute(query) + messages = result.scalars().all() + else: + # For representation tasks, get a batch based on token count. + # Always get at least the first message, then include additional messages + # as long as cumulative token count stays within limit. + # Join with messages table to get the actual token_count + + # Create CTE with row numbers and cumulative token counts + cte = ( + select( + models.QueueItem.id, + func.row_number() + .over(order_by=models.QueueItem.id) + .label("row_num"), + func.sum(models.Message.token_count) + .over(order_by=models.QueueItem.id) + .label("cumulative_token_count"), + ) + .select_from( + models.QueueItem.__table__.join( + models.Message.__table__, + func.cast( + models.QueueItem.payload["message_id"].astext, Integer + ) + == models.Message.id, + ) + ) + .where(models.QueueItem.work_unit_key == work_unit_key) + .where(~models.QueueItem.processed) + .order_by(models.QueueItem.id) + .cte() + ) + + # Select messages where either: + # 1. It's the first message (row_num = 1), OR + # 2. The cumulative token count is within the limit + query = ( + select(models.QueueItem) + .where( + models.QueueItem.id.in_( + select(cte.c.id).where( + (cte.c.row_num == 1) + | ( + cte.c.cumulative_token_count + <= settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + ) + ) + ) + ) + .order_by(models.QueueItem.id) + ) + + result = await db.execute(query) + messages = result.scalars().all() + # Important: commit to avoid tracked_db's rollback expiring the instance # We rely on expire_on_commit=False to keep attributes accessible post-close await db.commit() diff --git a/src/models.py b/src/models.py index 9b63d632..defe2110 100644 --- a/src/models.py +++ b/src/models.py @@ -376,7 +376,6 @@ class QueueItem(Base): task_type: Mapped[TaskType] = mapped_column(TEXT, nullable=False) payload: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False) processed: Mapped[bool] = mapped_column(Boolean, default=False) - token_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False) def __repr__(self) -> str: return f"QueueItem(id={self.id}, session_id={self.session_id}, work_unit_key={self.work_unit_key}, task_type={self.task_type}, payload={self.payload}, processed={self.processed})" diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index 8975d5f6..4f342261 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -118,22 +118,34 @@ class TestQueueProcessing: session, peers = sample_session_with_peers peer = peers[0] - payloads: list[Any] = [] - for i in range(3): - payloads.append( - create_queue_payload( # type: ignore[reportUnknownArgumentType] - message=models.Message( - id=i, - session_name=session.name, - workspace_name=session.workspace_name, - peer_name=peer.name, - content="hello", - ), # include id for payload builder - task_type="representation", - sender_name=peer.name, - target_name=peer.name, - ) + # Create and save messages to the database first + messages: list[models.Message] = [] + for _ in range(3): + message = models.Message( + session_name=session.name, + workspace_name=session.workspace_name, + peer_name=peer.name, + content="hello", + token_count=10, ) + db_session.add(message) + messages.append(message) + + await db_session.commit() + + # Refresh to get the actual IDs + for message in messages: + await db_session.refresh(message) + + payloads: list[Any] = [] + for message in messages: + payload = create_queue_payload( # type: ignore[reportUnknownArgumentType] + message=message, + task_type="representation", + sender_name=peer.name, + target_name=peer.name, + ) + payloads.append(payload) items = await add_queue_items(payloads, session.id) # Determine ascending order by DB id @@ -151,14 +163,20 @@ class TestQueueProcessing: first, second = ordered[0], ordered[1] qm = QueueManager() - batch = await qm.get_message_batch(first.work_unit_key, limit=1) + batch = await qm.get_message_batch( + first.work_unit_key, + task_type="representation", + ) nxt = batch[0] if batch else None assert nxt is not None and nxt.id == first.id # Mark first processed, next should be the second first.processed = True await db_session.commit() - batch2 = await qm.get_message_batch(first.work_unit_key, limit=1) + batch2 = await qm.get_message_batch( + first.work_unit_key, + task_type="representation", + ) nxt2 = batch2[0] if batch2 else None assert nxt2 is not None and nxt2.id == second.id @@ -253,17 +271,26 @@ class TestQueueProcessing: # Create messages with token counts that exceed batch limit limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS - messages = [ - models.Message( - id=i, + token_counts = [limit // 2, limit // 2, limit // 2] + + # Create and save messages to the database first + 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"Test message {i}", token_count=token_count, ) - for i, token_count in enumerate([limit // 2, limit // 2, limit // 2]) - ] + db_session.add(message) + messages.append(message) + + await db_session.commit() + + # Refresh to get the actual IDs + for message in messages: + await db_session.refresh(message) # Create queue items with token counts payloads = [ @@ -280,7 +307,7 @@ class TestQueueProcessing: from src.deriver.utils import get_work_unit_key queue_items: list[models.QueueItem] = [] - for payload, message in zip(payloads, messages, strict=False): + for payload in payloads: task_type = payload.get("task_type", "unknown") work_unit_key = get_work_unit_key(task_type, payload) @@ -290,7 +317,6 @@ class TestQueueProcessing: work_unit_key=work_unit_key, payload=payload, processed=False, - token_count=message.token_count, ) db_session.add(queue_item) queue_items.append(queue_item) @@ -325,68 +351,6 @@ class TestQueueProcessing: assert processed_batches[1]["payload_count"] == 1 assert all(b["task_type"] == "representation" for b in processed_batches) - @pytest.mark.asyncio - async def test_hard_batch_size_limit( - self, - db_session: AsyncSession, - sample_session_with_peers: tuple[models.Session, list[models.Peer]], - create_queue_payload: Callable[..., Any], - ) -> None: - """Test that get_message_batch respects the hard limit of 10 messages""" - session, peers = sample_session_with_peers - peer = peers[0] - - # Create 15 messages to test batch size limit - messages = [ - models.Message( - id=i, - session_name=session.name, - workspace_name=session.workspace_name, - peer_name=peer.name, - content=f"Test message {i}", - token_count=5, # Small tokens to avoid token-based batching - ) - for i in range(15) - ] - - payloads = [ - create_queue_payload(msg, "representation", peer.name, peer.name) - for msg in messages - ] - - # Create queue items - from src.deriver.utils import get_work_unit_key - - queue_items: list[models.QueueItem] = [] - for payload, message in zip(payloads, messages, strict=False): - task_type = payload.get("task_type", "unknown") - work_unit_key = get_work_unit_key(task_type, payload) - - queue_item = models.QueueItem( - session_id=session.id, - task_type=task_type, - work_unit_key=work_unit_key, - payload=payload, - processed=False, - token_count=message.token_count, - ) - db_session.add(queue_item) - queue_items.append(queue_item) - - await db_session.commit() - - # Test batch size limit - qm = QueueManager() - batch = await qm.get_message_batch(queue_items[0].work_unit_key, limit=10) - assert len(batch) == 10 # Should respect hard limit - - # Mark the first batch as processed - await qm.mark_messages_as_processed(batch, queue_items[0].work_unit_key) - - # Get the next batch - should return remaining 5 items - next_batch = await qm.get_message_batch(queue_items[0].work_unit_key, limit=10) - assert len(next_batch) == 5 # Should return remaining items - @pytest.mark.asyncio async def test_single_message_processing( self, @@ -401,6 +365,7 @@ class TestQueueProcessing: peer = peers[0] # Create two summary messages + token_counts = [500, 600] messages = [ models.Message( id=999, @@ -408,7 +373,6 @@ class TestQueueProcessing: workspace_name=session.workspace_name, peer_name=peer.name, content="First summary message", - token_count=500, ), models.Message( id=1000, @@ -416,7 +380,6 @@ class TestQueueProcessing: workspace_name=session.workspace_name, peer_name=peer.name, content="Second summary message", - token_count=600, ), ] @@ -426,6 +389,7 @@ class TestQueueProcessing: payload = create_queue_payload( message, "summary", message_seq_in_session=i + 1 ) + payload["token_count"] = token_counts[i] from src.deriver.utils import get_work_unit_key work_unit_key = get_work_unit_key("summary", payload) @@ -436,7 +400,6 @@ class TestQueueProcessing: work_unit_key=work_unit_key, payload=payload, processed=False, - token_count=message.token_count, ) db_session.add(queue_item) queue_items.append(queue_item) @@ -490,5 +453,207 @@ class TestQueueProcessing: # Optionally verify the items have the expected token counts from the messages expected_token_counts = [500, 600] # From the test messages - actual_token_counts = [item.token_count for item in processed_items] + actual_token_counts = [ + item.payload.get("token_count") or 0 for item in processed_items + ] assert sorted(actual_token_counts) == sorted(expected_token_counts) + + @pytest.mark.asyncio + async def test_first_message_exceeds_token_limit_still_included( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + ) -> None: + """Test that if the first message exceeds BATCH_MAX_TOKENS, it's still included alone""" + from unittest.mock import patch + + session, peers = sample_session_with_peers + peer = peers[0] + + # Create messages where first message exceeds the batch limit + limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + token_counts = [limit + 1000, 100, 200] # First message way over limit + + # Create and save messages to the database first + 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"Test message {i}", + token_count=token_count, + ) + db_session.add(message) + messages.append(message) + + await db_session.commit() + + # Refresh to get the actual IDs + for message in messages: + await db_session.refresh(message) + + # Create queue items + payloads = [ + create_queue_payload( # type: ignore[reportUnknownArgumentType] + message=msg, + task_type="representation", + sender_name=peer.name, + target_name=peer.name, + ) + for msg in messages + ] + + # Add items to queue + from src.deriver.utils import get_work_unit_key + + queue_items: list[models.QueueItem] = [] + for payload in payloads: + task_type = payload.get("task_type", "unknown") + work_unit_key = get_work_unit_key(task_type, payload) + + queue_item = models.QueueItem( + session_id=session.id, + task_type=task_type, + work_unit_key=work_unit_key, + payload=payload, + processed=False, + ) + db_session.add(queue_item) + queue_items.append(queue_item) + + await db_session.commit() + for item in queue_items: + await db_session.refresh(item) + + # Mock process_items to capture batches + processed_batches: list[dict[str, Any]] = [] + + async def mock_process_items( + task_type: str, queue_payloads: list[dict[str, Any]] + ) -> None: + processed_batches.append( + { + "task_type": task_type, + "payload_count": len(queue_payloads), + } + ) + + # Process work unit and verify batching + qm = QueueManager() + with patch( + "src.deriver.queue_manager.process_items", side_effect=mock_process_items + ): + await qm.process_work_unit(queue_items[0].work_unit_key) + + # Should create 2 batches: first large message alone, then second and third together + assert len(processed_batches) == 2 + assert ( + processed_batches[0]["payload_count"] == 1 + ) # First message (over limit) alone + assert processed_batches[1]["payload_count"] == 2 # Second and third messages + assert all(b["task_type"] == "representation" for b in processed_batches) + + @pytest.mark.asyncio + async def test_message_exactly_at_token_limit( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + ) -> None: + """Test boundary condition when cumulative sum exactly equals limit""" + from unittest.mock import patch + + session, peers = sample_session_with_peers + peer = peers[0] + + # Create messages that test the exact boundary + limit = settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS + token_counts = [ + limit // 2, + limit // 2, + 1, + ] # First two exactly at limit, third exceeds + + # Create and save messages to the database first + 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"Test message {i}", + token_count=token_count, + ) + db_session.add(message) + messages.append(message) + + await db_session.commit() + + # Refresh to get the actual IDs + for message in messages: + await db_session.refresh(message) + + # Create queue items + payloads = [ + create_queue_payload( # type: ignore[reportUnknownArgumentType] + message=msg, + task_type="representation", + sender_name=peer.name, + target_name=peer.name, + ) + for msg in messages + ] + + # Add items to queue + from src.deriver.utils import get_work_unit_key + + queue_items: list[models.QueueItem] = [] + for payload in payloads: + task_type = payload.get("task_type", "unknown") + work_unit_key = get_work_unit_key(task_type, payload) + + queue_item = models.QueueItem( + session_id=session.id, + task_type=task_type, + work_unit_key=work_unit_key, + payload=payload, + processed=False, + ) + db_session.add(queue_item) + queue_items.append(queue_item) + + await db_session.commit() + for item in queue_items: + await db_session.refresh(item) + + # Mock process_items to capture batches + processed_batches: list[dict[str, Any]] = [] + + async def mock_process_items( + task_type: str, queue_payloads: list[dict[str, Any]] + ) -> None: + processed_batches.append( + { + "task_type": task_type, + "payload_count": len(queue_payloads), + } + ) + + # Process work unit and verify batching + qm = QueueManager() + with patch( + "src.deriver.queue_manager.process_items", side_effect=mock_process_items + ): + await qm.process_work_unit(queue_items[0].work_unit_key) + + # Should create 2 batches: first two messages together (exactly at limit), third alone + assert len(processed_batches) == 2 + assert ( + processed_batches[0]["payload_count"] == 2 + ) # First two messages (exactly at limit) + assert ( + processed_batches[1]["payload_count"] == 1 + ) # Third message (exceeds limit) + assert all(b["task_type"] == "representation" for b in processed_batches)