Enhance RLM Playground with new features and UI improvements

- Added `custom_system_prompt` and `verbose` fields to `RunRequest` model for enhanced configuration.
- Implemented verbose output capturing in the `run_rlm` function, allowing detailed console output to be returned in `RunResponse`.
- Introduced `KeyValueEditor` component for dynamic key-value pair management in the frontend.
- Updated `PlaygroundForm` to support custom system prompts, environment kwargs, and other backends with their respective configurations.
- Enhanced `PlaygroundResults` to display verbose output if available, improving user feedback on execution details.

Author: lakshyaag

playground/models.py

@@ -32,6 +32,10 @@ class RunRequest(BaseModel):
     other_backend_kwargs: list[dict[str, Any]] | None = Field(
         None, description="Configuration for additional backends"
     )
+    custom_system_prompt: str | None = Field(
+        None, description="Custom system prompt to override the default"
+    )
+    verbose: bool = Field(False, description="Enable verbose console output")
     enable_logging: bool = Field(False, description="Whether to save logs to file")
 
 
@@ -43,4 +47,5 @@ class RunResponse(BaseModel):
     root_model: str | None = Field(None, description="Model name used")
     execution_time: float | None = Field(None, description="Total execution time in seconds")
     usage_summary: dict[str, Any] | None = Field(None, description="Token usage summary per model")
+    verbose_output: str | None = Field(None, description="Captured verbose console output")
     error: str | None = Field(None, description="Error message if completion failed")

playground/server.py

@@ -1,6 +1,8 @@
 """FastAPI server for RLM playground."""
 
+import io
 import os
+import sys
 
 from dotenv import load_dotenv
 from fastapi import FastAPI
@@ -56,25 +58,47 @@ def run_rlm(request: RunRequest) -> RunResponse:
             log_dir = os.getenv("RLM_LOG_DIR", "./logs")
             logger = RLMLogger(log_dir=log_dir)
 
-        # Create RLM instance
-        rlm = RLM(
-            backend=request.backend,
-            backend_kwargs=request.backend_kwargs or {},
-            environment=request.environment,
-            environment_kwargs=request.environment_kwargs or {},
-            max_depth=request.max_depth,
-            max_iterations=request.max_iterations,
-            other_backends=request.other_backends,
-            other_backend_kwargs=request.other_backend_kwargs,
-            logger=logger,
-            verbose=False,  # Disable console output in API
-        )
-
-        # Run completion
-        result = rlm.completion(
-            prompt=request.prompt,
-            root_prompt=request.root_prompt,
-        )
+        # Capture verbose output if enabled
+        verbose_output = None
+        if request.verbose:
+            stdout_capture = io.StringIO()
+            stderr_capture = io.StringIO()
+            old_stdout = sys.stdout
+            old_stderr = sys.stderr
+
+        try:
+            if request.verbose:
+                sys.stdout = stdout_capture
+                sys.stderr = stderr_capture
+
+            # Create RLM instance
+            rlm = RLM(
+                backend=request.backend,
+                backend_kwargs=request.backend_kwargs or {},
+                environment=request.environment,
+                environment_kwargs=request.environment_kwargs or {},
+                max_depth=request.max_depth,
+                max_iterations=request.max_iterations,
+                other_backends=request.other_backends,
+                other_backend_kwargs=request.other_backend_kwargs,
+                custom_system_prompt=request.custom_system_prompt,
+                logger=logger,
+                verbose=request.verbose,
+            )
+
+            # Run completion
+            result = rlm.completion(
+                prompt=request.prompt,
+                root_prompt=request.root_prompt,
+            )
+        finally:
+            # Restore stdout/stderr and capture output
+            if request.verbose:
+                sys.stdout = old_stdout
+                sys.stderr = old_stderr
+                stdout_text = stdout_capture.getvalue()
+                stderr_text = stderr_capture.getvalue()
+                verbose_output = stdout_text + stderr_text if stderr_text else stdout_text
 
         # Extract response - handle both string and tuple (type, content) formats
         response_text = result.response
@@ -104,6 +128,7 @@ def run_rlm(request: RunRequest) -> RunResponse:
             root_model=result.root_model,
             execution_time=result.execution_time,
             usage_summary=result.usage_summary.to_dict(),
+            verbose_output=verbose_output,
             error=None,
         )
 
@@ -115,5 +140,6 @@ def run_rlm(request: RunRequest) -> RunResponse:
             root_model=None,
             execution_time=None,
             usage_summary=None,
+            verbose_output=None,
             error=str(e),
         )

rlm/clients/openai.py

@@ -82,6 +82,9 @@ async def acompletion(
 
     def _track_cost(self, response: openai.ChatCompletion, model: str):
         self.model_call_counts[model] += 1
+        if not hasattr(response, "usage") or response.usage is None:
+            print("WARNING: No usage data received!")
+            return
         self.model_input_tokens[model] += response.usage.prompt_tokens
         self.model_output_tokens[model] += response.usage.completion_tokens
         self.model_total_tokens[model] += response.usage.total_tokens

visualizer/src/components/KeyValueEditor.tsx

@@ -0,0 +1,84 @@
+'use client';
+
+import { Button } from '@/components/ui/button';
+
+interface KeyValuePair {
+  key: string;
+  value: string;
+}
+
+interface KeyValueEditorProps {
+  pairs: KeyValuePair[];
+  onChange: (pairs: KeyValuePair[]) => void;
+  disabled?: boolean;
+  keyPlaceholder?: string;
+  valuePlaceholder?: string;
+}
+
+export function KeyValueEditor({
+  pairs,
+  onChange,
+  disabled = false,
+  keyPlaceholder = 'key',
+  valuePlaceholder = 'value',
+}: KeyValueEditorProps) {
+  const updatePair = (index: number, field: 'key' | 'value', newValue: string) => {
+    const newPairs = [...pairs];
+    newPairs[index] = { ...newPairs[index], [field]: newValue };
+    onChange(newPairs);
+  };
+
+  const addPair = () => {
+    onChange([...pairs, { key: '', value: '' }]);
+  };
+
+  const removePair = (index: number) => {
+    onChange(pairs.filter((_, i) => i !== index));
+  };
+
+  return (
+    <div className="space-y-2">
+      {pairs.map((pair, index) => (
+        <div key={index} className="flex gap-2 items-center">
+          <input
+            type="text"
+            value={pair.key}
+            onChange={(e) => updatePair(index, 'key', e.target.value)}
+            placeholder={keyPlaceholder}
+            className="flex-1 px-3 py-2 bg-background border border-input rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-ring"
+            disabled={disabled}
+          />
+          <input
+            type="text"
+            value={pair.value}
+            onChange={(e) => updatePair(index, 'value', e.target.value)}
+            placeholder={valuePlaceholder}
+            className="flex-1 px-3 py-2 bg-background border border-input rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-ring"
+            disabled={disabled}
+          />
+          <Button
+            type="button"
+            variant="outline"
+            size="sm"
+            onClick={() => removePair(index)}
+            disabled={disabled}
+            className="shrink-0"
+          >
+            Remove
+          </Button>
+        </div>
+      ))}
+      <Button
+        type="button"
+        variant="outline"
+        size="sm"
+        onClick={addPair}
+        disabled={disabled}
+        className="w-full"
+      >
+        + Add Key-Value Pair
+      </Button>
+    </div>
+  );
+}
+