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

feat: use token limit from all senders for rollup (#221) 

* feat: use token limit from all senders for rollup

* fix: send all messages to queue_manager

* test: fixe tests and move ownership check to separate query

* fix: if first sender message exceeds batch then we should return an empty batch

* feat: separate representation and non-representation processing

* fix: return msg object and queue items

* fix: rollup logic with interleaving messages

* fix: rm TODO

* Fix langfuse trace and add fallback for weird error in deriver

* fix: simplify query

* fix: comment out update_current_generation

* fix: BPR

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
Rajat Ahuja 2025-10-01 16:10:10 -04:00 committed by GitHub
parent ef84e8f99d
commit 3df044e1b7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 787 additions and 349 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

190
src/deriver/consumer.py
View File

@ -1,7 +1,6 @@
import logging
from typing import Any
import sentry_sdk
from langfuse import get_client
from pydantic import ValidationError
from rich.console import Console
@ -9,13 +8,12 @@ from rich.console import Console
from src.config import settings
from src.dependencies import tracked_db
from src.deriver.deriver import process_representation_tasks_batch
from src.models import Message
from src.utils import summarizer
from src.utils.logging import log_performance_metrics
from src.webhooks import webhook_delivery
from .queue_payload import (
RepresentationPayload,
RepresentationPayloads,
SummaryPayload,
WebhookPayload,
)
@ -28,100 +26,118 @@ console = Console(markup=True)
lf = get_client()
async def process_items(task_type: str, queue_payloads: list[dict[str, Any]]) -> None:
async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None:
"""Process a single item from the queue."""
if task_type == "webhook":
try:
validated = WebhookPayload(**queue_payload)
except ValidationError as e:
logger.error(
"Invalid webhook payload received: %s. Payload: %s",
str(e),
queue_payload,
)
raise ValueError(f"Invalid payload structure: {str(e)}") from e
async with tracked_db() as db:
await webhook_delivery.deliver_webhook(db, validated)
elif task_type == "summary":
try:
validated = SummaryPayload(**queue_payload)
except ValidationError as e:
logger.error(
"Invalid summary payload received: %s. Payload: %s",
str(e),
queue_payload,
)
raise ValueError(f"Invalid payload structure: {str(e)}") from e
if settings.LANGFUSE_PUBLIC_KEY:
with lf.start_as_current_span(
name="summary_processing",
input={
"workspace_name": validated.workspace_name,
"session_name": validated.session_name,
"message_id": validated.message_id,
},
metadata={
"summary_model": settings.SUMMARY.MODEL,
},
):
await summarizer.summarize_if_needed(
validated.workspace_name,
validated.session_name,
validated.message_id,
validated.message_seq_in_session,
)
log_performance_metrics(
f"summary_{validated.workspace_name}_{validated.message_id}"
)
else:
await summarizer.summarize_if_needed(
validated.workspace_name,
validated.session_name,
validated.message_id,
validated.message_seq_in_session,
)
log_performance_metrics(
f"summary_{validated.workspace_name}_{validated.message_id}"
)
else:
raise ValueError(f"Invalid task type: {task_type}")
async def process_representation_batch(
messages: list[Message],
sender_name: str | None,
target_name: str | None,
) -> None:
"""Validate incoming queue payloads and dispatch to the appropriate handler.
This function centralizes payload validation using a simple mapping from
task type to Pydantic model. After validation, routes the request to
the correct processor without repeating type checks elsewhere.
Args:
task_type: The type of task to process
queue_payloads: List of payload dictionaries to process
sender_name (optional): For representation tasks, the sender_name from work_unit_key
to identify which messages should be focused on
target_name (optional): For representation tasks, the target_name from work_unit_key
to identify which messages should be focused on
"""
if not queue_payloads or not queue_payloads[0]:
logger.debug("process_items received no payloads for task type %s", task_type)
if not messages or not messages[0]:
logger.debug("process_representation_batch received no payloads")
return
if sender_name is None or target_name is None:
raise ValueError(
"sender_name and target_name are required for representation tasks"
)
logger.debug(
"process_items received %s payloads for task type %s",
len(queue_payloads),
task_type,
"process_representation_batch received %s payloads",
len(messages),
)
if task_type == "webhook":
try:
validated = WebhookPayload(**queue_payloads[0])
except ValidationError as e:
logger.error(
"Invalid webhook payload received: %s. Payload: %s",
str(e),
queue_payloads[0],
)
raise ValueError(f"Invalid payload structure: {str(e)}") from e
await process_webhook(validated)
logger.debug("Finished processing webhook %s", validated.event_type)
elif task_type == "summary":
if settings.LANGFUSE_PUBLIC_KEY:
lf.update_current_trace( # type: ignore
metadata={
"critical_analysis_model": settings.DERIVER.MODEL,
}
)
try:
validated = SummaryPayload(**queue_payloads[0])
except ValidationError as e:
logger.error(
"Invalid summary payload received: %s. Payload: %s",
str(e),
queue_payloads[0],
)
raise ValueError(f"Invalid payload structure: {str(e)}") from e
await process_summary_task(validated)
elif task_type == "representation":
if settings.LANGFUSE_PUBLIC_KEY:
lf.update_current_trace(
metadata={
"critical_analysis_model": settings.DERIVER.MODEL,
}
)
try:
validated_payloads = RepresentationPayloads(
payloads=[
RepresentationPayload(**payload) for payload in queue_payloads
if settings.LANGFUSE_PUBLIC_KEY:
with lf.start_as_current_span(
name="representation_processing",
input={
"payloads": [
{
"message_id": msg.id,
"sender_name": sender_name,
"target_name": target_name,
"session_name": msg.session_name,
}
for msg in messages
]
)
except ValidationError as e:
logger.error(
"Invalid representation payloads received: %s. Payloads: %s",
str(e),
queue_payloads,
)
raise ValueError(f"Invalid payload structure: {str(e)}") from e
await process_representation_tasks_batch(validated_payloads.payloads)
},
metadata={
"critical_analysis_model": settings.DERIVER.MODEL,
},
):
await process_representation_tasks_batch(sender_name, target_name, messages)
else:
raise ValueError(f"Invalid task type: {task_type}")
@sentry_sdk.trace
async def process_webhook(
payload: WebhookPayload,
) -> None:
async with tracked_db() as db:
await webhook_delivery.deliver_webhook(db, payload)
@sentry_sdk.trace
async def process_summary_task(
payload: SummaryPayload,
) -> None:
"""
Process a summary task by generating summaries if needed.
"""
await summarizer.summarize_if_needed(
payload.workspace_name,
payload.session_name,
payload.message_id,
payload.message_seq_in_session,
)
log_performance_metrics(f"summary_{payload.workspace_name}_{payload.message_id}")
await process_representation_tasks_batch(sender_name, target_name, messages)

191
src/deriver/deriver.py
View File

@ -12,6 +12,7 @@ from src.config import settings
from src.crud.representation import GLOBAL_REPRESENTATION_COLLECTION_NAME
from src.dependencies import tracked_db
from src.deriver.utils import estimate_tokens
from src.models import Message
from src.utils import summarizer
from src.utils.clients import honcho_llm_call
from src.utils.embedding_store import EmbeddingStore
@ -26,7 +27,6 @@ from src.utils.formatting import (
from src.utils.logging import (
accumulate_metric,
conditional_observe,
format_reasoning_inputs_as_markdown,
format_reasoning_response_as_markdown,
log_observations_tree,
log_performance_metrics,
@ -46,9 +46,6 @@ from .prompts import (
estimate_base_prompt_tokens,
peer_card_prompt,
)
from .queue_payload import (
RepresentationPayload,
)
logger = logging.getLogger(__name__)
logging.getLogger("sqlalchemy.engine.Engine").disabled = True
@ -63,7 +60,7 @@ async def critical_analysis_call(
working_representation: str | None,
history: str,
new_turns: list[str],
) -> ReasoningResponse:
) -> ReasoningResponse | None:
prompt = critical_analysis_prompt(
peer_id=peer_id,
peer_card=peer_card,
@ -121,38 +118,40 @@ async def peer_card_call(
@sentry_sdk.trace
async def process_representation_tasks_batch(
payloads: list[RepresentationPayload],
sender_name: str,
target_name: str,
messages: list[Message],
) -> None:
"""
Process a batch of representation tasks by extracting insights and updating working representations.
"""
if not payloads or len(payloads) == 0:
if not messages or len(messages) == 0:
return
payloads.sort(key=lambda x: x.message_id)
messages.sort(key=lambda x: x.id)
latest_payload = payloads[-1]
earliest_payload = payloads[0]
latest_message = messages[-1]
earliest_message = messages[0]
# Start overall timing
overall_start = time.perf_counter()
logger.debug(
"Starting insight extraction for message batch starting with: %s",
earliest_payload.message_id,
earliest_message.id,
)
async with tracked_db("deriver.get_peer_card") as db:
speaker_peer_card: list[str] | None = await crud.get_peer_card(
db,
latest_payload.workspace_name,
latest_payload.sender_name,
latest_payload.target_name,
latest_message.workspace_name,
sender_name,
target_name,
)
if speaker_peer_card is None:
logger.warning("No peer card found for %s", latest_payload.sender_name)
logger.warning("No peer card found for %s", sender_name)
else:
logger.debug("Using peer card for %s", latest_payload.sender_name)
logger.debug("Using peer card for %s", sender_name)
# Get working representation data early for token estimation
async with tracked_db("deriver.get_working_representation_data") as db:
@ -160,10 +159,10 @@ async def process_representation_tasks_batch(
dict[str, Any] | str | None
) = await crud.get_working_representation_data(
db,
latest_payload.workspace_name,
latest_payload.target_name,
latest_payload.sender_name,
latest_payload.session_name,
latest_message.workspace_name,
target_name,
sender_name,
latest_message.session_name,
)
# Estimate tokens for deriver input
@ -173,8 +172,8 @@ async def process_representation_tasks_batch(
# Estimate tokens for new conversation turns
new_turns = [
format_new_turn_with_timestamp(p.content, p.created_at, p.sender_name)
for p in payloads
format_new_turn_with_timestamp(m.content, m.created_at, m.peer_name)
for m in messages
]
new_turns_tokens = estimate_tokens(new_turns)
@ -204,10 +203,10 @@ async def process_representation_tasks_batch(
async with tracked_db("deriver.get_session_context") as db:
formatted_history = await summarizer.get_session_context_formatted(
db,
latest_payload.workspace_name,
latest_payload.session_name,
latest_message.workspace_name,
latest_message.session_name,
token_limit=available_context_tokens,
cutoff=earliest_payload.message_id,
cutoff=earliest_message.id,
include_summary=True,
)
@ -216,10 +215,8 @@ async def process_representation_tasks_batch(
# otherwise, we're handling a directional representation task where the sender is
# being observed by the target.
collection_name = (
crud.construct_collection_name(
observer=latest_payload.target_name, observed=latest_payload.sender_name
)
if latest_payload.sender_name != latest_payload.target_name
crud.construct_collection_name(observer=target_name, observed=sender_name)
if sender_name != target_name
else GLOBAL_REPRESENTATION_COLLECTION_NAME
)
@ -227,21 +224,26 @@ async def process_representation_tasks_batch(
async with tracked_db("deriver.get_or_create_collection") as db:
collection = await crud.get_or_create_collection(
db,
latest_payload.workspace_name,
latest_message.workspace_name,
collection_name,
latest_payload.sender_name,
sender_name,
)
collection_name_loaded = collection.name
# Use the embedding store directly
embedding_store = EmbeddingStore(
workspace_name=latest_payload.workspace_name,
peer_name=latest_payload.sender_name,
workspace_name=latest_message.workspace_name,
peer_name=sender_name,
collection_name=collection_name_loaded,
)
# Create reasoner instance
reasoner = CertaintyReasoner(embedding_store=embedding_store, ctx=payloads)
reasoner = CertaintyReasoner(
embedding_store=embedding_store,
ctx=messages,
sender_name=sender_name,
target_name=target_name,
)
# Time context preparation
context_prep_start = time.perf_counter()
@ -273,7 +275,7 @@ async def process_representation_tasks_batch(
)
else:
# No existing working representation, use global search
query_text = [payload.content for payload in payloads]
query_text = [m.content for m in messages]
query_text = "\n".join(
query_text
) # TODO: consider a smarter strategy than concatenation
@ -323,15 +325,15 @@ async def process_representation_tasks_batch(
async with tracked_db("deriver.recalc_session_context") as db:
formatted_history = await summarizer.get_session_context_formatted(
db,
latest_payload.workspace_name,
latest_payload.session_name,
latest_message.workspace_name,
latest_message.session_name,
token_limit=available_context_tokens,
cutoff=earliest_payload.message_id,
cutoff=earliest_message.id,
include_summary=True,
)
context_prep_duration = (time.perf_counter() - context_prep_start) * 1000
accumulate_metric(
f"deriver_representation_{latest_payload.message_id}_{latest_payload.target_name}",
f"deriver_representation_{latest_message.id}_{target_name}",
"context_preparation",
context_prep_duration,
"ms",
@ -358,11 +360,13 @@ async def process_representation_tasks_batch(
log_observations_tree(final_obs_dict)
# Always save working representation to peer for dialectic access
await save_working_representation_to_peer(latest_payload, final_observations)
await save_working_representation_to_peer(
latest_message, final_observations, sender_name, target_name
)
# Calculate and log overall timing
overall_duration = (time.perf_counter() - overall_start) * 1000
accumulate_metric(
f"deriver_representation_{latest_payload.message_id}_{latest_payload.target_name}",
f"deriver_representation_{latest_message.id}_{target_name}",
"total_processing_time",
overall_duration,
"ms",
@ -371,14 +375,12 @@ async def process_representation_tasks_batch(
total_observations = sum(len(obs_list) for obs_list in final_obs_dict.values())
accumulate_metric(
f"deriver_representation_{latest_payload.message_id}_{latest_payload.target_name}",
f"deriver_representation_{latest_message.id}_{target_name}",
"final_observation_count",
total_observations,
"count",
)
log_performance_metrics(
f"deriver_representation_{latest_payload.message_id}_{latest_payload.target_name}"
)
log_performance_metrics(f"deriver_representation_{latest_message.id}_{target_name}")
if settings.LANGFUSE_PUBLIC_KEY:
lf.update_current_trace(
@ -386,24 +388,25 @@ async def process_representation_tasks_batch(
)
# The old function now just calls the batch processor with a single payload
async def process_representation_task(
payload: RepresentationPayload,
) -> None:
await process_representation_tasks_batch([payload])
class CertaintyReasoner:
"""Certainty reasoner for analyzing and deriving insights."""
embedding_store: EmbeddingStore
ctx: list[RepresentationPayload]
ctx: list[Message]
sender_name: str
target_name: str
def __init__(
self, embedding_store: EmbeddingStore, ctx: list[RepresentationPayload]
self,
embedding_store: EmbeddingStore,
ctx: list[Message],
sender_name: str,
target_name: str,
) -> None:
self.embedding_store = embedding_store
self.ctx = ctx
self.sender_name = sender_name
self.target_name = target_name
@conditional_observe
@sentry_sdk.trace
@ -417,21 +420,21 @@ class CertaintyReasoner:
Critically analyzes and revises understanding, returning structured observations.
"""
# For logging, we can just show the content of the last message
latest_payload = self.ctx[-1]
latest_message = self.ctx[-1]
if settings.LANGFUSE_PUBLIC_KEY:
lf.update_current_generation(
input=format_reasoning_inputs_as_markdown(
working_representation,
history,
latest_payload.content,
latest_payload.created_at,
)
)
# if settings.LANGFUSE_PUBLIC_KEY:
# lf.update_current_generation(
# input=format_reasoning_inputs_as_markdown(
# working_representation,
# history,
# latest_message.content,
# latest_message.created_at,
# )
# )
new_turns = [
format_new_turn_with_timestamp(p.content, p.created_at, p.sender_name)
for p in self.ctx
format_new_turn_with_timestamp(m.content, m.created_at, m.peer_name)
for m in self.ctx
]
formatted_working_representation = format_context_for_prompt(
@ -440,15 +443,15 @@ class CertaintyReasoner:
logger.debug(
"CRITICAL ANALYSIS: message_created_at='%s', new_turns_count=%s",
latest_payload.created_at,
latest_message.created_at,
len(new_turns),
)
try:
response_obj = await critical_analysis_call(
peer_id=latest_payload.sender_name,
peer_id=self.sender_name,
peer_card=speaker_peer_card,
message_created_at=latest_payload.created_at,
message_created_at=latest_message.created_at,
working_representation=formatted_working_representation,
history=history,
new_turns=new_turns,
@ -461,8 +464,12 @@ class CertaintyReasoner:
new_turns=new_turns,
) from e
# Handle None response from LLM
if response_obj is None:
logger.warning("LLM returned None response, using empty insights")
new_insights = ReasoningResponse(explicit=[], deductive=[])
# If response is a string, try to parse as JSON
if isinstance(response_obj, str):
elif isinstance(response_obj, str):
try:
response_data = json.loads(response_obj)
new_insights = ReasoningResponse(
@ -509,10 +516,10 @@ class CertaintyReasoner:
len(thinking) if thinking else 0,
)
if settings.LANGFUSE_PUBLIC_KEY:
lf.update_current_generation(
output=format_reasoning_response_as_markdown(response),
)
# if settings.LANGFUSE_PUBLIC_KEY:
# lf.update_current_generation(
# output=format_reasoning_response_as_markdown(response),
# )
return response
@ -528,7 +535,7 @@ class CertaintyReasoner:
Single-pass reasoning function that critically analyzes and derives insights.
Performs one analysis pass and returns the final observations.
"""
latest_payload = self.ctx[-1]
latest_message = self.ctx[-1]
analysis_start = time.perf_counter()
# Perform critical analysis to get observation lists
@ -543,7 +550,7 @@ class CertaintyReasoner:
analysis_duration_ms = (time.perf_counter() - analysis_start) * 1000
accumulate_metric(
f"deriver_representation_{latest_payload.message_id}_{latest_payload.target_name}",
f"deriver_representation_{latest_message.id}_{self.target_name}",
"critical_analysis_duration",
analysis_duration_ms,
"ms",
@ -554,13 +561,13 @@ class CertaintyReasoner:
new_observations_by_level: dict[
str, list[str]
] = await self._save_new_observations(
working_representation, reasoning_response, latest_payload
working_representation, reasoning_response, latest_message
)
save_observations_duration = (
time.perf_counter() - save_observations_start
) * 1000
accumulate_metric(
f"deriver_representation_{latest_payload.message_id}_{latest_payload.target_name}",
f"deriver_representation_{latest_message.id}_{self.target_name}",
"save_new_observations",
save_observations_duration,
"ms",
@ -579,7 +586,7 @@ class CertaintyReasoner:
time.perf_counter() - update_peer_card_start
) * 1000
accumulate_metric(
f"deriver_representation_{latest_payload.message_id}_{latest_payload.target_name}",
f"deriver_representation_{latest_message.id}_{self.target_name}",
"update_peer_card",
update_peer_card_duration,
"ms",
@ -594,7 +601,7 @@ class CertaintyReasoner:
original_working_representation: ReasoningResponse
| ReasoningResponseWithThinking,
revised_observations: ReasoningResponse | ReasoningResponseWithThinking,
latest_payload: RepresentationPayload,
latest_message: Message,
) -> dict[str, list[str]]:
"""Save only the observations that are new compared to the original context."""
# Use the utility function to find new observations
@ -642,9 +649,9 @@ class CertaintyReasoner:
if all_unified_observations:
await self.embedding_store.save_unified_observations(
all_unified_observations,
latest_payload.message_id,
latest_payload.session_name,
latest_payload.created_at,
latest_message.id,
latest_message.session_name,
latest_message.created_at,
)
else:
logger.debug("No new observations to save")
@ -679,8 +686,8 @@ class CertaintyReasoner:
await crud.set_peer_card(
db,
self.ctx[0].workspace_name,
self.ctx[0].sender_name,
self.ctx[0].target_name,
self.sender_name,
self.target_name,
new_peer_card,
)
except Exception as e:
@ -718,8 +725,10 @@ def observation_context_to_reasoning_response(
@sentry_sdk.trace
async def save_working_representation_to_peer(
payload: RepresentationPayload,
latest_message: Message,
final_observations: ReasoningResponseWithThinking,
sender_name: str,
target_name: str,
) -> None:
"""Save working representation to peer internal_metadata for dialectic access."""
@ -738,7 +747,7 @@ async def save_working_representation_to_peer(
working_rep_data = {
"final_observations": final_obs_dict,
"message_id": payload.message_id,
"message_id": latest_message.id,
"created_at": utc_now_iso(),
}
@ -746,10 +755,10 @@ async def save_working_representation_to_peer(
await crud.set_working_representation(
db,
working_rep_data,
payload.workspace_name,
payload.target_name,
payload.sender_name,
payload.session_name,
latest_message.workspace_name,
target_name,
sender_name,
latest_message.session_name,
)

16
src/deriver/prompts.py
View File

@ -37,9 +37,8 @@ def critical_analysis_prompt(
# Format the peer card as a string with newlines
peer_card_section = (
f"""
The user's known biographical information:
{peer_id}'s known biographical information:
<peer_card>
Peer ID: {peer_id}
{chr(10).join(peer_card)}
</peer_card>
"""
@ -49,7 +48,7 @@ Peer ID: {peer_id}
working_representation_section = (
f"""
The current user understanding:
Current understanding of {peer_id}:
<current_context>
{working_representation}
</current_context>
@ -64,12 +63,19 @@ The current user understanding:
f"""
You are an agent who critically analyzes user messages through rigorous logical reasoning to produce only conclusions about the user that are CERTAIN.
TARGET USER TO ANALYZE
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
You are analyzing: {peer_id}
The conversation may include messages from multiple participants, but you MUST focus ONLY on deriving conclusions about {peer_id}. Only use other participants' messages as context for understanding {peer_id}.
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
IMPORTANT NAMING RULES
• When you write a conclusion about the current user, always start the sentence with the user's name (e.g. "Anthony is 25 years old").
• When you write a conclusion about {peer_id}, always start the sentence with their name (e.g. "Anthony is 25 years old").
• NEVER start a conclusion with generic phrases like "The user …" unless the user name is not known.
• If you must reference a third person, use their explicit name, and add clarifiers such as "(third-party)" when confusion is possible.
Your goal is to IMPROVE understanding of the user through careful analysis. Your task is to arrive at truthful, factual conclusions via explicit and deductive reasoning.
Your goal is to IMPROVE understanding of {peer_id} through careful analysis. Your task is to arrive at truthful, factual conclusions via explicit and deductive reasoning.
Here are strict definitions for the reasoning modes you are to employ:

322
src/deriver/queue_manager.py
View File

@ -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 BigInteger, delete, select, update
from sqlalchemy import BigInteger, and_, delete, select, update
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql import func
@ -18,7 +18,10 @@ from sqlalchemy.sql import func
from src import models
from src.config import settings
from src.dependencies import tracked_db
from src.deriver.consumer import process_items
from src.deriver.consumer import (
process_item,
process_representation_batch,
)
from src.deriver.utils import parse_work_unit_key
from src.models import QueueItem
@ -294,6 +297,7 @@ class QueueManager:
)
async with self.semaphore:
message_count = 0
messages_to_process: list[QueueItem] = []
try:
parsed_key = parse_work_unit_key(work_unit_key)
task_type = parsed_key["task_type"]
@ -306,21 +310,55 @@ class QueueManager:
f"Worker {worker_id} lost ownership of work unit {work_unit_key}, stopping processing {work_unit_key}"
)
break
messages_to_process: list[QueueItem] = await self.get_message_batch(
task_type, work_unit_key, ownership.aqs_id
)
logger.debug(
f"Worker {worker_id} retrieved {len(messages_to_process)} messages for work unit {work_unit_key} (AQS ID: {ownership.aqs_id})"
)
if not messages_to_process:
logger.debug(
f"No more messages to process for work unit {work_unit_key} for worker {worker_id}"
)
break
try:
payloads = [msg.payload for msg in messages_to_process]
await process_items(task_type, payloads)
if task_type == "representation":
(
messages_context,
items_to_process,
) = await self.get_message_batch(
task_type, work_unit_key, ownership.aqs_id
)
logger.debug(
f"Worker {worker_id} retrieved {len(messages_context)} messages and {len(items_to_process)} queue items for work unit {work_unit_key} (AQS ID: {ownership.aqs_id})"
)
if not items_to_process:
logger.debug(
f"No more queue items to process for work unit {work_unit_key} for worker {worker_id}"
)
break
# Build payloads from the unique messages context window
sender_name = parsed_key["sender_name"]
target_name = parsed_key["target_name"]
await process_representation_batch(
messages_context,
sender_name=sender_name,
target_name=target_name,
)
await self.mark_messages_as_processed(
items_to_process, work_unit_key
)
message_count += len(items_to_process)
else:
messages_to_process = await self.get_next_message(
task_type, work_unit_key, ownership.aqs_id
)
if not messages_to_process:
logger.debug(
f"No more messages to process for work unit {work_unit_key} for worker {worker_id}"
)
break
await process_item(
task_type, messages_to_process[0].payload
)
await self.mark_messages_as_processed(
messages_to_process, work_unit_key
)
message_count += len(messages_to_process)
except Exception as e:
logger.error(
f"Error processing tasks for work unit {work_unit_key}: {e}",
@ -329,11 +367,6 @@ class QueueManager:
if settings.SENTRY.ENABLED:
sentry_sdk.capture_exception(e)
await self.mark_messages_as_processed(
messages_to_process, work_unit_key
)
message_count += len(messages_to_process)
# Check for shutdown after processing each batch
if self.shutdown_event.is_set():
logger.debug(
@ -387,113 +420,175 @@ class QueueManager:
)
@sentry_sdk.trace
async def get_message_batch(
async def get_next_message(
self, task_type: str, work_unit_key: str, aqs_id: str
) -> list[QueueItem]:
"""
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.
Args:
task_type: The type of task to process
work_unit_key: The key of the work unit to process
aqs_id: The ID of the active queue session to process
Returns:
A list of QueueItem objects
"""
async with tracked_db("get_message_batch") as db:
"""Get the next message to process for a specific work unit."""
if task_type == "representation":
raise ValueError(
"Representation tasks are not supported for get_next_message"
)
async with tracked_db("get_next_message") as db:
# ActiveQueueSession conditions for worker ownership verification
aqs_conditions = [
models.ActiveQueueSession.work_unit_key == work_unit_key,
models.ActiveQueueSession.id == aqs_id,
]
if task_type != "representation":
# For non-representation tasks, just get the next single message.
query = (
select(models.QueueItem)
.join(
models.ActiveQueueSession,
models.QueueItem.work_unit_key
== models.ActiveQueueSession.work_unit_key,
)
.where(models.QueueItem.work_unit_key == work_unit_key)
.where(~models.QueueItem.processed)
.where(*aqs_conditions)
.order_by(models.QueueItem.id)
.limit(1)
# For non-representation tasks, just get the next single message.
query = (
select(models.QueueItem)
.join(
models.ActiveQueueSession,
models.QueueItem.work_unit_key
== models.ActiveQueueSession.work_unit_key,
)
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,
BigInteger,
)
== 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
# Also ensure worker ownership verification by joining with ActiveQueueSession
query = (
select(models.QueueItem)
.join(
models.ActiveQueueSession,
models.QueueItem.work_unit_key
== models.ActiveQueueSession.work_unit_key,
)
.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
)
)
)
)
.where(*aqs_conditions)
.order_by(models.QueueItem.id)
)
result = await db.execute(query)
messages = result.scalars().all()
.where(models.QueueItem.work_unit_key == work_unit_key)
.where(~models.QueueItem.processed)
.where(*aqs_conditions)
.order_by(models.QueueItem.id)
.limit(1)
)
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()
return list(messages)
@sentry_sdk.trace
async def get_message_batch(
self,
task_type: str,
work_unit_key: str,
aqs_id: str,
) -> tuple[list[models.Message], list[QueueItem]]:
"""
Representation-only: returns a tuple of (messages_context, items_to_process).
- messages_context: unique Message rows (conversation turns) forming the context window
- items_to_process: QueueItems for the current work_unit_key within that window
"""
if task_type != "representation":
raise ValueError(
"Non-representation tasks are not supported for get_message_batch"
)
async with tracked_db("get_message_batch") as db:
# For representation tasks, get a batch based on token limit.
# Step 1: Parse work_unit_key to get session context and focused sender
parsed_key = parse_work_unit_key(work_unit_key)
session_name = parsed_key["session_name"]
workspace_name = parsed_key["workspace_name"]
# Verify worker still owns the work_unit_key
ownership_check = await db.execute(
select(models.ActiveQueueSession.id)
.where(models.ActiveQueueSession.work_unit_key == work_unit_key)
.where(models.ActiveQueueSession.id == aqs_id)
)
if not ownership_check.scalar_one_or_none():
# Worker lost ownership, return empty
await db.commit()
return [], []
# Step 2: Build a single SQL query that:
# 1. Finds the earliest unprocessed message for this work_unit_key
# 2. Gets ALL messages from that point forward (for conversational context)
# 3. Tracks cumulative tokens and focused sender position
# 4. Returns empty if focused sender is beyond token limit
# 5. Otherwise returns messages up to token limit + first focused sender message
# Find the minimum message_id with an unprocessed queue item across the session
min_unprocessed_message_id_subq = (
select(func.min(models.Message.id))
.select_from(models.QueueItem)
.join(
models.Message,
func.cast(models.QueueItem.payload["message_id"].astext, BigInteger)
== models.Message.id,
)
.where(~models.QueueItem.processed)
.where(models.Message.session_name == session_name)
.where(models.Message.workspace_name == workspace_name)
.where(models.QueueItem.work_unit_key == work_unit_key)
.scalar_subquery()
)
# Build CTE with ALL messages starting from the earliest unprocessed message
# This includes interleaving messages for conversational context
cte = (
select(
models.Message.id.label("message_id"),
models.Message.token_count.label("token_count"),
models.Message.peer_name.label("sender_name"),
func.sum(models.Message.token_count)
.over(order_by=models.Message.id)
.label("cumulative_token_count"),
)
.where(models.Message.session_name == session_name)
.where(models.Message.workspace_name == workspace_name)
.where(models.Message.id >= min_unprocessed_message_id_subq)
.order_by(models.Message.id)
.cte()
)
allowed_condition = (
(
cte.c.cumulative_token_count
<= settings.DERIVER.REPRESENTATION_BATCH_MAX_TOKENS
)
| (
cte.c.message_id == min_unprocessed_message_id_subq
) # always include the first unprocessed message
)
query = (
select(models.Message, models.QueueItem)
.select_from(cte)
.join(models.Message, models.Message.id == cte.c.message_id)
.outerjoin(
models.QueueItem,
and_(
models.QueueItem.work_unit_key == work_unit_key,
~models.QueueItem.processed,
func.cast(
models.QueueItem.payload["message_id"].astext, BigInteger
)
== models.Message.id,
),
)
.where(allowed_condition)
.order_by(models.Message.id, models.QueueItem.id)
)
result = await db.execute(query)
rows = result.all()
if not rows:
await db.commit()
return [], []
messages_context: list[models.Message] = []
items_to_process: list[QueueItem] = []
seen_messages: set[int] = set()
for m, qi in rows:
if m.id not in seen_messages:
messages_context.append(m)
seen_messages.add(m.id)
if qi is not None:
items_to_process.append(qi)
if items_to_process:
max_queue_item_message_id = max(
[qi.payload["message_id"] for qi in items_to_process]
)
messages_context = [ # remove any messages that are after the last message_id from queue items
m for m in messages_context if m.id <= max_queue_item_message_id
]
await db.commit()
return messages_context, items_to_process
async def mark_messages_as_processed(
self, messages: list[QueueItem], work_unit_key: str
) -> None:
@ -504,6 +599,7 @@ class QueueManager:
await db.execute(
update(models.QueueItem)
.where(models.QueueItem.id.in_(message_ids))
.where(models.QueueItem.work_unit_key == work_unit_key)
.values(processed=True)
)
await db.execute(

26
tests/deriver/test_deriver_processing.py
View File

@ -8,7 +8,6 @@ import pytest
from src import models
from src.deriver.deriver import process_representation_tasks_batch
from src.deriver.queue_payload import RepresentationPayload
from src.utils.shared_models import ReasoningResponseWithThinking
@ -168,27 +167,26 @@ class TestDeriverProcessing:
AsyncMock(),
)
# Create test payloads with different message IDs (earlier message has lower ID)
# Create test messages with different IDs (earlier message has lower ID)
now = datetime.now(timezone.utc)
payloads: list[RepresentationPayload] = []
messages: list[models.Message] = []
for i in range(8):
message_id = 100 + i # 100, 101, 102, ..., 107
payloads.append(
RepresentationPayload(
messages.append(
models.Message(
id=message_id,
workspace_name="test_workspace",
session_name="test_session",
message_id=message_id,
peer_name="alice",
content=f"message {message_id}",
sender_name="alice",
target_name="alice",
created_at=now
- timedelta(
minutes=7 - i
), # Earlier messages have earlier timestamps
token_count=0,
created_at=now - timedelta(minutes=7 - i),
)
)
await process_representation_tasks_batch(payloads)
await process_representation_tasks_batch(
sender_name="alice", target_name="alice", messages=messages
)
# Verify that the earliest message ID was used as the cutoff
assert captured_cutoffs == [payloads[0].message_id]
assert captured_cutoffs == [messages[0].id]

391
tests/deriver/test_queue_processing.py
View File

@ -1,5 +1,6 @@
from collections.abc import Callable
from typing import Any
from unittest.mock import patch
import pytest
from sqlalchemy import select
@ -8,6 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
from src import models
from src.config import settings
from src.deriver.queue_manager import QueueManager, WorkerOwnership
from src.deriver.utils import get_work_unit_key
@pytest.mark.asyncio
@ -169,23 +171,23 @@ class TestQueueProcessing:
await db_session.commit()
await db_session.refresh(aqs)
batch = await qm.get_message_batch(
_, items_to_process = await qm.get_message_batch(
task_type="representation",
work_unit_key=first.work_unit_key,
aqs_id=aqs.id,
)
nxt = batch[0] if batch else None
nxt = items_to_process[0] if items_to_process 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(
_, items_to_process2 = await qm.get_message_batch(
task_type="representation",
work_unit_key=first.work_unit_key,
aqs_id=aqs.id,
)
nxt2 = batch2[0] if batch2 else None
nxt2 = items_to_process2[0] if items_to_process2 else None
assert nxt2 is not None and nxt2.id == second.id
@pytest.mark.asyncio
@ -271,7 +273,6 @@ class TestQueueProcessing:
create_queue_payload: Callable[..., Any],
) -> None:
"""Test that representation tasks are batched based on token limits"""
from unittest.mock import patch
session, peers = sample_session_with_peers
peer = peers[0]
@ -310,9 +311,6 @@ class TestQueueProcessing:
for msg in messages
]
# Add items with token counts
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")
@ -335,13 +333,15 @@ class TestQueueProcessing:
# 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]]
async def mock_process_representation_batch(
messages: list[models.Message],
sender_name: str | None = None, # pyright: ignore[reportUnusedParameter]
target_name: str | None = None, # pyright: ignore[reportUnusedParameter]
) -> None:
processed_batches.append(
{
"task_type": task_type,
"payload_count": len(queue_payloads),
"task_type": "representation",
"payload_count": len(messages),
}
)
@ -358,8 +358,8 @@ class TestQueueProcessing:
)
with patch(
"src.deriver.queue_manager.process_items",
side_effect=mock_process_items,
"src.deriver.queue_manager.process_representation_batch",
side_effect=mock_process_representation_batch,
):
await qm.process_work_unit(work_unit_key, worker_id)
@ -369,6 +369,322 @@ 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_token_batching_filters_by_work_unit(
self,
db_session: AsyncSession,
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
create_queue_payload: Callable[..., Any],
) -> None:
"""Test that messages are batched chronologically across all peers, then filtered by work unit"""
session, peers = sample_session_with_peers
alice, bob, steve = peers[0], peers[1], peers[2]
# Create messages from different peers with specific token counts
# Message sequence: alice(250), bob(400), steve(300), alice(500), bob(500), alice(40), steve(100)
# Token limit: 2000 - should batch messages 1-6 (1990 tokens)
messages_data = [
(alice, 250),
(bob, 400),
(steve, 300),
(alice, 500),
(bob, 500),
(alice, 40),
(steve, 100),
]
messages: list[models.Message] = []
for peer, token_count in messages_data:
message = models.Message(
session_name=session.name,
workspace_name=session.workspace_name,
peer_name=peer.name,
content=f"Message from {peer.name}",
token_count=token_count,
)
db_session.add(message)
messages.append(message)
await db_session.commit()
for message in messages:
await db_session.refresh(message)
# Create queue items for alice's work unit (representation)
alice_queue_items: list[models.QueueItem] = []
bob_queue_items: list[models.QueueItem] = []
steve_queue_items: list[models.QueueItem] = []
for i, message in enumerate(messages):
peer = messages_data[i][0]
target = alice # All observing alice for simplicity
payload = create_queue_payload( # type: ignore[reportUnknownArgumentType]
message=message,
task_type="representation",
sender_name=peer.name,
target_name=target.name,
)
work_unit_key = get_work_unit_key("representation", payload)
queue_item = models.QueueItem(
session_id=session.id,
task_type="representation",
work_unit_key=work_unit_key,
payload=payload,
processed=False,
)
db_session.add(queue_item)
# Track items by peer
if peer == alice:
alice_queue_items.append(queue_item)
elif peer == bob:
bob_queue_items.append(queue_item)
else:
steve_queue_items.append(queue_item)
await db_session.commit()
for item in alice_queue_items + bob_queue_items + steve_queue_items:
await db_session.refresh(item)
qm = QueueManager()
# Mock the token limit to 2000 for this test
with patch.object(settings.DERIVER, "REPRESENTATION_BATCH_MAX_TOKENS", 2000):
# Test alice's work unit
alice_work_unit_key = alice_queue_items[0].work_unit_key
alice_aqs = models.ActiveQueueSession(work_unit_key=alice_work_unit_key)
db_session.add(alice_aqs)
await db_session.commit()
await db_session.refresh(alice_aqs)
alice_messages, alice_items = await qm.get_message_batch(
task_type="representation",
work_unit_key=alice_work_unit_key,
aqs_id=alice_aqs.id,
)
assert len(alice_messages) == 6
alice_message_ids: set[int] = {m.id for m in alice_messages}
expected_batch_ids = {
messages[0].id, # alice(250)
messages[1].id, # bob(400)
messages[2].id, # steve(300)
messages[3].id, # alice(500)
messages[4].id, # bob(500)
messages[5].id, # alice(40)
}
assert alice_message_ids == expected_batch_ids
# Ensure items are only for alice
assert all(
qi.payload.get("sender_name") == alice.name for qi in alice_items
)
# Test bob's work unit - starts at message 2 for per-work-unit anchoring
bob_work_unit_key = bob_queue_items[0].work_unit_key
bob_aqs = models.ActiveQueueSession(work_unit_key=bob_work_unit_key)
db_session.add(bob_aqs)
await db_session.commit()
await db_session.refresh(bob_aqs)
bob_messages, bob_items = await qm.get_message_batch(
task_type="representation",
work_unit_key=bob_work_unit_key,
aqs_id=bob_aqs.id,
)
# Bob should get 4 messages (2..5)
assert len(bob_messages) == 4
bob_message_ids: set[int] = {m.id for m in bob_messages}
expected_bob_ids = {
messages[1].id, # bob(400)
messages[2].id, # steve(300)
messages[3].id, # alice(500)
messages[4].id, # bob(500)
}
assert bob_message_ids == expected_bob_ids
# Ensure items are only for bob
assert all(qi.payload.get("sender_name") == bob.name for qi in bob_items)
# Test steve's work unit - starts at message 3 for per-work-unit anchoring
steve_work_unit_key = steve_queue_items[0].work_unit_key
steve_aqs = models.ActiveQueueSession(work_unit_key=steve_work_unit_key)
db_session.add(steve_aqs)
await db_session.commit()
await db_session.refresh(steve_aqs)
steve_messages, steve_items = await qm.get_message_batch(
task_type="representation",
work_unit_key=steve_work_unit_key,
aqs_id=steve_aqs.id,
)
# Steve should get 5 messages (3..7)
assert len(steve_messages) == 5
steve_message_ids: set[int] = {m.id for m in steve_messages}
expected_steve_ids = {
messages[2].id, # steve(300)
messages[3].id, # alice(500)
messages[4].id, # bob(500)
messages[5].id, # alice(40)
messages[6].id, # steve(100)
}
assert steve_message_ids == expected_steve_ids
# Ensure items are only for steve
assert all(
qi.payload.get("sender_name") == steve.name for qi in steve_items
)
@pytest.mark.asyncio
async def test_per_work_unit_anchoring_with_token_limits(
self,
db_session: AsyncSession,
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
create_queue_payload: Callable[..., Any],
) -> None:
"""Test that per-work-unit anchoring processes each sender's messages independently within token limits"""
session, peers = sample_session_with_peers
bob, steve, alice = peers[0], peers[1], peers[2]
# Create messages: bob(800), steve(800), alice(100), alice(200)
# Token limit: 1500
# With per-work-unit anchoring:
# - Bob's work unit: starts at message 1, includes only bob(800)
# - Steve's work unit: starts at message 2, includes only steve(800)
# - Alice's work unit: starts at message 3, includes alice(100) + alice(200)
messages_data = [
(bob, 800),
(steve, 800),
(alice, 100),
(alice, 200),
]
messages: list[models.Message] = []
for peer, token_count in messages_data:
message = models.Message(
session_name=session.name,
workspace_name=session.workspace_name,
peer_name=peer.name,
content=f"Message from {peer.name}",
token_count=token_count,
)
db_session.add(message)
messages.append(message)
await db_session.commit()
for message in messages:
await db_session.refresh(message)
# Create queue items
alice_queue_items: list[models.QueueItem] = []
bob_queue_items: list[models.QueueItem] = []
steve_queue_items: list[models.QueueItem] = []
for i, message in enumerate(messages):
peer = messages_data[i][0]
target = alice # All observing alice
payload = create_queue_payload( # type: ignore[reportUnknownArgumentType]
message=message,
task_type="representation",
sender_name=peer.name,
target_name=target.name,
)
work_unit_key = get_work_unit_key("representation", payload)
queue_item = models.QueueItem(
session_id=session.id,
task_type="representation",
work_unit_key=work_unit_key,
payload=payload,
processed=False,
)
db_session.add(queue_item)
if peer == alice:
alice_queue_items.append(queue_item)
elif peer == bob:
bob_queue_items.append(queue_item)
elif peer == steve:
steve_queue_items.append(queue_item)
await db_session.commit()
for item in alice_queue_items + bob_queue_items + steve_queue_items:
await db_session.refresh(item)
qm = QueueManager()
# Mock the token limit to 1500 for this test
with patch.object(settings.DERIVER, "REPRESENTATION_BATCH_MAX_TOKENS", 1500):
# Test alice's work unit
# With per-work-unit anchoring, Alice starts at her own first message (message 3)
# Alice's batch: alice(100) + alice(200) = 300 tokens, well under 1500 limit
if alice_queue_items:
alice_work_unit_key = alice_queue_items[0].work_unit_key
alice_aqs = models.ActiveQueueSession(work_unit_key=alice_work_unit_key)
db_session.add(alice_aqs)
await db_session.commit()
await db_session.refresh(alice_aqs)
alice_messages2, _ = await qm.get_message_batch(
task_type="representation",
work_unit_key=alice_work_unit_key,
aqs_id=alice_aqs.id,
)
# Per-work-unit anchoring: Alice starts at message 3 -> [3,4]
assert len(alice_messages2) == 2
assert [m.id for m in alice_messages2] == [
messages[2].id,
messages[3].id,
]
# Test bob's work unit
# With per-work-unit anchoring, Bob starts at his own first message (message 1)
# Bob's batch: bob(800) only, under 1500 limit
if bob_queue_items:
bob_work_unit_key = bob_queue_items[0].work_unit_key
bob_aqs = models.ActiveQueueSession(work_unit_key=bob_work_unit_key)
db_session.add(bob_aqs)
await db_session.commit()
await db_session.refresh(bob_aqs)
bob_messages2, _ = await qm.get_message_batch(
task_type="representation",
work_unit_key=bob_work_unit_key,
aqs_id=bob_aqs.id,
)
assert len(bob_messages2) == 1
assert bob_messages2[0].id == messages[0].id # bob only
# Test steve's work unit
# With per-work-unit anchoring, Steve starts at his own first message (message 2)
# Steve's batch: steve(800) only, under 1500 limit
if steve_queue_items:
steve_work_unit_key = steve_queue_items[0].work_unit_key
steve_aqs = models.ActiveQueueSession(work_unit_key=steve_work_unit_key)
db_session.add(steve_aqs)
await db_session.commit()
await db_session.refresh(steve_aqs)
steve_messages2, _ = await qm.get_message_batch(
task_type="representation",
work_unit_key=steve_work_unit_key,
aqs_id=steve_aqs.id,
)
# Per-work-unit anchoring: Steve starts at message 2 -> [2]
assert len(steve_messages2) == 1
assert [m.id for m in steve_messages2] == [
messages[1].id,
]
@pytest.mark.asyncio
async def test_single_message_processing(
self,
@ -377,7 +693,6 @@ class TestQueueProcessing:
create_queue_payload: Callable[..., Any],
) -> None:
"""Test that multiple summary messages in same work unit are processed separately"""
from unittest.mock import patch
session, peers = sample_session_with_peers
peer = peers[0]
@ -408,7 +723,6 @@ class TestQueueProcessing:
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)
@ -427,12 +741,11 @@ class TestQueueProcessing:
# Mock and process work unit
processed_batches: list[dict[str, Any]] = []
async def mock_process_items(
task_type: str, queue_payloads: list[dict[str, Any]]
async def mock_process_item(
task_type: str,
queue_payload: dict[str, Any], # pyright: ignore[reportUnusedParameter]
) -> None:
processed_batches.append(
{"task_type": task_type, "payload_count": len(queue_payloads)}
)
processed_batches.append({"task_type": task_type, "payload_count": 1})
qm = QueueManager()
work_unit_key = queue_items[0].work_unit_key
@ -446,8 +759,8 @@ class TestQueueProcessing:
)
with patch(
"src.deriver.queue_manager.process_items",
side_effect=mock_process_items,
"src.deriver.queue_manager.process_item",
side_effect=mock_process_item,
):
await qm.process_work_unit(work_unit_key, worker_id)
@ -491,7 +804,6 @@ class TestQueueProcessing:
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]
@ -531,7 +843,6 @@ class TestQueueProcessing:
]
# Add items to queue
from src.deriver.utils import get_work_unit_key
queue_items: list[models.QueueItem] = []
for payload in payloads:
@ -555,13 +866,15 @@ class TestQueueProcessing:
# 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]]
async def mock_process_representation_batch(
messages: list[models.Message],
sender_name: str | None = None, # pyright: ignore[reportUnusedParameter]
target_name: str | None = None, # pyright: ignore[reportUnusedParameter]
) -> None:
processed_batches.append(
{
"task_type": task_type,
"payload_count": len(queue_payloads),
"task_type": "representation",
"payload_count": len(messages),
}
)
@ -577,8 +890,8 @@ class TestQueueProcessing:
)
with patch(
"src.deriver.queue_manager.process_items",
side_effect=mock_process_items,
"src.deriver.queue_manager.process_representation_batch",
side_effect=mock_process_representation_batch,
):
await qm.process_work_unit(work_unit_key, worker_id)
@ -598,7 +911,6 @@ class TestQueueProcessing:
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]
@ -642,7 +954,6 @@ class TestQueueProcessing:
]
# Add items to queue
from src.deriver.utils import get_work_unit_key
queue_items: list[models.QueueItem] = []
for payload in payloads:
@ -666,13 +977,15 @@ class TestQueueProcessing:
# 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]]
async def mock_process_representation_batch(
messages: list[models.Message],
sender_name: str | None = None, # pyright: ignore[reportUnusedParameter]
target_name: str | None = None, # pyright: ignore[reportUnusedParameter]
) -> None:
processed_batches.append(
{
"task_type": task_type,
"payload_count": len(queue_payloads),
"task_type": "representation",
"payload_count": len(messages),
}
)
@ -688,8 +1001,8 @@ class TestQueueProcessing:
)
with patch(
"src.deriver.queue_manager.process_items",
side_effect=mock_process_items,
"src.deriver.queue_manager.process_representation_batch",
side_effect=mock_process_representation_batch,
):
await qm.process_work_unit(work_unit_key, worker_id)