Cognition Management in Thalaplan

Overview

Cognition management refers to the orchestrator's ability to cherry-pick which information to include or exclude from each LLM call (LLM_IN), in order to optimise the following objective:

Include the least possible of what is not needed; include the maximum possible of what is needed — all within input token limits.

This document tracks goals, current state, design decisions, and ideas related to cognition management in the Thalaplan pipeline.

Long-Term Goals

Goal 1 — Include the right content representation for loaded documents

When a document is loaded, the Thalaplan orchestrator currently always supplies the same representation to every agent. However, different agents have different needs:

Target behaviour: supply enhanced_text when the agent may need to edit the document; supply plain text otherwise.

Goal 2 — Include only the necessary loaded documents (or parts of them)

When multiple documents are loaded, an agent may only need: - A subset of the loaded documents - Specific chapters of a document, not the full content - A summary of a document rather than its full text

Target behaviour: dynamically select which documents (and at which granularity: full / chapter / summary) to include per agent call, guided by the task context and token budget.

Status: deferred — not implementing now.

Goal 3 — Dynamic token-budget-aware context assembly

As the number of loaded documents, RAG fragments, agent instructions, and conversation history grows, the total LLM_IN payload may exceed model limits. The orchestrator should: - Estimate token counts per context section - Prioritise sections by relevance to the current agent's task - Truncate, compress, or omit low-priority sections to stay within limits

Status: future — requires token-estimation infrastructure.

Current Situation

How Loaded Documents Enter the LLM Context

  1. Contextualizer goal (early pipeline step): calls _add_loaded_documents_to_goal_summary(), which in turn calls _get_loaded_documents_string(). This builds a <LOADED_DOCUMENTS> XML block and stores it inside goal.goal_summary — a pre-baked string that is computed once at context initialisation.

  2. build_messages() (called per agent): retrieves the pre-built goal_summary via get_summary() and inserts it into the LLM message list. It does not query Neo4j again for loaded documents.

  3. Result: every agent in the pipeline receives the exact same loaded-document representation that was frozen by the Contextualizer.

The Wrong Field Was Being Used

_get_loaded_documents_string() already calls load_documents_or_chapters(..., load_enhanced_text=True), meaning Neo4j returns both text and enhanced_text. However, the method was then selecting the text field for the output block:

# Before fix — plain text, no element IDs
{loaded_content['text']}

This meant that even agents like AgentToolPreparer, which need to generate DocumentEditor.SURROUND_WITH_TAG calls with element IDs, received content with no IDs at all. The LLM correctly identified that it could not proceed without them and returned "OK" (no tool call) instead of the desired edit operation.

Root-Cause Evidence

During diagnosis of Use Case 2 ("make character names bold and blue"), the full LLM_IN for goal 1.8.2.1 (AgentToolPreparer) was captured and inspected. The loaded document appeared as plain HTML without any id attributes on its elements. The model's reasoning trace explicitly stated:

"I realize I can't directly apply formatting without knowing specific element IDs… I'll explore potential strategies…"

Fix Applied (Goal 1 — first step)

The _get_loaded_documents_string() method now uses enhanced_text unconditionally, giving all agents (including those that only read) access to element IDs. This is a conservative first step: it increases context size slightly but unblocks document editing entirely.

# After fix — HTML with element IDs
{loaded_content.get('enhanced_text') or loaded_content['text']}

The fallback to text guards against documents that have not yet had enhanced_text generated.

Design Tension: Static Pre-baking vs. Dynamic Selection

The Contextualizer's pre-baking approach was designed for efficiency (Neo4j is queried once, result cached). However, it conflicts with dynamic cognition management because:

A future redesign should move loaded-document assembly out of the Contextualizer and into build_messages() (or a dedicated per-agent context builder), so that each agent can request exactly the content it needs. The pre-baking approach can be retained as a fast-path default.

Ideas