From f6520c1ffee01446465b6cf8339c06fb824b9280 Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Fri, 26 Sep 2025 11:37:40 -0400 Subject: [PATCH] fix: sum input tokens and set session context cutoff to MAX_INPUT_TOKENS - sum (#218) * fix: sum input tokens and set session context cutoff to MAX_INPUT_TOKENS - sum * fix: CR comments + remove context var * fix: CR comments 2 * fix: use cache decorator; rm logger * refactor: re-fetch formatted history after getting global working rep --- .env.template | 2 +- config.toml.example | 2 +- docs/v2/contributing/configuration.mdx | 3 - src/config.py | 11 +- src/deriver/deriver.py | 160 ++++++++++++++++++++----- src/deriver/prompts.py | 25 ++++ src/deriver/utils.py | 15 +++ src/utils/summarizer.py | 6 + 8 files changed, 185 insertions(+), 39 deletions(-) diff --git a/.env.template b/.env.template index 9cc57de0..2ddb6802 100644 --- a/.env.template +++ b/.env.template @@ -89,9 +89,9 @@ LLM_ANTHROPIC_API_KEY=your-anthropic-api-key-here # DERIVER_PEER_CARD_PROVIDER=openai # DERIVER_PEER_CARD_MODEL=gpt-5-nano-2025-08-07 # DERIVER_PEER_CARD_MAX_OUTPUT_TOKENS=2000 -# DERIVER_CONTEXT_TOKEN_LIMIT=30000 # DERIVER_WORKING_REPRESENTATION_MAX_OBSERVATIONS=100 # DERIVER_REPRESENTATION_BATCH_MAX_TOKENS=4096 +# DERIVER_MAX_INPUT_TOKENS=23000 # ============================================================================= # Dialectic Settings diff --git a/config.toml.example b/config.toml.example index 3be89786..6209f255 100644 --- a/config.toml.example +++ b/config.toml.example @@ -69,9 +69,9 @@ THINKING_BUDGET_TOKENS = 1024 # only applied when using Anthropic PEER_CARD_PROVIDER = "openai" PEER_CARD_MODEL = "gpt-5-nano-2025-08-07" PEER_CARD_MAX_OUTPUT_TOKENS = 2000 -CONTEXT_TOKEN_LIMIT = 30000 WORKING_REPRESENTATION_MAX_OBSERVATIONS = 100 REPRESENTATION_BATCH_MAX_TOKENS = 4096 +MAX_INPUT_TOKENS = 23000 # Dialectic settings [dialectic] diff --git a/docs/v2/contributing/configuration.mdx b/docs/v2/contributing/configuration.mdx index cf903c97..da5f2d6a 100644 --- a/docs/v2/contributing/configuration.mdx +++ b/docs/v2/contributing/configuration.mdx @@ -260,9 +260,6 @@ DERIVER_PEER_CARD_PROVIDER=openai DERIVER_PEER_CARD_MODEL=gpt-5-nano-2025-08-07 DERIVER_PEER_CARD_MAX_OUTPUT_TOKENS=2000 -# Context token limit for get_context method -DERIVER_CONTEXT_TOKEN_LIMIT=30000 - # Maximum number of observations to store in working representation # This is applied to both explicit and deductive observations DERIVER_WORKING_REPRESENTATION_MAX_OBSERVATIONS=100 diff --git a/src/config.py b/src/config.py index fec3b882..a8bed82e 100644 --- a/src/config.py +++ b/src/config.py @@ -201,11 +201,6 @@ class DeriverSettings(HonchoSettings): int, Field(default=4000, gt=1000, le=10_000) ] = 4000 - # Context token limit for get_context method - CONTEXT_TOKEN_LIMIT: Annotated[int, Field(default=30_000, gt=1000, le=100_000)] = ( - 30_000 - ) - # Maximum number of observations to store in working representation # This is applied to both explicit and deductive observations WORKING_REPRESENTATION_MAX_OBSERVATIONS: Annotated[ @@ -220,11 +215,13 @@ class DeriverSettings(HonchoSettings): ), ] = 4096 + MAX_INPUT_TOKENS: Annotated[int, Field(default=23000, gt=0, le=23000)] = 23000 + @model_validator(mode="after") def validate_batch_tokens_vs_context_limit(self): - if self.REPRESENTATION_BATCH_MAX_TOKENS > self.CONTEXT_TOKEN_LIMIT: + if self.REPRESENTATION_BATCH_MAX_TOKENS > self.MAX_INPUT_TOKENS: raise ValueError( - f"REPRESENTATION_BATCH_MAX_TOKENS ({self.REPRESENTATION_BATCH_MAX_TOKENS}) cannot exceed CONTEXT_TOKEN_LIMIT ({self.CONTEXT_TOKEN_LIMIT})" + f"REPRESENTATION_BATCH_MAX_TOKENS ({self.REPRESENTATION_BATCH_MAX_TOKENS}) cannot exceed max deriver input tokens ({self.MAX_INPUT_TOKENS})" ) return self diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py index 9b53d3e6..e6bf7024 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -11,6 +11,7 @@ from src import crud, exceptions 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.utils import summarizer from src.utils.clients import honcho_llm_call from src.utils.embedding_store import EmbeddingStore @@ -40,7 +41,11 @@ from src.utils.shared_models import ( UnifiedObservation, ) -from .prompts import critical_analysis_prompt, peer_card_prompt +from .prompts import ( + critical_analysis_prompt, + estimate_base_prompt_tokens, + peer_card_prompt, +) from .queue_payload import ( RepresentationPayload, ) @@ -137,13 +142,71 @@ async def process_representation_tasks_batch( earliest_payload.message_id, ) - # Use get_session_context_formatted with configurable token limit + 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, + ) + if speaker_peer_card is None: + logger.warning("No peer card found for %s", latest_payload.sender_name) + else: + logger.debug("Using peer card for %s", latest_payload.sender_name) + + # Get working representation data early for token estimation + async with tracked_db("deriver.get_working_representation_data") as db: + working_rep_data: ( + 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, + ) + + # Estimate tokens for deriver input + peer_card_tokens = estimate_tokens(speaker_peer_card) + working_rep_tokens = _estimate_working_representation_tokens(working_rep_data) + base_prompt_tokens = estimate_base_prompt_tokens() + + # 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 + ] + new_turns_tokens = estimate_tokens(new_turns) + + estimated_input_tokens = ( + peer_card_tokens + working_rep_tokens + base_prompt_tokens + new_turns_tokens + ) + + # Calculate available tokens for context + safety_buffer = 500 + available_context_tokens = max( + 0, + 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, + ) + + # Use get_session_context_formatted with dynamic token limit 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, - token_limit=settings.DERIVER.CONTEXT_TOKEN_LIMIT, + token_limit=available_context_tokens, cutoff=earliest_payload.message_id, include_summary=True, ) @@ -180,18 +243,6 @@ async def process_representation_tasks_batch( # Create reasoner instance reasoner = CertaintyReasoner(embedding_store=embedding_store, ctx=payloads) - # Check for existing working representation first, fall back to global search - async with tracked_db("deriver.get_working_representation_data") as db: - working_rep_data: ( - 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, - ) - # Time context preparation context_prep_start = time.perf_counter() if ( @@ -236,6 +287,48 @@ async def process_representation_tasks_batch( working_representation ) logger.info("No working representation found, using global semantic search") + + # Recalculate tokens now that we have a new working representation + new_working_rep_tokens = _estimate_working_representation_tokens( + { + "final_observations": { + "thinking": working_representation.thinking, + "explicit": working_representation.explicit, + "deductive": [ + { + "conclusion": obs.conclusion, + "premises": obs.premises, + } + for obs in working_representation.deductive + ], + } + } + ) + + # Update estimated input tokens with the new working representation + estimated_input_tokens = ( + peer_card_tokens + + new_working_rep_tokens + + base_prompt_tokens + + new_turns_tokens + ) + + # Recalculate available tokens for context + available_context_tokens = max( + 0, + settings.DERIVER.MAX_INPUT_TOKENS - estimated_input_tokens - safety_buffer, + ) + + # Recalculate formatted_history with updated token limit + 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, + token_limit=available_context_tokens, + cutoff=earliest_payload.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}", @@ -249,18 +342,6 @@ async def process_representation_tasks_batch( "REASONING: Running unified insight derivation across explicit and deductive reasoning levels" ) - 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, - ) - if speaker_peer_card is None: - logger.warning("No peer card found for %s", latest_payload.sender_name) - else: - logger.info("Using peer card: %s", speaker_peer_card) - # Run single-pass reasoning final_observations = await reasoner.reason( working_representation, @@ -670,3 +751,28 @@ async def save_working_representation_to_peer( payload.sender_name, payload.session_name, ) + + +def _estimate_working_representation_tokens( + working_rep_data: dict[str, Any] | str | None, +) -> int: + """Estimate tokens for working representation data.""" + if isinstance(working_rep_data, str): + return estimate_tokens(working_rep_data) + + if ( + not isinstance(working_rep_data, dict) + or "final_observations" not in working_rep_data + ): + return 0 + + final_obs = working_rep_data["final_observations"] + + explicit_tokens = estimate_tokens(final_obs.get("explicit")) + thinking_tokens = estimate_tokens(final_obs.get("thinking")) + deductive_tokens = sum( + estimate_tokens(d.get("conclusion")) + estimate_tokens(d.get("premises")) + for d in final_obs.get("deductive", []) + ) + + return explicit_tokens + thinking_tokens + deductive_tokens diff --git a/src/deriver/prompts.py b/src/deriver/prompts.py index 30009fbd..c57aa888 100644 --- a/src/deriver/prompts.py +++ b/src/deriver/prompts.py @@ -6,8 +6,11 @@ and reasoning tasks. """ import datetime +from functools import cache from inspect import cleandoc as c +from src.deriver.utils import estimate_tokens + def critical_analysis_prompt( peer_id: str, @@ -157,3 +160,25 @@ New observations: If there's no new key info, set "card" to null (or omit it) to signal no update. **NEVER** include notes or temporary information in the card itself, instead use the notes field. There are no mandatory fields -- if you can't find a value, just leave it out. **ONLY** include information that is **GIVEN**. """ # nosec B608 <-- this is a really dumb false positive ) + + +@cache +def estimate_base_prompt_tokens() -> int: + """Estimate base 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( + peer_id="", + peer_card=None, + message_created_at=datetime.datetime.now(datetime.timezone.utc), + working_representation=None, + history="", + new_turns=[], + ) + return estimate_tokens(base_prompt) + except Exception: + # Return a conservative estimate if estimation fails + return 500 diff --git a/src/deriver/utils.py b/src/deriver/utils.py index d29aa8cd..244dbd33 100644 --- a/src/deriver/utils.py +++ b/src/deriver/utils.py @@ -1,5 +1,8 @@ +import tiktoken from typing_extensions import Any, TypedDict +tokenizer = tiktoken.get_encoding("cl100k_base") + class ParsedWorkUnit(TypedDict): task_type: str @@ -65,3 +68,15 @@ def parse_work_unit_key(work_unit_key: str) -> ParsedWorkUnit: } raise ValueError(f"Invalid task type in work_unit_key: {task_type}") + + +def estimate_tokens(text: str | list[str] | None) -> int: + """Estimate token count using tiktoken for text or list of strings.""" + if not text: + return 0 + if isinstance(text, list): + text = "\n".join(text) + try: + return len(tokenizer.encode(text)) + except Exception: + return len(text) // 4 diff --git a/src/utils/summarizer.py b/src/utils/summarizer.py index 9e10fa2a..a40eea1d 100644 --- a/src/utils/summarizer.py +++ b/src/utils/summarizer.py @@ -580,6 +580,9 @@ async def get_session_context( Tuple of (summary, messages) where summary is a Summary pydantic model (or None) and messages is the list of message objects """ + if token_limit <= 0: + return None, [] + summary = None messages_tokens = token_limit messages_start_id = 0 @@ -658,6 +661,9 @@ async def get_session_context_formatted( This is a convenience wrapper around get_session_context that formats the output as a string. """ + if token_limit <= 0: + return "" + summary, messages = await get_session_context( db, workspace_name,