Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 85 additions & 0 deletions docs/getting-started.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
53 changes: 51 additions & 2 deletions docs/src/app/environments/page.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand All @@ -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>

Expand Down
133 changes: 133 additions & 0 deletions examples/custom_tools_example.py
Original file line number Diff line number Diff line change
@@ -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()
6 changes: 3 additions & 3 deletions examples/quickstart.py
Original file line number Diff line number Diff line change
Expand Up @@ -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={},
Expand Down
20 changes: 19 additions & 1 deletion rlm/core/rlm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down
Loading