diff options
| -rw-r--r-- | README.md | 14 | ||||
| -rwxr-xr-x | llm.py | 74 | ||||
| -rw-r--r-- | web_utils.py | 4 |
3 files changed, 84 insertions, 8 deletions
@@ -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 @@ -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 |