add verbose logging

Author: alexzhang13

examples/rlm_example.py

@@ -18,6 +18,7 @@
     environment_kwargs={},
     max_depth=1,
     logger=logger,
+    verbose=True,
 )
 
 result = rlm.completion("Using your code, solve 2^(2^(2^(2))). Show your work in Python.")

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "openai>=2.14.0",
     "portkey-ai>=2.1.0",
     "requests>=2.32.5",
+    "rich>=13.0.0",
 ]
 
 [project.optional-dependencies]

rlm/clients/openai.py

@@ -5,7 +5,8 @@
 import openai
 from dotenv import load_dotenv
 
-from rlm.clients.base_lm import BaseLM, ModelUsageSummary, UsageSummary
+from rlm.clients.base_lm import BaseLM
+from rlm.core.types import ModelUsageSummary, UsageSummary
 
 load_dotenv()
 

rlm/clients/portkey.py

@@ -4,7 +4,8 @@
 from portkey_ai import Portkey
 from portkey_ai.api_resources.types.chat_complete_type import ChatCompletions
 
-from rlm.clients.base_lm import BaseLM, ModelUsageSummary, UsageSummary
+from rlm.clients.base_lm import BaseLM
+from rlm.core.types import ModelUsageSummary, UsageSummary
 
 
 class PortkeyClient(BaseLM):

rlm/core/rlm.py

@@ -14,7 +14,7 @@
     RLMMetadata,
 )
 from rlm.environments import BaseEnv, get_environment
-from rlm.logger.rlm_logger import RLMLogger
+from rlm.logger import RLMLogger, VerbosePrinter
 from rlm.utils.parsing import (
     find_code_blocks,
     find_final_answer,
@@ -50,6 +50,7 @@ def __init__(
         other_backends: list[ClientBackend] | None = None,
         other_backend_kwargs: list[dict[str, Any]] | None = None,
         logger: RLMLogger | None = None,
+        verbose: bool = False,
     ):
         """
         Args:
@@ -58,12 +59,13 @@ def __init__(
             environment: The environment to use for the RLM.
             environment_kwargs: The kwargs to pass to the environment.
             depth: The current depth of the RLM (0-indexed).
-            max_depth: The maximum depth of the RLM.
+            max_depth: The maximum depth of the RLM. Currently, only depth 1 is supported.
             max_iterations: The maximum number of iterations of the RLM.
             custom_system_prompt: The custom system prompt to use for the RLM.
             other_backends: A list of other client backends that the environments can use to make sub-calls.
             other_backend_kwargs: The kwargs to pass to the other client backends (ordered to match other_backends).
             logger: The logger to use for the RLM.
+            verbose: Whether to print verbose output in rich to console.
         """
         # Store config for spawning per-completion
         self.backend = backend
@@ -78,9 +80,10 @@ def __init__(
         self.max_iterations = max_iterations
         self.system_prompt = custom_system_prompt if custom_system_prompt else RLM_SYSTEM_PROMPT
         self.logger = logger
+        self.verbose = VerbosePrinter(enabled=verbose)
 
         # Log metadata if logger is provided
-        if self.logger:
+        if self.logger or verbose:
             metadata = RLMMetadata(
                 root_model=backend_kwargs.get("model_name", "unknown"),
                 max_depth=max_depth,
@@ -91,7 +94,9 @@ def __init__(
                 environment_kwargs=filter_sensitive_keys(environment_kwargs),
                 other_backends=other_backends,
             )
-            self.logger.log_metadata(metadata)
+            if self.logger:
+                self.logger.log_metadata(metadata)
+            self.verbose.print_metadata(metadata)
 
     @contextmanager
     def _spawn_completion_context(self, prompt: str | dict[str, Any]):
@@ -156,6 +161,11 @@ def completion(
             A final answer as a string.
         """
         time_start = time.perf_counter()
+
+        # If we're at max depth, the RLM is an LM, so we fallback to the regular LM.
+        if self.depth >= self.max_depth:
+            return self._fallback_answer(prompt)
+
         with self._spawn_completion_context(prompt) as (lm_handler, environment):
             message_history = self._setup_prompt(prompt)
 
@@ -177,13 +187,19 @@ def completion(
                 if self.logger:
                     self.logger.log(iteration)
 
+                # Verbose output for this iteration
+                self.verbose.print_iteration(iteration, i + 1)
+
                 if final_answer:
                     time_end = time.perf_counter()
+                    usage = lm_handler.get_usage_summary()
+                    self.verbose.print_final_answer(final_answer)
+                    self.verbose.print_summary(i + 1, time_end - time_start, usage.to_dict())
                     return RLMChatCompletion(
                         root_model=self.backend_kwargs.get("model_name", "unknown"),
                         prompt=prompt,
                         response=final_answer,
-                        usage_summary=lm_handler.get_usage_summary(),
+                        usage_summary=usage,
                         execution_time=time_end - time_start,
                     )
 
@@ -194,8 +210,18 @@ def completion(
                 message_history.extend(new_messages)
 
             # Default behavior: we run out of iterations, provide one final answer
+            time_end = time.perf_counter()
             final_answer = self._default_answer(message_history, lm_handler)
-            return final_answer
+            usage = lm_handler.get_usage_summary()
+            self.verbose.print_final_answer(final_answer)
+            self.verbose.print_summary(self.max_iterations, time_end - time_start, usage.to_dict())
+            return RLMChatCompletion(
+                root_model=self.backend_kwargs.get("model_name", "unknown"),
+                prompt=prompt,
+                response=final_answer,
+                usage_summary=usage,
+                execution_time=time_end - time_start,
+            )
 
     def _completion_turn(
         self,
@@ -248,3 +274,11 @@ def _default_answer(self, message_history: list[dict[str, Any]], lm_handler: LMH
             )
 
         return response
+
+    def _fallback_answer(self, message: str | dict[str, Any]) -> str:
+        """
+        Fallback behavior if the RLM is actually at max depth, and should be treated as an LM.
+        """
+        client: BaseLM = get_client(self.backend, self.backend_kwargs)
+        response = client.completion(message)
+        return response

rlm/environments/__init__.py

@@ -5,21 +5,24 @@
 
 
 def get_environment(
-    environment: Literal["local", "modal"],
+    environment: Literal["local", "modal", "docker"],
     environment_kwargs: dict[str, Any],
 ) -> BaseEnv:
     """
     Routes a specific environment and the args (as a dict) to the appropriate environment if supported.
-    Currently supported environments: ['local', 'modal']
+    Currently supported environments: ['local', 'modal', 'docker']
     """
     if environment == "local":
         return LocalREPL(**environment_kwargs)
     elif environment == "modal":
-        # Lazy import to avoid requiring modal as a hard dependency
         from rlm.environments.modal_repl import ModalREPL
 
         return ModalREPL(**environment_kwargs)
+    elif environment == "docker":
+        from rlm.environments.docker_repl import DockerREPL
+
+        return DockerREPL(**environment_kwargs)
     else:
         raise ValueError(
-            f"Unknown environment: {environment}. Supported environments: ['local', 'modal']"
+            f"Unknown environment: {environment}. Supported: ['local', 'modal', 'docker']"
         )

rlm/logger/__init__.py

@@ -1,3 +1,4 @@
 from rlm.logger.rlm_logger import RLMLogger
+from rlm.logger.verbose import VerbosePrinter
 
-__all__ = ["RLMLogger"]
+__all__ = ["RLMLogger", "VerbosePrinter"]