add logger examples + extra logging func

Author: alexzhang13

README.md

@@ -135,17 +135,20 @@ If you use this code or repository in your research, please cite:
 }
 ```
 
+## Optional: Trajectory metadata and logging
+`RLMChatCompletion` has an optional `metadata` field (default empty) that can hold the full trajectory (run config + all iterations and sub-calls) so you can reconstruct the run. Pass an `RLMLogger` to capture it:
+
+- **In-memory only** (trajectory on `completion.metadata`): `logger=RLMLogger()` (no `log_dir`).
+- **Also save to disk** (JSONL for the visualizer): `logger=RLMLogger(log_dir="./logs")`.
+
 ## Optional Debugging: Visualizing RLM Trajectories
-We additionally provide a simple visualizer tool to examine and view the code, sub-LM, and root-LM calls of an RLM trajectory. To save log files (`.jsonl`) on every completion call that can be viewed in the visualizer, initialize the `RLMLogger` object and pass it into the `RLM` on initialization:
+We provide a simple visualizer to inspect code, sub-LM, and root-LM calls. Use `RLMLogger(log_dir="./logs")` so each completion writes a `.jsonl` file:
 ```python
 from rlm.logger import RLMLogger
 from rlm import RLM
 
 logger = RLMLogger(log_dir="./logs")
-rlm = RLM(
-    ...
-    logger=logger
-)
+rlm = RLM(..., logger=logger)
 ```
 
 To run the visualizer locally, we use Node.js and shadcn/ui:

examples/logger_example.py

@@ -0,0 +1,46 @@
+"""
+Logger example: same structure as quickstart, with a different prompt.
+Shows trajectory capture in completion.metadata (and optional save to disk).
+"""
+
+import os
+
+from dotenv import load_dotenv
+
+from rlm import RLM
+from rlm.logger import RLMLogger
+
+load_dotenv()
+
+# In-memory only: trajectory on completion.metadata (no disk write)
+logger = RLMLogger()
+# To also save JSONL for the visualizer: logger = RLMLogger(log_dir="./logs")
+
+rlm = RLM(
+    backend="portkey",
+    backend_kwargs={
+        "model_name": "@openai/gpt-5-nano",
+        "api_key": os.getenv("PORTKEY_API_KEY"),
+    },
+    environment="local",
+    environment_kwargs={},
+    max_depth=1,
+    logger=logger,
+    verbose=True,
+)
+
+prompt = "Compute the sum 1 + 2 + 3 + 4 + 5 and print the result using a REPL block, then return it with FINAL_VAR."
+result = rlm.completion(prompt)
+
+print("Response:", result.response)
+if result.metadata:
+    traj = result.metadata
+    print("Trajectory: run_metadata +", len(traj.get("iterations", [])), "iterations")
+    print("Trajectory:", traj)
+    if traj.get("iterations"):
+        first = traj["iterations"][0]
+        print("  First iteration keys:", list(first.keys()))
+else:
+    print("Trajectory: (none — no logger or metadata not captured)")
+
+print("Full response:", result)

rlm/core/rlm.py

@@ -229,6 +229,9 @@ def completion(
         if self.depth >= self.max_depth:
             return self._fallback_answer(prompt)
 
+        if self.logger:
+            self.logger.clear_iterations()
+
         with self._spawn_completion_context(prompt) as (lm_handler, environment):
             message_history = self._setup_prompt(prompt)
 
@@ -283,6 +286,7 @@ def completion(
                         response=final_answer,
                         usage_summary=usage,
                         execution_time=time_end - time_start,
+                        metadata=self.logger.get_trajectory() if self.logger else None,
                     )
 
                 # Format the iteration for the next prompt.
@@ -310,6 +314,7 @@ def completion(
                 response=final_answer,
                 usage_summary=usage,
                 execution_time=time_end - time_start,
+                metadata=self.logger.get_trajectory() if self.logger else None,
             )
 
     def _completion_turn(

rlm/core/types.py

@@ -96,15 +96,21 @@ class RLMChatCompletion:
     response: str
     usage_summary: UsageSummary
     execution_time: float
+    metadata: dict | None = (
+        None  # Full trajectory (run_metadata + iterations) when logger captures it
+    )
 
     def to_dict(self):
-        return {
+        out = {
             "root_model": self.root_model,
             "prompt": self.prompt,
             "response": self.response,
             "usage_summary": self.usage_summary.to_dict(),
             "execution_time": self.execution_time,
         }
+        if self.metadata is not None:
+            out["metadata"] = self.metadata
+        return out
 
     @classmethod
     def from_dict(cls, data: dict) -> "RLMChatCompletion":
@@ -114,6 +120,7 @@ def from_dict(cls, data: dict) -> "RLMChatCompletion":
             response=data.get("response"),
             usage_summary=UsageSummary.from_dict(data.get("usage_summary")),
             execution_time=data.get("execution_time"),
+            metadata=data.get("metadata"),
         )
 
 

rlm/logger/rlm_logger.py

@@ -1,7 +1,8 @@
 """
 Logger for RLM iterations.
 
-Writes RLMIteration data to JSON-lines files for analysis and debugging.
+Captures run metadata and iterations in memory so they can be attached to
+RLMChatCompletion.metadata. Optionally writes the same data to JSON-lines files.
 """
 
 import json
@@ -13,50 +14,77 @@
 
 
 class RLMLogger:
-    """Logger that writes RLMIteration data to a JSON-lines file."""
+    """
+    Captures trajectory (run metadata + iterations) for each completion.
+    By default only captures in memory; set log_dir to also save to disk.
 
-    def __init__(self, log_dir: str, file_name: str = "rlm"):
-        self.log_dir = log_dir
-        os.makedirs(log_dir, exist_ok=True)
-
-        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-        run_id = str(uuid.uuid4())[:8]
-        self.log_file_path = os.path.join(log_dir, f"{file_name}_{timestamp}_{run_id}.jsonl")
+    - log_dir=None: trajectory is available via get_trajectory() and can be
+      attached to RLMChatCompletion.metadata (no disk write).
+    - log_dir="path": same capture plus appends to a JSONL file per run.
+    """
 
+    def __init__(self, log_dir: str | None = None, file_name: str = "rlm"):
+        self._save_to_disk = log_dir is not None
+        self.log_dir = log_dir
+        self.log_file_path: str | None = None
+        if self._save_to_disk and log_dir:
+            os.makedirs(log_dir, exist_ok=True)
+            timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+            run_id = str(uuid.uuid4())[:8]
+            self.log_file_path = os.path.join(log_dir, f"{file_name}_{timestamp}_{run_id}.jsonl")
+
+        self._run_metadata: dict | None = None
+        self._iterations: list[dict] = []
         self._iteration_count = 0
         self._metadata_logged = False
 
-    def log_metadata(self, metadata: RLMMetadata):
-        """Log RLM metadata as the first entry in the file."""
+    def log_metadata(self, metadata: RLMMetadata) -> None:
+        """Capture run metadata (and optionally write to file)."""
         if self._metadata_logged:
             return
 
-        entry = {
-            "type": "metadata",
-            "timestamp": datetime.now().isoformat(),
-            **metadata.to_dict(),
-        }
-
-        with open(self.log_file_path, "a") as f:
-            json.dump(entry, f)
-            f.write("\n")
-
+        self._run_metadata = metadata.to_dict()
         self._metadata_logged = True
 
-    def log(self, iteration: RLMIteration):
-        """Log an RLMIteration to the file."""
+        if self._save_to_disk and self.log_file_path:
+            entry = {
+                "type": "metadata",
+                "timestamp": datetime.now().isoformat(),
+                **self._run_metadata,
+            }
+            with open(self.log_file_path, "a") as f:
+                json.dump(entry, f)
+                f.write("\n")
+
+    def log(self, iteration: RLMIteration) -> None:
+        """Capture one iteration (and optionally append to file)."""
         self._iteration_count += 1
-
         entry = {
             "type": "iteration",
             "iteration": self._iteration_count,
             "timestamp": datetime.now().isoformat(),
             **iteration.to_dict(),
         }
+        self._iterations.append(entry)
+
+        if self._save_to_disk and self.log_file_path:
+            with open(self.log_file_path, "a") as f:
+                json.dump(entry, f)
+                f.write("\n")
 
-        with open(self.log_file_path, "a") as f:
-            json.dump(entry, f)
-            f.write("\n")
+    def clear_iterations(self) -> None:
+        """Reset iterations for the next completion (trajectory is per completion)."""
+        self._iterations = []
+        self._iteration_count = 0
+
+    def get_trajectory(self) -> dict | None:
+        """Return captured run_metadata + iterations for the current completion, or None if no metadata yet."""
+        if self._run_metadata is None:
+            return None
+        return {
+            "run_metadata": self._run_metadata,
+            "iterations": list(self._iterations),
+        }
 
     @property
     def iteration_count(self) -> int:

tests/test_types.py

@@ -5,6 +5,7 @@
     ModelUsageSummary,
     QueryMetadata,
     REPLResult,
+    RLMChatCompletion,
     RLMIteration,
     RLMMetadata,
     UsageSummary,
@@ -156,6 +157,39 @@ def test_to_dict(self):
         assert d["iteration_time"] == 1.5
 
 
+class TestRLMChatCompletion:
+    """Tests for RLMChatCompletion."""
+
+    def test_metadata_default_none(self):
+        usage = UsageSummary(model_usage_summaries={})
+        c = RLMChatCompletion(
+            root_model="gpt-4",
+            prompt="hi",
+            response="hello",
+            usage_summary=usage,
+            execution_time=1.0,
+        )
+        assert c.metadata is None
+        d = c.to_dict()
+        assert "metadata" not in d
+
+    def test_metadata_roundtrip(self):
+        usage = UsageSummary(model_usage_summaries={})
+        trajectory = {"run_metadata": {"root_model": "gpt-4"}, "iterations": []}
+        c = RLMChatCompletion(
+            root_model="gpt-4",
+            prompt="hi",
+            response="hello",
+            usage_summary=usage,
+            execution_time=1.0,
+            metadata=trajectory,
+        )
+        d = c.to_dict()
+        assert d["metadata"] == trajectory
+        c2 = RLMChatCompletion.from_dict(d)
+        assert c2.metadata == trajectory
+
+
 class TestQueryMetadata:
     """Tests for QueryMetadata."""