fix: resolve FINAL_VAR answers + improve Python 3.11 compatibility for structured prompts (#8)

* fix: final var

* fix: types parsing

* refactored final var logic, and added unit test

---------

Co-authored-by: Alex Zhang <alex.lx.zhang@gmail.com>

Author: maximebrugel

rlm/core/rlm.py

@@ -186,7 +186,7 @@ def completion(
                 )
 
                 # Check if RLM is done and has a final answer.
-                final_answer = find_final_answer(iteration.response)
+                final_answer = find_final_answer(iteration.response, environment=environment)
                 iteration.final_answer = final_answer
 
                 # If logger is used, log the iteration.
@@ -196,7 +196,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)

rlm/core/types.py

@@ -10,11 +10,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"<module '{value.__name__}'>"
-    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()}
@@ -215,20 +215,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:

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
 
@@ -106,7 +126,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:
@@ -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):

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: