dspy-rlm-hooks 0.1.11

About

DSPy RLM Hooks injects lifecycle hooks into DSPy's internal RLM iteration loop, giving you full control over every stage of code generation, execution, and history tracking.

• Code Rewriting — Fix or augment LLM-generated code before it runs
• Variable Injection — Seed the interpreter with persistent variables and imports
• Result Auditing — Transform, validate, or retry on errors
• History Management — Inspect and modify the REPL history between iterations
• Sync & Async — Hooks work in either mode; coroutines are auto-detected
• Optional MLflow Tracing — Hook spans are added automatically when MLflow is installed
• PredictRLM Support — Same hook API works on PredictRLM instances

Requires DSPy 3.1+ and Pydantic 2+.

(back to top)

Architecture

RLM Hook Lifecycle

┌──────────────────────────┐
│    pre_iteration_hook    │
└────────────┬─────────────┘
             │
             │ inject vars,
             │ prepend code
             ▼
┌──────────────────────────┐
│      Generate Code       │
└────────────┬─────────────┘
             │
             ▼
┌──────────────────────────┐
│    pre_execution_hook    │
└────────────┬─────────────┘
             │
             │ rewrite code
             ▼
┌──────────────────────────┐
│       Execute Code       │
└────────────┬─────────────┘
             │
             ▼
┌──────────────────────────┐
│   post_execution_hook    │
└────────────┬─────────────┘
             │
             │ transform result
             ▼
┌──────────────────────────┐
│   post_iteration_hook    │
└──────────────────────────┘

Hooks fire at each stage of an RLM iteration, allowing inspection and modification of behaviour.

(back to top)

Quick Start

Install

Install dspy-rlm-hooks with uv (recommended):

uv add dspy-rlm-hooks

Or with pip:

pip install dspy-rlm-hooks

MLflow tracing is optional. Install the extra to record hook spans:

uv add "dspy-rlm-hooks[tracing]"
# or: pip install "dspy-rlm-hooks[tracing]"

The same enable_rlm_hooks(...) call is used either way. The package checks MLflow at runtime: if its tracing API is installed, hook inputs and outputs are recorded in MLflow spans; otherwise hooks run normally with no MLflow import requirement. To combine these spans with DSPy's native traces, configure mlflow.dspy.autolog() in your application.

Basic Usage

import dspy
from dspy_rlm_hooks import enable_rlm_hooks, PreIterationOutput

rlm = dspy.RLM(...)

def inject_math(iteration, variables, history, input_args):
    return PreIterationOutput(
        extra_vars={"tool": "calculator"},
        persistent_python_code="import math",
        prompt_context="Verify the calculation before answering.",
    )

enable_rlm_hooks(rlm, pre_iteration_hook=inject_math)

result = rlm(question="What is the square root of 1764?")

PreIterationOutput keeps code lifetimes explicit:

• python_code is prepended only to the current iteration's generated code.
• persistent_python_code=None keeps the current persistent prelude, a string replaces it, and "" clears it. Persistent code runs before every execution.
• prompt_context is shown to the action-generating LLM for the current iteration; it is not executed as Python.
• extra_vars are interpreter variables for the current iteration.

(back to top)

Usage

All Four Hooks

A realistic example showing how each hook can be used to build a safe, instrumented agent:

from dspy_rlm_hooks import (
    enable_rlm_hooks,
    PreIterationOutput,
    PreExecutionOutput,
    PostExecutionOutput,
    PostIterationOutput,
)
from dspy.primitives.repl_types import REPLHistory
import re

# ── Pre-iteration: seed interpreter with a regex toolkit ──

def pre_iteration(iteration, variables, history, input_args):
    """Inject a regex helper and seed variables before every iteration."""
    return PreIterationOutput(
        extra_vars={"search_pattern": r"TODO|FIXME|HACK"},
        persistent_python_code="""
import re

def grep(pattern, text):
    return re.findall(pattern, text)
""",
    )

# ── Pre-execution: block dangerous code ──

FORBIDDEN = re.compile(r"\b(eval|exec|compile|__import__)\b")

def pre_execution(iteration, code, variables, history, input_args):
    """Sanitise generated code before it reaches the interpreter."""
    if FORBIDDEN.search(code):
        safe_code = FORBIDDEN.sub("# BLOCKED", code)
        return PreExecutionOutput(code=safe_code)
    return PreExecutionOutput(code=code)

# ── Post-execution: retry on error ──

def post_execution(iteration, code, result, variables, history, input_args):
    """If execution raised an error, wrap a hint so the LLM retries next round."""
    if isinstance(result, str) and result.startswith("[Error]"):
        return PostExecutionOutput(
            result=f"{result}\n# Hint: the variable 'search_pattern' is already in scope."
        )
    return PostExecutionOutput(result=result)

# ── Post-iteration: enforce a price budget ──

MAX_COST_USD = 0.50

def _estimate_cost(pred):
    # In production, derive this from response.usage or similar.
    return 0.015

def make_budget_hook(max_cost=MAX_COST_USD):
    """Return a post_iteration hook with isolated, per-request state.

    Create a new hook for every RLM session so budgets don't leak
    across concurrent requests on a multi-threaded or async server.
    """
    accumulated_cost = 0.0

    def post_iteration(iteration, pred, code, result, history: REPLHistory):
        nonlocal accumulated_cost
        accumulated_cost += _estimate_cost(pred)
        if accumulated_cost >= max_cost:
            return PostIterationOutput(history=history, stop=True)
        return PostIterationOutput(history=history)

    return post_iteration

# ── Wire everything up ──

enable_rlm_hooks(
    rlm,
    pre_iteration_hook=pre_iteration,
    pre_execution_hook=pre_execution,
    post_execution_hook=post_execution,
    post_iteration_hook=make_budget_hook(max_cost=0.50),
)

result = rlm(question="Find all TODO comments in the codebase")

Async Hooks

Return a coroutine and the system handles it automatically:

async def fetch_context(iteration, variables, history, input_args):
    context = await remote_cache.get(input_args["question"])
    return PreIterationOutput(extra_vars={"cached_context": context})

enable_rlm_hooks(rlm, pre_iteration_hook=fetch_context)

MLflow Tracing

No tracing-specific enable function is needed:

import mlflow
from dspy_rlm_hooks import enable_rlm_hooks

mlflow.dspy.autolog()
enable_rlm_hooks(rlm, pre_iteration_hook=fetch_context)

When MLflow is available, every configured lifecycle hook gets a stable span name, rlm_hook/<hook_name>, for every iteration. The iteration number remains available in the span inputs. Python source in the trace's python_code, code, original_code, and modified_code fields is formatted as fenced Python Markdown for readable rendering. Every interpreter call also creates an rlm/execute_code span after action generation and any pre_execution hook. Its code input is the final source that actually ran, including persistent Python injected by pre_iteration, and its input_args include injected variables. With DSPy autologging enabled, these spans nest under DSPy's active trace. Without MLflow, the same enable_rlm_hooks call continues to run as regular untraced hooks. The older enable_rlm_hooks_with_tracing function remains available for callers that explicitly want an ImportError when MLflow is missing.

Disabling Hooks

from dspy_rlm_hooks import disable_rlm_hooks

disable_rlm_hooks(rlm)

Removes all monkey-patched overrides and reverts to original behaviour.

(back to top)

Speculative Execution

Concept

Speculative execution (sPTC, speculative programmatic tool calling) runs a shadow pre-pass over generated code and pre-dispatches independent tool calls so the real run claims the results instead of re-calling. By default the built-in sub-LLM tools llm_query and llm_query_batched are speculated.

In streaming mode (the default) the shadow feeds the model's streamed code output during generate_action — the RLM's underlying dspy.Predict — so sub-LLM tool calls overlap with main-context token generation. code deltas are streamed via dspy.streamify and fed into the speculation turn as they arrive; the real run then claims the pre-dispatched results.

In Lazy/JIT mode (streaming=False) the shadow runs a one-shot pass over the fully assembled code block (persistent prelude plus injected variables) ahead of real execution. If streaming is unavailable (non-streaming adapter or LM, cache hit) or fails, execution transparently falls back to Lazy/JIT.

Install

No extra dependency. Speculative execution ships in the same dspy-rlm-hooks package.

Quick Start

import dspy
from dspy_rlm_hooks import enable_rlm_speculation

rlm = dspy.RLM(...)
enable_rlm_speculation(rlm)

result = rlm(question="...")

Classification API

By default only the built-in LLM tools are speculated. To speculate a read-only user tool, mark it with speculate() and pass it through the tools mapping with speculate_user_tools=True:

from dspy_rlm_hooks import enable_rlm_speculation, speculate

def lookup_price(symbol: str) -> float:
    ...

speculate(lookup_price, speculatable=True, pure=True)

enable_rlm_speculation(
    rlm,
    tools={"lookup_price": lookup_price},
    speculate_user_tools=True,
)

speculate() folds a SpeculationPolicy into the tool's classification. speculatable=True requires pure=True: a tool with observable side effects must never run early. SpeculationPolicy also carries deterministic, latency_hint_ms, and an optional per-call gate predicate.

Budget and Timeout

enable_rlm_speculation accepts:

• max_inflight (default 8): max speculative executions in flight at once.
• max_dispatches_per_turn (default 2048): hard cap on speculative dispatches per RLM turn.
• timeout_s (default 5.0): how long to wait on the shadow pre-pass before falling back to real execution.
• streaming (default True): stream the code output during generate_action so tool calls overlap with token generation. Set streaming=False for the Lazy/JIT one-shot shadow over the assembled block.

Composition with Hooks

Speculation composes with enable_rlm_hooks. Call hooks first, then speculation, so both stay active:

from dspy_rlm_hooks import enable_rlm_hooks, enable_rlm_speculation

enable_rlm_hooks(rlm, pre_execution_hook=sanitize_code)
enable_rlm_speculation(rlm)

The reverse order leaves speculation inactive (hooks overwrite the wrapper), though hooks still work.

Limitations

• Streaming requires a streaming-capable adapter (ChatAdapter/XMLAdapter/JSONAdapter) and an LM that supports streaming. Otherwise execution falls back to Lazy/JIT.
• The shadow runs in a subprocess for side-effect safety, which adds spawn overhead.
• Only speculatable, pure tools are speculated. Unmarked tools are never run early.
• PredictRLM is not supported.

(back to top)

PredictRLM Support

enable_rlm_hooks works on both dspy.RLM and PredictRLM with the same API. The function auto-detects the RLM type and uses the appropriate mechanism.

Install with the predict-rlm extra:

uv add "dspy-rlm-hooks[predict-rlm]"

Quick Example

from predict_rlm import PredictRLM
from dspy_rlm_hooks import enable_rlm_hooks, PreExecutionOutput

rlm = PredictRLM("query -> answer")

def sanitize_code(iteration, code, variables, history, input_args):
    """Block dangerous code patterns."""
    if "os.system" in code:
        code = code.replace("os.system", "# BLOCKED")
    return PreExecutionOutput(code=code)

enable_rlm_hooks(rlm, pre_execution_hook=sanitize_code)
result = rlm(query="...")

(back to top)

Hook Reference

(back to top)

Development

Code Quality

This project uses several tools to maintain code quality:

• Ruff: Linting and formatting
• isort: Import sorting
• pytest: Testing framework
• ty: Type checking
• deptry: Dependency analysis

Available commands:

# Run all quality checks
uv run poe clean-full

# Individual checks
uv run poe lint          # Ruff linting
uv run poe format        # Ruff formatting
uv run poe sort          # Import sorting
uv run poe typecheck     # Type checking
uv run poe deptry        # Dependency analysis

Testing

Run tests using pytest:

# Run all tests
uv run pytest

# Run specific test
uv run pytest path/to/test.py::test_name

(back to top)

Contributing

Quick workflow:

1. Fork and branch: git checkout -b feature/name
2. Make changes
3. Run checks: uv run poe clean-full
4. Commit and push
5. Open a Pull Request

See the full contributing guide for detailed setup instructions, project structure, and style guidelines.

(back to top)

License

MIT (as declared in pyproject.toml).