Add deepagents_rl package sources

src/deepagents_rl/train/swe_env.py

@@ -2,6 +2,9 @@
 
 Kaggle has no Docker for instance images — `get_reward` POSTs the model patch
 to `SWE_GRADER_URL` (local `scripts/swe_grader_server.py` + tunnel).
+
+Each rollout opens a Langfuse agent trace; every tool call is a nested `tool`
+span; grading is an `evaluator` span with reward scores.
 """
 from __future__ import annotations
 
@@ -9,8 +12,11 @@ import json
 import os
 import shutil
 import tempfile
+import traceback
+import uuid
+from functools import wraps
 from pathlib import Path
-from typing import Any
+from typing import Any, Callable
 
 
 def _result_to_str(obj: Any) -> str:
@@ -47,6 +53,16 @@ def _as_str_list(val: Any) -> list[str]:
     return [str(val)]
 
 
+def _truncate(obj: Any, n: int = 4000) -> Any:
+    if isinstance(obj, str):
+        return obj[:n]
+    if isinstance(obj, dict):
+        return {k: _truncate(v, n) for k, v in list(obj.items())[:40]}
+    if isinstance(obj, list):
+        return [_truncate(v, n) for v in obj[:40]]
+    return obj
+
+
 def grade_patch_remote(
     *,
     instance_id: str,
@@ -66,7 +82,6 @@ def grade_patch_remote(
 
     headers = {
         "Content-Type": "application/json",
-        # ngrok free tier interstitial
         "ngrok-skip-browser-warning": "1",
     }
     token = os.environ.get("SWE_GRADER_TOKEN") or ""
@@ -87,6 +102,16 @@ def grade_patch_remote(
     return resp.json()
 
 
+def _traced_tool(method: Callable[..., str]) -> Callable[..., str]:
+    """Wrap a public tool method so each call becomes a Langfuse tool span."""
+
+    @wraps(method)
+    def wrapper(self: "DeepAgentSWEEnv", *args: Any, **kwargs: Any) -> str:
+        return self._call_tool(method.__name__, method, args, kwargs)
+
+    return wrapper
+
+
 class DeepAgentSWEEnv:
     """Per-rollout TRL environment: LocalShellBackend tools + remote Docker reward."""
 
@@ -97,17 +122,95 @@ class DeepAgentSWEEnv:
         self._task: dict[str, Any] = {}
         self._last_reward: float = 0.0
         self._last_report: dict[str, Any] = {}
+        self._rollout_id: str = ""
+        self._tool_calls: list[dict[str, Any]] = []
+        self._lf: Any = None
+        self._root_cm: Any = None
+        self._root: Any = None
+        self._trace_id: str = ""
+
+    def _langfuse_enabled(self) -> bool:
+        if os.getenv("DARL_LANGFUSE", "1").strip() in ("0", "false", "no"):
+            return False
+        return bool(os.getenv("LANGFUSE_PUBLIC_KEY") and os.getenv("LANGFUSE_SECRET_KEY"))
+
+    def _ensure_langfuse(self) -> Any | None:
+        if not self._langfuse_enabled():
+            return None
+        if self._lf is not None:
+            return self._lf
+        try:
+            from langfuse import Langfuse, get_client
+
+            Langfuse(
+                public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
+                secret_key=os.environ["LANGFUSE_SECRET_KEY"],
+                host=os.getenv("LANGFUSE_HOST") or os.getenv("LANGFUSE_BASE_URL"),
+            )
+            self._lf = get_client()
+            return self._lf
+        except Exception as e:  # noqa: BLE001
+            print(f"[langfuse] init failed: {e}")
+            return None
+
+    def _end_trace(self) -> None:
+        if self._root_cm is not None:
+            try:
+                self._root_cm.__exit__(None, None, None)
+            except Exception:
+                pass
+        self._root_cm = None
+        self._root = None
+        if self._lf is not None:
+            try:
+                self._lf.flush()
+            except Exception:
+                pass
+
+    def _call_tool(
+        self,
+        name: str,
+        method: Callable[..., str],
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> str:
+        payload = _truncate({"args": list(args), "kwargs": kwargs})
+        lf = self._lf
+        if lf is None or self._root is None:
+            out = method(self, *args, **kwargs)
+            self._tool_calls.append({"name": name, "input": payload, "output": _truncate(out)})
+            return out
+
+        with lf.start_as_current_observation(
+            name=name,
+            as_type="tool",
+            input=payload,
+            metadata={
+                "rollout_id": self._rollout_id,
+                "instance_id": self._task.get("instance_id"),
+                "harness": "deepagents-swe-grpo",
+            },
+        ) as span:
+            try:
+                out = method(self, *args, **kwargs)
+                span.update(output=_truncate(out))
+                self._tool_calls.append({"name": name, "input": payload, "output": _truncate(out)})
+                return out
+            except Exception as e:  # noqa: BLE001
+                err = f"TOOL_ERROR: {type(e).__name__}: {e}"
+                span.update(output=err, level="ERROR")
+                self._tool_calls.append({"name": name, "input": payload, "output": err})
+                return err
 
     def reset(self, **kwargs: Any) -> None:
         """Materialize repo sandbox from a SWE Lite row. Prompt stays on the dataset."""
         self._cleanup()
+        self._end_trace()
 
         instance_id = str(kwargs.get("instance_id") or kwargs.get("task_id") or "")
         repo = str(kwargs.get("repo") or "")
         base_commit = str(kwargs.get("base_commit") or "")
-        docker_image = str(
-            kwargs.get("docker_image") or kwargs.get("image") or ""
-        )
+        docker_image = str(kwargs.get("docker_image") or kwargs.get("image") or "")
         eval_script = str(kwargs.get("eval_script") or "")
         fail_to_pass = _as_str_list(
             kwargs.get("fail_to_pass")
@@ -131,6 +234,8 @@ class DeepAgentSWEEnv:
 
         from deepagents_rl.harness.repo import materialize_repo_at_commit
 
+        self._rollout_id = str(uuid.uuid4())
+        self._tool_calls = []
         self._task = {
             "instance_id": instance_id,
             "task_id": instance_id,
@@ -144,17 +249,66 @@ class DeepAgentSWEEnv:
         }
         self._last_reward = 0.0
         self._last_report = {}
+        self._trace_id = ""
+
+        lf = self._ensure_langfuse()
+        if lf is not None:
+            try:
+                self._root_cm = lf.start_as_current_observation(
+                    name="grpo-swe-rollout",
+                    as_type="agent",
+                    input={
+                        "instance_id": instance_id,
+                        "repo": repo,
+                        "base_commit": base_commit,
+                        "prompt_preview": _truncate(
+                            kwargs.get("prompt")
+                            if isinstance(kwargs.get("prompt"), str)
+                            else (
+                                kwargs.get("problem_statement")
+                                or (kwargs.get("prompt") or "")
+                            ),
+                            2000,
+                        ),
+                    },
+                    metadata={
+                        "rollout_id": self._rollout_id,
+                        "instance_id": instance_id,
+                        "docker_image": docker_image,
+                        "phase": "grpo_train_or_eval",
+                        "harness": "deepagents-swe-grpo",
+                        "tags": ["deepagents-rl", "swe-bench", "grpo", "tools"],
+                    },
+                )
+                self._root = self._root_cm.__enter__()
+                self._trace_id = lf.get_current_trace_id() or ""
+            except Exception as e:  # noqa: BLE001
+                print(f"[langfuse] start rollout failed: {e}")
+                self._root_cm = None
+                self._root = None
+
+        def _materialize() -> None:
+            self._sandbox_parent = Path(
+                tempfile.mkdtemp(prefix=f"darl_grpo_{instance_id[:40]}_")
+            )
+            self._sandbox = self._sandbox_parent / "repo"
+            materialize_repo_at_commit(repo=repo, base_commit=base_commit, dest=self._sandbox)
+            self._backend = LocalShellBackend(
+                root_dir=str(self._sandbox),
+                virtual_mode=True,
+                timeout=60,
+            )
 
-        self._sandbox_parent = Path(
-            tempfile.mkdtemp(prefix=f"darl_grpo_{instance_id[:40]}_")
-        )
-        self._sandbox = self._sandbox_parent / "repo"
-        materialize_repo_at_commit(repo=repo, base_commit=base_commit, dest=self._sandbox)
-        self._backend = LocalShellBackend(
-            root_dir=str(self._sandbox),
-            virtual_mode=True,
-            timeout=60,
-        )
+        if lf is not None and self._root is not None:
+            with lf.start_as_current_observation(
+                name="repo_materialize",
+                as_type="span",
+                input={"repo": repo, "base_commit": base_commit},
+            ) as span:
+                _materialize()
+                span.update(output={"sandbox": str(self._sandbox)})
+        else:
+            _materialize()
         return None
 
     def _cleanup(self) -> None:
@@ -169,6 +323,7 @@ class DeepAgentSWEEnv:
             raise RuntimeError("DeepAgentSWEEnv tool called before reset()")
         return self._backend
 
+    @_traced_tool
     def ls(self, path: str = ".") -> str:
         """List files in a directory inside the sandbox.
 
@@ -180,6 +335,7 @@ class DeepAgentSWEEnv:
         """
         return _result_to_str(self._require_backend().ls(path))
 
+    @_traced_tool
     def glob(self, pattern: str, path: str = ".") -> str:
         """Glob for files matching pattern under path.
 
@@ -192,6 +348,7 @@ class DeepAgentSWEEnv:
         """
         return _result_to_str(self._require_backend().glob(pattern, path))
 
+    @_traced_tool
     def grep(self, pattern: str, path: str = ".") -> str:
         """Search file contents for a regex pattern.
 
@@ -204,6 +361,7 @@ class DeepAgentSWEEnv:
         """
         return _result_to_str(self._require_backend().grep(pattern, path))
 
+    @_traced_tool
     def read_file(self, file_path: str, offset: int = 0, limit: int = 200) -> str:
         """Read a text file from the sandbox.
 
@@ -219,6 +377,7 @@ class DeepAgentSWEEnv:
             self._require_backend().read(file_path, offset=offset, limit=limit)
         )
 
+    @_traced_tool
     def write_file(self, file_path: str, content: str) -> str:
         """Write or create a text file in the sandbox.
 
@@ -231,6 +390,7 @@ class DeepAgentSWEEnv:
         """
         return _result_to_str(self._require_backend().write(file_path, content))
 
+    @_traced_tool
     def edit_file(self, file_path: str, old_string: str, new_string: str) -> str:
         """Replace old_string with new_string in an existing file.
 
@@ -246,6 +406,7 @@ class DeepAgentSWEEnv:
             self._require_backend().edit(file_path, old_string, new_string)
         )
 
+    @_traced_tool
     def execute(self, command: str) -> str:
         """Run a shell command in the sandbox.
 
@@ -259,41 +420,101 @@ class DeepAgentSWEEnv:
 
     def get_reward(self) -> float:
         """Export git diff and grade via remote Docker SWE evaluator."""
-        if self._sandbox is None:
-            self._last_reward = 0.0
-            return 0.0
-
-        from deepagents_rl.harness.repo import git_diff_patch
-
-        try:
-            model_patch = git_diff_patch(self._sandbox)
-        except Exception as e:  # noqa: BLE001
-            self._last_report = {"error": f"git_diff_failed: {e}"}
-            self._last_reward = 0.0
-            return 0.0
-
-        if not model_patch.strip():
-            self._last_report = {"error": "empty_model_patch", "reward_total": 0.0}
-            self._last_reward = 0.0
-            return 0.0
-
-        try:
-            report = grade_patch_remote(
-                instance_id=self._task["instance_id"],
-                docker_image=self._task["docker_image"],
-                eval_script=self._task["eval_script"],
-                fail_to_pass=list(self._task["fail_to_pass"]),
-                pass_to_pass=list(self._task["pass_to_pass"]),
-                model_patch=model_patch,
-            )
-        except Exception as e:  # noqa: BLE001
-            self._last_report = {"error": f"grader_failed: {e}"}
-            self._last_reward = 0.0
-            return 0.0
+        lf = self._lf
+        reward = 0.0
+        report: dict[str, Any] = {}
+
+        def _grade() -> float:
+            nonlocal report
+            if self._sandbox is None:
+                report = {"error": "no_sandbox", "reward_total": 0.0}
+                return 0.0
+
+            from deepagents_rl.harness.repo import git_diff_patch
+
+            try:
+                model_patch = git_diff_patch(self._sandbox)
+            except Exception as e:  # noqa: BLE001
+                report = {"error": f"git_diff_failed: {e}"}
+                return 0.0
+
+            if not model_patch.strip():
+                report = {"error": "empty_model_patch", "reward_total": 0.0}
+                return 0.0
+
+            try:
+                report = grade_patch_remote(
+                    instance_id=self._task["instance_id"],
+                    docker_image=self._task["docker_image"],
+                    eval_script=self._task["eval_script"],
+                    fail_to_pass=list(self._task["fail_to_pass"]),
+                    pass_to_pass=list(self._task["pass_to_pass"]),
+                    model_patch=model_patch,
+                )
+            except Exception as e:  # noqa: BLE001
+                report = {
+                    "error": f"grader_failed: {e}",
+                    "traceback": traceback.format_exc()[-1500:],
+                }
+                return 0.0
+
+            return float(report.get("reward_total") or 0.0)
+
+        if lf is not None and self._root is not None:
+            with lf.start_as_current_observation(
+                name="reward_grader",
+                as_type="evaluator",
+                input={
+                    "instance_id": self._task.get("instance_id"),
+                    "n_tool_calls": len(self._tool_calls),
+                    "tools_used": [c["name"] for c in self._tool_calls],
+                },
+            ) as span:
+                reward = _grade()
+                span.update(
+                    output=_truncate(
+                        {
+                            "reward_total": reward,
+                            "pass_frac": report.get("pass_frac"),
+                            "resolved": report.get("resolved"),
+                            "resolution": report.get("resolution"),
+                            "error": report.get("error"),
+                            "n_tool_calls": len(self._tool_calls),
+                        }
+                    )
+                )
+                try:
+                    lf.score_current_trace(name="reward_total", value=float(reward))
+                    lf.score_current_trace(
+                        name="pass_frac", value=float(report.get("pass_frac") or 0.0)
+                    )
+                    lf.score_current_trace(
+                        name="resolved", value=1.0 if report.get("resolved") else 0.0
+                    )
+                    lf.score_current_trace(
+                        name="n_tool_calls", value=float(len(self._tool_calls))
+                    )
+                except Exception:
+                    pass
+                if self._root is not None:
+                    try:
+                        self._root.update(
+                            output={
+                                "reward_total": reward,
+                                "n_tool_calls": len(self._tool_calls),
+                                "tools_used": [c["name"] for c in self._tool_calls],
+                                "resolution": report.get("resolution"),
+                            }
+                        )
+                    except Exception:
+                        pass
+        else:
+            reward = _grade()
 
         self._last_report = report
-        self._last_reward = float(report.get("reward_total") or 0.0)
-        return self._last_reward
+        self._last_reward = reward
+        self._end_trace()
+        return reward
 
 
 def prepare_swe_grpo_rows(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
@@ -306,7 +527,6 @@ def prepare_swe_grpo_rows(records: list[dict[str, Any]]) -> list[dict[str, Any]]
 
     rows: list[dict[str, Any]] = []
     for rec in records:
-        # Accept both HF records and already-normalized tasks
         if "fail_to_pass_json" in rec or (
             "prompt" in rec and isinstance(rec.get("prompt"), str) and "image" in rec
         ):