predict-rlm 0.2.2__tar.gz → 0.2.4__tar.gz

{predict_rlm-0.2.2 → predict_rlm-0.2.4}/.gitignore

@@ -11,5 +11,6 @@ node_modules/
 examples/*/output/
 .claude
 .env
+.env.*
 .DS_Store
 .vscode/

{predict_rlm-0.2.2 → predict_rlm-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: predict-rlm
-Version: 0.2.2
+Version: 0.2.4
 Summary: Production-grade RLMs (Recursive Language Models) with tool use, built on DSPy
 Project-URL: Homepage, https://www.trampoline.ai/
 Project-URL: Repository, https://github.com/Trampoline-AI/predict-rlm
@@ -25,7 +25,7 @@ Requires-Dist: pymupdf>=1.24.0; extra == 'examples'
 Description-Content-Type: text/markdown
 
 # predict-rlm
-Harness-less LM runtime built on top of [DSPy](https://dspy.ai). Define your inputs, outputs, and tools — the model handles its own control flow. Get fully interpretable trajectories and performance that scales directly with model improvements. Without context rot.
+Production focused Self-harnessed LM runtime (RLM) that allows the LM to call its sub-lm with [DSPy](https://dspy.ai) signatures. Define your inputs, outputs, and tools — the model handles its own control flow. Get fully interpretable trajectories and performance that scales directly with model improvements. Without context rot.
 
 Based on the [Recursive Language Models](https://arxiv.org/abs/2512.24601v1) paper by [Alex L. Zhang](https://x.com/a1zhang), [Tim Kraska](https://x.com/tim_kraska), and [Omar Khattab](https://x.com/lateinteraction) from the Stanford NLP lab.<br/>
 
@@ -68,7 +68,7 @@ uv add predict-rlm
 
 - **Multimodal** — process images, documents, audio, and video through sub-LM calls using native provider multimodal APIs. 
 - **Async tool calling** — native RLM async support in the WASM sandbox, enabling concurrent sub-LM invocations and tool calls
-- **Prompt-optimized skills & tools** — predic-rlm skills comes tested and optimized to ensure maximum LM interoperability and performance, bundling instructions, PyPI packages, and tools for domain-specific tasks
+- **Prompt-optimized skills & tools** — predict-rlm skills comes tested and optimized to ensure maximum LM interoperability and performance, bundling instructions, PyPI packages, and tools for domain-specific tasks
 - **Simple file I/O** — pass local or cloud files as typed inputs and outputs via `File`, keeping interop with your existing data pipelines straightforward. (S3 files support soon)
 - **Structured sub-LM calls** — native Pydantic and DSPy signature support for type-safe sub-LM invocations with structured outputs
 

{predict_rlm-0.2.2 → predict_rlm-0.2.4}/README.md

@@ -1,5 +1,5 @@
 # predict-rlm
-Harness-less LM runtime built on top of [DSPy](https://dspy.ai). Define your inputs, outputs, and tools — the model handles its own control flow. Get fully interpretable trajectories and performance that scales directly with model improvements. Without context rot.
+Production focused Self-harnessed LM runtime (RLM) that allows the LM to call its sub-lm with [DSPy](https://dspy.ai) signatures. Define your inputs, outputs, and tools — the model handles its own control flow. Get fully interpretable trajectories and performance that scales directly with model improvements. Without context rot.
 
 Based on the [Recursive Language Models](https://arxiv.org/abs/2512.24601v1) paper by [Alex L. Zhang](https://x.com/a1zhang), [Tim Kraska](https://x.com/tim_kraska), and [Omar Khattab](https://x.com/lateinteraction) from the Stanford NLP lab.<br/>
 
@@ -42,7 +42,7 @@ uv add predict-rlm
 
 - **Multimodal** — process images, documents, audio, and video through sub-LM calls using native provider multimodal APIs. 
 - **Async tool calling** — native RLM async support in the WASM sandbox, enabling concurrent sub-LM invocations and tool calls
-- **Prompt-optimized skills & tools** — predic-rlm skills comes tested and optimized to ensure maximum LM interoperability and performance, bundling instructions, PyPI packages, and tools for domain-specific tasks
+- **Prompt-optimized skills & tools** — predict-rlm skills comes tested and optimized to ensure maximum LM interoperability and performance, bundling instructions, PyPI packages, and tools for domain-specific tasks
 - **Simple file I/O** — pass local or cloud files as typed inputs and outputs via `File`, keeping interop with your existing data pipelines straightforward. (S3 files support soon)
 - **Structured sub-LM calls** — native Pydantic and DSPy signature support for type-safe sub-LM invocations with structured outputs
 

{predict_rlm-0.2.2 → predict_rlm-0.2.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "predict-rlm"
-version = "0.2.2"
+version = "0.2.4"
 description = "Production-grade RLMs (Recursive Language Models) with tool use, built on DSPy"
 authors = [{ name = "Trampoline AI" }]
 license = { text = "MIT" }

{predict_rlm-0.2.2 → predict_rlm-0.2.4}/src/predict_rlm/__init__.py

@@ -9,9 +9,10 @@ File I/O:
            (sync from sandbox). Use ``list[File]`` for multiple files.
 """
 
-from .files import File, LocalDir, LocalFile, OutputDir, OutputFile
+from .files import File, LocalDir, LocalFile, OutputDir, OutputFile, SyncedFile
 from .predict_rlm import PredictRLM
 from .rlm_skills import Skill
+from .trace import RunTrace
 
 __all__ = [
     "File",
@@ -20,5 +21,7 @@ __all__ = [
     "OutputDir",
     "OutputFile",
     "PredictRLM",
+    "RunTrace",
     "Skill",
+    "SyncedFile",
 ]

{predict_rlm-0.2.2 → predict_rlm-0.2.4}/src/predict_rlm/_shared.py

@@ -4,7 +4,9 @@ from __future__ import annotations
 
 import inspect
 import textwrap
-from typing import TYPE_CHECKING, Callable
+import typing
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated, Callable
 
 import dspy
 from dspy.adapters.utils import translate_field_type
@@ -29,20 +31,32 @@ def format_tool_docs_full(tools: dict[str, Callable]) -> str:
         # Get function signature with types
         try:
             sig = inspect.signature(func)
+            # Resolve string annotations (from `from __future__ import annotations`)
+            try:
+                resolved = typing.get_type_hints(func, include_extras=True)
+            except (TypeError, NameError):
+                resolved = {}
+
             params = []
             for p in sig.parameters.values():
-                if p.annotation != inspect.Parameter.empty:
-                    type_name = getattr(p.annotation, "__name__", str(p.annotation))
+                ann = resolved.get(p.name)
+                if ann is not None:
+                    # Unwrap Annotated[X, ...] → X (e.g. SyncedFile markers)
+                    if typing.get_origin(ann) is Annotated:
+                        ann = typing.get_args(ann)[0]
+                    # Show Path as str — the RLM passes sandbox paths as strings
+                    if ann is Path:
+                        ann = str
+                    type_name = getattr(ann, "__name__", str(ann))
                     params.append(f"{p.name}: {type_name}")
                 else:
                     params.append(p.name)
             params_str = ", ".join(params)
 
             # Get return type
-            if sig.return_annotation != inspect.Parameter.empty:
-                ret_type = getattr(
-                    sig.return_annotation, "__name__", str(sig.return_annotation)
-                )
+            ret_ann = resolved.get("return")
+            if ret_ann is not None:
+                ret_type = getattr(ret_ann, "__name__", str(ret_ann))
                 sig_str = f"{name}({params_str}) -> {ret_type}"
             else:
                 sig_str = f"{name}({params_str})"
@@ -128,7 +142,7 @@ def build_rlm_signatures(
     )
     action_sig = action_sig.append(
         "code",
-        dspy.OutputField(desc="Python code wrapped in ```repl blocks."),
+        dspy.OutputField(desc="Python code wrapped in ```python blocks."),
         type_=str,
     )
 

{predict_rlm-0.2.2 → predict_rlm-0.2.4}/src/predict_rlm/files.py

@@ -25,7 +25,8 @@ from __future__ import annotations
 
 import os
 import typing
-from typing import Any
+from dataclasses import dataclass
+from typing import Annotated, Any
 
 from pydantic import BaseModel, Field
 
@@ -64,6 +65,34 @@ OutputFile = File
 OutputDir = File
 
 
+@dataclass(frozen=True)
+class SyncedFile:
+    """Annotation marker for tool parameters that need sandbox-host file sync.
+
+    Use with ``typing.Annotated`` on tool function parameters to declare that
+    a parameter is a sandbox file path. The framework automatically syncs the
+    file from the sandbox to the host before calling the tool, and optionally
+    mounts the modified file back into the sandbox after the tool returns.
+
+    Example::
+
+        def recalculate(
+            workbook: Annotated[Path, SyncedFile(host_dir="/tmp/wb")],
+            reference: Annotated[Path, SyncedFile(writeback=False)],
+        ) -> str:
+            ...
+    """
+
+    writeback: bool = True
+    """If True (default), mount the file back into the sandbox after the tool
+    returns. Set to False for read-only access (skip the mount-after step)."""
+
+    host_dir: str | None = None
+    """Host directory for the synced file. If None, a temporary directory is
+    created and cleaned up after the call. If specified, the directory is used
+    as-is and not cleaned up."""
+
+
 def _unwrap_annotation(annotation: Any) -> Any:
     """Unwrap Optional/Annotated/list to get the inner file type."""
     origin = typing.get_origin(annotation)
@@ -263,3 +292,26 @@ def build_file_plan(
         "output_field_map": output_field_map,
         "instructions": instructions,
     }
+
+
+def get_synced_file_params(fn: Any) -> dict[str, SyncedFile]:
+    """Extract SyncedFile annotations from a tool function's type hints.
+
+    Returns a dict mapping parameter names to their ``SyncedFile`` marker
+    for all parameters annotated with ``Annotated[..., SyncedFile(...)]``.
+    """
+    try:
+        hints = typing.get_type_hints(fn, include_extras=True)
+    except (TypeError, NameError):
+        return {}
+
+    result: dict[str, SyncedFile] = {}
+    for name, hint in hints.items():
+        if name == "return":
+            continue
+        if typing.get_origin(hint) is Annotated:
+            for arg in typing.get_args(hint)[1:]:
+                if isinstance(arg, SyncedFile):
+                    result[name] = arg
+                    break
+    return result

{predict_rlm-0.2.2 → predict_rlm-0.2.4}/src/predict_rlm/interpreter.py

@@ -16,11 +16,15 @@ from __future__ import annotations
 import asyncio
 import concurrent.futures
 import functools
+import inspect
 import json
 import logging
 import os
 import re
 import select
+import shutil
+import tempfile
+import time
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -35,6 +39,7 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+
 # JSON-RPC 2.0 helpers (local to avoid coupling to dspy internals)
 JSONRPC_APP_ERRORS = {
     "SyntaxError": -32000,
@@ -223,6 +228,17 @@ class JspiInterpreter(PythonInterpreter):
         all_read_paths = list(enable_read_paths or []) + list(extra_read_paths or [])
         all_write_paths = list(enable_write_paths or []) + list(extra_write_paths or [])
 
+        # Scan tools for SyncedFile annotations with custom host_dir paths
+        # and add them to Deno permissions so the runner can write there.
+        if tools:
+            from predict_rlm.files import get_synced_file_params
+
+            for tool_fn in tools.values():
+                for sf in get_synced_file_params(tool_fn).values():
+                    if sf.host_dir is not None:
+                        all_write_paths.append(sf.host_dir)
+                        all_read_paths.append(sf.host_dir)
+
         # Build custom deno command if not provided
         if deno_command is None:
             deno_command = self._build_deno_command(
@@ -249,6 +265,9 @@ class JspiInterpreter(PythonInterpreter):
         # Per-interpreter thread pool for sync tool calls (avoids starving
         # the shared default executor when many interpreters run concurrently)
         self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+        # Pending file-sync operations requested by tools during execution.
+        # Maps request ID → asyncio.Future resolved by the execute loop.
+        self._pending_file_ops: dict[int, asyncio.Future] = {}
 
     def _ensure_deno_process(self) -> None:
         """Override to capture raw fds for non-blocking I/O."""
@@ -327,6 +346,11 @@ class JspiInterpreter(PythonInterpreter):
         allowed_read.extend(str(p) for p in read_paths)
         allowed_read.extend(str(p) for p in write_paths)
 
+        # Allow reading temp dirs so @file_sync tools can mount files back
+        import tempfile as _tempfile
+        allowed_read.append(_tempfile.gettempdir())
+        allowed_read.append("/tmp")
+
         if allowed_read:
             args.append(f"--allow-read={','.join(allowed_read)}")
 
@@ -384,7 +408,11 @@ class JspiInterpreter(PythonInterpreter):
             self._write_stdin(msg + "\n")
 
     def _send_request(self, method: str, params: dict, context: str) -> dict:
-        """Send a JSON-RPC request without blocking the OS pipe."""
+        """Send a JSON-RPC request without blocking the OS pipe.
+
+        Skips non-JSON lines (e.g. Pyodide package-loading messages) matching
+        the parent PythonInterpreter behaviour.
+        """
         self._request_id += 1
         request_id = self._request_id
         msg = json.dumps(
@@ -397,26 +425,40 @@ class JspiInterpreter(PythonInterpreter):
         )
         self._write_stdin(msg + "\n")
 
-        response_line = self._read_with_timeout(timeout=None)
-        if not response_line:
-            exit_code = self.deno_process.poll()
-            if exit_code is not None:
-                stderr = self.deno_process.stderr.read() if self.deno_process.stderr else ""
+        max_skip = 100
+        skipped = 0
+        while skipped <= max_skip:
+            response_line = self._read_with_timeout(timeout=None)
+            if not response_line:
+                exit_code = self.deno_process.poll()
+                if exit_code is not None:
+                    stderr = self.deno_process.stderr.read() if self.deno_process.stderr else ""
+                    raise CodeInterpreterError(
+                        f"Deno exited (code {exit_code}) {context}: {stderr}"
+                    )
+                raise CodeInterpreterError(f"No response {context}")
+
+            if not response_line.startswith("{"):
+                skipped += 1
+                continue
+
+            try:
+                response = json.loads(response_line)
+            except json.JSONDecodeError:
+                skipped += 1
+                continue
+
+            if response.get("id") != request_id:
                 raise CodeInterpreterError(
-                    f"Deno exited (code {exit_code}) {context}: {stderr}"
+                    f"Response ID mismatch {context}: expected {request_id}, got {response.get('id')}"
                 )
-            raise CodeInterpreterError(f"No response {context}")
+            if "error" in response:
+                raise CodeInterpreterError(
+                    f"Error {context}: {response['error'].get('message', 'Unknown error')}"
+                )
+            return response
 
-        response = json.loads(response_line)
-        if response.get("id") != request_id:
-            raise CodeInterpreterError(
-                f"Response ID mismatch {context}: expected {request_id}, got {response.get('id')}"
-            )
-        if "error" in response:
-            raise CodeInterpreterError(
-                f"Error {context}: {response['error'].get('message', 'Unknown error')}"
-            )
-        return response
+        raise CodeInterpreterError(f"Too many non-JSON lines ({skipped}) {context}")
 
     def _get_deno_dir(self) -> list[str]:
         """Get Deno cache directory paths (may have multiple on different platforms)."""
@@ -473,28 +515,16 @@ class JspiInterpreter(PythonInterpreter):
         self.deno_process.stdin.flush()
 
     def _strip_code_fences(self, code: str) -> str:
-        """Extract code from ```repl fences.
+        """Extract code from markdown fences.
 
-        Uses a specific ```repl tag (like the original RLM) to avoid ambiguity.
         The closing ``` must be on its own line (^```$) to handle:
         1. Code containing inline ``` (like in strings) - not on own line, won't match
         2. Double fences from model (```...```\\n```) - stops at first proper close
-        Falls back to generic fence matching for backwards compatibility.
 
-        Supports multiple ```repl blocks - all blocks are concatenated with newlines.
+        Supports multiple fenced blocks - all blocks are concatenated with newlines.
         """
-        # Primary: look for ```repl blocks
-        # Use MULTILINE so ^ matches start of line - closing ``` must be alone on a line
-        # Non-greedy .*? stops at FIRST ``` on its own line
-        # findall to get ALL blocks, not just the first
-        matches = re.findall(r"```repl\s*\n(.*?)^```\s*$", code, re.DOTALL | re.MULTILINE)
-        if matches:
-            # Join all blocks with double newlines to ensure separation
-            return "\n\n".join(block.rstrip() for block in matches)
-
-        # Fallback: try generic ```python or ``` blocks for backwards compatibility
         matches = re.findall(
-            r"```(?:python|py)?\s*\n(.*?)^```\s*$", code, re.DOTALL | re.MULTILINE
+            r"```(?:python|py|repl)?\s*\n(.*?)^```\s*$", code, re.DOTALL | re.MULTILINE
         )
         if matches:
             return "\n\n".join(block.rstrip() for block in matches)
@@ -681,6 +711,14 @@ class JspiInterpreter(PythonInterpreter):
                 logger.info(f"Skipping malformed JSON: {output_line[:100]}")
                 continue
 
+            # Route file-sync responses to pending futures (from _execute_tool_async)
+            resp_id = result.get("id")
+            if resp_id is not None and resp_id in self._pending_file_ops:
+                future = self._pending_file_ops.pop(resp_id)
+                if not future.done():
+                    future.set_result(result)
+                continue
+
             # JSON-RPC request from sandbox (tool call)
             if "method" in result:
                 if result["method"] == "tool_call":
@@ -823,8 +861,53 @@ class JspiInterpreter(PythonInterpreter):
         line, self._read_buf = self._read_buf.split("\n", 1)
         return line.strip()
 
+    async def _sync_file_during_tool(self, virtual_path: str, host_path: str) -> None:
+        """Sync a file from sandbox MEMFS to host during a tool call.
+
+        Sends a sync_file request to the Deno runner's responseReader (which
+        handles it during tool execution) and awaits the response via a Future
+        resolved by the _execute_async loop.
+        """
+        self._request_id += 1
+        req_id = self._request_id
+        loop = asyncio.get_running_loop()
+        future = loop.create_future()
+        self._pending_file_ops[req_id] = future
+        msg = json.dumps({
+            "jsonrpc": "2.0", "method": "sync_file",
+            "params": {"virtual_path": virtual_path, "host_path": host_path},
+            "id": req_id,
+        })
+        await self._write_stdin_async(msg + "\n")
+        result = await future
+        if "error" in result:
+            raise CodeInterpreterError(
+                f"sync_file failed: {result['error'].get('message', result['error'])}"
+            )
+
+    async def _mount_file_during_tool(self, host_path: str, virtual_path: str) -> None:
+        """Mount a file from host into sandbox MEMFS during a tool call."""
+        self._request_id += 1
+        req_id = self._request_id
+        loop = asyncio.get_running_loop()
+        future = loop.create_future()
+        self._pending_file_ops[req_id] = future
+        msg = json.dumps({
+            "jsonrpc": "2.0", "method": "mount_file",
+            "params": {"host_path": host_path, "virtual_path": virtual_path},
+            "id": req_id,
+        })
+        await self._write_stdin_async(msg + "\n")
+        result = await future
+        if "error" in result:
+            raise CodeInterpreterError(
+                f"mount_file failed: {result['error'].get('message', result['error'])}"
+            )
+
     async def _execute_tool_async(self, tool_name: str, call_args: dict) -> dict:
         """Execute a tool asynchronously and return the response dict."""
+        from .trace import ToolCall, ms_since, record_tool_call
+
         if self._debug:
             import sys
 
@@ -832,39 +915,116 @@ class JspiInterpreter(PythonInterpreter):
             print(
                 f"\n\033[33m── Tool: {tool_name}({kwargs_preview}) ──\033[0m", file=sys.stderr
             )
+
+        call_start = time.perf_counter()
+        # Copy to mutable containers so the SyncedFile handler below can
+        # rewrite sandbox paths to host paths before invoking the tool.
+        args = list(call_args.get("args", []))
+        kwargs = dict(call_args.get("kwargs", {}))
+        temp_dir: str | None = None
+
         try:
             if tool_name not in self.tools:
                 raise CodeInterpreterError(f"Unknown tool: {tool_name}")
 
             tool_fn = self.tools[tool_name]
-            args = call_args.get("args", [])
-            kwargs = call_args.get("kwargs", {})
 
             # Pass pydantic_schemas through to predict tool if present
             pydantic_schemas = call_args.get("pydantic_schemas")
             if pydantic_schemas and tool_name == "predict":
                 kwargs["pydantic_schemas"] = pydantic_schemas
 
+            # Handle SyncedFile-annotated tool parameters: sync sandbox files
+            # to host before calling, and mount modified files back after.
+            from predict_rlm.files import get_synced_file_params
+
+            synced_params = get_synced_file_params(tool_fn)
+            temp_dir = None
+            # (sandbox_path, host_path, writeback) for each synced param
+            synced_entries: list[tuple[str, str, bool]] = []
+
+            if synced_params:
+                sig = inspect.signature(tool_fn)
+                param_names = list(sig.parameters.keys())
+
+                for param_name, sf in synced_params.items():
+                    # Resolve the sandbox path from args or kwargs
+                    sandbox_path = kwargs.get(param_name)
+                    if sandbox_path is None and param_name in param_names:
+                        idx = param_names.index(param_name)
+                        if idx < len(args):
+                            sandbox_path = args[idx]
+                    if not sandbox_path or not isinstance(sandbox_path, str):
+                        continue
+
+                    # Determine host directory
+                    if sf.host_dir is not None:
+                        host_dir = sf.host_dir
+                        os.makedirs(host_dir, exist_ok=True)
+                    else:
+                        if temp_dir is None:
+                            temp_dir = tempfile.mkdtemp(prefix="tool-file-sync-")
+                        host_dir = temp_dir
+
+                    host_path = os.path.join(host_dir, os.path.basename(sandbox_path))
+                    await self._sync_file_during_tool(sandbox_path, host_path)
+                    synced_entries.append((sandbox_path, host_path, sf.writeback))
+
+                    # Replace the sandbox path with the host path in args/kwargs
+                    if param_name in kwargs:
+                        kwargs[param_name] = host_path
+                    elif param_name in param_names:
+                        idx = param_names.index(param_name)
+                        if idx < len(args):
+                            args[idx] = host_path
+
             # Check if tool is async or sync
             if asyncio.iscoroutinefunction(tool_fn):
                 result = await tool_fn(*args, **kwargs)
             else:
-                # Run sync function in per-interpreter thread pool (not the
-                # shared default pool) to prevent starvation when many
-                # interpreters run concurrently.
-                # loop.run_in_executor only accepts positional args, so wrap
-                # the call in functools.partial to bind **kwargs.
                 loop = asyncio.get_running_loop()
                 result = await loop.run_in_executor(
                     self._executor, functools.partial(tool_fn, *args, **kwargs)
                 )
 
+            # Mount modified files back into the sandbox (only for writeback params)
+            if synced_entries:
+                for sandbox_path, host_path, writeback in synced_entries:
+                    if writeback and os.path.isfile(host_path):
+                        await self._mount_file_during_tool(host_path, sandbox_path)
+                if temp_dir:
+                    shutil.rmtree(temp_dir, ignore_errors=True)
+
             is_json = isinstance(result, (list, dict))
-            return {
+            response = {
                 "value": json.dumps(result) if is_json else str(result or ""),
                 "type": "json" if is_json else "string",
             }
+
+            # Record non-predict tool calls (predict records itself with richer detail)
+            if tool_name != "predict":
+                record_tool_call(ToolCall(
+                    name=tool_name,
+                    args=args,
+                    kwargs={k: v for k, v in kwargs.items() if k != "pydantic_schemas"},
+                    result=result,
+                    duration_ms=ms_since(call_start),
+                ))
+
+            return response
         except Exception as e:
+            # Clean up any SyncedFile temp dir before returning
+            if temp_dir:
+                shutil.rmtree(temp_dir, ignore_errors=True)
+            if tool_name != "predict":
+                record_tool_call(ToolCall(
+                    name=tool_name,
+                    args=args,
+                    kwargs={k: v for k, v in kwargs.items() if k != "pydantic_schemas"},
+                    result=None,
+                    error=str(e),
+                    duration_ms=ms_since(call_start),
+                ))
             return {"error": str(e)}
 
     def _write_stdin(self, data: str) -> None: