Skip to main content

Recursive Language Model (RLM) toolset for Pydantic AI - handle extremely large contexts

Project description

Pydantic AI RLM

Handle Extremely Large Contexts with Any LLM Provider

GitHubExamples

Python 3.10+ License: MIT Pydantic AI

Switch Providers Instantly  •  Sandboxed Code Execution  •  Sub-Model Delegation  •  Fully Type-Safe


What is RLM?

RLM (Recursive Language Model) is a pattern for handling contexts that exceed a model's context window, introduced by Alex L. Zhang, Tim Kraska, and Omar Khattab in their paper Recursive Language Models. Instead of trying to fit everything into one prompt, the LLM writes Python code to programmatically explore and analyze the data.

The key insight: An LLM can write code to search through millions of lines in seconds, then use llm_query() to delegate semantic analysis of relevant chunks to a sub-model.

This library is an implementation inspired by the original minimal implementation.


Get Started in 60 Seconds

pip install pydantic-ai-rlm
from pydantic_ai_rlm import run_rlm_analysis

answer = await run_rlm_analysis(
    context=massive_document,  # Can be millions of characters
    query="Find the magic number hidden in the text",
    model="openai:gpt-5",
    sub_model="openai:gpt-5-mini",
)

That's it. Your agent can now:

  • Write Python code to analyze massive contexts
  • Use llm_query() to delegate semantic analysis to sub-models
  • Work with any Pydantic AI compatible provider

Why pydantic-ai-rlm?

Switch Providers Instantly

Built on Pydantic AI, you can test any model with a single line change:

# OpenAI
agent = create_rlm_agent(model="openai:gpt-5", sub_model="openai:gpt-5-mini")

# Anthropic
agent = create_rlm_agent(model="anthropic:claude-sonnet-4-5", sub_model="anthropic:claude-haiku-4-5")

agent = create_rlm_agent(model="anthropic:claude-sonnet-4-5", sub_model="openai:gpt-5-mini")

Reusable Toolset

The RLM toolset integrates with any pydantic-ai agent:

from pydantic_ai import Agent
from pydantic_ai_rlm import create_rlm_toolset, RLMDependencies

# Use the toolset in any agent
toolset = create_rlm_toolset(sub_model="openai:gpt-5-mini")
agent = Agent("openai:gpt-5", toolsets=[toolset])

How It Works

┌─────────────────────────────────────────────────────────────────┐
│                         pydantic-ai-rlm                         │
│                                                                 │
│   ┌─────────────┐         ┌─────────────────────────────────┐   │
│   │   Main LLM  │         │     Sandboxed REPL Environment  │   │
│   │   (gpt-5)   │────────>│                                 │   │
│   └─────────────┘         │   context = <your massive data> │   │
│         │                 │                                 │   │
│         │                 │   # LLM writes Python code:     │   │
│         │                 │   for line in context.split():  │   │
│         │                 │       if "magic" in line:       │   │
│         │                 │           result = llm_query(   │   │
│         │                 │               f"Analyze: {line}"│   │
│         │                 │           )                     │   │
│         │                 │                                 │   │
│         │                 └───────────────┬─────────────────┘   │
│         │                                 │                     │
│         │                                 ▼                     │
│         │                       ┌─────────────────┐             │
│         │                       │    Sub LLM      │             │
│         │                       │  (gpt-5-mini)   │             │
│         │                       └─────────────────┘             │
│         │                                                       │
│         ▼                                                       │
│   ┌─────────────┐                                               │
│   │   Answer    │                                               │
│   └─────────────┘                                               │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘
  1. Main LLM receives the query and writes Python code
  2. REPL Environment executes code with access to context variable
  3. llm_query() delegates semantic analysis to a cheaper/faster sub-model
  4. Main LLM synthesizes the final answer from code execution results

Examples

Needle in Haystack

Find specific information in massive text:

from pydantic_ai_rlm import run_rlm_analysis

# 1 million lines of text with a hidden number
massive_text = generate_haystack(num_lines=1_000_000)

answer = await run_rlm_analysis(
    context=massive_text,
    query="Find the magic number hidden in the text",
    model="openai:gpt-5",
    sub_model="openai:gpt-5-mini",
)

JSON Data Analysis

Works with structured data too:

from pydantic_ai_rlm import create_rlm_agent, RLMDependencies

agent = create_rlm_agent(model="openai:gpt-5")

deps = RLMDependencies(
    context={"users": [...], "transactions": [...], "logs": [...]},
)

result = await agent.run(
    "Find all users with suspicious transaction patterns",
    deps=deps,
)

API Reference

create_rlm_agent()

Create a Pydantic AI agent with RLM capabilities.

agent = create_rlm_agent(
    model="openai:gpt-5",           # Main model for orchestration
    sub_model="openai:gpt-5-mini",  # Model for llm_query() (optional)
    code_timeout=60.0,               # Timeout for code execution
    custom_instructions="...",       # Additional instructions
)

create_rlm_toolset()

Create a standalone RLM toolset for composition.

toolset = create_rlm_toolset(
    code_timeout=60.0,
    sub_model="openai:gpt-5-mini",
)

run_rlm_analysis() / run_rlm_analysis_sync()

Convenience functions for quick analysis.

# Async
answer = await run_rlm_analysis(context, query, model="openai:gpt-5")

# Sync
answer = run_rlm_analysis_sync(context, query, model="openai:gpt-5")

RLMDependencies

Dependencies for RLM agents.

deps = RLMDependencies(
    context="...",  # str, dict, or list
    config=RLMConfig(
        code_timeout=60.0,
        truncate_output_chars=50_000,
        sub_model="openai:gpt-5-mini",
    ),
)

configure_logging()

Enable verbose logging to see what the agent is doing in real-time.

from pydantic_ai_rlm import configure_logging, run_rlm_analysis

# Enable logging (uses rich if installed, falls back to plain text)
configure_logging(enabled=True)

# Now you'll see code executions and outputs in the terminal
answer = await run_rlm_analysis(
    context=massive_document,
    query="Find the magic number",
    model="openai:gpt-5",
)

# Disable logging when done
configure_logging(enabled=False)

Install with rich logging support for syntax highlighting and styled output:

pip install pydantic-ai-rlm[logging]

Or install rich separately:

pip install rich

When enabled, you'll see:

  • Syntax-highlighted code being executed (with rich)
  • Execution results with status indicators (SUCCESS/ERROR)
  • Execution time for each code block
  • Variables created during execution
  • LLM sub-queries and responses (when using llm_query())

Note: Logging works without rich installed - it will use plain text output instead of styled panels


REPL Environment

The sandboxed REPL provides:

Feature Description
context variable Your data loaded and ready to use
llm_query(prompt) Delegate to sub-model (if configured)
Safe built-ins print, len, range, etc.
Common imports json, re, collections, etc.
Persistent state Variables persist across executions
Output capture stdout/stderr returned to agent

Blocked for security: eval, exec, compile, open (outside temp dir)


Related Projects

  • rlm - Original RLM implementation by Alex L. Zhang, Tim Kraska, and Omar Khattab
  • rlm-minimal - Minal RLM implementation by Alex L. Zhang
  • pydantic-ai - The foundation: Agent framework by Pydantic
  • pydantic-deep - Full deep agent framework with planning, filesystem, and more

Contributing

git clone https://github.com/vstorm-co/pydantic-ai-rlm.git
cd pydantic-ai-rlm
pip install -e ".[dev]"
pytest

License

MIT — see LICENSE

Built with Pydantic AI by vstorm-co

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

pydantic_ai_rlm-0.1.0.tar.gz (20.0 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

pydantic_ai_rlm-0.1.0-py3-none-any.whl (20.0 kB view details)

Uploaded Python 3

File details

Details for the file pydantic_ai_rlm-0.1.0.tar.gz.

File metadata

  • Download URL: pydantic_ai_rlm-0.1.0.tar.gz
  • Upload date:
  • Size: 20.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.9.27 {"installer":{"name":"uv","version":"0.9.27","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}

File hashes

Hashes for pydantic_ai_rlm-0.1.0.tar.gz
Algorithm Hash digest
SHA256 c4d4d6143ecd918c3583c2fb25f2934266bbd39eaf7f4581bc9cfe85c743d932
MD5 a4aa5f92937b49ad5b742b21de8234c2
BLAKE2b-256 96a70b3bb6d444aa853a37880da30f64382a523a4cc79ffbd2833b1ad925eef9

See more details on using hashes here.

File details

Details for the file pydantic_ai_rlm-0.1.0-py3-none-any.whl.

File metadata

  • Download URL: pydantic_ai_rlm-0.1.0-py3-none-any.whl
  • Upload date:
  • Size: 20.0 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.9.27 {"installer":{"name":"uv","version":"0.9.27","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}

File hashes

Hashes for pydantic_ai_rlm-0.1.0-py3-none-any.whl
Algorithm Hash digest
SHA256 a2c7594339d781bf8557dc056fa2b6ea5807a60ac125074e4c7d7e63c0880df0
MD5 ecc9bef854143a259dd583c2f21ccde0
BLAKE2b-256 8a7eb5286ba56fc5c066c44a6a9edd2446e1ff0304f3ee19ad9a279ad9321cdb

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page