diff --git a/.env.example b/.env.example index 2fafeb8..d991f40 100644 --- a/.env.example +++ b/.env.example @@ -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) diff --git a/blueprints/conversation/agents.py b/blueprints/conversation/agents.py index c9079ac..445077c 100644 --- a/blueprints/conversation/agents.py +++ b/blueprints/conversation/agents.py @@ -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, ) diff --git a/blueprints/conversation/context.py b/blueprints/conversation/context.py index 1e0afb5..9babe95 100644 --- a/blueprints/conversation/context.py +++ b/blueprints/conversation/context.py @@ -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 diff --git a/blueprints/conversation/memory.py b/blueprints/conversation/memory.py index 5db37b5..06d4656 100644 --- a/blueprints/conversation/memory.py +++ b/blueprints/conversation/memory.py @@ -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}" diff --git a/docs/context-management.md b/docs/context-management.md new file mode 100644 index 0000000..9a60960 --- /dev/null +++ b/docs/context-management.md @@ -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. diff --git a/tests/unit/test_context_trim.py b/tests/unit/test_context_trim.py index 88fb7f8..6764133 100644 --- a/tests/unit/test_context_trim.py +++ b/tests/unit/test_context_trim.py @@ -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 diff --git a/tests/unit/test_tool_artifacts.py b/tests/unit/test_tool_artifacts.py index a381c04..a16196f 100644 --- a/tests/unit/test_tool_artifacts.py +++ b/tests/unit/test_tool_artifacts.py @@ -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