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

metrics for deriver / dialectic input + output tokens (#274) 

* feat: separate deriver input / output tokens

* fix: track dialectic input / output tokens

* feat: track dialectic output tokens in streaming API

* fix: add component to metric ad change critical_analysis -> representation

* feat: track summary metrics in prometheus

* test: fix client streaming mocks

* fix: update tokenizer

* fix: add helper method; instrument peer card

* fix: count previous summary if not fallback
This commit is contained in:
Rajat Ahuja 2025-11-19 21:20:34 -05:00 committed by GitHub
parent 2ffa7b7b30
commit 0f1e1dec20
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 383 additions and 95 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

74
src/deriver/deriver.py
View File

@ -20,12 +20,13 @@ from src.utils.logging import (
)
from src.utils.peer_card import PeerCardQuery
from src.utils.representation import PromptRepresentation, Representation
from src.utils.tokens import estimate_tokens
from src.utils.tokens import estimate_tokens, track_input_tokens
from src.utils.tracing import with_sentry_transaction
from .prompts import (
critical_analysis_prompt,
estimate_base_prompt_tokens,
estimate_critical_analysis_prompt_tokens,
estimate_peer_card_prompt_tokens,
peer_card_prompt,
)
@ -41,7 +42,6 @@ async def critical_analysis_call(
working_representation: Representation,
history: str,
new_turns: list[str],
estimated_input_tokens: int,
) -> PromptRepresentation:
prompt = critical_analysis_prompt(
peer_id=peer_id,
@ -70,7 +70,9 @@ async def critical_analysis_call(
prometheus.DERIVER_TOKENS_PROCESSED.labels(
task_type="representation",
).inc(response.output_tokens + estimated_input_tokens)
token_type="output", # nosec B106
component="total",
).inc(response.output_tokens)
return response.content
@ -101,6 +103,23 @@ async def peer_card_call(
retry_attempts=3,
)
# Track input tokens for peer_card task
track_input_tokens(
task_type="peer_card",
components={
"prompt": estimate_peer_card_prompt_tokens(),
"old_peer_card": estimate_tokens(old_peer_card),
"new_observations": estimate_tokens(new_observations.str_no_timestamps()),
},
)
# Track output tokens for peer_card task
prometheus.DERIVER_TOKENS_PROCESSED.labels(
task_type="peer_card",
token_type="output", # nosec B106
component="total",
).inc(response.output_tokens)
return response.content
@ -165,10 +184,12 @@ async def process_representation_tasks_batch(
# Estimate tokens for deriver input
peer_card_tokens = estimate_tokens(speaker_peer_card)
working_rep_tokens = estimate_tokens(
str(working_representation) if not working_representation.is_empty() else None
)
base_prompt_tokens = estimate_base_prompt_tokens()
prompt_tokens = estimate_critical_analysis_prompt_tokens()
# Estimate tokens for new conversation turns
new_turns = [
@ -178,7 +199,7 @@ async def process_representation_tasks_batch(
new_turns_tokens = estimate_tokens(new_turns)
estimated_input_tokens = (
peer_card_tokens + working_rep_tokens + base_prompt_tokens + new_turns_tokens
peer_card_tokens + working_rep_tokens + prompt_tokens + new_turns_tokens
)
# Calculate available tokens for context
@ -188,17 +209,6 @@ async def process_representation_tasks_batch(
settings.DERIVER.MAX_INPUT_TOKENS - estimated_input_tokens - safety_buffer,
)
logger.debug(
"Token estimation - Peer card: %d, Working rep: %d, Base prompt: %d, "
+ "New turns: %d, Total estimated: %d, Available for context: %d",
peer_card_tokens,
working_rep_tokens,
base_prompt_tokens,
new_turns_tokens,
estimated_input_tokens,
available_context_tokens,
)
async with tracked_db("deriver.get_session_context_formatted") as db:
formatted_history = await summarizer.get_session_context_formatted(
db,
@ -211,6 +221,32 @@ async def process_representation_tasks_batch(
session_context_tokens = estimate_tokens(formatted_history)
# Update total estimated input tokens with session context
estimated_input_tokens += session_context_tokens
logger.debug(
"Token estimation - Peer card: %d, Working rep: %d, Base prompt: %d, "
+ "New turns: %d, Session context: %d, Total estimated: %d",
peer_card_tokens,
working_rep_tokens,
prompt_tokens,
new_turns_tokens,
session_context_tokens,
estimated_input_tokens,
)
# Track all input token components
track_input_tokens(
task_type="representation",
components={
"peer_card": peer_card_tokens,
"working_representation": working_rep_tokens,
"prompt": prompt_tokens,
"new_turns": new_turns_tokens,
"session_context": session_context_tokens,
},
)
# got working representation and peer card, log timing
context_prep_duration = (time.perf_counter() - context_prep_start) * 1000
accumulate_metric(
@ -243,7 +279,6 @@ async def process_representation_tasks_batch(
ctx=messages,
observed=observed,
observer=observer,
estimated_input_tokens=estimated_input_tokens + session_context_tokens,
)
# Run single-pass reasoning
@ -294,13 +329,11 @@ class CertaintyReasoner:
*,
observed: str,
observer: str,
estimated_input_tokens: int,
) -> None:
self.representation_manager = representation_manager
self.ctx = ctx
self.observed = observed
self.observer = observer
self.estimated_input_tokens: int = estimated_input_tokens
@conditional_observe(name="Deriver")
@sentry_sdk.trace
@ -341,7 +374,6 @@ class CertaintyReasoner:
working_representation=working_representation,
history=history,
new_turns=new_turns,
estimated_input_tokens=self.estimated_input_tokens,
)
except Exception as e:
raise exceptions.LLMError(

26
src/deriver/prompts.py
View File

@ -213,14 +213,14 @@ If there's no new key info, set "card" to null (or omit it) to signal no update.
@cache
def estimate_base_prompt_tokens() -> int:
"""Estimate base prompt tokens by calling critical_analysis_prompt with empty values.
def estimate_critical_analysis_prompt_tokens() -> int:
"""Estimate critical analysis prompt tokens by calling critical_analysis_prompt with empty values.
This value is cached since it only changes on redeploys when the prompt template changes.
"""
try:
base_prompt = critical_analysis_prompt(
prompt = critical_analysis_prompt(
peer_id="",
peer_card=None,
message_created_at=datetime.datetime.now(datetime.timezone.utc),
@ -228,7 +228,25 @@ def estimate_base_prompt_tokens() -> int:
history="",
new_turns=[],
)
return estimate_tokens(base_prompt)
return estimate_tokens(prompt)
except Exception:
# Return a conservative estimate if estimation fails
return 500
@cache
def estimate_peer_card_prompt_tokens() -> int:
"""Estimate peer card prompt tokens by calling peer_card_prompt with empty values.
This value is cached since it only changes on redeploys when the prompt template changes.
"""
try:
prompt = peer_card_prompt(
old_peer_card=None,
new_observations="",
)
return estimate_tokens(prompt)
except Exception:
# Return a conservative estimate if estimation fails
return 400

55
src/dialectic/chat.py
View File

@ -13,7 +13,7 @@ from collections.abc import AsyncIterator
from dotenv import load_dotenv
from src import crud
from src import crud, prometheus
from src.config import settings
from src.dependencies import tracked_db
from src.utils import summarizer
@ -26,7 +26,7 @@ from src.utils.logging import (
from src.utils.representation import Representation
from src.utils.tokens import estimate_tokens
from .prompts import dialectic_prompt
from .prompts import dialectic_prompt, estimate_dialectic_prompt_tokens
# Configure logging
logger = logging.getLogger(__name__)
@ -60,6 +60,18 @@ async def dialectic_call(
Returns:
Model response
"""
# Estimate input tokens by concatenating all inputs
prompt_tokens = estimate_dialectic_prompt_tokens()
inputs = [
query,
working_representation,
recent_conversation_history or "",
"\n".join(peer_card) if peer_card else "",
"\n".join(observed_peer_card) if observed_peer_card else "",
]
contextual_tokens = estimate_tokens("".join(inputs))
estimated_input_tokens = prompt_tokens + contextual_tokens
# Generate the prompt and log it
prompt = dialectic_prompt(
query,
@ -87,6 +99,15 @@ async def dialectic_call(
logger.debug(prompt)
logger.debug("=== END DIALECTIC PROMPT ===")
# Track tokens in prometheus
prometheus.DIALECTIC_TOKENS_PROCESSED.labels(
token_type="input", # nosec B106
).inc(estimated_input_tokens)
prometheus.DIALECTIC_TOKENS_PROCESSED.labels(
token_type="output", # nosec B106
).inc(response.output_tokens)
return response.content
@ -115,6 +136,18 @@ async def dialectic_stream(
Returns:
Streaming model response
"""
# Estimate input tokens by concatenating all inputs
prompt_tokens = estimate_dialectic_prompt_tokens()
variable_inputs = [
query,
working_representation,
recent_conversation_history or "",
"\n".join(peer_card) if peer_card else "",
"\n".join(observed_peer_card) if observed_peer_card else "",
]
variable_tokens = estimate_tokens("".join(variable_inputs))
estimated_input_tokens = prompt_tokens + variable_tokens
# Generate the prompt and log it
prompt = dialectic_prompt(
query,
@ -143,7 +176,23 @@ async def dialectic_stream(
logger.debug(prompt)
logger.debug("=== END DIALECTIC PROMPT ===")
return response
# Track input tokens in prometheus
# Note: Output tokens are available in the final chunk of the stream (is_done=True)
prometheus.DIALECTIC_TOKENS_PROCESSED.labels(
token_type="input", # nosec B106
).inc(estimated_input_tokens)
# Wrap the response to log output tokens from final chunk
async def log_streaming_response():
async for chunk in response:
if chunk.is_done and chunk.output_tokens is not None:
# TODO: Currently not tracking output tokens for groq models
prometheus.DIALECTIC_TOKENS_PROCESSED.labels(
token_type="output", # nosec B106
).inc(chunk.output_tokens)
yield chunk
return log_streaming_response()
@conditional_observe(name="Dialectic")

26
src/dialectic/prompts.py
View File

@ -1,5 +1,8 @@
from functools import cache
from inspect import cleandoc as c
from src.utils.tokens import estimate_tokens
def dialectic_prompt(
query: str,
@ -141,6 +144,29 @@ Provide a natural language response that:
)
@cache
def estimate_dialectic_prompt_tokens() -> int:
"""Estimate base dialectic prompt tokens by calling dialectic_prompt with empty values.
This value is cached since it only changes on redeploys when the prompt template changes.
"""
try:
prompt = dialectic_prompt(
query="",
working_representation="",
recent_conversation_history=None,
observer_peer_card=None,
observed_peer_card=None,
observer="",
observed="",
)
return estimate_tokens(prompt)
except Exception:
# Return a conservative estimate if estimation fails
return 750
def query_generation_prompt(query: str, observed: str) -> str:
"""
Generate the prompt for semantic query expansion.

26
src/prometheus.py
View File

@ -86,7 +86,7 @@ DIALECTIC_CALLS = NamespacedCounter(
# Incremented in: src/deriver/queue_manager.py when queue items are processed
# Labels:
# - workspace_name: The workspace where items were processed
# - task_type: The type of task processed (e.g., "representation", "summary")
# - task_type: The type of task processed (e.g., "representation", "summary", "peer_card")
DERIVER_QUEUE_ITEMS_PROCESSED = NamespacedCounter(
"deriver_queue_items_processed_total",
"Total deriver queue items processed",
@ -97,13 +97,31 @@ DERIVER_QUEUE_ITEMS_PROCESSED = NamespacedCounter(
#
# Incremented in: src/deriver/deriver.py after the critical analysis call is made
# Labels:
# - task_type: The type of task that processed the tokens (e.g., "representation")
# - task_type: The type of task that processed the tokens (e.g., "representation", "summary")
# - token_type: The type of tokens ("input" or "output")
# - component: The component of the input (e.g., "peer_card", "working_representation", "prompt", "new_turns", "session_context")
DERIVER_TOKENS_PROCESSED = NamespacedCounter(
"tokens_processed_total",
"Total tokens processed",
"deriver_tokens_processed_total",
"Total tokens processed by the deriver",
[
"namespace",
"task_type",
"token_type",
"component",
],
)
# Tracks the total number of input and output tokens processed by the dialectic.
#
# Incremented in: src/dialectic/chat.py after the dialectic call is made
# Labels:
# - token_type: The type of tokens ("input" or "output")
DIALECTIC_TOKENS_PROCESSED = NamespacedCounter(
"dialectic_tokens_processed_total",
"Total tokens processed by the dialectic",
[
"namespace",
"token_type",
],
)

2
src/schemas.py
View File

@ -136,7 +136,7 @@ class MessageCreate(MessageBase):
@model_validator(mode="after")
def validate_and_set_token_count(self) -> Self:
encoding = tiktoken.get_encoding("cl100k_base")
encoding = tiktoken.get_encoding("o200k_base")
encoded_message = encoding.encode(self.content)
self._encoded_message = encoded_message

48
src/utils/clients.py
View File

@ -101,11 +101,13 @@ class HonchoLLMCallStreamChunk(BaseModel):
content: The text content for this chunk. Empty for chunks that only contain metadata.
is_done: Whether this is the final chunk in the stream.
finish_reasons: List of finish reasons if the stream is complete.
output_tokens: Number of tokens generated in the response. Only set on the final chunk.
"""
content: str
is_done: bool = False
finish_reasons: list[str] = Field(default_factory=list)
output_tokens: int | None = None
@overload
@ -756,12 +758,15 @@ async def handle_streaming_response(
text_content = getattr(chunk.delta, "text", "")
yield HonchoLLMCallStreamChunk(content=text_content)
final_message = await anthropic_stream.get_final_message()
usage = final_message.usage
output_tokens = usage.output_tokens if usage else None
yield HonchoLLMCallStreamChunk(
content="",
is_done=True,
finish_reasons=[final_message.stop_reason]
if final_message.stop_reason
else [],
output_tokens=output_tokens,
)
case AsyncOpenAI():
@ -769,6 +774,7 @@ async def handle_streaming_response(
"model": params["model"],
"messages": params["messages"],
"stream": True,
"stream_options": {"include_usage": True},
}
model_name = params["model"]
@ -787,18 +793,35 @@ async def handle_streaming_response(
openai_params["response_format"] = {"type": "json_object"}
openai_stream = await client.chat.completions.create(**openai_params) # pyright: ignore
finish_reason: str | None = None
usage_chunk_received = False
async for chunk in openai_stream: # pyright: ignore
chunk = cast(ChatCompletionChunk, chunk)
if chunk.choices and chunk.choices[0].delta.content:
yield HonchoLLMCallStreamChunk(
content=chunk.choices[0].delta.content
)
content = chunk.choices[0].delta.content
yield HonchoLLMCallStreamChunk(content=content)
# Track finish_reason when it appears (before usage chunk)
if chunk.choices and chunk.choices[0].finish_reason:
finish_reason = chunk.choices[0].finish_reason
# Check for usage info in chunk (with include_usage, this is a separate chunk with empty choices)
if hasattr(chunk, "usage") and chunk.usage:
yield HonchoLLMCallStreamChunk(
content="",
is_done=True,
finish_reasons=[chunk.choices[0].finish_reason],
finish_reasons=[finish_reason] if finish_reason else [],
output_tokens=chunk.usage.completion_tokens,
)
usage_chunk_received = True
# If stream ended without usage chunk (interrupted), still yield final chunk
if not usage_chunk_received and finish_reason:
logger.warning("OpenAI stream ended without usage chunk (interrupted)")
yield HonchoLLMCallStreamChunk(
content="",
is_done=True,
finish_reasons=[finish_reason],
output_tokens=None,
)
case genai.Client():
prompt_text = params["messages"][0]["content"] if params["messages"] else ""
@ -828,6 +851,7 @@ async def handle_streaming_response(
final_chunk = chunk
finish_reason = "stop" # Default fallback
gemini_output_tokens: int | None = None
if (
final_chunk
and hasattr(final_chunk, "candidates")
@ -837,8 +861,22 @@ async def handle_streaming_response(
):
finish_reason = final_chunk.candidates[0].finish_reason.name
# Extract output tokens from usage_metadata if available
if (
final_chunk
and hasattr(final_chunk, "usage_metadata")
and final_chunk.usage_metadata
and hasattr(final_chunk.usage_metadata, "candidates_token_count")
):
gemini_output_tokens = (
final_chunk.usage_metadata.candidates_token_count or None
)
yield HonchoLLMCallStreamChunk(
content="", is_done=True, finish_reasons=[finish_reason]
content="",
is_done=True,
finish_reasons=[finish_reason],
output_tokens=gemini_output_tokens,
)
case AsyncGroq():

180
src/utils/summarizer.py
View File

@ -2,19 +2,21 @@ import asyncio
import logging
import time
from enum import Enum
from functools import cache
from inspect import cleandoc as c
from typing import TypedDict
from sqlalchemy import update
from sqlalchemy.ext.asyncio import AsyncSession
from src import schemas
from src import prometheus, schemas
from src.config import settings
from src.dependencies import tracked_db
from src.exceptions import ResourceNotFoundException
from src.utils.clients import HonchoLLMCallResponse, honcho_llm_call
from src.utils.formatting import utc_now_iso
from src.utils.logging import accumulate_metric, conditional_observe
from src.utils.tokens import estimate_tokens, track_input_tokens
from .. import crud, models
@ -79,25 +81,13 @@ class SummaryType(Enum):
LONG = "honcho_chat_summary_long"
@conditional_observe(name="Create Short Summary")
async def create_short_summary(
def short_summary_prompt(
messages: list[models.Message],
input_tokens: int,
previous_summary: str | None = None,
) -> HonchoLLMCallResponse[str]:
# input_tokens indicates how many tokens the message list + previous summary take up
# we want to optimize short summaries to be smaller than the actual content being summarized
# so we ask the agent to produce a word count roughly equal to either the input, or the max
# size if the input is larger. the word/token ratio is roughly 4:3 so we multiply by 0.75.
# LLMs *seem* to respond better to getting asked for a word count but should workshop this.
output_words = int(min(input_tokens, settings.SUMMARY.MAX_TOKENS_SHORT) * 0.75)
if previous_summary:
previous_summary_text = previous_summary
else:
previous_summary_text = "There is no previous summary -- the messages are the beginning of the conversation."
prompt = c(f"""
output_words: int,
previous_summary_text: str,
) -> str:
"""Generate the short summary prompt."""
return c(f"""
You are a system that summarizes parts of a conversation to create a concise and accurate summary. Focus on capturing:
1. Key facts and information shared (**Capture as many explicit facts as possible**)
@ -122,28 +112,14 @@ Return only the summary without any explanation or meta-commentary.
Produce as thorough a summary as possible in {output_words} words or less.
""")
return await honcho_llm_call(
llm_settings=settings.SUMMARY,
prompt=prompt,
max_tokens=settings.SUMMARY.MAX_TOKENS_SHORT,
)
@conditional_observe(name="Create Long Summary")
async def create_long_summary(
def long_summary_prompt(
messages: list[models.Message],
previous_summary: str | None = None,
) -> HonchoLLMCallResponse[str]:
# the word/token ratio is roughly 4:3 so we multiply by 0.75.
# LLMs *seem* to respond better to getting asked for a word count but should workshop this.
output_words = int(settings.SUMMARY.MAX_TOKENS_LONG * 0.75)
if previous_summary:
previous_summary_text = previous_summary
else:
previous_summary_text = "There is no previous summary -- the messages are the beginning of the conversation."
prompt = c(f"""
output_words: int,
previous_summary_text: str,
) -> str:
"""Generate the long summary prompt."""
return c(f"""
You are a system that creates thorough, comprehensive summaries of conversations. Focus on capturing:
1. Key facts and information shared (**Capture as many explicit facts as possible**)
@ -170,6 +146,82 @@ Return only the summary without any explanation or meta-commentary.
Produce as thorough a summary as possible in {output_words} words or less.
""")
@cache
def estimate_short_summary_prompt_tokens() -> int:
"""Estimate tokens for the short summary prompt (without messages/previous_summary)."""
try:
return estimate_tokens(
short_summary_prompt(
messages=[],
output_words=0,
previous_summary_text="",
)
)
except Exception:
# Return a rough estimate if estimation fails
return 200
@cache
def estimate_long_summary_prompt_tokens() -> int:
"""Estimate tokens for the long summary prompt (without messages/previous_summary)."""
try:
return estimate_tokens(
long_summary_prompt(
messages=[],
output_words=0,
previous_summary_text="",
)
)
except Exception:
# Return a rough estimate if estimation fails
return 200
@conditional_observe(name="Create Short Summary")
async def create_short_summary(
messages: list[models.Message],
input_tokens: int,
previous_summary: str | None = None,
) -> HonchoLLMCallResponse[str]:
# input_tokens indicates how many tokens the message list + previous summary take up
# we want to optimize short summaries to be smaller than the actual content being summarized
# so we ask the agent to produce a word count roughly equal to either the input, or the max
# size if the input is larger. the word/token ratio is roughly 4:3 so we multiply by 0.75.
# LLMs *seem* to respond better to getting asked for a word count but should workshop this.
output_words = int(min(input_tokens, settings.SUMMARY.MAX_TOKENS_SHORT) * 0.75)
if previous_summary:
previous_summary_text = previous_summary
else:
previous_summary_text = "There is no previous summary -- the messages are the beginning of the conversation."
prompt = short_summary_prompt(messages, output_words, previous_summary_text)
return await honcho_llm_call(
llm_settings=settings.SUMMARY,
prompt=prompt,
max_tokens=settings.SUMMARY.MAX_TOKENS_SHORT,
)
@conditional_observe(name="Create Long Summary")
async def create_long_summary(
messages: list[models.Message],
previous_summary: str | None = None,
) -> HonchoLLMCallResponse[str]:
# the word/token ratio is roughly 4:3 so we multiply by 0.75.
# LLMs *seem* to respond better to getting asked for a word count but should workshop this.
output_words = int(settings.SUMMARY.MAX_TOKENS_LONG * 0.75)
if previous_summary:
previous_summary_text = previous_summary
else:
previous_summary_text = "There is no previous summary -- the messages are the beginning of the conversation."
prompt = long_summary_prompt(messages, output_words, previous_summary_text)
return await honcho_llm_call(
llm_settings=settings.SUMMARY,
prompt=prompt,
@ -311,7 +363,7 @@ async def _create_and_save_summary(
previous_summary_tokens = latest_summary["token_count"] if latest_summary else 0
input_tokens = messages_tokens + previous_summary_tokens
new_summary = await _create_summary(
new_summary, is_fallback = await _create_summary(
messages=messages,
previous_summary_text=previous_summary_text,
summary_type=summary_type,
@ -319,6 +371,30 @@ async def _create_and_save_summary(
message_public_id=message_public_id,
)
# Only track tokens if this was a real LLM call
if not is_fallback:
# Get base prompt tokens based on summary type
if summary_type == SummaryType.SHORT:
prompt_tokens = estimate_short_summary_prompt_tokens()
else:
prompt_tokens = estimate_long_summary_prompt_tokens()
track_input_tokens(
task_type="summary",
components={
"prompt": prompt_tokens,
"messages": messages_tokens,
"previous_summary": previous_summary_tokens,
},
)
# Track output tokens
prometheus.DERIVER_TOKENS_PROCESSED.labels(
task_type="summary",
token_type="output", # nosec B106
component="total",
).inc(new_summary["token_count"])
await _save_summary(
db,
new_summary,
@ -354,7 +430,7 @@ async def _create_summary(
summary_type: SummaryType,
input_tokens: int,
message_public_id: str,
) -> Summary:
) -> tuple[Summary, bool]:
"""
Generate a summary of the provided messages using an LLM.
@ -364,10 +440,12 @@ async def _create_summary(
summary_type: Type of summary to create ("short" or "long")
Returns:
A full summary of the conversation up to the last message
A tuple of (Summary, is_fallback) where is_fallback indicates if
the summary was generated using a fallback instead of an LLM call
"""
response: HonchoLLMCallResponse[str] | None = None
is_fallback = False
try:
if summary_type == SummaryType.SHORT:
response = await create_short_summary(
@ -393,14 +471,18 @@ async def _create_summary(
else ""
)
summary_tokens = 50
is_fallback = True
return Summary(
content=summary_text,
message_id=messages[-1].id if messages else 0,
summary_type=summary_type.value,
created_at=utc_now_iso(),
token_count=summary_tokens,
message_public_id=message_public_id,
return (
Summary(
content=summary_text,
message_id=messages[-1].id if messages else 0,
summary_type=summary_type.value,
created_at=utc_now_iso(),
token_count=summary_tokens,
message_public_id=message_public_id,
),
is_fallback,
)

20
src/utils/tokens.py
View File

@ -1,6 +1,8 @@
import tiktoken
tokenizer = tiktoken.get_encoding("cl100k_base")
from src import prometheus
tokenizer = tiktoken.get_encoding("o200k_base")
def estimate_tokens(text: str | list[str] | None) -> int:
@ -13,3 +15,19 @@ def estimate_tokens(text: str | list[str] | None) -> int:
return len(tokenizer.encode(text))
except Exception:
return len(text) // 4
def track_input_tokens(task_type: str, components: dict[str, int]) -> None:
"""
Helper method to track input token components for a given task type.
Args:
task_type: The type of task (e.g., "representation", "peer_card", "summary")
components: Dict mapping component names to token counts
"""
for component, token_count in components.items():
prometheus.DERIVER_TOKENS_PROCESSED.labels(
task_type=task_type,
token_type="input", # nosec B106
component=component,
).inc(token_count)

4
tests/test_llm_mock.py
View File

@ -8,6 +8,7 @@ from src.models import Message
from src.utils.representation import (
DeductiveObservation,
ExplicitObservation,
PromptRepresentation,
Representation,
)
@ -19,7 +20,7 @@ async def test_generic_honcho_llm_call_mock():
from src.deriver.deriver import critical_analysis_call
# Call the decorated function - this should use our mock
result = await critical_analysis_call(
result: PromptRepresentation = await critical_analysis_call(
peer_id="test_peer_id",
peer_card=["test_peer_card"],
message_created_at=datetime(2023, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
@ -44,7 +45,6 @@ async def test_generic_honcho_llm_call_mock():
),
history="test history",
new_turns=["test new turn"],
estimated_input_tokens=100,
)
# Verify that we get a mock result, not an actual LLM call

17
tests/utils/test_clients.py
View File

@ -235,8 +235,9 @@ class TestAnthropicClient:
# Set up the async iterator (same as working test_streaming_call)
mock_stream.__aiter__.return_value = iter(mock_chunks)
# Mock final message
mock_final_message = Mock(stop_reason="stop")
# Mock final message with usage tokens
mock_usage = Mock(output_tokens=42)
mock_final_message = Mock(stop_reason="stop", usage=mock_usage)
mock_stream.get_final_message.return_value = mock_final_message
mock_client.messages.stream.return_value = mock_stream
@ -632,10 +633,15 @@ class TestGoogleClient:
# Mock streaming chunks
mock_finish_reason = Mock()
mock_finish_reason.name = "STOP"
mock_usage_metadata = Mock(candidates_token_count=35)
mock_chunks = [
Mock(text="Hello"),
Mock(text=" world"),
Mock(text="", candidates=[Mock(finish_reason=mock_finish_reason)]),
Mock(
text="",
candidates=[Mock(finish_reason=mock_finish_reason)],
usage_metadata=mock_usage_metadata,
),
]
# Create async iterator for the chunks
@ -956,8 +962,9 @@ class TestMainLLMCallFunction:
mock_stream.__aenter__.return_value = mock_stream
mock_stream.__aiter__.return_value = iter(mock_chunks)
# Mock final message
mock_final_message = Mock(stop_reason="stop")
# Mock final message with usage tokens
mock_usage = Mock(output_tokens=28)
mock_final_message = Mock(stop_reason="stop", usage=mock_usage)
mock_stream.get_final_message.return_value = mock_final_message
mock_client.messages.stream.return_value = mock_stream