Merge pull request 'Budget conversation context and preserve complete tool exchanges' (#69) from feat/context-management into main
This commit was merged in pull request #69.
This commit is contained in:
@@ -27,6 +27,11 @@ LLAMA_MODEL_NAME=llama-3.1-8b-instruct
|
||||
# Abandon a turn when an in-flight model call produces no output (prefill included)
|
||||
# for this long. Thinking tokens count as output, so long reasoning is fine.
|
||||
TTFT_TIMEOUT_SECONDS=30
|
||||
# Context estimates include system instructions and tool definitions.
|
||||
# Set the window to the smallest backend context; leave room for generation.
|
||||
CONTEXT_WINDOW_TOKENS=32768
|
||||
PROMPT_BUDGET_TOKENS=16000
|
||||
RESPONSE_RESERVE_TOKENS=4096
|
||||
|
||||
# Embedding Server Configuration
|
||||
# If set, uses a custom OpenAI-compatible embedding server (e.g. llama-server)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import os
|
||||
from functools import partial
|
||||
from typing import cast
|
||||
|
||||
import httpx
|
||||
@@ -8,6 +9,7 @@ from langchain.tools import tool
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
from langchain_core.messages import SystemMessage
|
||||
from langchain_core.runnables import RunnableConfig, RunnableLambda
|
||||
from langchain_core.utils.function_calling import convert_to_openai_tool
|
||||
from langchain_openai import ChatOpenAI
|
||||
from tavily import AsyncTavilyClient
|
||||
|
||||
@@ -24,7 +26,7 @@ from blueprints.conversation.artifacts_util import (
|
||||
truncate,
|
||||
unique,
|
||||
)
|
||||
from blueprints.conversation.context import trim_history
|
||||
from blueprints.conversation.context import RESPONSE_RESERVE_TOKENS, build_context
|
||||
from blueprints.conversation.memory import save_memory
|
||||
from blueprints.conversation.prompts import SIMBA_SYSTEM_PROMPT
|
||||
from blueprints.rag.logic import query_vector_store
|
||||
@@ -57,11 +59,12 @@ if llama_url:
|
||||
model_kwargs={"reasoning_effort": "low"},
|
||||
# Ask llama-server for token usage on the final chunk (prompt size logging).
|
||||
stream_usage=True,
|
||||
max_tokens=RESPONSE_RESERVE_TOKENS,
|
||||
)
|
||||
else:
|
||||
llama_chat = None
|
||||
|
||||
openai_fallback = ChatOpenAI(model="gpt-5-mini")
|
||||
openai_fallback = ChatOpenAI(model="gpt-5-mini", max_tokens=RESPONSE_RESERVE_TOKENS)
|
||||
model_with_fallback = cast(
|
||||
BaseChatModel,
|
||||
llama_chat.with_fallbacks([openai_fallback]) if llama_chat else openai_fallback,
|
||||
@@ -468,14 +471,22 @@ async def obsidian_search_notes(query: str):
|
||||
|
||||
|
||||
@tool(response_format="content_and_artifact")
|
||||
async def obsidian_read_note(relative_path: str):
|
||||
"""Read the full content of a specific Obsidian note by path.
|
||||
async def obsidian_read_note(
|
||||
relative_path: str, start_line: int = 1, max_lines: int | None = None
|
||||
):
|
||||
"""Read an Obsidian note, optionally selecting a small range of content lines.
|
||||
|
||||
Use line ranges to retrieve omitted sections of a large note.
|
||||
|
||||
Args:
|
||||
relative_path: Path relative to vault root (e.g., "notes/my-note.md")
|
||||
start_line: First content line to read, starting at 1 (excludes frontmatter).
|
||||
max_lines: Number of lines to read, or None for all remaining lines.
|
||||
"""
|
||||
if not obsidian_enabled:
|
||||
return OBSIDIAN_NOT_CONFIGURED, None
|
||||
if start_line < 1 or (max_lines is not None and max_lines < 1):
|
||||
return "start_line and max_lines must be positive integers.", None
|
||||
|
||||
try:
|
||||
note = obsidian_service.read_note(relative_path)
|
||||
@@ -483,7 +494,25 @@ async def obsidian_read_note(relative_path: str):
|
||||
|
||||
result = f"File: {note['path']}\n\n"
|
||||
result += f"Frontmatter:\n{content_data['metadata']}\n\n"
|
||||
result += f"Content:\n{content_data['content']}\n\n"
|
||||
content = content_data["content"]
|
||||
if start_line != 1 or max_lines is not None:
|
||||
lines = content.splitlines()
|
||||
if start_line > len(lines):
|
||||
return (
|
||||
f"File: {note['path']} has {len(lines)} content lines; start_line is beyond the end.",
|
||||
None,
|
||||
)
|
||||
end_line = (
|
||||
min(len(lines), start_line - 1 + max_lines)
|
||||
if max_lines is not None
|
||||
else len(lines)
|
||||
)
|
||||
content = "\n".join(lines[start_line - 1 : end_line])
|
||||
result += (
|
||||
f"Content lines {start_line}-{end_line} of {len(lines)}:\n{content}\n\n"
|
||||
)
|
||||
else:
|
||||
result += f"Content:\n{content}\n\n"
|
||||
result += f"Tags: {', '.join(content_data['tags'])}\n"
|
||||
result += f"Contains {len(content_data['wikilinks'])} wikilinks and {len(content_data['embeds'])} embeds"
|
||||
|
||||
@@ -1100,8 +1129,7 @@ async def beli_search_restaurants(
|
||||
tools.append(beli_search_restaurants)
|
||||
|
||||
|
||||
@RunnableLambda
|
||||
def _build_prompt(state, config: RunnableConfig = None) -> list:
|
||||
def _build_prompt(state, config: RunnableConfig = None, *, tool_schemas=()) -> list:
|
||||
"""Inject the per-request system prompt, replacing any prior system messages."""
|
||||
system_prompt = (
|
||||
(config or {}).get("configurable", {}).get("system_prompt", SIMBA_SYSTEM_PROMPT)
|
||||
@@ -1109,8 +1137,7 @@ def _build_prompt(state, config: RunnableConfig = None) -> list:
|
||||
# LangGraph passes the full state dict {"messages": [...]}, not just the list
|
||||
messages = state["messages"] if isinstance(state, dict) else state
|
||||
non_system = [m for m in messages if not isinstance(m, SystemMessage)]
|
||||
# Window + trim earlier turns to stay within the 32k context window.
|
||||
return [SystemMessage(content=system_prompt)] + trim_history(non_system)
|
||||
return build_context(non_system, system_prompt, list(tool_schemas))
|
||||
|
||||
|
||||
def build_agent(checkpointer=None, extra_tools=()):
|
||||
@@ -1119,9 +1146,11 @@ def build_agent(checkpointer=None, extra_tools=()):
|
||||
``extra_tools`` are appended to the shared tool list, e.g. admin-only
|
||||
Homebox and BookStack tools for the admin agent.
|
||||
"""
|
||||
agent_tools = tools + list(extra_tools)
|
||||
tool_schemas = [convert_to_openai_tool(t) for t in agent_tools]
|
||||
return create_react_agent(
|
||||
model=model_with_fallback,
|
||||
tools=tools + list(extra_tools),
|
||||
prompt=_build_prompt,
|
||||
tools=agent_tools,
|
||||
prompt=RunnableLambda(partial(_build_prompt, tool_schemas=tool_schemas)),
|
||||
checkpointer=checkpointer,
|
||||
)
|
||||
|
||||
@@ -1,67 +1,286 @@
|
||||
"""Shape the message history replayed to the LLM so it fits the context window.
|
||||
"""Budget the model's view of history without changing checkpointed messages.
|
||||
|
||||
The LangGraph checkpoint keeps every message at full size: whole search hits,
|
||||
full notes, attached files, images. Replaying all of that each turn pushes the
|
||||
prompt toward the 32k-token llama-server slot, which makes every prefill slow.
|
||||
The current turn keeps everything; earlier turns are trimmed.
|
||||
Token estimates use UTF-8 bytes / 3 plus message overhead. They are deliberately
|
||||
more conservative than the usual chars / 4 heuristic, but are not a model
|
||||
specific tokenizer. A safety reserve covers estimation and chat-template costs.
|
||||
"""
|
||||
|
||||
from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
|
||||
# Matches TOOL_RESULT_MAX_CHARS in logic.py (what gets persisted to Postgres).
|
||||
from dotenv import load_dotenv
|
||||
from langchain_core.messages import (
|
||||
AIMessage,
|
||||
BaseMessage,
|
||||
HumanMessage,
|
||||
SystemMessage,
|
||||
ToolMessage,
|
||||
)
|
||||
|
||||
load_dotenv()
|
||||
|
||||
CONTEXT_WINDOW_TOKENS = int(os.getenv("CONTEXT_WINDOW_TOKENS", "32768"))
|
||||
PROMPT_BUDGET_TOKENS = int(os.getenv("PROMPT_BUDGET_TOKENS", "16000"))
|
||||
RESPONSE_RESERVE_TOKENS = int(os.getenv("RESPONSE_RESERVE_TOKENS", "4096"))
|
||||
SAFETY_RESERVE_TOKENS = 2048
|
||||
EARLIER_MESSAGE_MAX_CHARS = 4000
|
||||
TOOL_RESULT_MAX_CHARS = 12000
|
||||
TRUNCATION_NOTE = "\n[truncated: earlier context]"
|
||||
IMAGE_PLACEHOLDER = "[image from an earlier message]"
|
||||
# 12 ≈ 4-6 conversational turns (tool call/result pairs count as 2 messages).
|
||||
HISTORY_WINDOW = 12
|
||||
|
||||
|
||||
def _cap(text: str) -> str:
|
||||
if len(text) <= EARLIER_MESSAGE_MAX_CHARS:
|
||||
def estimate_tokens(value) -> int:
|
||||
text = (
|
||||
value
|
||||
if isinstance(value, str)
|
||||
else json.dumps(value, ensure_ascii=False, default=str)
|
||||
)
|
||||
return math.ceil(len(text.encode("utf-8")) / 3)
|
||||
|
||||
|
||||
def message_tokens(message: BaseMessage) -> int:
|
||||
# Count tool arguments but never artifact payloads (not sent to the model).
|
||||
content = message.content
|
||||
images = 0
|
||||
if isinstance(content, list):
|
||||
images = sum(
|
||||
isinstance(p, dict) and p.get("type") in ("image", "image_url")
|
||||
for p in content
|
||||
)
|
||||
content = [
|
||||
p
|
||||
for p in content
|
||||
if not (isinstance(p, dict) and p.get("type") in ("image", "image_url"))
|
||||
]
|
||||
return (
|
||||
16
|
||||
+ estimate_tokens(content)
|
||||
+ images * 4096
|
||||
+ estimate_tokens(getattr(message, "tool_calls", []))
|
||||
)
|
||||
|
||||
|
||||
def _cap(text: str, limit: int, note: str) -> str:
|
||||
if len(text) <= limit:
|
||||
return text
|
||||
return text[:EARLIER_MESSAGE_MAX_CHARS] + TRUNCATION_NOTE
|
||||
# Retain both ends: identifiers and totals often appear at the end of results.
|
||||
head = limit * 3 // 4
|
||||
return text[:head] + note + (text[-(limit - head) :] if limit > head else "")
|
||||
|
||||
|
||||
def _shrink(message: BaseMessage) -> BaseMessage:
|
||||
"""Cap text and drop images from a message that belongs to an earlier turn."""
|
||||
def _shrink(message: BaseMessage, limit: int, *, earlier: bool) -> BaseMessage:
|
||||
note = (
|
||||
TRUNCATION_NOTE
|
||||
if earlier
|
||||
else "\n[tool result excerpt; omitted content can be fetched again with the source tool]\n"
|
||||
)
|
||||
content = message.content
|
||||
if isinstance(content, str):
|
||||
new_content = _cap(content)
|
||||
new_content = _cap(content, limit, note)
|
||||
else:
|
||||
new_content = []
|
||||
remaining = limit
|
||||
for part in content:
|
||||
if isinstance(part, dict) and part.get("type") in ("image_url", "image"):
|
||||
if (
|
||||
isinstance(part, dict)
|
||||
and part.get("type") in ("image_url", "image")
|
||||
and earlier
|
||||
):
|
||||
new_content.append({"type": "text", "text": IMAGE_PLACEHOLDER})
|
||||
elif isinstance(part, dict) and isinstance(part.get("text"), str):
|
||||
new_content.append({**part, "text": _cap(part["text"])})
|
||||
text = _cap(part["text"], remaining, note)
|
||||
new_content.append({**part, "text": text})
|
||||
remaining = max(0, remaining - min(len(part["text"]), remaining))
|
||||
elif isinstance(part, str):
|
||||
new_content.append(_cap(part))
|
||||
new_content.append(_cap(part, remaining, note))
|
||||
remaining = max(0, remaining - min(len(part), remaining))
|
||||
else:
|
||||
new_content.append(part)
|
||||
if new_content == content:
|
||||
return message
|
||||
return message.model_copy(update={"content": new_content})
|
||||
|
||||
|
||||
def trim_history(messages: list[BaseMessage]) -> list[BaseMessage]:
|
||||
"""Window the history and trim tool results/user messages from earlier turns.
|
||||
|
||||
The current turn starts at the last HumanMessage; it and everything after it
|
||||
(this turn's tool calls and results) are left untouched.
|
||||
"""
|
||||
current_turn = max(
|
||||
(i for i, m in enumerate(messages) if isinstance(m, HumanMessage)),
|
||||
default=len(messages),
|
||||
return (
|
||||
message
|
||||
if new_content == content
|
||||
else message.model_copy(update={"content": new_content})
|
||||
)
|
||||
trimmed = [
|
||||
_shrink(m)
|
||||
if i < current_turn and isinstance(m, (ToolMessage, HumanMessage))
|
||||
|
||||
|
||||
def _turns(messages: list[BaseMessage]) -> list[list[BaseMessage]]:
|
||||
turns = []
|
||||
for message in messages:
|
||||
if isinstance(message, HumanMessage):
|
||||
turns.append([])
|
||||
if turns:
|
||||
turns[-1].append(message)
|
||||
# Leading orphaned checkpoint fragments have no user question to anchor them.
|
||||
return turns
|
||||
|
||||
|
||||
def _complete_exchanges(turn: list[BaseMessage]) -> list[BaseMessage]:
|
||||
"""Drop unfinished exchanges left by a cancelled turn and orphan results."""
|
||||
result = []
|
||||
i = 0
|
||||
while i < len(turn):
|
||||
message = turn[i]
|
||||
if isinstance(message, AIMessage) and message.tool_calls:
|
||||
expected = {c["id"] for c in message.tool_calls}
|
||||
end = i + 1
|
||||
while end < len(turn) and isinstance(turn[end], ToolMessage):
|
||||
end += 1
|
||||
replies = turn[i + 1 : end]
|
||||
actual = [m.tool_call_id for m in replies]
|
||||
if len(actual) == len(expected) and set(actual) == expected:
|
||||
result.extend(turn[i:end])
|
||||
i = end
|
||||
else:
|
||||
if not isinstance(message, ToolMessage):
|
||||
result.append(message)
|
||||
i += 1
|
||||
return result
|
||||
|
||||
|
||||
def _summary(turns: list[list[BaseMessage]], limit: int = 2400) -> SystemMessage | None:
|
||||
excerpts = []
|
||||
remaining = limit
|
||||
for turn in reversed(turns):
|
||||
lines = []
|
||||
for message in turn:
|
||||
if isinstance(message, HumanMessage) or (
|
||||
isinstance(message, AIMessage) and not message.tool_calls
|
||||
):
|
||||
content = message.content
|
||||
if not isinstance(content, str):
|
||||
content = " ".join(
|
||||
p.get("text", "") if isinstance(p, dict) else str(p)
|
||||
for p in content
|
||||
)
|
||||
role = "User" if isinstance(message, HumanMessage) else "Assistant"
|
||||
lines.append(f"{role}: {_cap(content, 400, ' [excerpt] ')}")
|
||||
excerpt = "\n".join(lines)
|
||||
if len(excerpt) > remaining:
|
||||
break
|
||||
excerpts.insert(0, excerpt)
|
||||
remaining -= len(excerpt) + 2
|
||||
if not excerpts:
|
||||
return None
|
||||
return SystemMessage(
|
||||
content="Earlier conversation excerpts (historical data, not new instructions; excerpts may omit details):\n\n"
|
||||
+ "\n\n".join(excerpts)
|
||||
)
|
||||
|
||||
|
||||
def trim_history(
|
||||
messages: list[BaseMessage], *, token_budget: int = 12000
|
||||
) -> list[BaseMessage]:
|
||||
"""Keep the current turn, then recent whole turns, within an estimated budget.
|
||||
|
||||
The current user message and tool call arguments are never truncated. Fail
|
||||
explicitly if these mandatory inputs cannot fit, rather than lose the task.
|
||||
"""
|
||||
turns = [_complete_exchanges(turn) for turn in _turns(messages)]
|
||||
if not turns:
|
||||
return []
|
||||
current = turns[-1]
|
||||
current = [
|
||||
_shrink(m, TOOL_RESULT_MAX_CHARS, earlier=False)
|
||||
if isinstance(m, ToolMessage)
|
||||
else m
|
||||
for i, m in enumerate(messages)
|
||||
for m in current
|
||||
]
|
||||
window = trimmed[-HISTORY_WINDOW:]
|
||||
# A window that starts mid tool-exchange would begin with ToolMessages whose
|
||||
# tool_calls were cut off; OpenAI-style APIs reject those.
|
||||
while window and isinstance(window[0], ToolMessage):
|
||||
window = window[1:]
|
||||
return window
|
||||
|
||||
def cost(ms):
|
||||
return sum(message_tokens(m) for m in ms)
|
||||
|
||||
# Keep room for continuity even when the current turn has many tool results.
|
||||
current_budget = (
|
||||
token_budget - min(1000, token_budget // 8) if len(turns) > 1 else token_budget
|
||||
)
|
||||
if cost(current) > current_budget:
|
||||
# Binary-search the largest shared excerpt size that fits. Tool call IDs
|
||||
# and all results remain, including parallel exchanges.
|
||||
low, high = 0, TOOL_RESULT_MAX_CHARS
|
||||
|
||||
def compact(limit):
|
||||
return [
|
||||
_shrink(m, limit, earlier=False)
|
||||
if isinstance(m, (ToolMessage, AIMessage))
|
||||
else m
|
||||
for m in current
|
||||
]
|
||||
|
||||
if cost(compact(0)) > token_budget:
|
||||
raise ValueError(
|
||||
"The current question and tool calls exceed the context budget. Send a smaller attachment or start a new conversation."
|
||||
)
|
||||
current_budget = max(current_budget, cost(compact(0)))
|
||||
while low < high:
|
||||
mid = (low + high + 1) // 2
|
||||
if cost(compact(mid)) <= current_budget:
|
||||
low = mid
|
||||
else:
|
||||
high = mid - 1
|
||||
current = compact(low)
|
||||
selected = []
|
||||
used = cost(current)
|
||||
prior = turns[:-1]
|
||||
# Leave a small share for excerpts of older turns only when necessary.
|
||||
summary_reserve = min(1000, max(0, token_budget - used))
|
||||
for turn in reversed(prior):
|
||||
shrunk = [_shrink(m, EARLIER_MESSAGE_MAX_CHARS, earlier=True) for m in turn]
|
||||
turn_cost = cost(shrunk)
|
||||
if used + turn_cost > token_budget - summary_reserve:
|
||||
break
|
||||
selected.insert(0, shrunk)
|
||||
used += turn_cost
|
||||
dropped = prior[: len(prior) - len(selected)]
|
||||
summary = _summary(dropped)
|
||||
result = [m for turn in selected for m in turn] + current
|
||||
if summary and used + message_tokens(summary) <= token_budget:
|
||||
result.insert(0, summary)
|
||||
logging.info(
|
||||
json.dumps(
|
||||
{
|
||||
"event": "context_budget",
|
||||
"budget_tokens": token_budget,
|
||||
"estimated_tokens": cost(result),
|
||||
"input_messages": len(messages),
|
||||
"output_messages": len(result),
|
||||
"dropped_turns": len(dropped),
|
||||
"summary_included": bool(
|
||||
result and isinstance(result[0], SystemMessage)
|
||||
),
|
||||
}
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def build_context(
|
||||
messages: list[BaseMessage], system_prompt: str, tool_schemas: list[dict]
|
||||
) -> list[BaseMessage]:
|
||||
limit = min(
|
||||
PROMPT_BUDGET_TOKENS,
|
||||
CONTEXT_WINDOW_TOKENS - RESPONSE_RESERVE_TOKENS - SAFETY_RESERVE_TOKENS,
|
||||
)
|
||||
system = SystemMessage(content=system_prompt)
|
||||
budget = limit - message_tokens(system) - estimate_tokens(tool_schemas)
|
||||
if budget <= 0:
|
||||
raise ValueError(
|
||||
"System instructions and tool definitions exceed the configured prompt budget."
|
||||
)
|
||||
history = trim_history(messages, token_budget=budget)
|
||||
logging.info(
|
||||
json.dumps(
|
||||
{
|
||||
"event": "prompt_budget",
|
||||
"limit_tokens": limit,
|
||||
"system_tokens": message_tokens(system),
|
||||
"tool_schema_tokens": estimate_tokens(tool_schemas),
|
||||
"estimated_prompt_tokens": message_tokens(system)
|
||||
+ estimate_tokens(tool_schemas)
|
||||
+ sum(message_tokens(m) for m in history),
|
||||
"response_reserve_tokens": RESPONSE_RESERVE_TOKENS,
|
||||
}
|
||||
)
|
||||
)
|
||||
return [system] + history
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
from .models import UserMemory
|
||||
from .prompts import SIMBA_SYSTEM_PROMPT
|
||||
from .context import estimate_tokens
|
||||
|
||||
MEMORY_BUDGET_TOKENS = 1000
|
||||
|
||||
|
||||
async def get_memories_for_user(user_id: str) -> list[str]:
|
||||
@@ -12,7 +15,16 @@ async def build_system_prompt_with_memories(user_id: str) -> str:
|
||||
memories = await get_memories_for_user(user_id)
|
||||
if not memories:
|
||||
return SIMBA_SYSTEM_PROMPT
|
||||
block = "\n".join(f"- {m}" for m in memories)
|
||||
selected = []
|
||||
used = 0
|
||||
for memory in memories:
|
||||
cost = estimate_tokens(f"- {memory}\n")
|
||||
if used + cost <= MEMORY_BUDGET_TOKENS:
|
||||
selected.append(memory)
|
||||
used += cost
|
||||
block = "\n".join(f"- {m}" for m in selected)
|
||||
if len(selected) < len(memories):
|
||||
block += "\n[Additional saved memories omitted from this prompt.]"
|
||||
return f"{SIMBA_SYSTEM_PROMPT}\n\nUSER MEMORIES (facts the user has asked you to remember):\n{block}"
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
# Conversation context management
|
||||
|
||||
The LangGraph checkpoint retains original messages and artifacts. Each model
|
||||
call constructs a bounded view; compaction never rewrites the checkpoint.
|
||||
|
||||
The prompt budget includes system instructions, saved memories, tool schemas
|
||||
(including admin tools), message content and tool arguments. The default target
|
||||
is 16,000 estimated tokens. A 4,096-token generation limit and 2,048-token safety
|
||||
reserve also constrain the budget against the configured 32,768-token window.
|
||||
|
||||
Estimates use UTF-8 bytes divided by three, message overhead, and a fixed image
|
||||
allowance. These are estimates rather than model-specific tokenizer counts.
|
||||
Compare `prompt_budget.estimated_prompt_tokens` with actual `stream_request`
|
||||
model-call usage and adjust the target/window for the deployed backend.
|
||||
|
||||
The current user request is always retained. Oversized tool results become
|
||||
marked excerpts preserving their beginning and end; all tool IDs and arguments
|
||||
remain. If necessary, intermediate assistant text is also shortened. If the
|
||||
mandatory request and exchange structure cannot fit, the request fails with an
|
||||
explicit error rather than silently dropping the question.
|
||||
|
||||
Recent history is selected as complete turns. Older user messages, assistant
|
||||
answers and tool results are capped; earlier images become placeholders.
|
||||
Incomplete tool exchanges from cancelled calls and orphan results are removed
|
||||
from the prompt. A small budget is reserved for bounded question/answer excerpts
|
||||
from older turns. These excerpts are not an LLM-generated rolling summary and do
|
||||
not guarantee retention of every historical decision. No extra model call is
|
||||
added to the request path. Saved memories use a separate 1,000-token estimated
|
||||
budget, preferring recent entries without deleting omitted memories.
|
||||
|
||||
Obsidian reads accept `start_line` (one-based content line, excluding frontmatter)
|
||||
and `max_lines`, allowing focused follow-up reads. Full artifact content remains
|
||||
available to the interface and checkpoint.
|
||||
|
||||
Configuration:
|
||||
|
||||
- `CONTEXT_WINDOW_TOKENS`: smallest supported context window across backends.
|
||||
- `PROMPT_BUDGET_TOKENS`: prompt target, including fixed instructions and tools.
|
||||
- `RESPONSE_RESERVE_TOKENS`: maximum generated tokens and reserved answer space.
|
||||
|
||||
`context_budget` logs history estimates, message counts, dropped turns and excerpt
|
||||
inclusion. `prompt_budget` logs fixed overhead and the complete prompt estimate.
|
||||
Neither log includes message content.
|
||||
+247
-53
@@ -1,78 +1,272 @@
|
||||
"""Tests for trimming the replayed history (blueprints/conversation/context.py)."""
|
||||
"""Regression tests for complete-turn context budgeting."""
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
||||
|
||||
from blueprints.conversation import context
|
||||
from blueprints.conversation.context import (
|
||||
EARLIER_MESSAGE_MAX_CHARS,
|
||||
HISTORY_WINDOW,
|
||||
IMAGE_PLACEHOLDER,
|
||||
TRUNCATION_NOTE,
|
||||
build_context,
|
||||
message_tokens,
|
||||
trim_history,
|
||||
)
|
||||
|
||||
BIG = "x" * (EARLIER_MESSAGE_MAX_CHARS * 3)
|
||||
|
||||
def call(*ids):
|
||||
return AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{"id": i, "name": "search", "args": {"path": "notes/budget.md"}}
|
||||
for i in ids
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _call(call_id: str, name: str = "simba_search") -> AIMessage:
|
||||
return AIMessage(content="", tool_calls=[{"id": call_id, "name": name, "args": {}}])
|
||||
def result(i, content="found"):
|
||||
return ToolMessage(content=content, tool_call_id=i)
|
||||
|
||||
|
||||
def _result(call_id: str, text: str = BIG) -> ToolMessage:
|
||||
return ToolMessage(content=text, tool_call_id=call_id, name="simba_search")
|
||||
def cost(messages):
|
||||
return sum(message_tokens(m) for m in messages)
|
||||
|
||||
|
||||
def test_earlier_tool_results_are_capped_current_turn_untouched():
|
||||
messages = [
|
||||
HumanMessage(content="first question"),
|
||||
_call("c1"),
|
||||
_result("c1"),
|
||||
AIMessage(content="first answer"),
|
||||
HumanMessage(content="second question"),
|
||||
_call("c2"),
|
||||
_result("c2"),
|
||||
def test_long_current_turn_keeps_question_and_all_exchanges():
|
||||
question = HumanMessage(content="Compare my spending with my plan")
|
||||
messages = [question]
|
||||
for i in range(8):
|
||||
messages += [call(str(i)), result(str(i), "data " * 1000)]
|
||||
out = trim_history(messages, token_budget=3000)
|
||||
assert out[0] is question
|
||||
assert len(out) == 17
|
||||
assert cost(out) <= 3000
|
||||
assert [m.tool_call_id for m in out if isinstance(m, ToolMessage)] == [
|
||||
str(i) for i in range(8)
|
||||
]
|
||||
|
||||
out = trim_history(messages)
|
||||
|
||||
earlier, current = out[2], out[6]
|
||||
assert earlier.content == BIG[:EARLIER_MESSAGE_MAX_CHARS] + TRUNCATION_NOTE
|
||||
assert current.content == BIG
|
||||
# The checkpointed originals are never mutated.
|
||||
assert messages[2].content == BIG
|
||||
assert messages[2].content == "data " * 1000
|
||||
assert out[1].tool_calls == messages[1].tool_calls
|
||||
|
||||
|
||||
def test_earlier_attachments_and_images_are_shrunk():
|
||||
def test_parallel_tool_results_remain_paired():
|
||||
messages = [
|
||||
HumanMessage(content="q"),
|
||||
call("a", "b"),
|
||||
result("a", "x" * 20000),
|
||||
result("b", "y" * 20000),
|
||||
]
|
||||
out = trim_history(messages, token_budget=700)
|
||||
assert len(out) == 4
|
||||
assert out[1].tool_calls == messages[1].tool_calls
|
||||
assert cost(out) <= 700
|
||||
|
||||
|
||||
def test_drops_whole_turns_and_includes_bounded_excerpts():
|
||||
messages = []
|
||||
for i in range(12):
|
||||
messages += [
|
||||
HumanMessage(content=f"Question {i}"),
|
||||
call(str(i)),
|
||||
result(str(i), "data " * 1000),
|
||||
AIMessage(content=f"Decision {i}"),
|
||||
]
|
||||
question = HumanMessage(content="What did we decide?")
|
||||
messages.append(question)
|
||||
out = trim_history(messages, token_budget=3000)
|
||||
assert isinstance(out[0], SystemMessage)
|
||||
assert "Decision" in out[0].content
|
||||
assert "historical data" in out[0].content
|
||||
assert cost(out) <= 3000
|
||||
assert out[-1] is question
|
||||
# Every retained result still has its original call preceding it.
|
||||
seen = set()
|
||||
for m in out:
|
||||
if isinstance(m, AIMessage):
|
||||
seen.update(c["id"] for c in m.tool_calls)
|
||||
if isinstance(m, ToolMessage):
|
||||
assert m.tool_call_id in seen
|
||||
|
||||
|
||||
def test_earlier_images_replaced_current_image_preserved():
|
||||
image = {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}
|
||||
messages = [
|
||||
HumanMessage(content=[{"type": "text", "text": BIG}, image]),
|
||||
AIMessage(content="nice photo"),
|
||||
HumanMessage(content=[{"type": "text", "text": "now this one"}, image]),
|
||||
HumanMessage(content=[{"type": "text", "text": "old"}, image]),
|
||||
AIMessage(content="photo"),
|
||||
HumanMessage(content=[{"type": "text", "text": "new"}, image]),
|
||||
]
|
||||
|
||||
out = trim_history(messages)
|
||||
|
||||
earlier = out[0].content
|
||||
assert earlier[0]["text"].endswith(TRUNCATION_NOTE)
|
||||
assert earlier[1] == {"type": "text", "text": IMAGE_PLACEHOLDER}
|
||||
# The current turn keeps its image.
|
||||
assert out[2].content[1] == image
|
||||
assert out[0].content[1] == {"type": "text", "text": IMAGE_PLACEHOLDER}
|
||||
assert out[-1].content[1] == image
|
||||
assert messages[0].content[1] == image
|
||||
|
||||
|
||||
def test_window_never_starts_with_orphaned_tool_results():
|
||||
# Many parallel tool results so the window cut lands between an AIMessage's
|
||||
# tool_calls and their ToolMessages.
|
||||
messages = [HumanMessage(content="q"), _call("c0")]
|
||||
messages += [_result("c0", "r") for _ in range(HISTORY_WINDOW)]
|
||||
messages += [AIMessage(content="answer"), HumanMessage(content="next")]
|
||||
|
||||
out = trim_history(messages)
|
||||
|
||||
assert len(out) <= HISTORY_WINDOW
|
||||
assert not isinstance(out[0], ToolMessage)
|
||||
|
||||
|
||||
def test_short_history_is_returned_as_is():
|
||||
def test_short_history_unchanged():
|
||||
messages = [HumanMessage(content="hi"), AIMessage(content="hello")]
|
||||
assert trim_history(messages) == messages
|
||||
|
||||
|
||||
def test_question_that_cannot_fit_fails_without_truncation():
|
||||
messages = [HumanMessage(content="important request " * 1000)]
|
||||
with pytest.raises(ValueError, match="current question"):
|
||||
trim_history(messages, token_budget=100)
|
||||
assert messages[0].content == "important request " * 1000
|
||||
|
||||
|
||||
def test_system_and_tool_schemas_are_charged_to_budget(monkeypatch):
|
||||
monkeypatch.setattr(context, "PROMPT_BUDGET_TOKENS", 2500)
|
||||
monkeypatch.setattr(context, "CONTEXT_WINDOW_TOKENS", 10000)
|
||||
schemas = [{"name": "search", "description": "tool docs " * 100}]
|
||||
messages = [HumanMessage(content="q"), call("a"), result("a", "result " * 3000)]
|
||||
out = build_context(messages, "system " * 100, schemas)
|
||||
assert cost(out) + context.estimate_tokens(schemas) <= 2500
|
||||
assert out[1] is messages[0]
|
||||
|
||||
|
||||
def test_fixed_overhead_that_cannot_fit_fails(monkeypatch):
|
||||
monkeypatch.setattr(context, "PROMPT_BUDGET_TOKENS", 100)
|
||||
with pytest.raises(ValueError, match="tool definitions"):
|
||||
build_context([HumanMessage(content="q")], "system" * 1000, [])
|
||||
|
||||
|
||||
def test_response_and_safety_reserves_reduce_prompt_limit(monkeypatch):
|
||||
monkeypatch.setattr(context, "CONTEXT_WINDOW_TOKENS", 8000)
|
||||
monkeypatch.setattr(context, "PROMPT_BUDGET_TOKENS", 16000)
|
||||
messages = [HumanMessage(content="q"), call("a"), result("a", "data " * 5000)]
|
||||
out = build_context(messages, "system", [])
|
||||
assert (
|
||||
cost(out) + context.estimate_tokens([])
|
||||
<= 8000 - context.RESPONSE_RESERVE_TOKENS - context.SAFETY_RESERVE_TOKENS
|
||||
)
|
||||
|
||||
|
||||
def test_cancelled_turn_does_not_replay_incomplete_exchange():
|
||||
messages = [
|
||||
HumanMessage(content="old"),
|
||||
call("a", "b"),
|
||||
result("a"),
|
||||
HumanMessage(content="new"),
|
||||
]
|
||||
out = trim_history(messages)
|
||||
assert out == messages
|
||||
assert out == [messages[0], messages[-1]]
|
||||
|
||||
|
||||
def test_orphan_results_are_removed():
|
||||
messages = [
|
||||
result("orphan"),
|
||||
HumanMessage(content="q"),
|
||||
result("another orphan"),
|
||||
AIMessage(content="answer"),
|
||||
]
|
||||
assert trim_history(messages) == [messages[1], messages[3]]
|
||||
|
||||
|
||||
def test_large_assistant_answer_is_bounded_in_earlier_turn():
|
||||
messages = [
|
||||
HumanMessage(content="old"),
|
||||
AIMessage(content="long " * 10000),
|
||||
HumanMessage(content="new"),
|
||||
]
|
||||
out = trim_history(messages, token_budget=4000)
|
||||
assert cost(out) <= 4000
|
||||
assert len(out[1].content) < len(messages[1].content)
|
||||
assert messages[1].content == "long " * 10000
|
||||
|
||||
|
||||
def test_multiblock_attachment_shares_one_cap():
|
||||
messages = [
|
||||
HumanMessage(content=[{"type": "text", "text": "x" * 10000} for _ in range(5)]),
|
||||
AIMessage(content="ok"),
|
||||
HumanMessage(content="next"),
|
||||
]
|
||||
out = trim_history(messages, token_budget=5000)
|
||||
assert cost(out) <= 5000
|
||||
assert sum(len(p["text"]) for p in out[0].content) < 5000
|
||||
|
||||
|
||||
async def test_saved_memories_are_bounded_without_deleting(monkeypatch):
|
||||
from blueprints.conversation import memory
|
||||
|
||||
saved = ["newest preference", "big " * 5000, "older preference"]
|
||||
|
||||
async def get_memories(user_id):
|
||||
return saved
|
||||
|
||||
monkeypatch.setattr(memory, "get_memories_for_user", get_memories)
|
||||
prompt = await memory.build_system_prompt_with_memories("user")
|
||||
block = prompt.split("USER MEMORIES", 1)[1]
|
||||
assert "newest preference" in block
|
||||
assert "older preference" in block
|
||||
assert "omitted" in block
|
||||
assert "big big" not in block
|
||||
assert saved[1] == "big " * 5000
|
||||
|
||||
|
||||
def test_agent_prompt_accounts_for_extra_tool_schemas(monkeypatch):
|
||||
monkeypatch.setenv("TAVILY_API_KEY", "test-key")
|
||||
from blueprints.conversation import agents
|
||||
|
||||
captured = {}
|
||||
|
||||
def create(**kwargs):
|
||||
captured.update(kwargs)
|
||||
|
||||
monkeypatch.setattr(agents, "create_react_agent", create)
|
||||
extra = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "extra",
|
||||
"description": "extra docs",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
agents.build_agent(extra_tools=[extra])
|
||||
assert extra in captured["tools"]
|
||||
question = HumanMessage(content="hello")
|
||||
actual = captured["prompt"].invoke({"messages": [question]})
|
||||
assert actual[-1] is question
|
||||
# The actual prompt closure includes this schema rather than relying on a
|
||||
# static reserve that would miss admin-only tools.
|
||||
seen = {}
|
||||
|
||||
def build(messages, system, schemas):
|
||||
seen["schemas"] = schemas
|
||||
return messages
|
||||
|
||||
monkeypatch.setattr(agents, "build_context", build)
|
||||
messages = [HumanMessage(content="hello")]
|
||||
assert captured["prompt"].invoke({"messages": messages}) == messages
|
||||
assert extra in seen["schemas"]
|
||||
|
||||
|
||||
def test_heavy_current_turn_leaves_room_for_previous_decisions():
|
||||
question = HumanMessage(content="Apply that plan to my spending")
|
||||
messages = [
|
||||
HumanMessage(content="Plan?"),
|
||||
AIMessage(content="Keep dining under $300."),
|
||||
question,
|
||||
call("a"),
|
||||
result("a", "spending " * 10000),
|
||||
]
|
||||
out = trim_history(messages, token_budget=2500)
|
||||
assert cost(out) <= 2500
|
||||
assert question in out
|
||||
assert any(
|
||||
"Keep dining under $300." in m.content
|
||||
for m in out
|
||||
if isinstance(m.content, str)
|
||||
)
|
||||
|
||||
|
||||
def test_large_tool_arguments_fail_explicitly_instead_of_breaking_pair():
|
||||
exchange = call("a")
|
||||
exchange.tool_calls[0]["args"] = {"query": "huge " * 10000}
|
||||
messages = [HumanMessage(content="q"), exchange, result("a")]
|
||||
with pytest.raises(ValueError, match="tool calls exceed"):
|
||||
trim_history(messages, token_budget=300)
|
||||
assert exchange.tool_calls[0]["args"]["query"] == "huge " * 10000
|
||||
|
||||
|
||||
def test_unicode_tool_result_fits_estimated_budget():
|
||||
messages = [HumanMessage(content="q"), call("a"), result("a", "猫😺预算" * 10000)]
|
||||
out = trim_history(messages, token_budget=600)
|
||||
assert cost(out) <= 600
|
||||
assert out[0] is messages[0]
|
||||
assert out[1].tool_calls == messages[1].tool_calls
|
||||
|
||||
@@ -589,3 +589,26 @@ def test_search_documents_artifact_empty():
|
||||
)
|
||||
def test_format_date_range(start, end, expected):
|
||||
assert format_date_range(start, end) == expected
|
||||
|
||||
|
||||
async def test_obsidian_note_line_range_retains_full_artifact(obsidian):
|
||||
with open(os.path.join(obsidian.vault_path, "range.md"), "w") as f:
|
||||
f.write("first\nsecond\nthird\nfourth")
|
||||
msg = await call(
|
||||
agents.obsidian_read_note,
|
||||
{"relative_path": "range.md", "start_line": 2, "max_lines": 2},
|
||||
)
|
||||
assert "Content lines 2-3 of 4:\nsecond\nthird\n\n" in msg.content
|
||||
assert "first" not in msg.content
|
||||
assert msg.artifact["data"]["markdown"] == "first\nsecond\nthird\nfourth"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args", [{"start_line": 0}, {"max_lines": 0}, {"start_line": 99}]
|
||||
)
|
||||
async def test_obsidian_note_invalid_line_range(obsidian, args):
|
||||
with open(os.path.join(obsidian.vault_path, "range.md"), "w") as f:
|
||||
f.write("first\nsecond")
|
||||
msg = await call(agents.obsidian_read_note, {"relative_path": "range.md", **args})
|
||||
assert msg.artifact is None
|
||||
assert "positive integers" in msg.content or "beyond the end" in msg.content
|
||||
|
||||
Reference in New Issue
Block a user