Skip to content

Commit 7c38465

Browse files
authored
Merge branch 'main' into sid/add-e2b-sandboxes
2 parents 8426806 + 015061e commit 7c38465

19 files changed

Lines changed: 1289 additions & 52 deletions

‎.github/workflows/publish.yml‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
name: Publish to PyPI
2+
3+
on:
4+
release:
5+
types: [published]
6+
7+
permissions:
8+
id-token: write
9+
10+
jobs:
11+
publish:
12+
name: Build and publish to PyPI
13+
runs-on: ubuntu-latest
14+
environment:
15+
name: pypi
16+
url: https://pypi.org/project/rlms/
17+
steps:
18+
- uses: actions/checkout@v4
19+
20+
- name: Set up Python
21+
uses: actions/setup-python@v5
22+
with:
23+
python-version: "3.12"
24+
25+
- name: Install build dependencies
26+
run: pip install build
27+
28+
- name: Build package
29+
run: python -m build
30+
31+
- name: Publish to PyPI
32+
uses: pypa/gh-action-pypi-publish@release/v1

‎MANIFEST.IN‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
include README.md
22
include LICENSE
33
include pyproject.toml
4-
recursive-include verifiers *.py
5-
recursive-include configs *.yaml *.yml *.json
4+
recursive-include rlm *.py
65
global-exclude *.pyc
76
global-exclude __pycache__
87
global-exclude .DS_Store

‎README.md‎

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -35,17 +35,28 @@ This repository provides an extensible inference engine for using RLMs around st
3535
> [!NOTE]
3636
> This repository contains inference code for RLMs with support for various sandbox environments. Open-source contributions are welcome. This repository is maintained by the authors of the paper from the MIT OASYS lab.
3737
38-
<!-- ## Installation
39-
```
40-
pip install rlm
41-
```
42-
To install the latest from `main`:
38+
## Quick Setup
39+
You can try out RLMs quickly by installing from PyPi:
40+
```bash
41+
pip install rlms
4342
```
44-
pip install git+https://github.com/alexzhang13/rlm.git
43+
44+
The default RLM client uses a REPL environment that runs on the host process through Python `exec` calls. It uses the same virtual environment as the host process (i.e. it will have access to the same dependencies), but with some limitations in its available global modules. As an example, we can call RLM completions using GPT-5-nano:
45+
```python
46+
from rlm import RLM
47+
48+
rlm = RLM(
49+
backend="openai",
50+
backend_kwargs={"model_name": "gpt-5-nano"},
51+
verbose=True, # For printing to console with rich, disabled by default.
52+
)
53+
54+
print(rlm.completion("Print me the first 100 powers of two, each on a newline.").response)
4555
```
46-
``` -->
4756

48-
## Quick Setup
57+
<details>
58+
<summary><b>Manual Setup</b></summary>
59+
4960
Set up the dependencies with `uv` (or your virtual environment of choice):
5061
```bash
5162
curl -LsSf https://astral.sh/uv/install.sh | sh
@@ -63,18 +74,7 @@ To run a quick test, the following will run an RLM query with the OpenAI client
6374
make quickstart
6475
```
6576

66-
The default RLM client uses a REPL environment that runs on the host process through Python `exec` calls. It uses the same virtual environment as the host process (i.e. it will have access to the same dependencies), but with some limitations in its available global modules. As an example, we can call RLM completions using GPT-5-nano:
67-
```python
68-
from rlm import RLM
69-
70-
rlm = RLM(
71-
backend="openai",
72-
backend_kwargs={"model_name": "gpt-5-nano"},
73-
verbose=True, # For printing to console with rich, disabled by default.
74-
)
75-
76-
print(rlm.completion("Print me the first 100 powers of two, each on a newline.").response)
77-
```
77+
</details>
7878

7979
## REPL Environments
8080
We support two types of REPL environments -- isolated, and non-isolated. Non-isolated environments (default) run code execution on the same machine as the RLM (e.g. through `exec`), which is pretty reasonable for some local low-risk tasks, like simple benchmarking, but can be problematic if the prompts or tool calls can interact with malicious users. Fully isolated environments used Cloud-based sandboxes (e.g. Prime Sandboxes, [Modal Sandboxes](https://modal.com/docs/guide/sandboxes)) to run code generated by the RLM, ensuring completely isolation from the host process. Environments can be added, but we natively support the following: `local` (default), `modal`, `prime`.

‎examples/daytona_repl_example.py‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
"""
2+
Example usage of Daytona REPL with code execution and LLM queries.
3+
4+
Run with: python -m examples.daytona_repl_example
5+
"""
6+
7+
from rlm.clients.base_lm import BaseLM
8+
from rlm.core.lm_handler import LMHandler
9+
from rlm.core.types import ModelUsageSummary, UsageSummary
10+
from rlm.environments.daytona_repl import DaytonaREPL
11+
12+
13+
class MockLM(BaseLM):
14+
"""Simple mock LM that echoes prompts."""
15+
16+
def __init__(self):
17+
super().__init__(model_name="mock-model")
18+
19+
def completion(self, prompt):
20+
return f"Mock response to: {prompt[:50]}"
21+
22+
async def acompletion(self, prompt):
23+
return self.completion(prompt)
24+
25+
def get_usage_summary(self):
26+
return UsageSummary(
27+
model_usage_summaries={
28+
"mock-model": ModelUsageSummary(
29+
total_calls=1, total_input_tokens=10, total_output_tokens=10
30+
)
31+
}
32+
)
33+
34+
def get_last_usage(self):
35+
return self.get_usage_summary()
36+
37+
38+
def main():
39+
print("=" * 60)
40+
print("Daytona REPL Example")
41+
print("=" * 60)
42+
43+
# Note: Requires DAYTONA_API_KEY environment variable to be set
44+
# or passed explicitly to DaytonaREPL(api_key="...")
45+
46+
# Example 1: Basic code execution
47+
print("\n[1] Basic code execution (no LLM handler)")
48+
print("-" * 40)
49+
50+
try:
51+
with DaytonaREPL(name="rlm-example") as repl:
52+
print(f"Daytona sandbox started with ID: {repl.sandbox.id}")
53+
result = repl.execute_code("x = 1 + 2")
54+
print("Executed: x = 1 + 2")
55+
print(f"Locals: {result.locals}")
56+
57+
result = repl.execute_code("print(x * 2)")
58+
print("Executed: print(x * 2)")
59+
print(f"Stdout: {result.stdout.strip()}")
60+
61+
result = repl.execute_code("answer = 42")
62+
result = repl.execute_code('print(FINAL_VAR("answer"))')
63+
print(f"FINAL_VAR('answer'): {result.stdout.strip()}")
64+
65+
# Example 2: With LLM handler
66+
print("\n[2] Code execution with LLM handler")
67+
print("-" * 40)
68+
69+
mock_client = MockLM()
70+
71+
with LMHandler(client=mock_client) as handler:
72+
print(f"LM Handler started at {handler.address}")
73+
74+
with DaytonaREPL(
75+
name="rlm-example-handler",
76+
lm_handler_address=handler.address,
77+
) as repl:
78+
print(f"Daytona sandbox started with ID: {repl.sandbox.id}")
79+
print(f"Broker URL: {repl.broker_url}")
80+
81+
# Single LLM query
82+
result = repl.execute_code('response = llm_query("What is 2+2?")')
83+
print("Executed: response = llm_query('What is 2+2?')")
84+
if result.stderr:
85+
print(f"Stderr: {result.stderr}")
86+
87+
result = repl.execute_code("print(response)")
88+
print(f"Response: {result.stdout.strip()}")
89+
90+
# Batched LLM query
91+
result = repl.execute_code(
92+
'responses = llm_query_batched(["Question 1", "Question 2", "Question 3"])'
93+
)
94+
print("\nExecuted: responses = llm_query_batched([...])")
95+
96+
result = repl.execute_code("print(f'Got {len(responses)} responses')")
97+
print(f"Result: {result.stdout.strip()}")
98+
99+
result = repl.execute_code("print(responses[0])")
100+
print(f"First response: {result.stdout.strip()}")
101+
102+
except Exception as e:
103+
print(f"Error: {e}")
104+
print("\nMake sure Daytona is configured correctly and DAYTONA_API_KEY is set.")
105+
106+
print("\n" + "=" * 60)
107+
print("Done!")
108+
print("=" * 60)
109+
110+
111+
if __name__ == "__main__":
112+
main()

‎examples/quickstart.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,13 @@
1515
"model_name": "gpt-5-nano",
1616
"api_key": os.getenv("OPENAI_API_KEY"),
1717
},
18-
environment="local",
18+
environment="docker",
1919
environment_kwargs={},
2020
max_depth=1,
2121
logger=logger,
2222
verbose=True, # For printing to console with rich, disabled by default.
2323
)
2424

25-
result = rlm.completion("Print me the first 100 powers of two, each on a newline.")
25+
result = rlm.completion("Print me the first 5 powers of two, each on a newline.")
2626

2727
print(result)

‎pyproject.toml‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,23 @@
11
[project]
2-
name = "rlm"
2+
name = "rlms"
33
authors = [
44
{name = "Alex Zhang", email = "altzhang@mit.edu"},
55
]
66
version = "0.1.0"
77
description = "Recursive Language Models."
88
readme = "README.md"
9+
license = "MIT"
910
requires-python = ">=3.11"
11+
classifiers = [
12+
"Development Status :: 4 - Beta",
13+
"Intended Audience :: Developers",
14+
"Programming Language :: Python :: 3",
15+
"Programming Language :: Python :: 3.11",
16+
"Programming Language :: Python :: 3.12",
17+
"Programming Language :: Python :: 3.13",
18+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
19+
]
20+
1021
dependencies = [
1122
"anthropic>=0.75.0",
1223
"google-genai>=1.56.0",
@@ -18,9 +29,15 @@ dependencies = [
1829
"rich>=13.0.0",
1930
]
2031

32+
[project.urls]
33+
Homepage = "https://github.com/alexzhang13/rlm"
34+
Repository = "https://github.com/alexzhang13/rlm"
35+
Issues = "https://github.com/alexzhang13/rlm/issues"
36+
2137
[project.optional-dependencies]
2238
modal = ["modal>=0.73.0", "dill>=0.3.7"]
2339
e2b = ["e2b-code-interpreter>=0.0.11", "dill>=0.3.7"]
40+
daytona = ["daytona>=0.128.1", "dill>=0.3.7"]
2441
prime = ["prime-sandboxes>=0.2.0", "dill>=0.3.7"]
2542

2643
[build-system]

‎rlm/clients/base_lm.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from abc import ABC, abstractmethod
22
from typing import Any
33

4-
from rlm.core.types import UsageSummary
4+
from rlm.core.types import ModelUsageSummary, UsageSummary
55

66

77
class BaseLM(ABC):
@@ -28,6 +28,6 @@ def get_usage_summary(self) -> UsageSummary:
2828
raise NotImplementedError
2929

3030
@abstractmethod
31-
def get_last_usage(self) -> UsageSummary:
31+
def get_last_usage(self) -> ModelUsageSummary:
3232
"""Get the last cost summary of the model."""
3333
raise NotImplementedError

‎rlm/core/lm_handler.py‎

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ def handle(self):
2222
request_data = socket_recv(self.connection)
2323
if not isinstance(request_data, dict):
2424
response = LMResponse.error_response("Request must be a JSON object")
25-
socket_send(self.connection, response.to_dict())
25+
self._safe_send(response)
2626
return
2727

2828
request = LMRequest.from_dict(request_data)
@@ -37,11 +37,26 @@ def handle(self):
3737
else:
3838
response = LMResponse.error_response("Missing 'prompt' or 'prompts' in request.")
3939

40-
socket_send(self.connection, response.to_dict())
40+
self._safe_send(response)
41+
42+
except (BrokenPipeError, ConnectionError, ConnectionResetError, OSError):
43+
# Client disconnected - this is expected during parallel execution
44+
# when workers complete and close their sockets. Silently ignore.
45+
pass
4146

4247
except Exception as e:
48+
# Try to send error response, but don't fail if socket is broken
4349
response = LMResponse.error_response(str(e))
50+
self._safe_send(response)
51+
52+
def _safe_send(self, response: LMResponse) -> bool:
53+
"""Send response, returning False if the socket is broken."""
54+
try:
4455
socket_send(self.connection, response.to_dict())
56+
return True
57+
except (BrokenPipeError, ConnectionError, ConnectionResetError, OSError):
58+
# Client disconnected - silently ignore
59+
return False
4560

4661
def _handle_single(self, request: LMRequest, handler: "LMHandler") -> LMResponse:
4762
"""Handle a single prompt request."""
@@ -51,10 +66,12 @@ def _handle_single(self, request: LMRequest, handler: "LMHandler") -> LMResponse
5166
content = client.completion(request.prompt)
5267
end_time = time.perf_counter()
5368

54-
usage_summary = client.get_last_usage()
69+
model_usage = client.get_last_usage()
70+
root_model = request.model or client.model_name
71+
usage_summary = UsageSummary(model_usage_summaries={root_model: model_usage})
5572
return LMResponse.success_response(
5673
chat_completion=RLMChatCompletion(
57-
root_model=request.model or client.model_name,
74+
root_model=root_model,
5875
prompt=request.prompt,
5976
response=content,
6077
usage_summary=usage_summary,
@@ -76,11 +93,13 @@ async def run_all():
7693
end_time = time.perf_counter()
7794

7895
total_time = end_time - start_time
79-
usage_summary = client.get_last_usage()
96+
model_usage = client.get_last_usage()
97+
root_model = request.model or client.model_name
98+
usage_summary = UsageSummary(model_usage_summaries={root_model: model_usage})
8099

81100
chat_completions = [
82101
RLMChatCompletion(
83-
root_model=request.model or client.model_name,
102+
root_model=root_model,
84103
prompt=prompt,
85104
response=content,
86105
usage_summary=usage_summary,

‎rlm/core/types.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
"azure_openai",
1414
"gemini",
1515
]
16-
EnvironmentType = Literal["local", "docker", "modal", "prime", "e2b"]
16+
EnvironmentType = Literal["local", "docker", "modal", "prime", "daytona", "e2b"]
1717

1818

1919
def _serialize_value(value: Any) -> Any:

‎rlm/environments/__init__.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,12 @@
77

88

99
def get_environment(
10-
environment: Literal["local", "modal", "docker", "prime", "e2b"],
10+
environment: Literal["local", "modal", "docker", "daytona", "prime", "e2b"],
1111
environment_kwargs: dict[str, Any],
1212
) -> BaseEnv:
1313
"""
1414
Routes a specific environment and the args (as a dict) to the appropriate environment if supported.
15-
Currently supported environments: ['local', 'modal', 'docker', 'prime', 'e2b']
15+
Currently supported environments: ['local', 'modal', 'docker', 'daytona', 'prime', 'e2b']
1616
"""
1717
if environment == "local":
1818
return LocalREPL(**environment_kwargs)
@@ -24,6 +24,10 @@ def get_environment(
2424
from rlm.environments.docker_repl import DockerREPL
2525

2626
return DockerREPL(**environment_kwargs)
27+
elif environment == "daytona":
28+
from rlm.environments.daytona_repl import DaytonaREPL
29+
30+
return DaytonaREPL(**environment_kwargs)
2731
elif environment == "prime":
2832
from rlm.environments.prime_repl import PrimeREPL
2933

@@ -34,5 +38,5 @@ def get_environment(
3438
return E2BREPL(**environment_kwargs)
3539
else:
3640
raise ValueError(
37-
f"Unknown environment: {environment}. Supported: ['local', 'modal', 'docker', 'prime', 'e2b']"
41+
f"Unknown environment: {environment}. Supported: ['local', 'modal', 'docker', 'daytona', 'prime', 'e2b']"
3842
)

0 commit comments

Comments
 (0)