Update

[ghstack-poisoned]

Author: vmoens

test/llm/test_data.py

@@ -13,7 +13,7 @@
 import torch
 from tensordict import lazy_stack, set_list_to_stack, TensorDict
 
-from torchrl import torchrl_logger
+from torchrl import logger as torchrl_logger
 
 from torchrl.data import (
     History,

test/llm/test_envs.py

@@ -375,7 +375,7 @@ def test_python_interpreter_single_batch(self):
             "```python\n"
             "print(1 + 1)\n"
             "```<|im_end|>\n"
-            "    <|im_start|>user\n"
+            "    <|im_start|>tool\n"
             "<tool_response>\n"
             "Code block 1 executed successfully:\n"
             "2\n"
@@ -395,7 +395,7 @@ def test_python_interpreter_single_batch(self):
                         content="Here is a python code to execute:\n```python\nprint(1 + 1)\n```",
                     ),
                     History(
-                        role="user",
+                        role="tool",
                         content="<tool_response>\nCode block 1 executed successfully:\n2\n\n</tool_response>",
                         tool_responses=["Code block 1 executed successfully:\n2\n"],
                     ),
@@ -478,7 +478,7 @@ def test_python_interpreter_persistent(self):
             "```python\n"
             "a=1\n"
             "```<|im_end|>\n"
-            "    <|im_start|>user\n"
+            "    <|im_start|>tool\n"
             "<tool_response>\n"
             "Code block 1 executed successfully:\n"
             "\n"
@@ -489,7 +489,7 @@ def test_python_interpreter_persistent(self):
             "a+=1\n"
             "assert a == 2\n"
             "```<|im_end|>\n"
-            "    <|im_start|>user\n"
+            "    <|im_start|>tool\n"
             "<tool_response>\n"
             "Code block 1 executed successfully:\n"
             "\n"

torchrl/envs/llm/transforms/tools.py

@@ -530,7 +530,10 @@ def _step(
 
         procs = []
         # Iterate over env batch-size
-        for i, t in enumerate(local_history.content):
+        content = local_history.content
+        if isinstance(content, str):
+            content = [content]
+        for i, t in enumerate(content):
             results = self._process_llm_response(t, i)
             if len(results) == 0:
                 procs.append(None)