refactor: remove build_context_code, keep only DataFrame-specific codegen

Remove the unified build_context_code() which handled strings, dicts,
and DataFrames with fragile escaping. Keep only build_dataframe_context_code()
for DataFrame-specific Parquet-over-base64 codegen. Restore original
per-environment string/dict handling in Modal, Prime, and Daytona REPLs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: kmad

rlm/environments/daytona_repl.py

@@ -610,9 +610,23 @@ def _handle_llm_request(self, req_data: dict) -> dict:
 
     def load_context(self, context_payload: ContextPayload):
         """Load context into the sandbox environment."""
-        from rlm.utils.dataframe_utils import build_context_code
+        from rlm.utils.dataframe_utils import (
+            build_dataframe_context_code,
+            dataframe_to_parquet_b64,
+            get_dataframe_type,
+        )
 
-        self.execute_code(build_context_code(context_payload))
+        df_type = get_dataframe_type(context_payload)
+        if df_type is not None:
+            parquet_b64, df_type = dataframe_to_parquet_b64(context_payload)
+            self.execute_code(build_dataframe_context_code(parquet_b64, df_type))
+        elif isinstance(context_payload, str):
+            escaped = context_payload.replace("\\", "\\\\").replace('"""', '\\"\\"\\"')
+            self.execute_code(f'context = """{escaped}"""')
+        else:
+            context_json = json.dumps(context_payload)
+            escaped_json = context_json.replace("\\", "\\\\").replace("'", "\\'")
+            self.execute_code(f"import json; context = json.loads('{escaped_json}')")
 
     def execute_code(self, code: str) -> REPLResult:
         """Execute code in the Daytona sandbox and return result."""

rlm/environments/modal_repl.py

@@ -438,9 +438,23 @@ def _handle_llm_request(self, req_data: dict) -> dict:
 
     def load_context(self, context_payload: ContextPayload):
         """Load context into the sandbox environment."""
-        from rlm.utils.dataframe_utils import build_context_code
+        from rlm.utils.dataframe_utils import (
+            build_dataframe_context_code,
+            dataframe_to_parquet_b64,
+            get_dataframe_type,
+        )
 
-        self.execute_code(build_context_code(context_payload))
+        df_type = get_dataframe_type(context_payload)
+        if df_type is not None:
+            parquet_b64, df_type = dataframe_to_parquet_b64(context_payload)
+            self.execute_code(build_dataframe_context_code(parquet_b64, df_type))
+        elif isinstance(context_payload, str):
+            escaped = context_payload.replace("\\", "\\\\").replace('"""', '\\"\\"\\"')
+            self.execute_code(f'context = """{escaped}"""')
+        else:
+            context_json = json.dumps(context_payload)
+            escaped_json = context_json.replace("\\", "\\\\").replace("'", "\\'")
+            self.execute_code(f"import json; context = json.loads('{escaped_json}')")
 
     def execute_code(self, code: str) -> REPLResult:
         """Execute code in the Modal sandbox and return result."""

rlm/environments/prime_repl.py

@@ -504,9 +504,23 @@ def _handle_llm_request(self, req_data: dict) -> dict:
 
     def load_context(self, context_payload: ContextPayload):
         """Load context into the sandbox environment."""
-        from rlm.utils.dataframe_utils import build_context_code
+        from rlm.utils.dataframe_utils import (
+            build_dataframe_context_code,
+            dataframe_to_parquet_b64,
+            get_dataframe_type,
+        )
 
-        self.execute_code(build_context_code(context_payload))
+        df_type = get_dataframe_type(context_payload)
+        if df_type is not None:
+            parquet_b64, df_type = dataframe_to_parquet_b64(context_payload)
+            self.execute_code(build_dataframe_context_code(parquet_b64, df_type))
+        elif isinstance(context_payload, str):
+            escaped = context_payload.replace("\\", "\\\\").replace('"""', '\\"\\"\\"')
+            self.execute_code(f'context = """{escaped}"""')
+        else:
+            context_json = json.dumps(context_payload)
+            escaped_json = context_json.replace("\\", "\\\\").replace("'", "\\'")
+            self.execute_code(f"import json; context = json.loads('{escaped_json}')")
 
     def execute_code(self, code: str) -> REPLResult:
         """Execute code in the Prime sandbox and return result."""

rlm/utils/dataframe_utils.py

@@ -1,6 +1,5 @@
 import base64
 import io
-import json
 from typing import Any, Literal
 
 DataFrameType = Literal["pandas"]
@@ -83,25 +82,3 @@ def build_dataframe_context_code(
         f"_parquet_bytes = base64.b64decode('{parquet_b64}')\n"
         f"{var_name} = pd.read_parquet(io.BytesIO(_parquet_bytes))"
     )
-
-
-def build_context_code(context_payload: Any, var_name: str = "context") -> str:
-    """Build Python code that reconstructs *context_payload* as a variable.
-
-    Handles three payload shapes:
-    - pandas DataFrame  → Parquet-over-base64 round-trip
-    - str               → triple-quoted string literal
-    - other (dict/list) → ``json.loads`` from a JSON string literal
-    """
-    df_type = get_dataframe_type(context_payload)
-    if df_type is not None:
-        parquet_b64, df_type = dataframe_to_parquet_b64(context_payload)
-        return build_dataframe_context_code(parquet_b64, df_type, var_name=var_name)
-
-    if isinstance(context_payload, str):
-        escaped = context_payload.replace("\\", "\\\\").replace('"""', '\\"\\"\\"')
-        return f'{var_name} = """{escaped}"""'
-
-    context_json = json.dumps(context_payload)
-    escaped_json = context_json.replace("\\", "\\\\").replace("'", "\\'")
-    return f"import json; {var_name} = json.loads('{escaped_json}')"

tests/test_dataframe_utils.py

@@ -2,12 +2,10 @@
 
 import base64
 import io
-import json
 
 import pytest
 
 from rlm.utils.dataframe_utils import (
-    build_context_code,
     build_dataframe_context_code,
     dataframe_metadata,
     dataframe_to_parquet_b64,
@@ -157,81 +155,3 @@ def test_custom_var_name(self):
         exec(code, ns)
         assert "my_df" in ns
         assert ns["my_df"]["v"].tolist() == [1]
-
-
-# ---------------------------------------------------------------------------
-# Unified dispatch — build_context_code()
-# ---------------------------------------------------------------------------
-
-
-class TestBuildContextCode:
-    """Tests for the unified build_context_code() dispatcher."""
-
-    def test_string_payload(self):
-        code = build_context_code("hello world")
-        ns = {}
-        exec(code, ns)
-        assert ns["context"] == "hello world"
-
-    def test_dict_payload(self):
-        payload = {"key": "value", "num": 42}
-        code = build_context_code(payload)
-        ns = {}
-        exec(code, ns)
-        assert ns["context"] == payload
-
-    def test_list_payload(self):
-        payload = [1, 2, "three"]
-        code = build_context_code(payload)
-        ns = {}
-        exec(code, ns)
-        assert ns["context"] == payload
-
-    def test_dataframe_payload(self):
-        pd = pytest.importorskip("pandas")
-        df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
-        code = build_context_code(df)
-        ns = {}
-        exec(code, ns)
-        assert ns["context"]["a"].tolist() == [1, 2]
-        assert ns["context"]["b"].tolist() == [3, 4]
-
-    def test_empty_string(self):
-        code = build_context_code("")
-        ns = {}
-        exec(code, ns)
-        assert ns["context"] == ""
-
-    def test_empty_dict(self):
-        code = build_context_code({})
-        ns = {}
-        exec(code, ns)
-        assert ns["context"] == {}
-
-    def test_empty_list(self):
-        code = build_context_code([])
-        ns = {}
-        exec(code, ns)
-        assert ns["context"] == []
-
-    def test_special_chars_in_string(self):
-        payload = 'He said "hello" and she said \'goodbye\'\nNew line\ttab'
-        code = build_context_code(payload)
-        ns = {}
-        exec(code, ns)
-        assert ns["context"] == payload
-
-    def test_custom_var_name(self):
-        code = build_context_code("data", var_name="my_var")
-        ns = {}
-        exec(code, ns)
-        assert ns["my_var"] == "data"
-        assert "context" not in ns
-
-    def test_nested_dict(self):
-        payload = {"outer": {"inner": [1, 2, 3]}, "flag": True}
-        code = build_context_code(payload)
-        ns = {}
-        exec(code, ns)
-        assert ns["context"]["outer"]["inner"] == [1, 2, 3]
-        assert ns["context"]["flag"] is True