pydantic-ai-rlm 0.1.0__py3-none-any.whl

pydantic_ai_rlm/__init__.py

@@ -0,0 +1,32 @@
+from .agent import create_rlm_agent, run_rlm_analysis, run_rlm_analysis_sync
+from .dependencies import ContextType, RLMConfig, RLMDependencies
+from .logging import configure_logging
+from .prompts import (
+    LLM_QUERY_INSTRUCTIONS,
+    RLM_INSTRUCTIONS,
+    build_rlm_instructions,
+)
+from .repl import REPLEnvironment, REPLResult
+from .toolset import (
+    cleanup_repl_environments,
+    create_rlm_toolset,
+)
+
+__all__ = [
+    "LLM_QUERY_INSTRUCTIONS",
+    "RLM_INSTRUCTIONS",
+    "ContextType",
+    "REPLEnvironment",
+    "REPLResult",
+    "RLMConfig",
+    "RLMDependencies",
+    "build_rlm_instructions",
+    "cleanup_repl_environments",
+    "configure_logging",
+    "create_rlm_agent",
+    "create_rlm_toolset",
+    "run_rlm_analysis",
+    "run_rlm_analysis_sync",
+]
+
+__version__ = "0.1.0"

pydantic_ai_rlm/agent.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic_ai import Agent, UsageLimits
+
+from .dependencies import ContextType, RLMConfig, RLMDependencies
+from .prompts import build_rlm_instructions
+from .toolset import create_rlm_toolset
+
+
+def create_rlm_agent(
+    model: str = "openai:gpt-5",
+    sub_model: str | None = None,
+    code_timeout: float = 60.0,
+    include_example_instructions: bool = True,
+    custom_instructions: str | None = None,
+) -> Agent[RLMDependencies, str]:
+    """
+    Create a Pydantic AI agent with REPL code execution capabilities.
+
+    Args:
+        model: Model to use for the main agent
+        sub_model: Model to use for llm_query() within the REPL environment.
+            If provided, a `llm_query(prompt: str) -> str` function becomes
+            available in the REPL, allowing the agent to delegate sub-queries.
+            Example: "openai:gpt-5-mini" or "anthropic:claude-3-haiku-20240307"
+        code_timeout: Timeout for code execution in seconds
+        include_example_instructions: Include detailed examples in instructions
+        custom_instructions: Additional instructions to append
+
+    Returns:
+        Configured Agent instance
+
+    Example:
+        ```python
+        from pydantic_ai_rlm import create_rlm_agent, RLMDependencies, RLMConfig
+
+        # Create agent with sub-model for llm_query
+        agent = create_rlm_agent(
+            model="openai:gpt-5",
+            sub_model="openai:gpt-5-mini",
+        )
+
+        deps = RLMDependencies(
+            context=very_large_document,
+            config=RLMConfig(sub_model="openai:gpt-5-mini"),
+        )
+        result = await agent.run("What are the main themes?", deps=deps)
+        print(result.output)
+        ```
+    """
+    toolset = create_rlm_toolset(code_timeout=code_timeout, sub_model=sub_model)
+
+    instructions = build_rlm_instructions(
+        include_llm_query=sub_model is not None,
+        custom_suffix=custom_instructions,
+    )
+
+    agent: Agent[RLMDependencies, str] = Agent(
+        model,
+        deps_type=RLMDependencies,
+        output_type=str,
+        toolsets=[toolset],
+        instructions=instructions,
+    )
+
+    return agent
+
+
+async def run_rlm_analysis(
+    context: ContextType,
+    query: str,
+    model: str = "openai:gpt-5",
+    sub_model: str | None = None,
+    config: RLMConfig | None = None,
+    max_tool_calls: int = 50,
+    **agent_kwargs: Any,
+) -> str:
+    """
+    Convenience function to run RLM analysis on a context.
+
+    Args:
+        context: The large context to analyze (string, dict, or list)
+        query: The question to answer about the context
+        model: Model to use for the main agent
+        sub_model: Model to use for llm_query() within the REPL environment.
+            If provided, a `llm_query(prompt: str) -> str` function becomes
+            available in the REPL, allowing the agent to delegate sub-queries.
+        config: Optional RLMConfig for customization
+        max_tool_calls: Maximum tool calls allowed
+        **agent_kwargs: Additional arguments passed to create_rlm_agent()
+
+    Returns:
+        The agent's final answer as a string
+
+    Example:
+        ```python
+        from pydantic_ai_rlm import run_rlm_analysis
+
+        # With sub-model for llm_query
+        answer = await run_rlm_analysis(
+            context=huge_document,
+            query="Find the magic number hidden in the text",
+            sub_model="openai:gpt-5-mini",
+        )
+        print(answer)
+        ```
+    """
+    agent = create_rlm_agent(model=model, sub_model=sub_model, **agent_kwargs)
+
+    effective_config = config or RLMConfig()
+    if sub_model and not effective_config.sub_model:
+        effective_config.sub_model = sub_model
+
+    deps = RLMDependencies(
+        context=context,
+        config=effective_config,
+    )
+
+    result = await agent.run(
+        query,
+        deps=deps,
+        usage_limits=UsageLimits(tool_calls_limit=max_tool_calls),
+    )
+
+    return result.output
+
+
+def run_rlm_analysis_sync(
+    context: ContextType,
+    query: str,
+    model: str = "openai:gpt-5",
+    sub_model: str | None = None,
+    config: RLMConfig | None = None,
+    max_tool_calls: int = 50,
+    **agent_kwargs: Any,
+) -> str:
+    """
+    Synchronous version of run_rlm_analysis.
+
+    See run_rlm_analysis() for full documentation.
+    """
+    agent = create_rlm_agent(model=model, sub_model=sub_model, **agent_kwargs)
+
+    effective_config = config or RLMConfig()
+    if sub_model and not effective_config.sub_model:
+        effective_config.sub_model = sub_model
+
+    deps = RLMDependencies(
+        context=context,
+        config=effective_config,
+    )
+
+    result = agent.run_sync(
+        query,
+        deps=deps,
+        usage_limits=UsageLimits(tool_calls_limit=max_tool_calls),
+    )
+
+    return result.output

pydantic_ai_rlm/dependencies.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+ContextType = str | dict[str, Any] | list[Any]
+
+
+@dataclass
+class RLMConfig:
+    """Configuration for RLM behavior."""
+
+    code_timeout: float = 60.0
+    """Timeout in seconds for code execution."""
+
+    truncate_output_chars: int = 50_000
+    """Maximum characters to return from code execution output."""
+
+    sub_model: str | None = None
+    """
+    Model to use for llm_query() within the REPL environment.
+
+    If set, a `llm_query(prompt: str) -> str` function becomes available
+    in the REPL environment, allowing the main LLM to delegate sub-queries
+    to another model. This is useful for processing large contexts in chunks.
+    """
+
+
+@dataclass
+class RLMDependencies:
+    """
+    Dependencies injected into RLM tools via RunContext.
+
+    This holds the context data and configuration that
+    the RLM toolset needs to operate.
+    """
+
+    context: ContextType
+    """The context to analyze (string, dict, or list)."""
+
+    config: RLMConfig = field(default_factory=RLMConfig)
+    """RLM configuration options."""
+
+    def __post_init__(self):
+        """Validate dependencies after initialization."""
+        if self.context is None:
+            raise ValueError("context cannot be None")

pydantic_ai_rlm/logging.py

@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .repl import REPLResult
+
+# Check if rich is available
+try:
+    from rich.console import Console
+    from rich.panel import Panel
+    from rich.syntax import Syntax
+    from rich.text import Text
+
+    RICH_AVAILABLE = True
+except ImportError:
+    RICH_AVAILABLE = False
+
+
+class RLMLogger:
+    """
+    Pretty logger for RLM code execution.
+
+    Uses rich for fancy terminal output with syntax highlighting and styled panels.
+    Falls back to plain text if rich is not installed.
+    """
+
+    def __init__(self, enabled: bool = True):
+        self.enabled = enabled
+        if RICH_AVAILABLE:
+            self.console = Console()
+        else:
+            self.console = None
+
+    def log_code_execution(self, code: str) -> None:
+        """Log the code being executed."""
+        if not self.enabled:
+            return
+
+        if RICH_AVAILABLE and self.console:
+            syntax = Syntax(code, "python", theme="monokai", line_numbers=True)
+            panel = Panel(
+                syntax,
+                title="[bold cyan]Code Execution[/bold cyan]",
+                border_style="cyan",
+                padding=(0, 1),
+            )
+            self.console.print(panel)
+        else:
+            print(f"\n{'='*50}")
+            print("CODE EXECUTION")
+            print("=" * 50)
+            print(code)
+            print("=" * 50)
+
+    def log_result(self, result: REPLResult) -> None:
+        """Log the execution result."""
+        if not self.enabled:
+            return
+
+        if RICH_AVAILABLE and self.console:
+            self._log_result_rich(result)
+        else:
+            self._log_result_plain(result)
+
+    def _log_result_rich(self, result: REPLResult) -> None:
+        """Log result using rich formatting."""
+        status, border_style = self._get_status_style(result.success)
+        content_parts = self._build_content_parts(result)
+        user_vars = self._get_user_vars(result.locals)
+
+        self._print_result_panel(content_parts, status, border_style, user_vars)
+
+    def _get_status_style(self, success: bool) -> tuple:
+        """Get status text and border style based on success."""
+        if success:
+            return Text("SUCCESS", style="bold green"), "green"
+        return Text("ERROR", style="bold red"), "red"
+
+    def _build_content_parts(self, result: REPLResult) -> list:
+        """Build content parts for the result panel."""
+        parts = [Text(f"Executed in {result.execution_time:.3f}s", style="dim")]
+
+        if result.stdout.strip():
+            stdout = result.stdout.strip()
+            if len(stdout) > 2000:
+                stdout = stdout[:2000] + "\n... (truncated)"
+            parts.extend([Text("\n"), Text("Output:", style="bold yellow"), Text("\n"), Text(stdout, style="white")])
+
+        if result.stderr.strip():
+            stderr = result.stderr.strip()
+            if len(stderr) > 1000:
+                stderr = stderr[:1000] + "\n... (truncated)"
+            parts.extend([Text("\n"), Text("Errors:", style="bold red"), Text("\n"), Text(stderr, style="red")])
+
+        return parts
+
+    def _get_user_vars(self, locals_dict: dict) -> dict:
+        """Extract user-defined variables from locals."""
+        excluded = ("context", "json", "re", "os", "collections", "math")
+        return {k: v for k, v in locals_dict.items() if not k.startswith("_") and k not in excluded}
+
+    def _print_result_panel(self, content_parts: list, status, border_style: str, user_vars: dict) -> None:
+        """Print the result panel and optional variables table."""
+        from rich.table import Table
+
+        if user_vars:
+            content_parts.extend([Text("\n"), Text("Variables:", style="bold magenta"), Text("\n")])
+            if len(user_vars) > 10:
+                content_parts.append(Text(f"  ... and {len(user_vars) - 10} more variables\n", style="dim"))
+
+        combined = Text()
+        for part in content_parts:
+            combined.append(part)
+
+        panel = Panel(combined, title=f"[bold]Result: {status}[/bold]", border_style=border_style, padding=(0, 1))
+        self.console.print(panel)
+
+        if user_vars:
+            var_table = Table(show_header=True, header_style="bold", box=None, padding=(0, 1))
+            var_table.add_column("Name", style="cyan")
+            var_table.add_column("Type", style="yellow")
+            var_table.add_column("Value", style="white", max_width=60)
+
+            for name, value in list(user_vars.items())[:10]:
+                value_str = self._format_var_value(value)
+                var_table.add_row(name, type(value).__name__, value_str)
+
+            self.console.print(var_table)
+
+    def _format_var_value(self, value) -> str:
+        """Format a variable value for display."""
+        try:
+            value_str = repr(value)
+            if len(value_str) > 60:
+                return value_str[:57] + "..."
+            return value_str
+        except Exception:
+            return "<unable to repr>"
+
+    def _log_result_plain(self, result: REPLResult) -> None:
+        """Log result using plain text."""
+        status = "SUCCESS" if result.success else "ERROR"
+        print(f"\n{'='*50}")
+        print(f"RESULT: {status} (executed in {result.execution_time:.3f}s)")
+        print("=" * 50)
+
+        if result.stdout.strip():
+            print("\nOutput:")
+            stdout = result.stdout.strip()
+            if len(stdout) > 2000:
+                stdout = stdout[:2000] + "\n... (truncated)"
+            print(stdout)
+
+        if result.stderr.strip():
+            print("\nErrors:")
+            stderr = result.stderr.strip()
+            if len(stderr) > 1000:
+                stderr = stderr[:1000] + "\n... (truncated)"
+            print(stderr)
+
+        user_vars = {
+            k: v
+            for k, v in result.locals.items()
+            if not k.startswith("_") and k not in ("context", "json", "re", "os")
+        }
+        if user_vars:
+            print("\nVariables:")
+            for name, value in list(user_vars.items())[:10]:
+                try:
+                    value_str = repr(value)
+                    if len(value_str) > 60:
+                        value_str = value_str[:57] + "..."
+                except Exception:
+                    value_str = "<unable to repr>"
+                print(f"  {name} ({type(value).__name__}): {value_str}")
+            if len(user_vars) > 10:
+                print(f"  ... and {len(user_vars) - 10} more variables")
+
+        print("=" * 50)
+
+    def log_llm_query(self, prompt: str) -> None:
+        """Log an llm_query call."""
+        if not self.enabled:
+            return
+
+        if RICH_AVAILABLE and self.console:
+            # Truncate long prompts
+            display_prompt = prompt
+            if len(display_prompt) > 500:
+                display_prompt = display_prompt[:500] + "..."
+
+            panel = Panel(
+                Text(display_prompt, style="white"),
+                title="[bold blue]LLM Query[/bold blue]",
+                border_style="blue",
+                padding=(0, 1),
+            )
+            self.console.print(panel)
+        else:
+            print(f"\n{'='*50}")
+            print("LLM QUERY")
+            print("=" * 50)
+            display_prompt = prompt
+            if len(display_prompt) > 500:
+                display_prompt = display_prompt[:500] + "..."
+            print(display_prompt)
+            print("=" * 50)
+
+    def log_llm_response(self, response: str) -> None:
+        """Log an llm_query response."""
+        if not self.enabled:
+            return
+
+        if RICH_AVAILABLE and self.console:
+            # Truncate long responses
+            display_response = response
+            if len(display_response) > 500:
+                display_response = display_response[:500] + "..."
+
+            panel = Panel(
+                Text(display_response, style="white"),
+                title="[bold blue]LLM Response[/bold blue]",
+                border_style="blue",
+                padding=(0, 1),
+            )
+            self.console.print(panel)
+        else:
+            print(f"\n{'='*50}")
+            print("LLM RESPONSE")
+            print("=" * 50)
+            display_response = response
+            if len(display_response) > 500:
+                display_response = display_response[:500] + "..."
+            print(display_response)
+            print("=" * 50)
+
+
+# Global logger instance
+_logger: RLMLogger | None = None
+
+
+def get_logger() -> RLMLogger:
+    """Get the global RLM logger instance."""
+    global _logger
+    if _logger is None:
+        _logger = RLMLogger(enabled=False)  # Disabled by default
+    return _logger
+
+
+def configure_logging(enabled: bool = True) -> RLMLogger:
+    """
+    Configure RLM logging.
+
+    Args:
+        enabled: Whether to enable logging output
+
+    Returns:
+        The configured logger instance
+
+    Example:
+        ```python
+        from pydantic_ai_rlm import configure_logging
+
+        # Enable fancy logging
+        configure_logging(enabled=True)
+
+        # Run your analysis - you'll see code and output in the terminal
+        result = await run_rlm_analysis(context, query)
+        ```
+    """
+    global _logger
+    _logger = RLMLogger(enabled=enabled)
+    return _logger

pydantic_ai_rlm/prompts.py

@@ -0,0 +1,118 @@
+RLM_INSTRUCTIONS = """You are an AI assistant that analyzes data using Python code execution. You have access to a REPL environment where code persists between executions.
+
+## REPL Environment
+
+The REPL environment provides:
+1. A `context` variable containing your data (string, dict, or list)
+2. Common modules available via import: `re`, `json`, `collections`, etc.
+3. Variables persist between code executions
+
+## Strategy for Large Contexts
+
+### Step 1: Explore the Context Structure
+```python
+print(f"Context type: {type(context)}")
+print(f"Context length: {len(context)}")
+if isinstance(context, str):
+    print(f"First 500 chars: {context[:500]}")
+```
+
+### Step 2: Process the Data
+For structured data:
+```python
+import re
+sections = re.split(r'### (.+)', context)
+for i in range(1, len(sections), 2):
+    header = sections[i]
+    content = sections[i+1][:200]
+    print(f"{header}: {content}...")
+```
+
+For raw text - search patterns:
+```python
+import re
+matches = re.findall(r'\\d{4}-\\d{2}-\\d{2}', context)
+print(f"Found {len(matches)} dates: {matches[:10]}")
+```
+
+### Step 3: Build Your Answer
+```python
+results = []
+# ... process data ...
+print(f"Final answer: {results}")
+```
+
+## Guidelines
+
+1. **Always explore first** - Check context type and size before processing
+2. **Use print() liberally** - See intermediate results
+3. **Store results in variables** - Build up your answer incrementally
+4. **Be thorough** - For needle-in-haystack, search the entire context
+"""
+
+LLM_QUERY_INSTRUCTIONS = """
+
+## Sub-LLM Queries
+
+You also have access to `llm_query(prompt: str) -> str` function that allows you to query another LLM from within your REPL code. This is extremely useful for:
+- **Semantic analysis** - Understanding meaning, not just text patterns
+- **Summarization** - Condensing large sections of context
+- **Chunked processing** - Analyzing context in manageable pieces
+- **Complex reasoning** - Delegating sub-tasks that require language understanding
+
+### Example: Chunked Analysis
+```python
+# Split context into chunks and analyze each with llm_query
+chunk_size = 50000
+chunks = [context[i:i+chunk_size] for i in range(0, len(context), chunk_size)]
+
+summaries = []
+for i, chunk in enumerate(chunks):
+    summary = llm_query(f"Summarize this section:\\n{chunk}")
+    summaries.append(f"Chunk {i+1}: {summary}")
+    print(f"Processed chunk {i+1}/{len(chunks)}")
+
+# Combine summaries for final answer
+final = llm_query(f"Based on these summaries, answer: What are the main themes?\\n" + "\\n".join(summaries))
+print(final)
+```
+
+### Example: Semantic Search
+```python
+# Use llm_query for semantic understanding
+result = llm_query(f"Find any mentions of 'magic number' in this text and return the value:\\n{context[:100000]}")
+print(result)
+```
+
+**Tips:**
+- The sub-LLM can handle ~500K characters per query
+- Use it for semantic analysis that regex/string operations can't do
+- Store sub-LLM results in variables to build up your answer
+"""
+
+
+def build_rlm_instructions(
+    include_llm_query: bool = False,
+    custom_suffix: str | None = None,
+) -> str:
+    """
+    Build RLM instructions with optional customization.
+
+    Args:
+        include_examples: Whether to include detailed examples
+        include_llm_query: Whether to include llm_query() documentation
+        custom_suffix: Additional instructions to append
+
+    Returns:
+        Complete instructions string
+    """
+    base = RLM_INSTRUCTIONS
+
+    if include_llm_query:
+        llm_docs = LLM_QUERY_INSTRUCTIONS
+        base = f"{base}{llm_docs}"
+
+    if custom_suffix:
+        base = f"{base}\n\n## Additional Instructions\n\n{custom_suffix}"
+
+    return base