feat: use messages from queue items for rep completed token count (#350)
This commit is contained in:
parent
bd68c19a07
commit
110787cdca
|
|
@ -158,7 +158,7 @@ async def process_representation_batch(
|
|||
*,
|
||||
observers: list[str] | None,
|
||||
observed: str | None,
|
||||
queue_items_count: int,
|
||||
queue_item_message_ids: list[int],
|
||||
) -> None:
|
||||
"""
|
||||
Prepares and processes a batch of messages for representation tasks.
|
||||
|
|
@ -168,6 +168,7 @@ async def process_representation_batch(
|
|||
message_level_configuration: Resolved configuration for this batch
|
||||
observers: List of observers for the messages
|
||||
observed: The observed of the messages
|
||||
queue_item_message_ids: Message IDs from queue items
|
||||
"""
|
||||
if not messages or not messages[0]:
|
||||
logger.debug("process_representation_batch received no messages")
|
||||
|
|
@ -181,7 +182,7 @@ async def process_representation_batch(
|
|||
message_level_configuration,
|
||||
observers=observers,
|
||||
observed=observed,
|
||||
queue_items_count=queue_items_count,
|
||||
queue_item_message_ids=queue_item_message_ids,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ from src.utils.clients import honcho_llm_call
|
|||
from src.utils.config_helpers import get_configuration
|
||||
from src.utils.formatting import format_new_turn_with_timestamp
|
||||
from src.utils.representation import PromptRepresentation, Representation
|
||||
from src.utils.tokens import estimate_tokens, track_deriver_input_tokens
|
||||
from src.utils.tokens import track_deriver_input_tokens
|
||||
|
||||
from .prompts import estimate_minimal_deriver_prompt_tokens, minimal_deriver_prompt
|
||||
|
||||
|
|
@ -34,7 +34,7 @@ async def process_representation_tasks_batch(
|
|||
*,
|
||||
observers: list[str],
|
||||
observed: str,
|
||||
queue_items_count: int,
|
||||
queue_item_message_ids: list[int],
|
||||
) -> None:
|
||||
"""
|
||||
Process messages with minimal overhead - single LLM call, save to multiple collections.
|
||||
|
|
@ -44,7 +44,7 @@ async def process_representation_tasks_batch(
|
|||
message_level_configuration: Optional configuration override.
|
||||
observers: List of observer peer IDs (collections to save to).
|
||||
observed: The observed peer ID.
|
||||
queue_items_count: Number of QueueItem records being processed in this batch.
|
||||
queue_item_message_ids: Message IDs from queue items being processed
|
||||
"""
|
||||
if not messages:
|
||||
return
|
||||
|
|
@ -93,9 +93,12 @@ async def process_representation_tasks_batch(
|
|||
for msg in messages
|
||||
)
|
||||
|
||||
# Track token usage
|
||||
# Track token usage - count only tokens from messages being processed
|
||||
prompt_tokens = estimate_minimal_deriver_prompt_tokens()
|
||||
messages_tokens = estimate_tokens(formatted_messages)
|
||||
queue_item_message_ids_set = set(queue_item_message_ids)
|
||||
messages_tokens = sum(
|
||||
msg.token_count for msg in messages if msg.id in queue_item_message_ids_set
|
||||
)
|
||||
track_deriver_input_tokens(
|
||||
task_type=DeriverTaskTypes.INGESTION,
|
||||
components={
|
||||
|
|
@ -235,7 +238,7 @@ async def process_representation_tasks_batch(
|
|||
workspace_name=latest_message.workspace_name,
|
||||
session_name=latest_message.session_name,
|
||||
observed=observed,
|
||||
queue_items_processed=queue_items_count,
|
||||
queue_items_processed=len(queue_item_message_ids),
|
||||
earliest_message_id=earliest_message.public_id,
|
||||
latest_message_id=latest_message.public_id,
|
||||
message_count=len(messages),
|
||||
|
|
|
|||
|
|
@ -450,12 +450,17 @@ class QueueManager:
|
|||
else:
|
||||
observers = []
|
||||
|
||||
queue_item_message_ids = [
|
||||
item.message_id
|
||||
for item in items_to_process
|
||||
if item.message_id is not None
|
||||
]
|
||||
await process_representation_batch(
|
||||
messages_context,
|
||||
message_level_configuration,
|
||||
observers=observers,
|
||||
observed=work_unit.observed,
|
||||
queue_items_count=len(items_to_process),
|
||||
queue_item_message_ids=queue_item_message_ids,
|
||||
)
|
||||
await self.mark_queue_items_as_processed(
|
||||
items_to_process, work_unit_key
|
||||
|
|
|
|||
|
|
@ -347,7 +347,7 @@ class TestQueueProcessing:
|
|||
*,
|
||||
observed: str | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
observers: list[str] | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
queue_items_count: int | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
queue_item_message_ids: list[int] | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
) -> None:
|
||||
processed_batches.append(
|
||||
{
|
||||
|
|
@ -910,7 +910,7 @@ class TestQueueProcessing:
|
|||
*,
|
||||
observed: str | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
observers: list[str] | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
queue_items_count: int | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
queue_item_message_ids: list[int] | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
) -> None:
|
||||
processed_batches.append(
|
||||
{
|
||||
|
|
@ -1029,7 +1029,7 @@ class TestQueueProcessing:
|
|||
*,
|
||||
observed: str | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
observers: list[str] | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
queue_items_count: int | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
queue_item_message_ids: list[int] | None = None, # pyright: ignore[reportUnusedParameter]
|
||||
) -> None:
|
||||
processed_batches.append(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -261,7 +261,7 @@ class TestDeriverIngestionMetrics:
|
|||
message_level_configuration=create_test_configuration(),
|
||||
observers=[peer.name],
|
||||
observed=peer.name,
|
||||
queue_items_count=len(messages),
|
||||
queue_item_message_ids=[m.id for m in messages],
|
||||
)
|
||||
|
||||
# Verify output tokens metric
|
||||
|
|
@ -318,7 +318,7 @@ class TestDeriverIngestionMetrics:
|
|||
message_level_configuration=create_test_configuration(),
|
||||
observers=[peer.name],
|
||||
observed=peer.name,
|
||||
queue_items_count=len(messages),
|
||||
queue_item_message_ids=[m.id for m in messages],
|
||||
)
|
||||
|
||||
metric_checker.assert_delta(
|
||||
|
|
@ -375,7 +375,7 @@ class TestDeriverIngestionMetrics:
|
|||
message_level_configuration=create_test_configuration(),
|
||||
observers=[peer.name],
|
||||
observed=peer.name,
|
||||
queue_items_count=len(messages),
|
||||
queue_item_message_ids=[m.id for m in messages],
|
||||
)
|
||||
|
||||
# Verify messages tokens were tracked (should be > 0)
|
||||
|
|
|
|||
Loading…
Reference in New Issue