diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py
index c564722f..0692db47 100644
--- a/src/deriver/consumer.py
+++ b/src/deriver/consumer.py
@@ -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)
diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py
index e6bf7024..d3373b9c 100644
--- a/src/deriver/deriver.py
+++ b/src/deriver/deriver.py
@@ -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,
)
diff --git a/src/deriver/prompts.py b/src/deriver/prompts.py
index c57aa888..1371fb65 100644
--- a/src/deriver/prompts.py
+++ b/src/deriver/prompts.py
@@ -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 ID: {peer_id}
{chr(10).join(peer_card)}
"""
@@ -49,7 +48,7 @@ Peer ID: {peer_id}
working_representation_section = (
f"""
-The current user understanding:
+Current understanding of {peer_id}:
{working_representation}
@@ -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:
diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py
index bfa1b7cf..1ba9fcfe 100644
--- a/src/deriver/queue_manager.py
+++ b/src/deriver/queue_manager.py
@@ -10,7 +10,7 @@ import sentry_sdk
from dotenv import load_dotenv
from nanoid import generate as generate_nanoid
from sentry_sdk.integrations.asyncio import AsyncioIntegration
-from sqlalchemy import 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(
diff --git a/tests/deriver/test_deriver_processing.py b/tests/deriver/test_deriver_processing.py
index 020d9429..14636185 100644
--- a/tests/deriver/test_deriver_processing.py
+++ b/tests/deriver/test_deriver_processing.py
@@ -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]
diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py
index c65f7333..9e4ae473 100644
--- a/tests/deriver/test_queue_processing.py
+++ b/tests/deriver/test_queue_processing.py
@@ -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)