diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py index 701bbc50..3b6ab521 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -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( diff --git a/src/deriver/prompts.py b/src/deriver/prompts.py index 4a414f88..0cdf9ccb 100644 --- a/src/deriver/prompts.py +++ b/src/deriver/prompts.py @@ -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 diff --git a/src/dialectic/chat.py b/src/dialectic/chat.py index db76201c..14e10a40 100644 --- a/src/dialectic/chat.py +++ b/src/dialectic/chat.py @@ -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") diff --git a/src/dialectic/prompts.py b/src/dialectic/prompts.py index 7d1fbb55..132e06b3 100644 --- a/src/dialectic/prompts.py +++ b/src/dialectic/prompts.py @@ -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. diff --git a/src/prometheus.py b/src/prometheus.py index f19d30f3..0ad90b45 100644 --- a/src/prometheus.py +++ b/src/prometheus.py @@ -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", ], ) diff --git a/src/schemas.py b/src/schemas.py index 8b6068da..32a88bc5 100644 --- a/src/schemas.py +++ b/src/schemas.py @@ -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 diff --git a/src/utils/clients.py b/src/utils/clients.py index 83e78bb0..deb371b4 100644 --- a/src/utils/clients.py +++ b/src/utils/clients.py @@ -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(): diff --git a/src/utils/summarizer.py b/src/utils/summarizer.py index 569a5e8f..084672ed 100644 --- a/src/utils/summarizer.py +++ b/src/utils/summarizer.py @@ -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, ) diff --git a/src/utils/tokens.py b/src/utils/tokens.py index 2af439c9..9bdb9efb 100644 --- a/src/utils/tokens.py +++ b/src/utils/tokens.py @@ -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) diff --git a/tests/test_llm_mock.py b/tests/test_llm_mock.py index 92292655..28841737 100644 --- a/tests/test_llm_mock.py +++ b/tests/test_llm_mock.py @@ -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 diff --git a/tests/utils/test_clients.py b/tests/utils/test_clients.py index f4cb9073..1cc0b155 100644 --- a/tests/utils/test_clients.py +++ b/tests/utils/test_clients.py @@ -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