From 1cd5230cd3975b50119228a9d9b121ab33bcb331 Mon Sep 17 00:00:00 2001 From: Maxime Brugel Date: Sun, 4 Jan 2026 16:07:31 +0100 Subject: [PATCH 1/3] fix: final var --- rlm/core/rlm.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index ca3f8ae5b..5f2584504 100644 --- a/rlm/core/rlm.py +++ b/rlm/core/rlm.py @@ -180,7 +180,19 @@ def completion( ) # Check if RLM is done and has a final answer. - final_answer = find_final_answer(iteration.response) + final_answer_match = find_final_answer(iteration.response) + final_answer: str | None = None + if final_answer_match is not None: + answer_type, content = final_answer_match + if answer_type == "FINAL": + final_answer = content + elif answer_type == "FINAL_VAR": + variable_name = content.strip().strip('"').strip("'") + result = environment.execute_code(f"print(FINAL_VAR({variable_name!r}))") + final_answer = result.stdout.strip() + if final_answer == "": + final_answer = result.stderr.strip() or "" + iteration.final_answer = final_answer # If logger is used, log the iteration. @@ -190,7 +202,7 @@ def completion( # Verbose output for this iteration self.verbose.print_iteration(iteration, i + 1) - if final_answer: + if final_answer is not None: time_end = time.perf_counter() usage = lm_handler.get_usage_summary() self.verbose.print_final_answer(final_answer) From 39799f0e5a87d112f1e1049774e986565920f96d Mon Sep 17 00:00:00 2001 From: Maxime Brugel Date: Sun, 4 Jan 2026 16:20:24 +0100 Subject: [PATCH 2/3] fix: types parsing --- rlm/core/types.py | 36 ++++++++++++++++++++++++++++-------- rlm/utils/parsing.py | 2 +- 2 files changed, 29 insertions(+), 9 deletions(-) diff --git a/rlm/core/types.py b/rlm/core/types.py index ad1676284..585ad3c2b 100644 --- a/rlm/core/types.py +++ b/rlm/core/types.py @@ -8,11 +8,11 @@ def _serialize_value(value: Any) -> Any: """Convert a value to a JSON-serializable representation.""" - if value is None or isinstance(value, bool | int | float | str): + if value is None or isinstance(value, (bool, int, float, str)): return value if isinstance(value, ModuleType): return f"" - if isinstance(value, list | tuple): + if isinstance(value, (list, tuple)): return [_serialize_value(v) for v in value] if isinstance(value, dict): return {str(k): _serialize_value(v) for k, v in value.items()} @@ -213,20 +213,40 @@ class QueryMetadata: context_total_length: int context_type: str - def __init__(self, prompt: str | list[str] | dict[Any, str] | list[dict[Any, str]]): + def __init__(self, prompt: str | list[str] | dict[Any, Any] | list[dict[Any, Any]]): if isinstance(prompt, str): self.context_lengths = [len(prompt)] self.context_type = "str" - elif isinstance(prompt, dict[Any, str]): - self.context_lengths = [len(chunk) for chunk in prompt.values()] + elif isinstance(prompt, dict): + self.context_type = "dict" + self.context_lengths = [] + for chunk in prompt.values(): + if isinstance(chunk, str): + self.context_lengths.append(len(chunk)) + continue + try: + import json + + self.context_lengths.append(len(json.dumps(chunk, default=str))) + except Exception: + self.context_lengths.append(len(repr(chunk))) self.context_type = "dict" elif isinstance(prompt, list): self.context_type = "list" - if isinstance(prompt[0], dict): + if len(prompt) == 0: + self.context_lengths = [0] + elif isinstance(prompt[0], dict): if "content" in prompt[0]: - self.context_lengths = [len(chunk["content"]) for chunk in prompt] + self.context_lengths = [len(str(chunk.get("content", ""))) for chunk in prompt] else: - self.context_lengths = [len(chunk) for chunk in prompt] + self.context_lengths = [] + for chunk in prompt: + try: + import json + + self.context_lengths.append(len(json.dumps(chunk, default=str))) + except Exception: + self.context_lengths.append(len(repr(chunk))) else: self.context_lengths = [len(chunk) for chunk in prompt] else: diff --git a/rlm/utils/parsing.py b/rlm/utils/parsing.py index 0f871cf7a..94a7d4931 100644 --- a/rlm/utils/parsing.py +++ b/rlm/utils/parsing.py @@ -106,7 +106,7 @@ def format_execution_result(result: REPLResult) -> str: "__doc__", ]: # Only show simple types or short representations - if isinstance(value, str | int | float | bool | list | dict | tuple): + if isinstance(value, (str, int, float, bool, list, dict, tuple)): important_vars[key] = "" if important_vars: From b6221367889a873d7c42b853960960422ebb7220 Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Mon, 5 Jan 2026 17:16:12 -0500 Subject: [PATCH 3/3] refactored final var logic, and added unit test --- rlm/core/rlm.py | 14 +---- rlm/utils/parsing.py | 59 ++++++++---------- tests/test_parsing.py | 137 ++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 159 insertions(+), 51 deletions(-) diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index 5b9ad2d11..55a3cf58d 100644 --- a/rlm/core/rlm.py +++ b/rlm/core/rlm.py @@ -186,19 +186,7 @@ def completion( ) # Check if RLM is done and has a final answer. - final_answer_match = find_final_answer(iteration.response) - final_answer: str | None = None - if final_answer_match is not None: - answer_type, content = final_answer_match - if answer_type == "FINAL": - final_answer = content - elif answer_type == "FINAL_VAR": - variable_name = content.strip().strip('"').strip("'") - result = environment.execute_code(f"print(FINAL_VAR({variable_name!r}))") - final_answer = result.stdout.strip() - if final_answer == "": - final_answer = result.stderr.strip() or "" - + final_answer = find_final_answer(iteration.response, environment=environment) iteration.final_answer = final_answer # If logger is used, log the iteration. diff --git a/rlm/utils/parsing.py b/rlm/utils/parsing.py index 94a7d4931..e4c235068 100644 --- a/rlm/utils/parsing.py +++ b/rlm/utils/parsing.py @@ -3,9 +3,13 @@ """ import re +from typing import TYPE_CHECKING from rlm.core.types import REPLResult, RLMIteration +if TYPE_CHECKING: + from rlm.environments.base_env import BaseEnv + def find_code_blocks(text: str) -> list[str]: """ @@ -22,22 +26,38 @@ def find_code_blocks(text: str) -> list[str]: return results -def find_final_answer(text: str) -> tuple[str, str] | None: +def find_final_answer(text: str, environment: "BaseEnv | None" = None) -> str | None: """ - Find FINAL(...) or FINAL_VAR(...) statement in response and return (type, content). + Find FINAL(...) or FINAL_VAR(...) statement in response and return the final answer string. + + If FINAL_VAR is found and an environment is provided, executes code to retrieve the variable value. Returns None if neither pattern is found. + + Args: + text: The response text to parse + environment: Optional environment to execute code for FINAL_VAR retrieval + + Returns: + The final answer string, or None if no final answer pattern is found """ # Check for FINAL_VAR pattern first - must be at start of line final_var_pattern = r"^\s*FINAL_VAR\((.*?)\)" match = re.search(final_var_pattern, text, re.MULTILINE | re.DOTALL) if match: - return ("FINAL_VAR", match.group(1).strip()) + variable_name = match.group(1).strip().strip('"').strip("'") + if environment is not None: + result = environment.execute_code(f"print(FINAL_VAR({variable_name!r}))") + final_answer = result.stdout.strip() + if final_answer == "": + final_answer = result.stderr.strip() or "" + return final_answer + return None # Check for FINAL pattern - must be at start of line final_pattern = r"^\s*FINAL\((.*?)\)" match = re.search(final_pattern, text, re.MULTILINE | re.DOTALL) if match: - return ("FINAL", match.group(1).strip()) + return match.group(1).strip() return None @@ -117,35 +137,8 @@ def format_execution_result(result: REPLResult) -> str: def check_for_final_answer(response: str, repl_env, logger) -> str | None: """Check if response contains a final answer.""" - result = find_final_answer(response) - if result is None: - return None - - answer_type, content = result - - if answer_type == "FINAL": - return content - elif answer_type == "FINAL_VAR": - # Get the variable directly from the REPL environment - try: - # Strip spaces, quotes, and newlines from variable name - variable_name = content.strip().strip('"').strip("'").strip("\n").strip("\r") - - # Check if variable exists in the REPL environment's locals - if variable_name in repl_env.locals: - variable_value = repl_env.locals[variable_name] - return str(variable_value) - else: - error_msg = f"Variable '{variable_name}' not found in REPL environment" - logger.log_tool_execution("FINAL_VAR", error_msg) - return None - except Exception as e: - error_msg = f"Error retrieving variable '{variable_name}': {str(e)}" - print("ERROR MESSAGE", error_msg) - logger.log_tool_execution("FINAL_VAR", error_msg) - return None - - return None + # Use the new find_final_answer function which handles both FINAL and FINAL_VAR + return find_final_answer(response, environment=repl_env) def convert_context_for_repl(context): diff --git a/tests/test_parsing.py b/tests/test_parsing.py index 1871e3382..a42065aa1 100644 --- a/tests/test_parsing.py +++ b/tests/test_parsing.py @@ -1,6 +1,9 @@ """Tests for parsing utilities.""" +from unittest.mock import Mock + from rlm.core.types import CodeBlock, REPLResult, RLMIteration +from rlm.environments.local_repl import LocalREPL from rlm.utils.parsing import ( convert_context_for_repl, find_code_blocks, @@ -81,12 +84,25 @@ class TestFindFinalAnswer: def test_final_answer(self): text = "The answer is:\nFINAL(42)" result = find_final_answer(text) - assert result == ("FINAL", "42") + assert result == "42" def test_final_var_answer(self): text = "Check the variable:\nFINAL_VAR(result)" + # Create a mock environment that returns the variable value + mock_env = Mock() + mock_env.execute_code.return_value = REPLResult(stdout="42", stderr="", locals={}) + result = find_final_answer(text, environment=mock_env) + assert result == "42" + # Verify execute_code was called with the correct code + mock_env.execute_code.assert_called_once() + call_args = mock_env.execute_code.call_args[0][0] + assert "FINAL_VAR('result')" in call_args or 'FINAL_VAR("result")' in call_args + + def test_final_var_without_environment(self): + text = "Check the variable:\nFINAL_VAR(result)" + # Without environment, FINAL_VAR should return None result = find_final_answer(text) - assert result == ("FINAL_VAR", "result") + assert result is None def test_no_final_answer(self): text = "Still working on the problem..." @@ -98,8 +114,8 @@ def test_final_with_multiline_content(self): multiline answer)""" result = find_final_answer(text) assert result is not None - assert result[0] == "FINAL" - assert "multiline" in result[1] + assert "multiline" in result + assert "This is a" in result def test_final_must_be_at_start_of_line(self): # FINAL not at start of line should not match @@ -110,7 +126,118 @@ def test_final_must_be_at_start_of_line(self): def test_final_with_whitespace(self): text = " FINAL(answer with spaces)" result = find_final_answer(text) - assert result == ("FINAL", "answer with spaces") + assert result == "answer with spaces" + + def test_final_and_final_var_parsing(self): + """Test that both FINAL and FINAL_VAR patterns are parsed correctly.""" + # Test FINAL with various content types + test_cases_final = [ + ("FINAL(42)", "42"), + ("FINAL('hello world')", "'hello world'"), + ('FINAL("test")', '"test"'), + ("FINAL(123.45)", "123.45"), + ("FINAL([1, 2, 3])", "[1, 2, 3]"), + ] + + for text, expected in test_cases_final: + result = find_final_answer(text) + assert result == expected, f"Failed for text: {text}" + + # Test FINAL_VAR with environment + mock_env = Mock() + mock_env.execute_code.return_value = REPLResult( + stdout="computed_result", stderr="", locals={} + ) + + test_cases_final_var = [ + ("FINAL_VAR(result)", "result"), + ("FINAL_VAR('my_var')", "my_var"), + ('FINAL_VAR("another_var")', "another_var"), + ("FINAL_VAR(answer)", "answer"), + ] + + for text, var_name in test_cases_final_var: + result = find_final_answer(text, environment=mock_env) + assert result == "computed_result", f"Failed for text: {text}" + # Verify the variable name was extracted correctly + call_args = mock_env.execute_code.call_args[0][0] + assert ( + var_name in call_args + or f"'{var_name}'" in call_args + or f'"{var_name}"' in call_args + ) + + def test_final_var_takes_precedence_over_final(self): + """Test that FINAL_VAR is checked first and takes precedence over FINAL.""" + mock_env = Mock() + mock_env.execute_code.return_value = REPLResult(stdout="var_value", stderr="", locals={}) + + # If both appear, FINAL_VAR should be found first (checked first in the function) + text = "FINAL_VAR(result)\nFINAL(direct_answer)" + result = find_final_answer(text, environment=mock_env) + assert result == "var_value" # FINAL_VAR should be returned + + # Without environment, FINAL_VAR pattern is found but returns None (no fallback to FINAL) + # This is expected behavior - FINAL_VAR takes precedence when found + result = find_final_answer(text) + assert result is None # FINAL_VAR found but no environment, so returns None + + # Test that FINAL alone works when FINAL_VAR is not present + text_final_only = "FINAL(direct_answer)" + result = find_final_answer(text_final_only) + assert result == "direct_answer" + + def test_final_var_retrieves_actual_variables_from_environment(self): + """Test that FINAL_VAR actually retrieves variables from a real code environment.""" + # Create a real LocalREPL environment + env = LocalREPL() + + try: + # Execute code to create variables with different types + env.execute_code("x = 42") + env.execute_code("result = 'hello world'") + env.execute_code("answer = [1, 2, 3, 4, 5]") + env.execute_code("computed = 10 * 5 + 2") + env.execute_code("nested = {'key': 'value', 'num': 123}") + + # Test retrieving integer variable + text = "FINAL_VAR(x)" + result = find_final_answer(text, environment=env) + assert result == "42", f"Expected '42', got '{result}'" + + # Test retrieving string variable + text = "FINAL_VAR(result)" + result = find_final_answer(text, environment=env) + assert result == "hello world", f"Expected 'hello world', got '{result}'" + + # Test retrieving list variable + text = "FINAL_VAR(answer)" + result = find_final_answer(text, environment=env) + assert result == "[1, 2, 3, 4, 5]", f"Expected '[1, 2, 3, 4, 5]', got '{result}'" + + # Test retrieving computed variable + text = "FINAL_VAR(computed)" + result = find_final_answer(text, environment=env) + assert result == "52", f"Expected '52', got '{result}'" + + # Test retrieving dictionary variable + text = "FINAL_VAR(nested)" + result = find_final_answer(text, environment=env) + assert "'key': 'value'" in result or '"key": "value"' in result + assert "'num': 123" in result or '"num": 123' in result + + # Test that variable updates are reflected + env.execute_code("x = 100") + text = "FINAL_VAR(x)" + result = find_final_answer(text, environment=env) + assert result == "100", f"Expected '100', got '{result}'" + + # Test that non-existent variable returns error message + text = "FINAL_VAR(nonexistent)" + result = find_final_answer(text, environment=env) + assert "Error" in result or "not found" in result.lower() + finally: + env.cleanup() class TestFormatExecutionResult: