feat: history compactingHEADmain

3 files changed, 84 insertions, 8 deletions

diff --git a/README.md b/README.md
index 9c30af7..0293439 100644
--- a/README.md
+++ b/README.md
@@ -14,6 +14,7 @@ This project implements a simple autonomous coding agent that can interact with
   - `mypy`: Run static type checking on a Python file.
   - `webgrep`: Search for keywords within the markdownified content of a URL.
 - **Reasoning Loop**: Uses a recursive loop to process tasks, analyze results, and determine the next step until a final response is reached.
+- **Context Management**: Estimates the token count of the conversation and automatically squashes (summarizes) older history when it approaches the model's context size limit, keeping the most recent messages intact.
 - **Logging**: Logs interactions to `thoughts.log` for debugging and analysis.
 
 ## Requirements
@@ -30,6 +31,12 @@ This project implements a simple autonomous coding agent that can interact with
    ```bash
    export LLM_API_KEY="your_api_key_here"
    ```
+   Optionally, you can also configure the endpoint, model, and context size via environment variables:
+   ```bash
+   export LLM_API_ENDPOINT="https://llm-gateway.k8s.aalto.fi/api/v1"
+   export LLM_MODEL="zai-org/GLM-5.3-Flash"
+   export LLM_CONTEXT_SIZE=1048576
+   ```
 
 2. **Run the Agent**:
    Execute the script using Python:
@@ -43,9 +50,11 @@ This project implements a simple autonomous coding agent that can interact with
   ```bash
   python3 llm.py --workdir /path/to/your/project
   ```
-- `--endpoint URL`: The API endpoint for the LLM (defaults to `https://llm-gateway.k8s.aalto.fi/api/v1`).
+- `--endpoint URL`: The API endpoint for the LLM (defaults to `LLM_API_ENDPOINT`, or `https://llm-gateway.k8s.aalto.fi/api/v1`).
 - `--key KEY`: The API key for the LLM endpoint (defaults to `LLM_API_KEY` environment variable).
-- `--model MODEL`: The model name to use (defaults to `RedHatAI/gemma-4-31B-it-FP8-Dynamic`).
+- `--model MODEL`: The model name to use (defaults to `LLM_MODEL`, or `zai-org/GLM-5.3-Flash`).
+- `--min-history N`: Minimum amount of most recent messages to keep intact when summarising the history to keep the context short enough (default: 3).
+- `--model-context-size N`: Context size supported by the model (defaults to `LLM_CONTEXT_SIZE`, or 1048576).
 - `--trust-level LEVEL`: The trust level of the agent (defaults to 2). Determines which actions require manual approval.
 
 ## How it Works
@@ -57,6 +66,7 @@ This project implements a simple autonomous coding agent that can interact with
    - The script executes the action (interacting with the file system, shell, or web).
    - The result of the action is fed back into the LLM as a new observation.
    - This repeats until the agent provides a final natural language response.
+3. **History Squashing**: When the estimated token count of the conversation exceeds a threshold of the model's context size, the older history is summarized by the LLM into a compact "Previous history summary" message, while the most recent `--min-history` messages are preserved.
 
 ## Security and Permissions
 
diff --git a/llm.py b/llm.py
index 7577f9f..ca61b48 100755
--- a/llm.py
+++ b/llm.py
@@ -42,7 +42,19 @@ class AgentConfig:
     search_padding: int = 200
 
 
-def mk_direct_query(client: OpenAI, model_name: str) -> QueryMethod:
+def estimate_tokens(state: DiscussionState) -> int:
+    """Simple heuristic to estimate tokens in the conversation state."""
+    text = "".join(str(m["content"]) for m in state)
+    return len(text) // 3.5
+
+
+def mk_direct_query(
+    client: OpenAI,
+    model_name: str,
+    max_token_count: int = 27500,
+    squash_threshold: float = 0.75,
+    squasher: Optional[Callable[[DiscussionState], DiscussionState]] = None,
+) -> QueryMethod:
     def direct_query(
         query: str,
         state: DiscussionState,
@@ -52,6 +64,13 @@ def mk_direct_query(client: OpenAI, model_name: str) -> QueryMethod:
 
         state = state + ({"role": role, "content": query},)
 
+        if (
+            squasher is not None
+            and estimate_tokens(state) > max_token_count * squash_threshold
+        ):
+            logger.info("Squashing history")
+            state = squasher(state)
+
         res = client.responses.create(model=model_name, input=list(state))
 
         logger.info(f"receive: {res.output_text}")
@@ -90,7 +109,12 @@ def query_agent(query: str, config: AgentConfig, state=()):
         except (json.JSONDecodeError, KeyError, TypeError):
             continue
 
-        f = first(config.actions, lambda f: f.name == action_name and all(k in action for k in f.required_keys))
+        f = first(
+            config.actions,
+            lambda f: (
+                f.name == action_name and all(k in action for k in f.required_keys)
+            ),
+        )
         if f is not None:
             desc = f.formatter(action)
             if (
@@ -142,8 +166,20 @@ def main():
     parser.add_argument(
         "--model",
         type=str,
-        default="RedHatAI/gemma-4-31B-it-FP8-Dynamic",
-        help="The model name to use",
+        default=os.environ.get("LLM_MODEL") or "zai-org/GLM-5.3-Flash",
+        help="The model to use",
+    )
+    parser.add_argument(
+        "--min-history",
+        type=int,
+        default=3,
+        help="Minimum amount of most recent messages to keep intact when summarising the history to keep the context short enough",
+    )
+    parser.add_argument(
+        "--model-context-size",
+        type=int,
+        default=os.environ.get("LLM_CONTEXT_SIZE") or 1048576,
+        help="Context size supported by the model.",
     )
     parser.add_argument(
         "--trust-level",
@@ -153,6 +189,31 @@ def main():
     )
     args = parser.parse_args()
 
+    client = OpenAI(base_url=args.endpoint, api_key=args.key)
+
+    # We need a version of direct_query that doesn't trigger squashing to use inside the squasher
+    # to avoid infinite recursion.
+    raw_direct_query = mk_direct_query(client, args.model, squasher=None)
+
+    def squash_history(state: DiscussionState) -> DiscussionState:
+        if len(state) <= 1 + args.min_history:
+            return (
+                state[0],
+                {"role": "system", "content": "(history removed)"},
+                state[-1],
+            )
+        else:
+            (summary, _) = raw_direct_query(
+                "Please summarize the preceding conversation history concisely, preserving all key technical decisions, progress made, and current goals.",
+                state[1 : -args.min_history],
+                role="system",
+            )
+            summary_msg: Message = {
+                "role": "system",
+                "content": f"Previous history summary: {summary}",
+            }
+            return (state[0], summary_msg) + state[-args.min_history :]
+
     config = AgentConfig(
         workdir=args.workdir,
         actions=[ls, read, write, tldr, shell, mypy, webgrep],
@@ -167,7 +228,10 @@ def main():
         },
         trustlevel=args.trust_level,
         query_method=mk_direct_query(
-            OpenAI(base_url=args.endpoint, api_key=args.key), args.model
+            client,
+            args.model,
+            squasher=squash_history,
+            max_token_count=args.model_context_size,
         ),
     )
 
diff --git a/web_utils.py b/web_utils.py
index 6fecb47..89c6632 100644
--- a/web_utils.py
+++ b/web_utils.py
@@ -50,7 +50,9 @@ def fetch_and_markdownify(url: str) -> FetchResult:
         return _cache[url]
 
     try:
-        with httpx2.Client(headers={"User-Agent": "Mozilla/5.0"}, timeout=10.0) as client:
+        with httpx2.Client(
+            headers={"User-Agent": "Mozilla/5.0"}, timeout=10.0
+        ) as client:
             response = client.get(url)
             response.raise_for_status()
             html = response.text