From 43d2ade8e065c27e2513d91606599865a65aaa47 Mon Sep 17 00:00:00 2001 From: Bubba Date: Mon, 30 Mar 2026 12:56:20 -0400 Subject: [PATCH] feat: add OpenRouter reasoning support with multi-turn pass-back Send reasoning: {enabled: true} in the request payload for OpenRouter models so they return reasoning tokens. Preserve reasoning_details in assistant messages for multi-turn conversations. Add --no-reasoning flag to disable this when reasoning adds unwanted overhead. Localhost llama.cpp calls are unaffected (reasoning_enabled defaults to False in LLMClient, only set True via create_client for openrouter://). Co-Authored-By: Claude Opus 4.6 (1M context) --- amnesia_bench.py | 33 ++++++++++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/amnesia_bench.py b/amnesia_bench.py index 49d8129..889eee5 100644 --- a/amnesia_bench.py +++ b/amnesia_bench.py @@ -209,11 +209,12 @@ def with_exponential_backoff(fn, max_retries=20, base_delay=2.0, max_delay=120.0 class LLMClient: """Wrapper for llama.cpp or any OpenAI-compatible /v1/chat/completions endpoint.""" - def __init__(self, server_url: str = SERVER_URL, temperature: float = TEMPERATURE, api_key: str = None, model_name: str = None): + def __init__(self, server_url: str = SERVER_URL, temperature: float = TEMPERATURE, api_key: str = None, model_name: str = None, reasoning_enabled: bool = False): self.server_url = server_url.rstrip("/") self.temperature = temperature self.model_name = model_name # passed to API as model field (required by OpenRouter) self.auth_header = {"Authorization": f"Bearer {api_key}"} if api_key else {} + self.reasoning_enabled = reasoning_enabled # send reasoning: {enabled: true} for OpenRouter def generate(self, messages: list[dict], max_tokens: int) -> dict: """ @@ -227,6 +228,8 @@ def generate(self, messages: list[dict], max_tokens: int) -> dict: "temperature": self.temperature, "stream": True, } + if self.reasoning_enabled: + payload["reasoning"] = {"enabled": True} def _do_request(): if self.model_name: @@ -293,9 +296,15 @@ def _do_request(): else: full_content = content + # Build reasoning_details for multi-turn pass-back (OpenRouter format) + reasoning_details = None + if reasoning and self.reasoning_enabled: + reasoning_details = [{"type": "thinking", "thinking": reasoning}] + return { "content": full_content, "reasoning_content": reasoning, + "reasoning_details": reasoning_details, "final_content": content, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, @@ -423,6 +432,7 @@ def _do_request(): return { "content": content, "reasoning_content": "", + "reasoning_details": None, "final_content": content, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, @@ -449,6 +459,7 @@ def create_client( api_key: str = None, model_name: str = None, temperature: float = TEMPERATURE, + reasoning_enabled: bool = True, ) -> Union[LLMClient, GeminiClient]: """ Create appropriate client based on server_url scheme. @@ -482,6 +493,7 @@ def create_client( temperature=temperature, api_key=api_key, model_name=or_model, + reasoning_enabled=reasoning_enabled, ) elif server_url.startswith("http"): return LLMClient(server_url=server_url, temperature=temperature) @@ -946,6 +958,7 @@ def run_trial( finish = "truncated" content = resp["content"] + reasoning_details = resp.get("reasoning_details") total_now = resp["total_tokens"] peak_tokens = max(peak_tokens, total_now) @@ -1007,11 +1020,17 @@ def run_trial( code_output=combined_output, ) conversation.append(code_turn) - messages.append({"role": "assistant", "content": content}) + asst_msg = {"role": "assistant", "content": content} + if reasoning_details: + asst_msg["reasoning_details"] = reasoning_details + messages.append(asst_msg) messages.append({"role": "user", "content": f"Code output:\n{combined_output}"}) continue - messages.append({"role": "assistant", "content": content}) + asst_msg = {"role": "assistant", "content": content} + if reasoning_details: + asst_msg["reasoning_details"] = reasoning_details + messages.append(asst_msg) messages.append({"role": "user", "content": "Continue solving."}) conversation.append(Turn(role="user", content="Continue solving.")) @@ -1341,6 +1360,7 @@ def run_all_models( trials: int = TRIALS_PER_WINDOW, temperature: float = TEMPERATURE, cli_api_key: str = None, + reasoning_enabled: bool = True, ): """ Iterate over all models in models.json, run all problems for each model. @@ -1379,6 +1399,7 @@ def run_all_models( api_key=api_key, model_name=mname, temperature=temperature, + reasoning_enabled=reasoning_enabled, ) except ValueError as e: print(f" ERROR: Could not create client for {mname}: {e} — skipping") @@ -1481,6 +1502,8 @@ def main(): parser.add_argument("--temperature", type=float, default=TEMPERATURE) parser.add_argument("--config", type=str, default=None, help="Run specific config only: NoTIR_HardCut, TIR_HardCut, NoTIR_Compact, TIR_Compact") + parser.add_argument("--no-reasoning", action="store_true", + help="Disable reasoning tokens for OpenRouter models (saves cost/latency)") parser.add_argument("--results-dir", type=str, default=None, help="Results directory for --scores / --analyze (default: ./results)") @@ -1536,6 +1559,8 @@ def main(): else: api_key = os.environ.get("GEMINI_API_KEY") + reasoning_enabled = not args.no_reasoning + # Multi-model mode if args.run_all_models: run_all_models( @@ -1546,6 +1571,7 @@ def main(): trials=trials_per_window, temperature=args.temperature, cli_api_key=api_key, + reasoning_enabled=reasoning_enabled, ) return @@ -1559,6 +1585,7 @@ def main(): api_key=api_key, model_name=model_name, temperature=args.temperature, + reasoning_enabled=reasoning_enabled, ) except ValueError as e: print(f"ERROR: {e}")