diff options
Diffstat (limited to 'llm.py')
| -rwxr-xr-x | llm.py | 74 |
1 files changed, 69 insertions, 5 deletions
@@ -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, ), ) |