Skip to content

Commit 1384031

Browse files
petroslambpetroslamb
committed
spec02: jupyter support
1 parent cde7ee1 commit 1384031

14 files changed

Lines changed: 2089 additions & 46 deletions

File tree

‎.gitignore‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
# MacOS
22
.DS_Store
33

4+
# Local agent instructions
5+
AGENTS.md
6+
47
# Logs
58
logs/
69
!visualizer/src/app/api/logs/

‎00-RLM-Notebook-Integration.ipynb‎

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"id": "5f12c999-c348-4247-a76b-35f8feb1b14e",
7+
"metadata": {
8+
"editable": true,
9+
"slideshow": {
10+
"slide_type": ""
11+
},
12+
"tags": []
13+
},
14+
"outputs": [],
15+
"source": [
16+
"import sys\n",
17+
"print(sys.executable)\n"
18+
]
19+
},
20+
{
21+
"cell_type": "code",
22+
"execution_count": null,
23+
"id": "5462657d-27e1-44a2-b46a-9097b28a6ed2",
24+
"metadata": {},
25+
"outputs": [],
26+
"source": [
27+
"from rlm import RLM\n",
28+
"from rlm.utils.notebook import render_completion\n",
29+
"from rlm.utils.trace_notebook import enable_trace_notebook\n",
30+
"\n",
31+
"logger, trace_handle = enable_trace_notebook(log_dir=\"logs\")\n",
32+
"rlm = RLM(environment=\"jupyter\", logger=logger)"
33+
]
34+
},
35+
{
36+
"cell_type": "code",
37+
"execution_count": null,
38+
"id": "a2d39c00-a9e7-457b-bffa-dd31b418d6c3",
39+
"metadata": {
40+
"jp-MarkdownHeadingCollapsed": true
41+
},
42+
"outputs": [],
43+
"source": [
44+
"rlm = RLM(\n",
45+
" backend=\"openai\",\n",
46+
" backend_kwargs={\"model_name\": \"gpt-5.2\"},\n",
47+
" environment=\"jupyter\",\n",
48+
" environment_kwargs={\"sync_to_user_ns\": True},\n",
49+
" max_iterations=15,\n",
50+
" persistent=True,\n",
51+
" logger=logger,\n",
52+
")"
53+
]
54+
},
55+
{
56+
"cell_type": "code",
57+
"execution_count": null,
58+
"id": "2cd91eff-71e1-4d0c-8c0b-5f33ddeb426c",
59+
"metadata": {
60+
"scrolled": true
61+
},
62+
"outputs": [],
63+
"source": [
64+
"session = rlm.start_session()\n",
65+
"result = session.chat(\"Create a Python variable x=42 and print it. Use a repl code block.\")\n",
66+
"result = session.chat(\"make that variable 24\")\n",
67+
"result.render_trace() # cumulative"
68+
]
69+
},
70+
{
71+
"cell_type": "code",
72+
"execution_count": null,
73+
"id": "bdc14eb2-8d4b-488b-8938-f48d57c61bb6",
74+
"metadata": {},
75+
"outputs": [],
76+
"source": [
77+
"x # noqa: F821"
78+
]
79+
},
80+
{
81+
"cell_type": "code",
82+
"execution_count": null,
83+
"id": "87b752f3-8002-4f1b-8b59-346ac2f1c63d",
84+
"metadata": {},
85+
"outputs": [],
86+
"source": [
87+
"session_context_1\n"
88+
]
89+
},
90+
{
91+
"cell_type": "code",
92+
"execution_count": null,
93+
"id": "210ec55a-2ef9-4e19-af11-19112cf0a734",
94+
"metadata": {},
95+
"outputs": [],
96+
"source": [
97+
"result = rlm.completion(\"Say hello in two lines.\")\n",
98+
"render_completion(result)"
99+
]
100+
},
101+
{
102+
"cell_type": "code",
103+
"execution_count": null,
104+
"id": "ee0b2e7b-301d-4b76-b8c4-238251e3ffbc",
105+
"metadata": {
106+
"scrolled": true
107+
},
108+
"outputs": [],
109+
"source": [
110+
"result.render_trace()"
111+
]
112+
},
113+
{
114+
"cell_type": "code",
115+
"execution_count": null,
116+
"id": "fdbfe41a-eb1e-42a9-8655-20913fc92631",
117+
"metadata": {},
118+
"outputs": [],
119+
"source": [
120+
"result = rlm.completion(\"Create a Python variable x=42 and print it. Use a repl code block.\")"
121+
]
122+
},
123+
{
124+
"cell_type": "code",
125+
"execution_count": null,
126+
"id": "de5d8459-36d6-4940-aa30-78bfc57b6168",
127+
"metadata": {},
128+
"outputs": [],
129+
"source": [
130+
"render_completion(result)"
131+
]
132+
},
133+
{
134+
"cell_type": "code",
135+
"execution_count": null,
136+
"id": "11e65233-c0dd-428a-bdbc-677755d5db33",
137+
"metadata": {},
138+
"outputs": [],
139+
"source": [
140+
"result.render_trace()"
141+
]
142+
},
143+
{
144+
"cell_type": "code",
145+
"execution_count": null,
146+
"id": "149e0dce-55b6-47ac-9035-f04d017f8300",
147+
"metadata": {},
148+
"outputs": [],
149+
"source": [
150+
"x"
151+
]
152+
},
153+
{
154+
"cell_type": "code",
155+
"execution_count": null,
156+
"id": "5f05357d-be81-4431-b992-865b969dbcd4",
157+
"metadata": {},
158+
"outputs": [],
159+
"source": [
160+
"completion_context\n"
161+
]
162+
},
163+
{
164+
"cell_type": "code",
165+
"execution_count": null,
166+
"id": "1f7c4f59-1d86-4451-a1ff-1cd24787704b",
167+
"metadata": {},
168+
"outputs": [],
169+
"source": [
170+
"from rlm.clients.openai import OpenAIClient\n",
171+
"\n",
172+
"client = OpenAIClient() # uses OPENAI_API_KEY by default\n",
173+
"models = client.client.models.list()\n",
174+
"[m.id for m in models.data]"
175+
]
176+
},
177+
{
178+
"cell_type": "code",
179+
"execution_count": null,
180+
"id": "1ba95705-1e0a-4ad3-a583-ec92033304e9",
181+
"metadata": {},
182+
"outputs": [],
183+
"source": []
184+
}
185+
],
186+
"metadata": {
187+
"kernelspec": {
188+
"display_name": "rlm-venv",
189+
"language": "python",
190+
"name": "rlm-venv"
191+
},
192+
"language_info": {
193+
"codemirror_mode": {
194+
"name": "ipython",
195+
"version": 3
196+
},
197+
"file_extension": ".py",
198+
"mimetype": "text/x-python",
199+
"name": "python",
200+
"nbconvert_exporter": "python",
201+
"pygments_lexer": "ipython3",
202+
"version": "3.11.9"
203+
},
204+
"rlm_trace_enabled": true
205+
},
206+
"nbformat": 4,
207+
"nbformat_minor": 5
208+
}

‎AGENTS.md‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ uv run pre-commit install
8282
- Small, focused diffs
8383
- One change per PR
8484
- Backward compatibility is only desirable if it can be done without introducing excessive maintenance burden
85+
- For POCs, prefer breaking changes over compatibility shims; remove obsolete args instead of accepting them
8586
- Delete dead code (don't guard it)
8687

8788
### Checklist
@@ -162,13 +163,14 @@ Environment implementations live in `rlm/environments/`. Choose the appropriate
162163

163164
### Key Implementation Details
164165
- `setup()`: Initialize globals, locals, and helper functions
165-
- `load_context()`: Make context available as `context` variable
166+
- `load_context()`: Make context available as `context_0` and set `context` to the latest prompt
166167
- `execute_code()`: Execute code, capture stdout/stderr, return `REPLResult`
167168
- Always provide `llm_query` and `llm_query_batched` functions in environment globals
168169

169170
### State Management
170171
Environments must provide these globals to executed code:
171-
- `context`: The loaded context payload
172+
- `context_0`: The loaded context payload
173+
- `context`: The latest prompt payload (updates as new contexts are added)
172174
- `llm_query(prompt, model=None)`: For sub-LM calls
173175
- `llm_query_batched(prompts, model=None)`: For batched sub-LM calls
174176
- `FINAL_VAR(variable_name)`: For returning final answers
@@ -191,7 +193,7 @@ class MyEnvironment(NonIsolatedEnv):
191193
# Initialize execution namespace
192194

193195
def load_context(self, context_payload: dict | list | str):
194-
# Make context available to executed code
196+
# Make context available as context_0 and set context to the latest prompt
195197

196198
def execute_code(self, code: str) -> REPLResult:
197199
# Execute code and return REPLResult
@@ -317,4 +319,3 @@ When building a new isolated environment (e.g., for a new cloud provider):
317319
5. **Handle state** - Serialize/deserialize execution state between code blocks
318320

319321
See `rlm/environments/modal_repl.py` as the canonical reference implementation.
320-

‎README.md‎

Lines changed: 84 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,14 +81,97 @@ We support two types of REPL environments -- isolated, and non-isolated. Non-iso
8181

8282
```python
8383
rlm = RLM(
84-
environment="...", # "local", "docker", "modal", "prime"
84+
environment="...", # "local", "docker", "modal", "prime", "jupyter"
8585
environment_kwargs={...},
8686
)
8787
```
8888

8989
### Local Environments
9090
The default `local` environment `LocalREPL` runs in the same process as the RLM itself, with specified global and local namespaces for minimal security. Using this REPL is generally safe, but should not be used for production settings. It also shares the same virtual environment (e.g. Conda or uv) as the host process.
9191

92+
#### Jupyter Notebook
93+
We support a `jupyter` environment that runs code blocks directly in the notebook kernel with stdout/stderr capture.
94+
```python
95+
rlm = RLM(
96+
environment="jupyter",
97+
environment_kwargs={"workdir": ".", "allow_builtin_imports": True},
98+
)
99+
```
100+
To keep REPL state across `completion()` and `session.chat()` calls, enable persistence (supported for `jupyter`):
101+
```python
102+
rlm = RLM(
103+
environment="jupyter",
104+
persistent=True,
105+
)
106+
```
107+
Persistent sessions expose prompt payloads as `session_context_0`, `session_context_1`, etc. The latest prompt is `session_context_{n-1}`. A `context_history` list contains all session contexts (overwritten each call), and `session_history` stores per-call message histories. Session prompts identify the current prompt as `session_context_{n-1}`; completion prompts only mention `completion_context`. Call `rlm.close()` when done.
108+
109+
Install notebook extras if you haven't already:
110+
```bash
111+
uv pip install -e ".[notebook]"
112+
```
113+
To make variables from executed code blocks available in the notebook namespace, enable sync:
114+
```python
115+
rlm = RLM(
116+
environment="jupyter",
117+
environment_kwargs={
118+
"workdir": ".",
119+
"allow_builtin_imports": True,
120+
"sync_to_user_ns": True,
121+
},
122+
)
123+
```
124+
For notebook-friendly rendering:
125+
```python
126+
from rlm.utils.notebook import render_completion
127+
result = rlm.completion("...")
128+
render_completion(result)
129+
```
130+
131+
To write a trace notebook per run:
132+
```python
133+
from rlm.utils.trace_notebook import enable_trace_notebook
134+
135+
logger, trace_handle = enable_trace_notebook(log_dir="logs")
136+
rlm = RLM(environment="jupyter", logger=logger)
137+
rlm.completion("...")
138+
139+
# stop_trace_writer(trace_handle) # optional
140+
```
141+
Trace notebooks include `llm_query` replay from recorded subcalls. Prompts must match exactly.
142+
143+
To access the markdown trace on the result:
144+
```python
145+
result = rlm.completion("...")
146+
result.trace_markdown
147+
```
148+
For multi-turn sessions with cumulative traces:
149+
```python
150+
session = rlm.start_session()
151+
result = session.chat("...")
152+
result.trace_markdown
153+
```
154+
For multi-turn sessions, the REPL exposes all prompt payloads as `session_context_0`, `session_context_1`, etc. The latest prompt is `session_context_{n-1}`; `context_history` tracks all session contexts (overwritten each call), and `session_history` stores per-call message histories. Session prompts identify the current prompt as `session_context_{n-1}`; completion prompts only mention `completion_context`.
155+
Assistant responses are always strings; `FINAL_VAR(...)` values are coerced to string.
156+
157+
To keep variables available in your notebook (after execution), enable sync:
158+
```python
159+
rlm = RLM(
160+
environment="jupyter",
161+
environment_kwargs={"sync_to_user_ns": True},
162+
)
163+
```
164+
165+
To add custom tools for a single session, use `setup_code`:
166+
```python
167+
rlm = RLM(
168+
environment="jupyter",
169+
environment_kwargs={
170+
"setup_code": "def my_tool(x): return x * 2",
171+
},
172+
)
173+
```
174+
92175
#### Docker <img src="https://github.com/docker.png" alt="Docker" height="20" style="vertical-align: middle;"/> (*requires [Docker installed](https://docs.docker.com/desktop/setup/install/)*)
93176
We also support a Docker-based environment called `DockerREPL` that launches the REPL environment as a Docker image. By default, we use the `python:3.11-slim` image, but the user can specify custom images as well.
94177

0 commit comments

Comments
 (0)