Update metadata_search/research_synthesis.py

metadata_search/research_synthesis.py

@@ -214,40 +214,40 @@ def _research_generate_answer(
         max_new_tokens,
     )
 
-# _generate_answer normally calls an imported agent_answer function.
-# Clone agent_answer and its run_agent dependency so the research-only
-# token override reaches the actual /agent_run call.
-original_agent_answer = generate_globals.get("agent_answer")
-
-if isinstance(original_agent_answer, types.FunctionType):
-    agent_globals = _override_token_globals(
-        original_agent_answer.__globals__,
-        max_new_tokens,
-    )
-
-    # agent_answer() calls run_agent(), and run_agent() is where
-    # WATERSCOPE_MAX_TOKENS is actually read.
-    original_run_agent = agent_globals.get("run_agent")
-
-    if isinstance(original_run_agent, types.FunctionType):
-        run_agent_globals = _override_token_globals(
-            original_run_agent.__globals__,
+    # _generate_answer normally calls an imported agent_answer function.
+    # Clone agent_answer and its run_agent dependency so the research-only
+    # token override reaches the actual /agent_run call.
+    original_agent_answer = generate_globals.get("agent_answer")
+    
+    if isinstance(original_agent_answer, types.FunctionType):
+        agent_globals = _override_token_globals(
+            original_agent_answer.__globals__,
             max_new_tokens,
         )
-
-        cloned_run_agent = _copy_function(
-            original_run_agent,
-            run_agent_globals,
+    
+        # agent_answer() calls run_agent(), and run_agent() is where
+        # WATERSCOPE_MAX_TOKENS is actually read.
+        original_run_agent = agent_globals.get("run_agent")
+    
+        if isinstance(original_run_agent, types.FunctionType):
+            run_agent_globals = _override_token_globals(
+                original_run_agent.__globals__,
+                max_new_tokens,
+            )
+    
+            cloned_run_agent = _copy_function(
+                original_run_agent,
+                run_agent_globals,
+            )
+    
+            agent_globals["run_agent"] = cloned_run_agent
+    
+        cloned_agent_answer = _copy_function(
+            original_agent_answer,
+            agent_globals,
         )
-
-        agent_globals["run_agent"] = cloned_run_agent
-
-    cloned_agent_answer = _copy_function(
-        original_agent_answer,
-        agent_globals,
-    )
-
-    generate_globals["agent_answer"] = cloned_agent_answer
+    
+        generate_globals["agent_answer"] = cloned_agent_answer
     
 
     cloned_generate = _copy_function(