lm-deluge 0.0.67__py3-none-any.whl → 0.0.88__py3-none-any.whl

lm_deluge/tool/prefab/rlm/__init__.py

@@ -0,0 +1,296 @@
+"""
+RLM (Recursive Language Model) for lm-deluge.
+
+Enables models to process long contexts through a REPL environment
+with recursive LM calls, based on the RLM paper:
+https://alexzhang13.github.io/blog/2025/rlm/
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from lm_deluge.prompt import Conversation
+from lm_deluge.tool import Tool
+
+from .executor import RLMExecutionError, RLMExecutor
+from .parse import RLMSecurityError
+
+if TYPE_CHECKING:
+    from lm_deluge.api_requests.base import APIResponse
+    from lm_deluge.client import _LLMClient
+
+
+RLM_SYSTEM_PROMPT = """You have access to a long context stored in the variable `{context_var}`.
+You can write Python code to analyze this context using the `execute` tool.
+
+IMPORTANT RULES:
+1. You MUST use print() to see output. Bare expressions produce NO output.
+2. You MUST call final(answer) when you have the answer. This is required!
+
+Available in your code environment:
+- `{context_var}`: The full context as a string ({context_len:,} characters)
+- `lm(prompt)`: Make a recursive LLM call (runs in parallel when possible)
+- `final(answer)`: Signal completion with the given answer - YOU MUST CALL THIS!
+- `final_var(varname)`: Signal completion with a variable's value
+- Modules: `re`, `math`, `collections`, `json` (imports are allowed but optional)
+- From collections: `Counter`, `defaultdict`, `deque`, `namedtuple`, `OrderedDict`
+- Standard builtins: `len`, `str`, `int`, `list`, `dict`, `sum`, `sorted`, `map`, `filter`, etc.
+
+Example - count word occurrences:
+```python
+count = len(re.findall(r'\\bword\\b', {context_var}))
+print(f"Found {{count}} occurrences")
+final(count)
+```
+
+Example - use Counter:
+```python
+words = {context_var}.split()
+counts = Counter(words)
+print(counts.most_common(10))
+final(counts.most_common(10))
+```
+
+Example - analyze with lm() calls:
+```python
+chunks = [{context_var}[i:i+2000] for i in range(0, len({context_var}), 2000)][:3]
+summaries = [lm(f"Summarize: {{chunk}}") for chunk in chunks]
+combined = "\\n".join(str(s) for s in summaries)
+final(f"Summary:\\n{{combined}}")
+```
+
+Variables persist between execute() calls. Always call final() when you have the answer!
+"""
+
+
+class RLMManager:
+    """Manages RLM execution for a long context.
+
+    The RLMManager exposes a REPL-like interface as tools that allow an LLM
+    to analyze a long context by writing Python code.
+
+    Example:
+        >>> manager = RLMManager(
+        ...     context=long_document,
+        ...     client=LLMClient("gpt-4.1-mini"),  # For lm() calls
+        ... )
+        >>> main_client = LLMClient("gpt-4.1")
+        >>> conv = Conversation.system(manager.get_system_prompt())
+        >>> conv = conv.user("What are the main themes in this document?")
+        >>> conv, resp = await main_client.run_agent_loop(
+        ...     conv,
+        ...     tools=manager.get_tools(),
+        ... )
+        >>> if manager.is_complete:
+        ...     print(manager.final_answer)
+    """
+
+    def __init__(
+        self,
+        context: str,
+        client: _LLMClient,
+        context_var_name: str = "CONTEXT",
+        max_lm_calls_per_execution: int = 20,
+    ):
+        """Initialize the RLMManager.
+
+        Args:
+            context: The long context string to analyze
+            client: LLMClient for making recursive lm() calls
+            context_var_name: Variable name for the context (default: "CONTEXT")
+            max_lm_calls_per_execution: Maximum lm() calls allowed per execute() call
+        """
+        self.context = context
+        self.client = client
+        self.context_var_name = context_var_name
+        self.max_lm_calls_per_execution = max_lm_calls_per_execution
+
+        self.executor = RLMExecutor(
+            context=context,
+            client=client,
+            context_var_name=context_var_name,
+            max_lm_calls_per_execution=max_lm_calls_per_execution,
+        )
+
+        self._final_answer: str | None = None
+        self._tools: list[Tool] | None = None
+
+    async def _execute(self, code: str) -> str:
+        """Execute code against the context."""
+        try:
+            answer, is_final = await self.executor.execute(code)
+            if is_final:
+                self._final_answer = answer
+                # Truncate for display but keep full answer stored
+                display = answer[:1000] + "..." if len(answer) > 1000 else answer
+                return f"[FINAL ANSWER SET]\n{display}"
+            return answer
+        except RLMSecurityError as e:
+            return f"Security error: {e}"
+        except RLMExecutionError as e:
+            return f"Execution error: {e}"
+        except Exception as e:
+            return f"Unexpected error: {type(e).__name__}: {e}"
+
+    def get_system_prompt(self) -> str:
+        """Get the system prompt explaining the RLM environment."""
+        return RLM_SYSTEM_PROMPT.format(
+            context_var=self.context_var_name,
+            context_len=len(self.context),
+        )
+
+    def get_tools(self) -> list[Tool]:
+        """Get the tools for RLM execution."""
+        if self._tools is not None:
+            return self._tools
+
+        self._tools = [
+            Tool(
+                name="execute",
+                description=(
+                    f"Execute Python code to analyze the context. "
+                    f"The context ({len(self.context):,} chars) is available as `{self.context_var_name}`. "
+                    f"Use `lm(prompt)` for recursive LLM calls (parallel when possible), and "
+                    f"`final(answer)` or `final_var(varname)` to signal completion. "
+                    f"Variables persist between calls. "
+                    f"Modules available without import: re, math, collections, json."
+                ),
+                run=self._execute,
+                parameters={
+                    "code": {
+                        "type": "string",
+                        "description": "Python code to execute",
+                    }
+                },
+                required=["code"],
+            )
+        ]
+        return self._tools
+
+    @property
+    def is_complete(self) -> bool:
+        """Check if FINAL() was called."""
+        return self._final_answer is not None
+
+    @property
+    def final_answer(self) -> str | None:
+        """Get the final answer if set."""
+        return self._final_answer
+
+    def reset(self) -> None:
+        """Reset the RLM state."""
+        self.executor.reset()
+        self._final_answer = None
+
+
+@dataclass
+class RLMResult:
+    """Result from RLMPipeline."""
+
+    answer: str
+    conversation: Conversation
+    rounds_used: int
+    final_response: APIResponse
+
+
+class RLMPipeline:
+    """High-level pipeline for RLM processing.
+
+    A thin wrapper that takes a long context and question, sets up an RLMManager,
+    runs an agent loop until final() is called, and returns the result.
+
+    Example:
+        >>> pipeline = RLMPipeline(
+        ...     context=long_document,
+        ...     client=LLMClient("gpt-4.1"),  # Smart orchestrator
+        ...     lm_client=LLMClient("gpt-4.1-mini"),  # Cheaper model for lm() calls
+        ...     question="What are the main themes in this document?",
+        ... )
+        >>> result = await pipeline.run()
+        >>> print(result.answer)
+    """
+
+    def __init__(
+        self,
+        context: str,
+        client: _LLMClient,
+        question: str,
+        *,
+        lm_client: _LLMClient | None = None,
+        context_var_name: str = "CONTEXT",
+        max_rounds: int = 15,
+        max_lm_calls_per_execution: int = 20,
+    ):
+        """Initialize the RLMPipeline.
+
+        Args:
+            context: The long context string to analyze
+            client: LLMClient for the main agent (runs the execute loop)
+            question: The question to answer about the context
+            lm_client: LLMClient for lm() calls (defaults to same as client)
+            context_var_name: Variable name for the context (default: "CONTEXT")
+            max_rounds: Maximum agent loop rounds (default: 15)
+            max_lm_calls_per_execution: Maximum lm() calls per execute() call
+        """
+        self.context = context
+        self.client = client
+        self.lm_client = lm_client or client
+        self.question = question
+        self.context_var_name = context_var_name
+        self.max_rounds = max_rounds
+        self.max_lm_calls_per_execution = max_lm_calls_per_execution
+
+    async def run(self) -> RLMResult:
+        """Run the RLM pipeline until completion."""
+        manager = RLMManager(
+            context=self.context,
+            client=self.lm_client,
+            context_var_name=self.context_var_name,
+            max_lm_calls_per_execution=self.max_lm_calls_per_execution,
+        )
+
+        # Build conversation with system prompt and question
+        conv = Conversation.system(manager.get_system_prompt())
+        conv = conv.user(
+            f"Question to answer about the context:\n\n{self.question}\n\n"
+            "Use the execute tool to analyze the context and find the answer. "
+            "Start by peeking at the context structure, then use appropriate "
+            "techniques (regex, chunking, lm() calls) to find the answer. "
+            "Call final(answer) when you have the answer."
+        )
+
+        # Run agent loop
+        conv, resp = await self.client.run_agent_loop(
+            conv,
+            tools=manager.get_tools(),
+            max_rounds=self.max_rounds,
+        )
+
+        # Extract answer
+        if manager.is_complete:
+            answer = manager.final_answer or "No answer produced"
+        else:
+            # Model stopped without calling final() - use last response
+            answer = resp.completion or "No answer produced (final not called)"
+
+        # Count rounds used
+        rounds_used = sum(1 for m in conv.messages if m.role == "assistant")
+
+        return RLMResult(
+            answer=answer,
+            conversation=conv,
+            rounds_used=rounds_used,
+            final_response=resp,
+        )
+
+
+__all__ = [
+    "RLMManager",
+    "RLMPipeline",
+    "RLMResult",
+    "RLMExecutor",
+    "RLMExecutionError",
+    "RLMSecurityError",
+]

lm_deluge/tool/prefab/rlm/executor.py

@@ -0,0 +1,349 @@
+"""
+RLM (Recursive Language Model) code executor.
+
+Executes Python code with access to a context variable and lm() function
+for recursive language model calls.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any, Callable
+
+from .parse import (
+    RLM_MODULES,
+    RLM_SAFE_BUILTINS,
+    RLMExecutionError,
+    validate_rlm_code,
+)
+
+if TYPE_CHECKING:
+    from lm_deluge.client import _LLMClient
+
+
+class OutputCapture:
+    """Captures print() output during execution."""
+
+    def __init__(self):
+        self.outputs: list[str] = []
+
+    def print(self, *args, **kwargs):
+        """Replacement print function that captures output."""
+        output = " ".join(str(arg) for arg in args)
+        self.outputs.append(output)
+
+    def get_output(self) -> str:
+        return "\n".join(self.outputs)
+
+
+class PendingLMResult:
+    """Placeholder for an lm() call result that hasn't completed yet."""
+
+    def __init__(self, call_id: int, results: dict[int, str]):
+        self._call_id = call_id
+        self._results = results
+
+    def _require_result(self) -> str:
+        if self._call_id not in self._results:
+            raise RuntimeError(f"LM result for call {self._call_id} not yet available")
+        return self._results[self._call_id]
+
+    def is_ready(self) -> bool:
+        return self._call_id in self._results
+
+    def __repr__(self) -> str:
+        return repr(self._require_result())
+
+    def __str__(self) -> str:
+        return str(self._require_result())
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._require_result(), name)
+
+    def __getitem__(self, key: Any) -> Any:
+        return self._require_result()[key]
+
+    def __iter__(self):
+        return iter(self._require_result())
+
+    def __len__(self) -> int:
+        return len(self._require_result())
+
+    def __bool__(self) -> bool:
+        return bool(self._require_result())
+
+    def __add__(self, other):
+        return self._require_result() + other
+
+    def __radd__(self, other):
+        return other + self._require_result()
+
+    def __contains__(self, item):
+        return item in self._require_result()
+
+
+class FinalAnswer(Exception):
+    """Raised when FINAL() or FINAL_VAR() is called to signal completion."""
+
+    def __init__(self, answer: Any):
+        self.answer = answer
+        super().__init__("Final answer signaled")
+
+
+def _resolve_value(value: Any, results: dict[int, str]) -> Any:
+    """Recursively resolve any PendingLMResult placeholders in a value."""
+    if isinstance(value, PendingLMResult):
+        return value._require_result()
+    if isinstance(value, list):
+        return [_resolve_value(v, results) for v in value]
+    if isinstance(value, tuple):
+        return tuple(_resolve_value(v, results) for v in value)
+    if isinstance(value, dict):
+        return {k: _resolve_value(v, results) for k, v in value.items()}
+    if isinstance(value, set):
+        return {_resolve_value(v, results) for v in value}
+    return value
+
+
+def _contains_unresolved(value: Any) -> bool:
+    """Check if a value contains any unresolved PendingLMResult."""
+    if isinstance(value, PendingLMResult):
+        return not value.is_ready()
+    if isinstance(value, (list, tuple, set)):
+        return any(_contains_unresolved(item) for item in value)
+    if isinstance(value, dict):
+        return any(_contains_unresolved(v) for v in value.values())
+    return False
+
+
+class RLMExecutor:
+    """Executes RLM code with access to context and lm() calls."""
+
+    def __init__(
+        self,
+        context: str,
+        client: _LLMClient,
+        context_var_name: str = "CONTEXT",
+        max_lm_calls_per_execution: int = 20,
+    ):
+        """Initialize the RLM executor.
+
+        Args:
+            context: The long context string to analyze
+            client: LLMClient for making recursive lm() calls
+            context_var_name: Variable name for the context (default: "CONTEXT")
+            max_lm_calls_per_execution: Maximum lm() calls allowed per execute() call
+        """
+        self.context = context
+        self.client = client
+        self.context_var_name = context_var_name
+        self.max_lm_calls_per_execution = max_lm_calls_per_execution
+
+        # Persistent state across execute() calls
+        self._persistent_locals: dict[str, Any] = {}
+
+    def _make_lm_wrapper(
+        self,
+        pending_lm_calls: list[dict],
+        lm_results: dict[int, str],
+        call_state: dict[str, int],
+        pending_call_ids: set[int],
+    ) -> Callable[[str], PendingLMResult]:
+        """Create the lm(prompt) wrapper function."""
+
+        def lm_call(prompt: str) -> PendingLMResult:
+            # Check for unresolved dependencies in the prompt
+            if _contains_unresolved(prompt):
+                raise RuntimeError("LM result for call dependency not yet available")
+
+            call_id = call_state["next_lm_id"]
+            call_state["next_lm_id"] += 1
+
+            # Only queue if not already completed or pending
+            if call_id not in lm_results and call_id not in pending_call_ids:
+                if len(pending_lm_calls) >= self.max_lm_calls_per_execution:
+                    raise RuntimeError(
+                        f"Too many lm() calls in single execution "
+                        f"(max {self.max_lm_calls_per_execution})"
+                    )
+                pending_call_ids.add(call_id)
+                pending_lm_calls.append(
+                    {
+                        "id": call_id,
+                        "prompt": str(prompt),
+                    }
+                )
+
+            return PendingLMResult(call_id, lm_results)
+
+        return lm_call
+
+    def _make_final_func(
+        self, exec_namespace: dict[str, Any], lm_results: dict[int, str]
+    ) -> Callable[[Any], None]:
+        """Create final(answer) function."""
+
+        def final_func(answer: Any) -> None:
+            resolved = _resolve_value(answer, lm_results)
+            raise FinalAnswer(resolved)
+
+        return final_func
+
+    def _make_final_var_func(
+        self, exec_namespace: dict[str, Any], lm_results: dict[int, str]
+    ) -> Callable[[str], None]:
+        """Create final_var(varname) function."""
+
+        def final_var_func(varname: str) -> None:
+            if varname not in exec_namespace:
+                raise RuntimeError(f"Variable '{varname}' not found")
+            value = exec_namespace[varname]
+            resolved = _resolve_value(value, lm_results)
+            raise FinalAnswer(resolved)
+
+        return final_var_func
+
+    async def _execute_pending_lm_calls(
+        self,
+        pending_calls: list[dict],
+        results: dict[int, str],
+    ) -> None:
+        """Execute all pending lm() calls in parallel."""
+        if not pending_calls:
+            return
+
+        from lm_deluge.prompt import Conversation
+
+        # Start all calls in parallel using start_nowait
+        task_mapping: list[tuple[int, int]] = []  # (call_id, task_id)
+        for call in pending_calls:
+            conv = Conversation.user(call["prompt"])
+            task_id = self.client.start_nowait(conv)
+            task_mapping.append((call["id"], task_id))
+
+        # Wait for all to complete
+        for call_id, task_id in task_mapping:
+            try:
+                response = await self.client.wait_for(task_id)
+                results[call_id] = response.completion or "(no response)"
+            except Exception as e:
+                results[call_id] = f"Error: {e}"
+
+        # Clear the pending list
+        pending_calls.clear()
+
+    def _format_answer(self, value: Any) -> str:
+        """Format the final answer as a string."""
+        if isinstance(value, str):
+            return value
+        try:
+            return json.dumps(value, default=str, indent=2)
+        except Exception:
+            return str(value)
+
+    async def execute(self, code: str) -> tuple[str, bool]:
+        """Execute RLM code.
+
+        Args:
+            code: Python code to execute
+
+        Returns:
+            Tuple of (output_string, is_final) where is_final indicates
+            whether FINAL()/FINAL_VAR() was called.
+        """
+        # Validate the code
+        tree = validate_rlm_code(code)
+
+        # Set up execution environment
+        pending_lm_calls: list[dict] = []
+        lm_results: dict[int, str] = {}
+        pending_call_ids: set[int] = set()
+        call_state = {"next_lm_id": 0}
+        output_capture = OutputCapture()
+
+        # Create the lm() wrapper
+        lm_wrapper = self._make_lm_wrapper(
+            pending_lm_calls, lm_results, call_state, pending_call_ids
+        )
+
+        # Build a single namespace for execution
+        # Using a single dict for both globals and locals ensures that
+        # variables are visible inside nested scopes (list comprehensions, etc.)
+        exec_namespace: dict[str, Any] = {
+            "__builtins__": {**RLM_SAFE_BUILTINS, "print": output_capture.print},
+            self.context_var_name: self.context,
+            "lm": lm_wrapper,
+            "json": json,  # Explicitly include json
+            **RLM_MODULES,
+            # Include persistent state from previous calls
+            **self._persistent_locals,
+        }
+
+        # Add final and final_var (they need access to exec_namespace for final_var)
+        exec_namespace["final"] = self._make_final_func(exec_namespace, lm_results)
+        exec_namespace["final_var"] = self._make_final_var_func(
+            exec_namespace, lm_results
+        )
+
+        # Track which keys are "system" keys that shouldn't be persisted
+        system_keys = set(exec_namespace.keys())
+
+        # Execute with retry loop for deferred lm() resolution
+        max_iterations = 50
+        compiled = compile(tree, "<rlm>", "exec")
+
+        for iteration in range(max_iterations):
+            # Reset call sequencing for this pass
+            call_state["next_lm_id"] = 0
+            pending_call_ids.clear()
+
+            try:
+                exec(compiled, exec_namespace)
+
+                # Execution completed - run any remaining pending calls
+                await self._execute_pending_lm_calls(pending_lm_calls, lm_results)
+
+                # Update persistent locals (exclude system keys)
+                for key, value in exec_namespace.items():
+                    if key not in system_keys:
+                        self._persistent_locals[key] = value
+
+                break
+
+            except FinalAnswer as fa:
+                # FINAL() or FINAL_VAR() was called
+                for key, value in exec_namespace.items():
+                    if key not in system_keys:
+                        self._persistent_locals[key] = value
+                return (self._format_answer(fa.answer), True)
+
+            except RuntimeError as e:
+                if "not yet available" in str(e):
+                    # Need to resolve pending lm() calls and retry
+                    await self._execute_pending_lm_calls(pending_lm_calls, lm_results)
+                    pending_call_ids.clear()
+                    # Continue to retry
+                else:
+                    raise RLMExecutionError(f"Runtime error: {e}")
+
+            except Exception as e:
+                raise RLMExecutionError(f"Execution error: {type(e).__name__}: {e}")
+
+        else:
+            raise RLMExecutionError(
+                f"Execution exceeded maximum iterations ({max_iterations})"
+            )
+
+        # Get output
+        output = output_capture.get_output()
+
+        # If no print output, check for result variable
+        if not output and "result" in exec_namespace:
+            result_value = _resolve_value(exec_namespace["result"], lm_results)
+            output = self._format_answer(result_value)
+
+        return (output or "Execution completed with no output", False)
+
+    def reset(self) -> None:
+        """Reset the persistent state."""
+        self._persistent_locals.clear()