Add custom tool call passing (#106)

* add custom tool calling

* remove subcall tools for now, until we add in recursive depth >1 and better agents with their own REPLs

Author: alexzhang13

docs/getting-started.md

@@ -138,6 +138,7 @@ This will display:
 | `other_backend_kwargs` | `list` | `None` | Configs for additional backends |
 | `logger` | `RLMLogger` | `None` | Logger for trajectory tracking |
 | `verbose` | `bool` | `False` | Enable console output |
+| `custom_tools` | `dict` | `None` | Custom functions/data available in REPL |
 
 ### The `completion()` Method
 
@@ -285,6 +286,90 @@ rlm = RLM(
 
 ---
 
+## Custom Tools
+
+You can provide custom functions and data that the RLM can use in its REPL environment. This allows you to give the model access to domain-specific tools, APIs, or helper functions.
+
+### Basic Usage
+
+```python
+def fetch_weather(city: str) -> str:
+    """Fetch weather data for a city."""
+    # Your API call here
+    return f"Weather in {city}: Sunny, 72°F"
+
+def calculate_shipping(weight: float, distance: float) -> float:
+    """Calculate shipping cost."""
+    return weight * 0.5 + distance * 0.1
+
+rlm = RLM(
+    backend="openai",
+    backend_kwargs={"model_name": "gpt-4o"},
+    custom_tools={
+        "fetch_weather": fetch_weather,
+        "calculate_shipping": calculate_shipping,
+        "API_KEY": "your-api-key",  # Non-callable values become variables
+    },
+)
+
+result = rlm.completion("What's the weather in Tokyo and calculate shipping for 10kg over 500km?")
+```
+
+Inside the REPL, the model can now call:
+```python
+weather = fetch_weather("Tokyo")
+cost = calculate_shipping(10, 500)
+```
+
+### Tool Descriptions
+
+You can provide descriptions for your tools that will be included in the system prompt, helping the model understand what each tool does:
+
+```python
+rlm = RLM(
+    backend="openai",
+    backend_kwargs={"model_name": "gpt-4o"},
+    custom_tools={
+        # Dict format: {"tool": callable_or_value, "description": "..."}
+        "fetch_weather": {"tool": fetch_weather, "description": "Fetch current weather data for a city name"},
+        "calculate_shipping": {"tool": calculate_shipping, "description": "Calculate shipping cost given weight (kg) and distance (km)"},
+        "API_KEY": {"tool": "your-api-key", "description": "API key for the weather service"},
+    },
+)
+```
+
+The descriptions are automatically added to the system prompt:
+```
+6. Custom tools and data available in the REPL:
+- `fetch_weather`: Fetch current weather data for a city name
+- `calculate_shipping`: Calculate shipping cost given weight (kg) and distance (km)
+- `API_KEY`: API key for the weather service
+```
+
+### Isolated Environments (Modal, Daytona)
+
+For isolated environments, custom tools must be serializable. You can provide:
+
+1. **Code strings** - Python code that defines the function:
+```python
+custom_tools = {
+    "helper": '''
+def helper(x):
+    return x * 2
+''',
+}
+```
+
+2. **Serializable data** - JSON-compatible values (strings, numbers, dicts, lists):
+```python
+custom_tools = {
+    "CONFIG": {"api_url": "https://api.example.com", "timeout": 30},
+    "ALLOWED_CITIES": ["Tokyo", "London", "New York"],
+}
+```
+
+---
+
 ## Logging and Debugging
 
 ### Enable Logging

docs/src/app/environments/page.tsx

@@ -99,16 +99,20 @@ export default function EnvironmentsPage() {
               ],
               [
                 <code key="2" className="text-sm font-semibold">llm_query(prompt, model=None)</code>, 
-                "Query a sub-LM from within the REPL. Returns the completion string."
+                "Single LM completion call. Returns the completion string. Does not have tool access."
               ],
               [
                 <code key="3" className="text-sm font-semibold">llm_query_batched(prompts, model=None)</code>, 
-                "Concurrent sub-LM queries. Returns a list of completion strings."
+                "Concurrent single LM completion calls. Returns a list of completion strings. Does not have tool access."
               ],
               [
                 <code key="4" className="text-sm font-semibold">FINAL_VAR(var_name)</code>, 
                 "Mark a variable as the final answer to return from the RLM"
               ],
+              [
+                <code key="5" className="text-sm font-semibold">custom_tools</code>, 
+                "Any custom functions or data you provide via the custom_tools parameter"
+              ],
             ]}
           />
         </div>
@@ -119,13 +123,58 @@ context = "Your input here"
 # Query a sub-LM
 result = llm_query("Summarize the context", model="gpt-5-mini")
 
+# Use a custom tool (if provided)
+data = fetch_data(context["url"])  # Custom function
+
 # Process the result
 summary = process(result)
 
 # Return final answer
 FINAL_VAR(summary)`} />
       </div>
 
+      <div className="my-12">
+        <h2 className="text-2xl font-bold mb-4">Custom Tools</h2>
+        
+        <p className="text-muted-foreground mb-6 leading-relaxed">
+          You can provide custom functions and data that the RLM can use in its REPL environment 
+          via the <code className="px-1.5 py-0.5 rounded bg-muted text-foreground text-sm font-semibold">custom_tools</code> parameter:
+        </p>
+
+        <CodeBlock code={`from rlm import RLM
+
+def fetch_weather(city: str) -> str:
+    """Fetch weather data for a city."""
+    return f"Weather in {city}: Sunny, 72°F"
+
+rlm = RLM(
+    backend="openai",
+    backend_kwargs={"model_name": "gpt-4o"},
+    custom_tools={
+        # Plain format (no description)
+        "fetch_weather": fetch_weather,
+        
+        # Dict format with description (recommended)
+        "calculate_tip": {"tool": lambda x: x * 0.2, "description": "Calculate 20% tip for a bill amount"},
+        "API_KEY": {"tool": "your-key", "description": "API key for external services"},
+    },
+)
+
+# The model can now call fetch_weather() in its REPL code`} />
+
+        <p className="text-muted-foreground mt-6 leading-relaxed">
+          <strong className="text-foreground">Tool descriptions:</strong> Use the dict format 
+          <code className="px-1.5 py-0.5 rounded bg-muted text-foreground text-sm font-semibold">{`{"tool": value, "description": "..."}`}</code> to 
+          provide descriptions that help the model understand what each tool does. Descriptions are automatically 
+          included in the system prompt.
+        </p>
+
+        <p className="text-muted-foreground mt-4 leading-relaxed">
+          <strong className="text-foreground">Note:</strong> <code className="px-1.5 py-0.5 rounded bg-muted text-foreground text-sm font-semibold">llm_query()</code> calls 
+          are single LM completions and do not have access to custom tools. Only the main RLM execution context has tool access.
+        </p>
+      </div>
+
       <div className="my-12">
         <h2 className="text-2xl font-bold mb-4">Architecture</h2>
 

examples/custom_tools_example.py

@@ -0,0 +1,133 @@
+"""
+Example demonstrating custom tools in RLM.
+
+Custom tools allow you to provide domain-specific functions and data
+that the RLM can use in its REPL environment.
+
+Run with: python -m examples.custom_tools_example
+"""
+
+import os
+from typing import Any
+
+from dotenv import load_dotenv
+
+from rlm import RLM
+
+load_dotenv()
+
+
+# =============================================================================
+# Define Custom Tools
+# =============================================================================
+
+
+def fetch_stock_price(symbol: str) -> dict[str, Any]:
+    """
+    Fetch stock price data for a given symbol.
+    In a real application, this would call a financial API.
+    """
+    # Mock data for demonstration
+    prices = {
+        "AAPL": {"price": 178.50, "change": 2.3, "volume": 52_000_000},
+        "GOOGL": {"price": 141.25, "change": -0.8, "volume": 21_000_000},
+        "MSFT": {"price": 378.90, "change": 1.5, "volume": 18_000_000},
+        "AMZN": {"price": 178.75, "change": 0.2, "volume": 35_000_000},
+    }
+    return prices.get(symbol.upper(), {"error": f"Symbol {symbol} not found"})
+
+
+def calculate_portfolio_value(holdings: dict[str, int]) -> float:
+    """
+    Calculate total portfolio value given holdings.
+    holdings: dict mapping symbol to number of shares
+    """
+    total = 0.0
+    for symbol, shares in holdings.items():
+        data = fetch_stock_price(symbol)
+        if "price" in data:
+            total += data["price"] * shares
+    return total
+
+
+def format_currency(amount: float) -> str:
+    """Format a number as USD currency."""
+    return f"${amount:,.2f}"
+
+
+# Configuration data (non-callable values become variables)
+MARKET_CONFIG = {
+    "trading_hours": {"open": "09:30", "close": "16:00"},
+    "currency": "USD",
+    "exchange": "NYSE",
+}
+
+
+# =============================================================================
+# Example 1: Basic Custom Tools
+# =============================================================================
+
+
+def example_basic_tools():
+    """Demonstrate basic custom tools usage with descriptions."""
+    print("=" * 60)
+    print("Example 1: Basic Custom Tools with Descriptions")
+    print("=" * 60)
+
+    # Tools can be provided with descriptions using dict format:
+    # {"name": {"tool": callable_or_value, "description": "..."}}
+    # The description will be included in the system prompt so the model
+    # knows what each tool does.
+
+    rlm = RLM(
+        backend="portkey",
+        backend_kwargs={
+            "model_name": "@openai/gpt-5-nano",
+            "api_key": os.getenv("PORTKEY_API_KEY"),
+        },
+        environment="local",
+        custom_tools={
+            # Callable functions with descriptions (dict format)
+            "fetch_stock_price": {
+                "tool": fetch_stock_price,
+                "description": "Fetch current stock price data for a symbol (AAPL, GOOGL, MSFT, AMZN)",
+            },
+            "calculate_portfolio_value": {
+                "tool": calculate_portfolio_value,
+                "description": "Calculate total portfolio value from a dict of {symbol: shares}",
+            },
+            "format_currency": {
+                "tool": format_currency,
+                "description": "Format a number as USD currency string",
+            },
+            # Data values with descriptions
+            "MARKET_CONFIG": {
+                "tool": MARKET_CONFIG,
+                "description": "Market configuration including trading hours and exchange info",
+            },
+        },
+        verbose=True,
+    )
+
+    # The model can now use these tools to answer questions
+    result = rlm.completion(
+        "What's the current price of AAPL and GOOGL? "
+        "Calculate the total value of a portfolio with 100 shares of each. "
+        "Format the result as currency."
+    )
+
+    print(f"\nFinal Answer: {result.response}")
+
+
+# =============================================================================
+# Main
+# =============================================================================
+
+
+if __name__ == "__main__":
+    if not os.getenv("OPENAI_API_KEY"):
+        print("Please set OPENAI_API_KEY environment variable")
+        print("You can also modify this example to use a different backend")
+        exit(1)
+
+    example_basic_tools()

examples/quickstart.py

@@ -10,10 +10,10 @@
 logger = RLMLogger(log_dir="./logs")
 
 rlm = RLM(
-    backend="openai",  # or "portkey", etc.
+    backend="portkey",  # or "openai", etc.
     backend_kwargs={
-        "model_name": "gpt-5-nano",
-        "api_key": os.getenv("OPENAI_API_KEY"),
+        "model_name": "@openai/gpt-5-nano",
+        "api_key": os.getenv("PORTKEY_API_KEY"),
     },
     environment="local",
     environment_kwargs={},

rlm/core/rlm.py

@@ -52,6 +52,8 @@ def __init__(
         logger: RLMLogger | None = None,
         verbose: bool = False,
         persistent: bool = False,
+        custom_tools: dict[str, Any] | None = None,
+        custom_sub_tools: dict[str, Any] | None = None,
     ):
         """
         Args:
@@ -68,6 +70,10 @@ def __init__(
             logger: The logger to use for the RLM.
             verbose: Whether to print verbose output in rich to console.
             persistent: If True, reuse the environment across completion() calls for multi-turn conversations.
+            custom_tools: Dict of custom functions/tools available in the REPL. Keys are function names,
+                values are callable functions. These are injected into the REPL globals.
+            custom_sub_tools: Dict of custom tools for sub-agents (llm_query calls). If None, inherits
+                from custom_tools. Pass an empty dict {} to disable tools for sub-agents.
         """
         # Store config for spawning per-completion
         self.backend = backend
@@ -87,6 +93,11 @@ def __init__(
         self.other_backends = other_backends
         self.other_backend_kwargs = other_backend_kwargs
 
+        # Custom tools: functions available in the REPL environment
+        self.custom_tools = custom_tools
+        # Sub-tools: if None, inherit from custom_tools; if {}, no tools for sub-agents
+        self.custom_sub_tools = custom_sub_tools if custom_sub_tools is not None else custom_tools
+
         self.depth = depth
         self.max_depth = max_depth
         self.max_iterations = max_iterations
@@ -165,6 +176,11 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
             env_kwargs["lm_handler_address"] = (lm_handler.host, lm_handler.port)
             env_kwargs["context_payload"] = prompt
             env_kwargs["depth"] = self.depth + 1  # Environment depth is RLM depth + 1
+            # Pass custom tools to the environment
+            if self.custom_tools is not None:
+                env_kwargs["custom_tools"] = self.custom_tools
+            if self.custom_sub_tools is not None:
+                env_kwargs["custom_sub_tools"] = self.custom_sub_tools
             environment: BaseEnv = get_environment(self.environment_type, env_kwargs)
 
             if self.persistent:
@@ -184,7 +200,9 @@ def _setup_prompt(self, prompt: str | dict[str, Any]) -> list[dict[str, Any]]:
         """
         metadata = QueryMetadata(prompt)
         message_history = build_rlm_system_prompt(
-            system_prompt=self.system_prompt, query_metadata=metadata
+            system_prompt=self.system_prompt,
+            query_metadata=metadata,
+            custom_tools=self.custom_tools,
         )
 
         return message_history