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
This commit is contained in:
Rajat Ahuja 2025-09-26 11:37:40 -04:00 committed by GitHub
parent 6e01af2005
commit f6520c1ffe
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 185 additions and 39 deletions

View File

@ -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

View File

@ -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]

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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,