diff --git a/ex_app/lib/agent.py b/ex_app/lib/agent.py index 80141d1..62d21d0 100644 --- a/ex_app/lib/agent.py +++ b/ex_app/lib/agent.py @@ -6,18 +6,22 @@ from time import monotonic from typing import Any, cast -from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, SystemMessage, ToolMessage +from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, HumanMessage, RemoveMessage, SystemMessage, ToolMessage from langchain_core.runnables import RunnableConfig +from langgraph.constants import TAG_NOSTREAM +from langgraph.graph.message import REMOVE_ALL_MESSAGES from nc_py_api import AsyncNextcloudApp -from nc_py_api.ex_app import persistent_storage +from nc_py_api.ex_app import LogLvl, persistent_storage import pytz from ex_app.lib.all_tools.lib.tz import user_timezone_async from ex_app.lib.all_tools.nextcloud_links import get_absolute_base_url from ex_app.lib.all_tools.skills import list_skills_metadata +from ex_app.lib.context_window import CHARS_PER_TOKEN, render_transcript, select_context from ex_app.lib.graph import CONFIRM_TOOLS_NODE, AgentState, get_graph from ex_app.lib.jsonplus import JsonPlusSerializer +from ex_app.lib.logger import log from ex_app.lib.memorysaver import MemorySaver from ex_app.lib.nc_model import ( MULTIMODAL_INTERACTION, @@ -27,7 +31,7 @@ extract_text_content, ) from ex_app.lib.signature import add_signature, load_legacy_key, load_or_create_key, verify_signature -from ex_app.lib.tools import get_confirmation_policy, get_tools +from ex_app.lib.tools import get_confirmation_policy, get_context_budget, get_tools graph_config = { # Dummy thread id as we return the whole state @@ -113,6 +117,23 @@ def format_user_instructions(user_instructions: str | None) -> str: + user_instructions.strip() + "\n" ) +# Instructions for summarizing the turns that no longer fit into the context window +SUMMARY_PROMPT = ( + "You summarize a conversation between a user and an AI assistant that uses tools, so the assistant" + " can continue the conversation without seeing these messages anymore." + " Keep what the user wants to achieve, their preferences, the decisions made, the open tasks" + " and the facts learned from tool results, including names, ids, file paths, links, dates and numbers" + " the assistant may need again. Leave out small talk and details that will not matter anymore." + " If a summary of the conversation so far is given, fold the new messages into it." + " Write in the language of the conversation, in under 400 words, and reply with the summary only.\n" +) + +# Precedes the summary in the message that stands in for the summarized turns +SUMMARY_MESSAGE_PREFIX = ( + "[The earlier part of this conversation is no longer shown in full. This is an automatically generated" + " summary of it, not a new message of the user.]\n\n" +) + # Value of the optional "output_mode" input when the agent's reply is converted to speech OUTPUT_MODE_SPEECH = "speech" @@ -139,6 +160,7 @@ async def react( tools = await get_tools(nc) policy = await get_confirmation_policy(nc) + context_budget = await get_context_budget(nc) tz_name = user_tz = None try: tz_name = await user_timezone_async(nc) @@ -156,6 +178,23 @@ def tool_enabled(tool_name): return True return False + async def summarize(previous_summary: str, messages: list[BaseMessage]) -> str: + """Fold the turns that no longer fit into the context window into the summary""" + content = '' + if previous_summary: + content += "Summary of the conversation so far:\n\n" + previous_summary + "\n\n" + content += "The conversation continued as follows:\n\n" + render_transcript(messages, context_budget * CHARS_PER_TOKEN) + summarizer = ChatWithNextcloud(nc=nc) + # The summary is no reply to the user, so it is not streamed to them + response = await summarizer.ainvoke( + [SystemMessage(SUMMARY_PROMPT), HumanMessage(content=content)], + {"tags": [TAG_NOSTREAM]}, + ) + summary = extract_text_content(response.content).strip() + if not summary: + raise ValueError("The language model returned an empty summary") + return summary + async def call_model( state: AgentState, config: RunnableConfig, @@ -228,15 +267,34 @@ async def call_model( system_prompt_text += format_output_mode(task['input'].get('output_mode')) + selection = select_context(state["messages"], context_budget) + summary = state.get("summary") or "" + update: dict[str, Any] = {"messages": []} + if selection.dropped: + try: + summary = await summarize(summary, selection.dropped) + update["summary"] = summary + # Replace the conversation with the turns that are still sent in full + update["messages"] = [RemoveMessage(id=REMOVE_ALL_MESSAGES), *selection.kept] + except Exception as e: + # Keep the dropped turns in the state, so summarizing them is tried again on the next step + await log(nc, LogLvl.WARNING, f"Failed to summarize the earlier part of the conversation: {e}") + model_messages = selection.model_messages + if summary: + # The summary is made from what users and tools wrote, so it is sent as the first message of + # the conversation and not as a part of the instructions. It is not kept in the state. + model_messages = [HumanMessage(content=SUMMARY_MESSAGE_PREFIX + summary)] + model_messages + # this is similar to customizing the create_react_agent with state_modifier, but is a lot more flexible system_prompt = SystemMessage( system_prompt_text.replace("{CURRENT_DATE}", current_date).replace("{DATE_SUFFIX}", date_suffix) + format_user_instructions(task['input'].get('user_instructions')) ) - response = await bound_model.ainvoke([system_prompt] + state["messages"], config) + response = await bound_model.ainvoke([system_prompt] + model_messages, config) # We return a list, because this will get added to the existing list - return {"messages": [response]} + update["messages"].append(response) + return update try: # Try to load state using the new conversation_token type @@ -277,7 +335,8 @@ async def call_model( last_message: AIMessage = AIMessage("") if len(snapshot_messages) > 0: last_message = cast(AIMessage, snapshot_messages[-1]) - previous_message_count = len(snapshot_messages) + # The earlier messages can be replaced by a summary during the run, so new messages are told apart by id + previous_message_ids = {message.id for message in snapshot_messages} source_list: list[str] = [] known_sources: set[str] = set() streamed_output = '' @@ -319,7 +378,9 @@ async def report_stream_state(force: bool = False): event = payload last_message = event['messages'][-1] - for message in event['messages'][previous_message_count:]: + for message in event['messages']: + if message.id in previous_message_ids: + continue if isinstance(message, AIMessage) and message.tool_calls: for tool_call in message.tool_calls: tool_name = tool_call['name'] diff --git a/ex_app/lib/context_window.py b/ex_app/lib/context_window.py new file mode 100644 index 0000000..ce7ba2d --- /dev/null +++ b/ex_app/lib/context_window.py @@ -0,0 +1,251 @@ +# SPDX-FileCopyrightText: 2026 Nextcloud GmbH and Nextcloud contributors +# SPDX-License-Identifier: AGPL-3.0-or-later +"""Fitting a conversation into the language model's context window + +The context window of the model behind the TaskProcessing provider is unknown, so the +conversation is fitted into a budget the admin can configure. Sizes are estimated from +characters, as there is no tokenizer for the unknown model either. + +- The current turn (from the user's latest message on) is sent in full, as the model + needs every tool result of it to finish its work, except that a single tool result + gets no more than half the budget. Only if the turn on its own outgrows the budget + are its earlier tool results shortened and the budget is shared out among the tool + results the model has not answered yet. +- Tool results of earlier turns are shortened, keeping their beginning and their end. +- If the conversation still does not fit, whole turns are dropped, oldest first, so a + tool call is never separated from its result. The caller summarizes them. +""" +import json +from dataclasses import dataclass + +from langchain_core.messages import AIMessage, BaseMessage + +from ex_app.lib.nc_model import extract_file_ids, extract_text_content + +# Rough average for the languages and JSON the model gets to see +CHARS_PER_TOKEN = 4 + +CONTEXT_BUDGET_SETTING_ID = 'context_budget' +# In tokens +DEFAULT_CONTEXT_BUDGET = 16000 +MIN_CONTEXT_BUDGET = 2000 + +# Once the budget is exceeded, the conversation is trimmed down to this share of it. +# Trimming further than needed leaves room for the next turns, so the conversation +# is not summarized again on every turn. +TRIM_TARGET = 0.6 + +# Tool results of earlier turns are shortened to this many characters +OLD_TOOL_OUTPUT_MAX_LENGTH = 2000 + +# Room reserved for the marker that replaces the omitted part +MARKER_RESERVE = 160 + + +def _omission_hint(tool_name: str | None, can_refetch: bool) -> str: + if not can_refetch: + # Calling the tool again would give the same result, shortened the same way + return '; the result is too large to be shown in full' + if not tool_name: + return '' + return f'; call {tool_name} again if you need the omitted part' + + +def _truncate_json_list(text: str, max_chars: int, tool_name: str | None, can_refetch: bool) -> str | None: + """Shorten a JSON list by leaving out items, so the result is still valid JSON + + Returns None if the text is no JSON list or not even a single item fits. + """ + try: + data = json.loads(text) + except (ValueError, TypeError): + return None + if not isinstance(data, list) or len(data) < 2: + return None + + items = [json.dumps(item, ensure_ascii=False) for item in data] + compact = '[' + ', '.join(items) + ']' + if len(compact) <= max_chars: + return compact + + budget = max_chars - MARKER_RESERVE + head: list[str] = [] + tail: list[str] = [] + used = 2 + i, j = 0, len(items) - 1 + while i <= j and used + len(items[i]) + 2 <= budget * 2 // 3: + head.append(items[i]) + used += len(items[i]) + 2 + i += 1 + while j >= i and used + len(items[j]) + 2 <= budget: + tail.insert(0, items[j]) + used += len(items[j]) + 2 + j -= 1 + if not head and not tail: + return None + + omitted = j - i + 1 + marker = json.dumps(f'… {omitted} of {len(items)} items omitted{_omission_hint(tool_name, can_refetch)} …', ensure_ascii=False) + return '[' + ', '.join(head + [marker] + tail) + ']' + + +def truncate_text(text: str, max_chars: int, tool_name: str | None = None, can_refetch: bool = True) -> str: + """Shorten a text to about max_chars, keeping its beginning and its end + + The end is often where the answer is, e.g. the latest entries of a list. JSON lists + are shortened by leaving out whole items. The marker says how much was left out and, + given the tool's name, that the model can call the tool again to see all of it. + With can_refetch=False it says that the result is too large instead, for results + that would come back shortened the same way. + """ + if len(text) <= max_chars: + return text + + truncated = _truncate_json_list(text, max_chars, tool_name, can_refetch) + if truncated is not None: + return truncated + + keep = max(max_chars - MARKER_RESERVE, 0) + head_len = keep * 2 // 3 + tail_len = keep - head_len + head = text[:head_len] + tail = text[len(text) - tail_len:] if tail_len > 0 else '' + # Cut at line boundaries if there is one close by + newline = head.rfind('\n') + if newline >= head_len * 0.8: + head = head[:newline + 1] + newline = tail.find('\n') + if 0 <= newline < tail_len * 0.2: + tail = tail[newline + 1:] + + omitted = len(text) - len(head) - len(tail) + return f'{head}\n[… {omitted} of {len(text)} characters omitted{_omission_hint(tool_name, can_refetch)}]\n{tail}' + + +def message_size(message: BaseMessage) -> int: + """The size of a message in characters, as it is sent to the model""" + size = len(json.dumps(message.content, ensure_ascii=False)) + if isinstance(message, AIMessage) and message.tool_calls: + size += len(json.dumps(message.tool_calls, ensure_ascii=False, default=str)) + return size + + +def _turn_size(turn: list[BaseMessage]) -> int: + return sum(message_size(message) for message in turn) + + +def split_turns(messages: list[BaseMessage]) -> list[list[BaseMessage]]: + """Split a conversation into turns, each starting with a message of the user""" + turns: list[list[BaseMessage]] = [] + for message in messages: + if message.type == 'human' or not turns: + turns.append([]) + turns[-1].append(message) + return turns + + +def _truncate_tool_message(message: BaseMessage, max_chars: int, can_refetch: bool = True) -> BaseMessage: + if message.type != 'tool' or not isinstance(message.content, str) or len(message.content) <= max_chars: + return message + content = truncate_text(message.content, max_chars, getattr(message, 'name', None), can_refetch) + return message.model_copy(update={'content': content}) + + +def _compact_old_turn(turn: list[BaseMessage]) -> list[BaseMessage]: + return [_truncate_tool_message(message, OLD_TOOL_OUTPUT_MAX_LENGTH) for message in turn] + + +def _share_out(sizes: list[int], total: int, floor: int, ceiling: int) -> list[int]: + """Share out a number of characters among texts of the given sizes + + What a small text does not need of its share goes to the larger ones. No text gets + less than floor, so many texts together can still exceed the total. + """ + shares = [0] * len(sizes) + remaining = total + smallest_first = sorted(range(len(sizes)), key=lambda i: sizes[i]) + for n, i in enumerate(smallest_first): + shares[i] = min(max(remaining // (len(sizes) - n), floor), ceiling) + remaining -= min(sizes[i], shares[i]) + return shares + + +def _compact_current_turn(turn: list[BaseMessage], budget_chars: int) -> list[BaseMessage]: + # A single tool result must not take up the whole budget + max_result_chars = budget_chars // 2 + capped = [_truncate_tool_message(message, max_result_chars, can_refetch=False) for message in turn] + if _turn_size(capped) <= budget_chars: + return capped + # The earlier tool results of the turn, which the model has already seen, are shortened + last_ai = max((i for i, message in enumerate(turn) if message.type == 'ai'), default=-1) + compacted = [ + _truncate_tool_message(message, OLD_TOOL_OUTPUT_MAX_LENGTH) if i < last_ai else message + for i, message in enumerate(turn) + ] + # The tool results the model has not answered yet share what is left of the budget + pending = [i for i, message in enumerate(compacted) if i > last_ai and message.type == 'tool'] + available = budget_chars - _turn_size([message for i, message in enumerate(compacted) if i not in pending]) + shares = _share_out( + [message_size(compacted[i]) for i in pending], max(available, 0), OLD_TOOL_OUTPUT_MAX_LENGTH, max_result_chars, + ) + for i, share in zip(pending, shares): + compacted[i] = _truncate_tool_message(compacted[i], share, can_refetch=False) + return compacted + + +@dataclass +class ContextSelection: + # Turns that do not fit anymore, to be summarized + dropped: list[BaseMessage] + # The rest of the conversation, unchanged, to be kept in the state + kept: list[BaseMessage] + # The rest of the conversation as it is sent to the model, with shortened tool results + model_messages: list[BaseMessage] + + +def select_context(messages: list[BaseMessage], budget_tokens: int) -> ContextSelection: + """Pick the part of the conversation that is sent to the model""" + budget_chars = max(budget_tokens, MIN_CONTEXT_BUDGET) * CHARS_PER_TOKEN + turns = split_turns(list(messages)) + if not turns: + return ContextSelection([], [], []) + + older, current = turns[:-1], turns[-1] + older_compacted = [_compact_old_turn(turn) for turn in older] + current_compacted = _compact_current_turn(current, budget_chars) + + sizes = [_turn_size(turn) for turn in older_compacted] + total = sum(sizes) + _turn_size(current_compacted) + first_kept = 0 + if total > budget_chars: + while first_kept < len(older) and total > budget_chars * TRIM_TARGET: + total -= sizes[first_kept] + first_kept += 1 + + return ContextSelection( + dropped=[message for turn in older[:first_kept] for message in turn], + kept=[message for turn in older[first_kept:] for message in turn] + current, + model_messages=[message for turn in older_compacted[first_kept:] for message in turn] + current_compacted, + ) + + +def render_transcript(messages: list[BaseMessage], max_chars: int) -> str: + """Write out a part of the conversation as plain text, for it to be summarized""" + lines = [] + for message in messages: + if message.type == 'human': + line = 'User: ' + extract_text_content(message.content) + file_ids = extract_file_ids(message.content) + if file_ids: + line += f' [attached files with the ids {", ".join(str(file_id) for file_id in file_ids)}]' + lines.append(line) + elif message.type == 'ai': + text = extract_text_content(message.content) + if text: + lines.append('Assistant: ' + text) + for tool_call in getattr(message, 'tool_calls', None) or []: + lines.append(f'Assistant called the tool {tool_call["name"]} with {json.dumps(tool_call.get("args"), ensure_ascii=False, default=str)}') + elif message.type == 'tool': + content = message.content if isinstance(message.content, str) else extract_text_content(message.content) + lines.append(f'Result of {getattr(message, "name", None) or "the tool"}: ' + truncate_text(content, OLD_TOOL_OUTPUT_MAX_LENGTH)) + return truncate_text('\n\n'.join(lines), max_chars) diff --git a/ex_app/lib/graph.py b/ex_app/lib/graph.py index 685db0d..26f2fde 100644 --- a/ex_app/lib/graph.py +++ b/ex_app/lib/graph.py @@ -10,7 +10,7 @@ from langgraph.graph.message import add_messages from langgraph.graph.state import CompiledStateGraph from langgraph.prebuilt import ToolNode, tools_condition -from typing_extensions import TypedDict, Annotated +from typing_extensions import TypedDict, Annotated, NotRequired from ex_app.lib.all_tools.lib.impulse import ( ALLOWED_CALL, @@ -39,6 +39,8 @@ class AgentState(TypedDict): # add_messages is a reducer # See https://langchain-ai.github.io/langgraph/concepts/low_level/#reducers messages: Annotated[Sequence[BaseMessage], add_messages] + # Summary of the earlier turns that no longer fit into the context window, see context_window.py + summary: NotRequired[str] class CallClassifications: diff --git a/ex_app/lib/main.py b/ex_app/lib/main.py index 2f80e98..64f9b53 100644 --- a/ex_app/lib/main.py +++ b/ex_app/lib/main.py @@ -43,6 +43,7 @@ ImpulseRadius, Reversibility, ) +from ex_app.lib.context_window import CONTEXT_BUDGET_SETTING_ID, DEFAULT_CONTEXT_BUDGET from ex_app.lib.tools import get_categories PROVIDERS = [provider, multimodal_provider] @@ -251,6 +252,18 @@ async def exapp_lifespan(app: FastAPI): type=SettingsFieldType.NUMBER, default=DEFAULT_ITEM_COUNT_THRESHOLD, ), + SettingsField( + id=CONTEXT_BUDGET_SETTING_ID, + title=_("Maximum length of the conversation sent to the language model, in tokens"), + description=_( + "When a conversation grows longer, the oldest messages are replaced by a summary and" + " long tool results are shortened. Set this well below the context window of your" + " language model, as the instructions and the tool descriptions come on top." + " Roughly four characters make one token." + ), + type=SettingsFieldType.NUMBER, + default=DEFAULT_CONTEXT_BUDGET, + ), SettingsField( id="here_api", title=_("API Key HERE"), diff --git a/ex_app/lib/nc_model.py b/ex_app/lib/nc_model.py index 6e92de7..212258c 100644 --- a/ex_app/lib/nc_model.py +++ b/ex_app/lib/nc_model.py @@ -98,9 +98,6 @@ class ChatWithNextcloud(BaseChatModel): tools: Sequence[ Union[typing.Dict[str, Any], type, Callable, BaseTool]] = [] TIMEOUT: int = 60 * 30 # 30 minutes - MAX_MESSAGE_HISTORY: int = 42 - TOOL_OUTPUT_TRUNCATE_AFTER: int = 10 - TOOL_OUTPUT_MAX_LENGTH: int = 2000 POLL_WAIT_TIME: int = 5 STREAMING_POLL_WAIT_TIME: int = 1 multimodal: bool = False @@ -115,20 +112,7 @@ def _build_task_input(self, messages: list[BaseMessage]) -> dict[str, typing.Any task_input['system_prompt'] = messages[0].content messages = messages[1:] - # Impose a history limit on non-tool messages to avoid token intake exploding. - # Tool messages don't count toward the limit so tool_call/tool_result pairs - # stay intact even in tool-heavy conversations. - non_tool_count = 0 - cutoff_idx = 0 - # idx=0 is the first message, idx=len-1 is the most recent, so we walk backward - for i in range(len(messages) - 1, -1, -1): - if messages[i].type != 'tool': - non_tool_count += 1 - if non_tool_count > self.MAX_MESSAGE_HISTORY: - cutoff_idx = i + 1 - break - messages = messages[cutoff_idx:] - + # The messages were already fitted into the context window, see context_window.py # first message cannot be a tool message while len(messages) > 0 and messages[0].type == 'tool': messages = messages[1:] @@ -157,9 +141,6 @@ def _build_task_input(self, messages: list[BaseMessage]) -> dict[str, typing.Any task_input['input_attachments'] = extract_file_ids(message.content) elif message.type == 'tool': content = message.content - age = len(messages) - 1 - i - if age > self.TOOL_OUTPUT_TRUNCATE_AFTER and isinstance(content, str) and len(content) > self.TOOL_OUTPUT_MAX_LENGTH: - content = content[:self.TOOL_OUTPUT_MAX_LENGTH] + "…[truncated]" if len(messages)-1 != i: history.append(json.dumps({"role": "tool", "content": content, "name": message.name, "tool_call_id": message.tool_call_id})) else: diff --git a/ex_app/lib/tools.py b/ex_app/lib/tools.py index bbaf7b4..f1f4748 100644 --- a/ex_app/lib/tools.py +++ b/ex_app/lib/tools.py @@ -28,6 +28,7 @@ parse_impulse_radius, parse_reversibility, ) +from ex_app.lib.context_window import CONTEXT_BUDGET_SETTING_ID, DEFAULT_CONTEXT_BUDGET, MIN_CONTEXT_BUDGET async def _get_threshold(nc: AsyncNextcloudApp, setting_id: str, default: ImpulseRadius) -> ImpulseRadius: @@ -50,6 +51,12 @@ async def _get_count(nc: AsyncNextcloudApp, setting_id: str, default: int) -> in return default +@timed_memoize(1*60) +async def get_context_budget(nc: AsyncNextcloudApp) -> int: + """The configured size of the conversation sent to the model, in tokens""" + return max(await _get_count(nc, CONTEXT_BUDGET_SETTING_ID, DEFAULT_CONTEXT_BUDGET), MIN_CONTEXT_BUDGET) + + # Cached like get_tools below, and for the same reason: every turn reads these, an # admin changes them once in a while. @timed_memoize(1*60) diff --git a/tests/test_concurrent_users.py b/tests/test_concurrent_users.py index a5eee3a..c347b51 100644 --- a/tests/test_concurrent_users.py +++ b/tests/test_concurrent_users.py @@ -32,6 +32,7 @@ from ex_app.lib import agent # noqa: E402 from ex_app.lib.all_tools.lib.impulse import ConfirmationPolicy # noqa: E402 +from ex_app.lib.context_window import DEFAULT_CONTEXT_BUDGET # noqa: E402 from ex_app.lib.nc_model import MULTIMODAL_INTERACTION, ChatWithNextcloud, Task # noqa: E402 @@ -69,6 +70,9 @@ async def get_tools(nc): async def get_confirmation_policy(nc): return ConfirmationPolicy() + async def get_context_budget(nc): + return DEFAULT_CONTEXT_BUDGET + async def user_timezone_async(nc): return None @@ -85,6 +89,7 @@ async def schedule_task(self, task_input, prefer_streaming=False): patches = [ (agent, "get_tools", get_tools), (agent, "get_confirmation_policy", get_confirmation_policy), + (agent, "get_context_budget", get_context_budget), (agent, "user_timezone_async", user_timezone_async), (ChatWithNextcloud, "_schedule_task", schedule_task), ] diff --git a/tests/test_context_window.py b/tests/test_context_window.py new file mode 100644 index 0000000..5d8e23b --- /dev/null +++ b/tests/test_context_window.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: 2026 Nextcloud GmbH and Nextcloud contributors +# SPDX-License-Identifier: AGPL-3.0-or-later +"""Fitting the conversation into the context window, see ex_app/lib/context_window.py""" +import asyncio +import json +import os +import sys +from pathlib import Path + +os.environ.setdefault("APP_ID", "context_agent") +os.environ.setdefault("APP_HOST", "127.0.0.1") +os.environ.setdefault("APP_PORT", "9081") +os.environ.setdefault("APP_SECRET", "test-secret") +os.environ.setdefault("APP_VERSION", "2.9.1") +os.environ.setdefault("NEXTCLOUD_URL", "http://127.0.0.1:8080") +os.environ.setdefault("APP_PERSISTENT_STORAGE", str(Path(__file__).parent / ".tmp_persistent")) + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, ToolMessage # noqa: E402 +from langgraph.graph.message import REMOVE_ALL_MESSAGES # noqa: E402 +from nc_py_api import AsyncNextcloudApp # noqa: E402 + +from ex_app.lib.context_window import ( # noqa: E402 + CHARS_PER_TOKEN, + MIN_CONTEXT_BUDGET, + OLD_TOOL_OUTPUT_MAX_LENGTH, + render_transcript, + select_context, + split_turns, + truncate_text, +) +from ex_app.lib.graph import get_graph # noqa: E402 +from ex_app.lib.jsonplus import JsonPlusSerializer # noqa: E402 +from ex_app.lib.memorysaver import MemorySaver # noqa: E402 + + +def tool_turn(n: int, output: str, calls: int = 1) -> list: + """A turn in which the user asks, the model calls tools and answers""" + call_ids = [f"c{n}-{i}" for i in range(calls)] + return [ + HumanMessage(content=f"question {n}", id=f"h{n}"), + AIMessage(content="", tool_calls=[{"name": "lookup", "args": {"i": i}, "id": cid} for i, cid in enumerate(call_ids)], id=f"a{n}"), + *[ToolMessage(content=output, name="lookup", tool_call_id=cid, id=f"t{cid}") for cid in call_ids], + AIMessage(content=f"answer {n}", id=f"r{n}"), + ] + + +def test_truncate_text_keeps_head_and_tail(): + text = "".join(f"line {i}\n" for i in range(2000)) + truncated = truncate_text(text, 1000, "lookup") + assert len(truncated) <= 1000 + 50 + assert truncated.startswith("line 0\n") + assert truncated.endswith("line 1999\n") + assert "characters omitted; call lookup again" in truncated + + +def test_truncate_text_leaves_short_text_alone(): + assert truncate_text("short", 100) == "short" + + +def test_truncate_json_list_stays_valid(): + data = [{"id": i, "title": f"item {i}"} for i in range(500)] + truncated = truncate_text(json.dumps(data, indent=2), 2000, "lookup") + assert len(truncated) <= 2000 + parsed = json.loads(truncated) + assert parsed[0] == data[0] + assert parsed[-1] == data[-1] + marker = next(item for item in parsed if isinstance(item, str)) + assert "items omitted" in marker + + +def test_split_turns(): + messages = tool_turn(1, "x") + tool_turn(2, "y") + turns = split_turns(messages) + assert [turn[0].id for turn in turns] == ["h1", "h2"] + assert sum(len(turn) for turn in turns) == len(messages) + + +def test_current_turn_tool_outputs_are_not_truncated(): + # Many parallel calls: the first results are far more than 10 messages from the end + big = "x" * (OLD_TOOL_OUTPUT_MAX_LENGTH * 3) + messages = tool_turn(1, big) + tool_turn(2, big, calls=12) + selection = select_context(messages, 100_000) + assert selection.dropped == [] + assert selection.kept == messages + old_tool, *current_tools = [m for m in selection.model_messages if m.type == "tool"] + assert len(old_tool.content) < len(big) + assert "omitted" in old_tool.content + assert all(m.content == big for m in current_tools) + + +def test_old_turns_are_dropped_whole_when_over_budget(): + budget = MIN_CONTEXT_BUDGET + messages = [] + for n in range(30): + messages += tool_turn(n, "y" * 1500) + selection = select_context(messages, budget) + assert selection.dropped + assert selection.dropped + selection.kept == messages + assert selection.kept[0].type == "human" + # Trimmed below the budget, leaving room for the next turns + assert sum(len(json.dumps(m.content)) for m in selection.model_messages) < budget * CHARS_PER_TOKEN + # The current turn is always kept + assert selection.kept[-len(tool_turn(0, "")):][0].id == "h29" + + +def test_current_turn_alone_over_budget(): + budget = MIN_CONTEXT_BUDGET + budget_chars = budget * CHARS_PER_TOKEN + huge = "z" * (budget_chars * 2) + messages = [ + # Long enough that halving the two results alone does not make the turn fit + HumanMessage(content="do it " + "d" * 1000, id="h"), + AIMessage(content="", tool_calls=[{"name": "lookup", "args": {}, "id": "c1"}], id="a1"), + ToolMessage(content=huge, name="lookup", tool_call_id="c1", id="t1"), + AIMessage(content="", tool_calls=[{"name": "lookup", "args": {}, "id": "c2"}], id="a2"), + ToolMessage(content=huge, name="lookup", tool_call_id="c2", id="t2"), + ] + selection = select_context(messages, budget) + assert selection.dropped == [] + earlier, latest = [m for m in selection.model_messages if m.type == "tool"] + # The result the model has not seen yet gets the most room, but no more than half the budget + assert OLD_TOOL_OUTPUT_MAX_LENGTH < len(latest.content) <= budget_chars // 2 + 200 + assert len(earlier.content) <= OLD_TOOL_OUTPUT_MAX_LENGTH + 200 + # The state keeps the full results + assert selection.kept == messages + + +def test_pending_tool_results_share_the_budget(): + budget = MIN_CONTEXT_BUDGET + budget_chars = budget * CHARS_PER_TOKEN + big = "z" * budget_chars + call_ids = [f"c{i}" for i in range(3)] + messages = [ + HumanMessage(content="do it", id="h"), + AIMessage(content="", tool_calls=[{"name": "lookup", "args": {}, "id": cid} for cid in call_ids + ["small"]], id="a"), + *[ToolMessage(content=big, name="lookup", tool_call_id=cid, id=f"t{cid}") for cid in call_ids], + ToolMessage(content="tiny", name="lookup", tool_call_id="small", id="tsmall"), + ] + selection = select_context(messages, budget) + *large, small = [m for m in selection.model_messages if m.type == "tool"] + # Together the results fit, and the room the small one does not need goes to the large ones + assert sum(len(json.dumps(m.content)) for m in selection.model_messages) <= budget_chars + assert small.content == "tiny" + assert all(len(m.content) > budget_chars // 4 for m in large) + assert selection.kept == messages + + +def test_current_turn_marker_does_not_suggest_calling_again(): + budget = MIN_CONTEXT_BUDGET + huge = "z" * (budget * CHARS_PER_TOKEN * 2) + messages = tool_turn(1, huge) + tool_turn(2, huge)[:-1] + selection = select_context(messages, budget) + old, current = [m for m in selection.model_messages if m.type == "tool"] + # Calling the tool again gets an old result back in full, but not one of the current turn + assert "call lookup again" in old.content + assert "call lookup again" not in current.content + assert "too large to be shown in full" in current.content + + +def test_render_transcript(): + transcript = render_transcript(tool_turn(1, "result " * 1000), 100_000) + assert "User: question 1" in transcript + assert 'Assistant called the tool lookup with {"i": 0}' in transcript + assert "Result of lookup:" in transcript + assert "Assistant: answer 1" in transcript + assert len(transcript) < 3000 + + +def test_graph_replaces_dropped_messages_and_keeps_summary(): + """The update call_model returns once it has summarized the dropped turns""" + # Without the final answer, which call_model is about to give + messages = (tool_turn(1, "a") + tool_turn(2, "b"))[:-1] + kept = tool_turn(2, "b")[:-1] + + async def call_model(state, config): + if len(state["messages"]) == len(messages): + return {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES), *kept, AIMessage(content="done")], "summary": "earlier: 1"} + return {"messages": [AIMessage(content="again")]} + + async def run(): + checkpointer = MemorySaver(serde=JsonPlusSerializer()) + graph, _ = await get_graph(call_model, [], checkpointer) + config = {"configurable": {"thread_id": "t"}} + await graph.ainvoke({"messages": messages}, config) + state = graph.get_state(config).values + assert state["summary"] == "earlier: 1" + assert [m.id for m in state["messages"][:-1]] == [m.id for m in kept] + assert state["messages"][-1].content == "done" + # The summary stays when later steps do not touch it + await graph.ainvoke({"messages": [HumanMessage(content="more")]}, config) + assert graph.get_state(config).values["summary"] == "earlier: 1" + + asyncio.run(run()) + + +def test_react_summarizes_dropped_turns_without_streaming_the_summary(): + from ex_app.lib import agent + from ex_app.lib.all_tools.lib.impulse import ConfirmationPolicy + from ex_app.lib.nc_model import ChatWithNextcloud, Task, extract_text_content + + scheduled = [] + + async def get_tools(nc): + return [] + + async def get_confirmation_policy(nc): + return ConfirmationPolicy() + + async def get_context_budget(nc): + return MIN_CONTEXT_BUDGET + + async def user_timezone_async(nc): + return None + + async def schedule_task(self, task_input, prefer_streaming=False): + scheduled.append(task_input) + if task_input["system_prompt"].startswith("You summarize"): + output = "SUMMARY-TEXT" + else: + output = "final answer" + return Task(id=len(scheduled), status="STATUS_SUCCESSFUL", output={"output": output}) + + async def answer(state, config): + return {"messages": [AIMessage(content="ok " + "w" * 1500)]} + + async def build_token(): + checkpointer = MemorySaver(serde=JsonPlusSerializer()) + graph, _ = await get_graph(answer, [], checkpointer) + config = {"configurable": {"thread_id": "thread-1"}} + for n in range(20): + await graph.ainvoke({"messages": [HumanMessage(content=f"question {n} " + "q" * 1500)]}, config) + return agent.export_conversation(checkpointer) + + streamed = [] + + async def stream_output(state): + streamed.append(state["output"]) + + patches = [ + (agent, "get_tools", get_tools), + (agent, "get_confirmation_policy", get_confirmation_policy), + (agent, "get_context_budget", get_context_budget), + (agent, "user_timezone_async", user_timezone_async), + (ChatWithNextcloud, "_schedule_task", schedule_task), + ] + originals = [(obj, name, getattr(obj, name)) for obj, name, _ in patches] + for obj, name, value in patches: + setattr(obj, name, value) + try: + token = asyncio.run(build_token()) + task = { + "type": "core:contextagent:interaction", + "preferStreaming": True, + "input": {"input": "latest question", "confirmation": 1, "conversation_token": token}, + } + result = asyncio.run(agent.react(task, AsyncNextcloudApp(), stream_output)) + finally: + for obj, name, value in originals: + setattr(obj, name, value) + + summary_calls = [t for t in scheduled if t["system_prompt"].startswith("You summarize")] + model_calls = [t for t in scheduled if t not in summary_calls] + assert len(summary_calls) == 1 + assert "question 0" in summary_calls[0]["input"] + assert len(model_calls) == 1 + # The summary leads the history and is kept out of the instructions + assert "SUMMARY-TEXT" not in model_calls[0]["system_prompt"] + leading = json.loads(model_calls[0]["history"][0]) + assert leading["role"] == "human" + assert leading["content"] == agent.SUMMARY_MESSAGE_PREFIX + "SUMMARY-TEXT" + assert not any("SUMMARY-TEXT" in entry for entry in model_calls[0]["history"][1:]) + assert model_calls[0]["input"] == "latest question" + assert len(model_calls[0]["history"]) < 40 + assert result["output"] == "final answer" + assert not any("SUMMARY-TEXT" in output for output in streamed) + + async def load_state(): + graph, _ = await get_graph(answer, [], agent.load_conversation(result["conversation_token"])) + return graph.get_state({"configurable": {"thread_id": "thread-1"}}).values + + state = asyncio.run(load_state()) + assert state["summary"] == "SUMMARY-TEXT" + assert state["messages"][0].type == "human" + # The summary message is not kept in the state: the history without it, the latest question and the answer + assert not any("SUMMARY-TEXT" in extract_text_content(m.content) for m in state["messages"]) + assert len(state["messages"]) == len(model_calls[0]["history"]) - 1 + 2