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

GitHubPyPIExamples

Python 3.10+ License: MIT Pydantic AI PyPI version

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.1.tar.gz (20.2 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.1-py3-none-any.whl (20.0 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: pydantic_ai_rlm-0.1.1.tar.gz
  • Upload date:
  • Size: 20.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for pydantic_ai_rlm-0.1.1.tar.gz
Algorithm Hash digest
SHA256 e02406611ae28a1d8d2d70ef4993bfe638e752e707dbf996f0f0c3f9a6b4b5ba
MD5 af68e2ed46725e5b5ffb3bdd0fcd5637
BLAKE2b-256 0444947b6ee93ebeaafb5025e1854a9283a8c85f2fd2626bc3355202b541cc6f

See more details on using hashes here.

Provenance

The following attestation bundles were made for pydantic_ai_rlm-0.1.1.tar.gz:

Publisher: publish.yml on vstorm-co/pydantic-ai-rlm

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

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

File metadata

File hashes

Hashes for pydantic_ai_rlm-0.1.1-py3-none-any.whl
Algorithm Hash digest
SHA256 565d73c34693ee8c82fe81d5f08126f93b33d09bc1e0a6a862493067be1ff271
MD5 d17b1a80c63e14fc008457e6a0cb907a
BLAKE2b-256 371444f9edd518163da6623433ba4cc2e210a3f2274cdc7e312d012c4ec890f1

See more details on using hashes here.

Provenance

The following attestation bundles were made for pydantic_ai_rlm-0.1.1-py3-none-any.whl:

Publisher: publish.yml on vstorm-co/pydantic-ai-rlm

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

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