RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

perf: use join on messages table instead of storing token_count in queue payload

This commit is contained in:
Rajat Ahuja 2025-09-24 11:45:04 -04:00
parent 0b58472ab1
commit 1cc3d9aaaf
5 changed files with 337 additions and 276 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

134
migrations/versions/394e7c39362d_add_token_count_to_queueitem_table.py
View File

@ -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)

2
src/deriver/enqueue.py
View File

@ -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"),
}

131
src/deriver/queue_manager.py
View File

@ -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()

1
src/models.py
View File

@ -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})"

345
tests/deriver/test_queue_processing.py
View File

@ -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)