Metadata-Version: 2.4
Name: dspy-rlm-hooks
Version: 0.1.10
Summary: Lifecycle instrumentation for DSPy's RLM (Recursive Language Model).
Keywords: dspy,ai,recursive language model,hooks,instrumentation,rlm
Author: Edward Boswell
Author-email: Edward Boswell <thememium@gmail.com>
License-Expression: MIT
License-File: LICENSE
Classifier: Operating System :: OS Independent
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Dist: dspy>=3.1.0
Requires-Dist: pydantic>=2.0.0
Requires-Dist: predict-rlm>=0.1.0 ; extra == 'predict-rlm'
Requires-Dist: mlflow>=2.14.0 ; extra == 'tracing'
Requires-Python: >=3.12
Project-URL: Homepage, https://github.com/thememium/dspy-rlm-hooks
Project-URL: Documentation, https://github.com/thememium/dspy-rlm-hooks
Project-URL: Repository, https://github.com/thememium/dspy-rlm-hooks.git
Project-URL: Issues, https://github.com/thememium/dspy-rlm-hooks/issues
Project-URL: Changelog, https://github.com/thememium/dspy-rlm-hooks/blob/master/CHANGELOG.md
Provides-Extra: predict-rlm
Provides-Extra: tracing
Description-Content-Type: text/markdown

<a name="readme-top"></a>

<div align="center">
  <h3 align="center">DSPy RLM Hooks</h3>

  <p align="center">
    Lifecycle instrumentation for DSPy's RLM (Recursive Language Model).
    <br />
    <a href="#table-of-contents"><strong>Explore the Documentation »</strong></a>
    <br />
    <a href="https://github.com/thememium/dspy-rlm-hooks/issues">Report Bug</a>
    ·
    <a href="https://github.com/thememium/dspy-rlm-hooks/issues">Request Feature</a>
  </p>
</div>

<!-- TABLE OF CONTENTS -->

<a name="table-of-contents"></a>

<details>
  <summary>Table of Contents</summary>
  <ol>
    <li><a href="#about">About</a></li>
    <li><a href="#quick-start">Quick Start</a></li>
    <li><a href="#usage">Usage</a></li>
    <li><a href="#speculative-execution">Speculative Execution</a></li>
    <li><a href="#predictrlm-support">PredictRLM Support</a></li>
    <li><a href="#development">Development</a></li>
    <li><a href="#contributing">Contributing</a></li>
    <li><a href="#license">License</a></li>
  </ol>
</details>

<!-- ABOUT -->

## 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](https://github.com/Trampoline-AI/predict-rlm) instances

Requires **DSPy 3.1+** and **Pydantic 2+**.

<p align="right">(<a href="#readme-top">back to top</a>)</p>

<!-- ARCHITECTURE -->

## 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.

<p align="right">(<a href="#readme-top">back to top</a>)</p>

<!-- QUICK START -->

## Quick Start

### Install

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

```bash
uv add dspy-rlm-hooks
```

Or with pip:

```bash
pip install dspy-rlm-hooks
```

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

```bash
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

```python
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.

<p align="right">(<a href="#readme-top">back to top</a>)</p>

<!-- USAGE -->

## Usage

### All Four Hooks

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

```python
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:

```python
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:

```python
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

```python
from dspy_rlm_hooks import disable_rlm_hooks

disable_rlm_hooks(rlm)
```

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

<p align="right">(<a href="#readme-top">back to top</a>)</p>

<!-- SPECULATIVE EXECUTION -->

## 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

```python
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`:

```python
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:

```python
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.

<p align="right">(<a href="#readme-top">back to top</a>)</p>

<!-- PREDICTRLM SUPPORT -->

## PredictRLM Support

`enable_rlm_hooks` works on both `dspy.RLM` and
[PredictRLM](https://github.com/Trampoline-AI/predict-rlm) with the same API.
The function auto-detects the RLM type and uses the appropriate mechanism.

Install with the `predict-rlm` extra:

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

### Quick Example

```python
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="...")
```

<p align="right">(<a href="#readme-top">back to top</a>)</p>

## Hook Reference

| Hook | When it fires | What it can do |
| --- | --- | --- |
| **PreIteration** | Before action generation | Inject current execution code (`python_code`), replace/clear persistent code (`persistent_python_code`), add interpreter variables (`extra_vars`), or steer action generation (`prompt_context`) |
| **PreExecution** | After code generation, before running | Rewrite or sanitise the generated `code` string |
| **PostExecution** | After code runs, before history processing | Transform, audit, or replace the raw `result` |
| **PostIteration** | After result is folded into history | Save learnings, trigger side effects, modify `history`, or set `stop=True` to force final extraction |

<p align="right">(<a href="#readme-top">back to top</a>)</p>

<!-- DEVELOPMENT -->

## 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:**

```sh
# 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:

```sh
# Run all tests
uv run pytest

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

<p align="right">(<a href="#readme-top">back to top</a>)</p>

<!-- CONTRIBUTING -->

## 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](https://github.com/thememium/dspy-rlm-hooks/blob/master/.github/contributing.md) for detailed setup instructions, project structure, and style guidelines.

<p align="right">(<a href="#readme-top">back to top</a>)</p>

<!-- LICENSE -->

## License

MIT (as declared in `pyproject.toml`).

---

<div align="center">
  <p>
    <sub>Built by <a href="https://github.com/thememium">thememium</a></sub>
  </p>
</div>
