Add deepagents_rl package sources

src/deepagents_rl/data/swe_bench_lite.py

@@ -46,10 +46,16 @@ def _as_str_list(val: Any) -> list[str]:
 def build_agent_prompt(problem_statement: str, *, instance_id: str) -> str:
     return (
         "You are a software engineering agent inside a full git checkout of the repo.\n"
-        "Resolve the GitHub issue by exploring and editing files with your tools "
-        "(ls, glob, grep, read_file, write_file, edit_file, execute).\n"
-        "Prefer minimal correct edits. Do not modify tests unless necessary.\n"
-        "When finished, stop calling tools and reply with DONE plus a short summary.\n\n"
+        "You MUST resolve the GitHub issue by calling tools via the function-calling "
+        "interface (not by writing a plan or code only in plain text).\n"
+        "Available tools: ls, glob, grep, read_file, write_file, edit_file, execute.\n"
+        "Required workflow:\n"
+        "1) Call ls or glob to inspect the repo.\n"
+        "2) Call read_file / grep to find the bug.\n"
+        "3) Call edit_file or write_file to apply a minimal fix.\n"
+        "4) Optionally execute a quick check command.\n"
+        "5) Stop calling tools and reply DONE with a short summary.\n"
+        "Prefer minimal correct edits. Do not modify tests unless necessary.\n\n"
         f"Instance: {instance_id}\n\n"
         f"Issue:\n{problem_statement.strip()}\n"
     )

src/deepagents_rl/train/swe_env.py

@@ -256,8 +256,8 @@ class DeepAgentSWEEnv:
             )
             return out
 
-    def reset(self, **kwargs: Any) -> None:
-        """Materialize repo sandbox from a SWE Lite row. Prompt stays on the dataset."""
+    def reset(self, **kwargs: Any) -> str:
+        """Materialize repo sandbox; return tool-forcing text TRL appends to the user message."""
         self._cleanup()
         self._end_trace()
 
@@ -317,13 +317,17 @@ class DeepAgentSWEEnv:
         _vlog("env.reset.langfuse", enabled=lf is not None)
         if lf is not None:
             try:
+                # Fresh root trace per rollout (avoid nesting everything under one parent).
+                self._trace_id = lf.create_trace_id()
                 self._root_cm = lf.start_as_current_observation(
-                    name="grpo-swe-rollout",
+                    name=f"grpo-swe-rollout:{instance_id}",
                     as_type="agent",
+                    trace_context={"trace_id": self._trace_id},
                     input={
                         "instance_id": instance_id,
                         "repo": repo,
                         "base_commit": base_commit,
+                        "rollout_id": self._rollout_id,
                         "prompt_preview": _truncate(
                             kwargs.get("prompt")
                             if isinstance(kwargs.get("prompt"), str)
@@ -340,16 +344,18 @@ class DeepAgentSWEEnv:
                         "docker_image": docker_image,
                         "phase": "grpo_train_or_eval",
                         "harness": "deepagents-swe-grpo",
+                        "session_id": self._rollout_id,
                         "tags": ["deepagents-rl", "swe-bench", "grpo", "tools"],
                     },
+                    end_on_exit=True,
                 )
                 self._root = self._root_cm.__enter__()
-                self._trace_id = lf.get_current_trace_id() or ""
                 _vlog("env.reset.langfuse_trace", trace_id=self._trace_id)
             except Exception as e:  # noqa: BLE001
                 _vlog("env.reset.langfuse_failed", error=str(e))
                 self._root_cm = None
                 self._root = None
+                self._trace_id = ""
 
         def _materialize() -> None:
             _vlog("env.reset.materialize", repo=repo, base_commit=base_commit[:12])
@@ -381,8 +387,15 @@ class DeepAgentSWEEnv:
             rollout_id=self._rollout_id,
             n_f2p=len(fail_to_pass),
             n_p2p=len(pass_to_pass),
+            langfuse_trace_id=self._trace_id,
+        )
+        # TRL appends this string to the last user message — force tool use.
+        return (
+            "\n\nCRITICAL: Use structured tool/function calls only. "
+            "Your first action MUST be a tool call to `ls` or `glob` on the repo. "
+            "Then `read_file`/`grep`, then `edit_file`/`write_file` to fix the issue. "
+            "Do not answer with a text-only plan. When the fix is applied, stop tools and say DONE."
         )
-        return None
 
     def _cleanup(self) -> None:
         self._backend = None
@@ -647,7 +660,7 @@ def prepare_swe_grpo_rows(records: list[dict[str, Any]]) -> list[dict[str, Any]]
     Renames `image` -> `docker_image` so TRL does not treat it as multimodal.
     Converts `prompt` to conversational chat messages (required for tool calling).
     """
-    from deepagents_rl.data.swe_bench_lite import from_hf_record
+    from deepagents_rl.data.swe_bench_lite import build_agent_prompt, from_hf_record
 
     rows: list[dict[str, Any]] = []
     for rec in records:
@@ -658,11 +671,16 @@ def prepare_swe_grpo_rows(records: list[dict[str, Any]]) -> list[dict[str, Any]]
         else:
             task = dict(rec)
 
-        prompt = task.get("prompt") or ""
-        if isinstance(prompt, list):
-            messages = prompt
-        else:
-            messages = [{"role": "user", "content": str(prompt)}]
+        # Rebuild prompt locally so tool-forcing text stays current (HF rows may be stale).
+        messages = [
+            {
+                "role": "user",
+                "content": build_agent_prompt(
+                    str(task.get("problem_statement") or ""),
+                    instance_id=str(task.get("instance_id") or task.get("task_id") or ""),
+                ),
+            }
+        ]
 
         docker_image = str(task.get("docker_image") or task.get("image") or "")
         rows.append(