Skip to content

Commit e2783d5

Browse files
alexzhang13alexzhang13
committed
add AGENTS md
1 parent 5d6b9ee commit e2783d5

1 file changed

Lines changed: 317 additions & 0 deletions

File tree

‎AGENTS.md‎

Lines changed: 317 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,317 @@
1+
# AGENTS.md
2+
3+
This guide covers best practices for contributing to the core Recursive Language Models `rlm` library and developing new environments (in `rlm/environments/`) and LM clients (in `rlm/clients/`).
4+
5+
## Setup
6+
7+
We use `uv` for developing `rlm`.
8+
```bash
9+
# Install uv (first time)
10+
curl -LsSf https://astral.sh/uv/install.sh | sh
11+
12+
# Setup blank project if needed
13+
uv init && uv venv --python 3.12
14+
source .venv/bin/activate
15+
16+
# Install in editable mode
17+
uv pip install -e .
18+
19+
# For Modal sandbox support
20+
uv pip install -e ".[modal]"
21+
```
22+
23+
## General Guidelines
24+
25+
### Code Style & Typing
26+
- **Formatting**: Strict `ruff` enforcement. All PRs must pass `ruff check --fix .`
27+
- **Typing**: Explicit types preferred
28+
- **OK**: `cast(...)`, `assert ...` for type narrowing
29+
- **SOMETIMES OK**: Untyped args for simple cases (e.g., prompt handlers)
30+
- **NOT OK**: `# type: ignore` without strong justification
31+
32+
### Naming Conventions
33+
- **Methods**: snake_case
34+
- **Classes**: PascalCase (e.g., `LocalREPL`, `PortkeyClient`)
35+
- **Variables**: snake_case
36+
- **Constants**: UPPER_CASE (e.g., `_SAFE_BUILTINS`, `RLM_SYSTEM_PROMPT`)
37+
38+
Do NOT use `_` prefix for private methods unless explicitly requested.
39+
40+
### Error Handling Philosophy
41+
- **Fail fast, fail loud** - No defensive programming or silent fallbacks
42+
- **Minimize branching** - Prefer single code paths; every `if`/`try` needs justification
43+
- **Example**: Missing API key → immediate `ValueError`, not graceful fallback
44+
45+
## Core Repository Development
46+
47+
For PRs to `rlm` core:
48+
```bash
49+
git clone https://github.com/alexzhang13/rlm.git
50+
cd rlm
51+
52+
# Standard development:
53+
uv sync
54+
55+
# Install dev + test dependencies:
56+
uv sync --group dev --group test
57+
58+
# Install pre-commit hooks:
59+
uv run pre-commit install
60+
```
61+
62+
### Dependencies
63+
- Avoid new core dependencies
64+
- Use optional extras for non-essential features (e.g., `modal` extra)
65+
- Exception: tiny deps that simplify widely-used code
66+
67+
### Testing
68+
- `uv run pytest` with discovery under `tests/`
69+
- Write simple, deterministic unit tests
70+
- Update tests when changing functionality
71+
- For isolated environments, mock external services
72+
73+
### Documentation
74+
- Keep concise and actionable
75+
- Update README when behavior changes
76+
- Avoid content duplication
77+
78+
### Scope
79+
- Small, focused diffs
80+
- One change per PR
81+
- Backward compatibility is only desirable if it can be done without introducing excessive maintenance burden
82+
- Delete dead code (don't guard it)
83+
84+
### Checklist
85+
86+
Before a PR:
87+
88+
```bash
89+
# Run style + lint checks:
90+
uv run ruff check --fix .
91+
uv run ruff format .
92+
uv run pre-commit run --all-files
93+
94+
# Run tests:
95+
uv run pytest
96+
```
97+
98+
Ensure docs and tests are updated if necessary, and dead code is deleted. Strive for minimal, surgical diffs.
99+
100+
## Developing LM Clients
101+
102+
LM client implementations live in `rlm/clients/`. All clients must inherit from `BaseLM`.
103+
104+
### Client Pattern
105+
106+
| Base Class | When to Use | Key Methods |
107+
|------------|-------------|-------------|
108+
| `BaseLM` | All LM integrations | `completion`, `acompletion`, `get_usage_summary`, `get_last_usage` |
109+
110+
### Requirements
111+
- Inherit from `BaseLM` in `rlm/clients/base_lm.py`
112+
- Implement all abstract methods: `completion`, `acompletion`, `get_usage_summary`, `get_last_usage`
113+
- Track per-model usage (calls, input/output tokens)
114+
- Handle both string and message list prompts
115+
- Register client in `rlm/clients/__init__.py`
116+
117+
### Example Structure
118+
```python
119+
from rlm.clients.base_lm import BaseLM
120+
from rlm.core.types import ModelUsageSummary, UsageSummary
121+
122+
class MyClient(BaseLM):
123+
def __init__(self, api_key: str, model_name: str, **kwargs):
124+
super().__init__(model_name=model_name, **kwargs)
125+
# Initialize your client
126+
127+
def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
128+
# Handle both str and message list formats
129+
# Track usage with _track_cost()
130+
# Return response string
131+
132+
def get_usage_summary(self) -> UsageSummary:
133+
# Return aggregated usage across all calls
134+
```
135+
136+
### Configuration Guidelines
137+
- **Environment variables**: ONLY for API keys (document in README)
138+
- **Hardcode**: Default base URLs, reasonable defaults
139+
- **Arguments**: Essential customization via `__init__()`
140+
141+
## Developing Environments
142+
143+
Environment implementations live in `rlm/environments/`. Choose the appropriate base class.
144+
145+
### Environment Pattern
146+
147+
| Pattern | Base Class | When to Use | Key Methods |
148+
|---------|------------|-------------|-------------|
149+
| **Non-isolated** | `NonIsolatedEnv` | Local execution, same machine | `setup`, `load_context`, `execute_code` |
150+
| **Isolated** | `IsolatedEnv` | Cloud sandboxes (Modal, Prime) | `setup`, `load_context`, `execute_code` |
151+
152+
### Requirements
153+
- Inherit from `NonIsolatedEnv` or `IsolatedEnv` in `rlm/environments/base_env.py`
154+
- Implement all abstract methods: `setup`, `load_context`, `execute_code`
155+
- Return `REPLResult` from `execute_code`
156+
- Handle `lm_handler_address` for sub-LM calls via `llm_query()`
157+
- Implement `cleanup()` for resource management
158+
- Register environment in `rlm/environments/__init__.py`
159+
160+
### Key Implementation Details
161+
- `setup()`: Initialize globals, locals, and helper functions
162+
- `load_context()`: Make context available as `context` variable
163+
- `execute_code()`: Execute code, capture stdout/stderr, return `REPLResult`
164+
- Always provide `llm_query` and `llm_query_batched` functions in environment globals
165+
166+
### State Management
167+
Environments must provide these globals to executed code:
168+
- `context`: The loaded context payload
169+
- `llm_query(prompt, model=None)`: For sub-LM calls
170+
- `llm_query_batched(prompts, model=None)`: For batched sub-LM calls
171+
- `FINAL_VAR(variable_name)`: For returning final answers
172+
173+
### Example Structure
174+
```python
175+
from rlm.environments.base_env import NonIsolatedEnv
176+
from rlm.core.types import REPLResult
177+
178+
class MyEnvironment(NonIsolatedEnv):
179+
def __init__(self, lm_handler_address: tuple[str, int] | None = None,
180+
context_payload: dict | list | str | None = None, **kwargs):
181+
super().__init__(**kwargs)
182+
self.lm_handler_address = lm_handler_address
183+
self.setup()
184+
if context_payload:
185+
self.load_context(context_payload)
186+
187+
def setup(self):
188+
# Initialize execution namespace
189+
190+
def load_context(self, context_payload: dict | list | str):
191+
# Make context available to executed code
192+
193+
def execute_code(self, code: str) -> REPLResult:
194+
# Execute code and return REPLResult
195+
196+
def cleanup(self):
197+
# Clean up resources
198+
```
199+
200+
### Checklist
201+
- Guidelines here are followed
202+
- Environment works with basic RLM completion calls
203+
- `cleanup()` properly releases all resources
204+
- Sub-LM calls work via `llm_query()`
205+
206+
## Architecture: Environment ↔ LM Handler Communication
207+
208+
Understanding how environments communicate with the LM Handler is essential for developing new environments.
209+
210+
### Overview
211+
212+
```
213+
┌─────────────────────────────────────────────────────────────────────┐
214+
│ Host Machine │
215+
│ ┌─────────────┐ Socket (TCP) ┌──────────────────────┐ │
216+
│ │ RLM │◄──────────────────────────► LMHandler │ │
217+
│ │ (main) │ │ (ThreadingTCPServer)│ │
218+
│ └─────────────┘ └──────────────────────┘ │
219+
│ │ ▲ │
220+
│ ▼ │ │
221+
│ ┌─────────────┐ Socket (TCP) │ │
222+
│ │ LocalREPL │────────────────────────────────────┘ │
223+
│ │ (exec code) │ llm_query() → send_lm_request() │
224+
│ └─────────────┘ │
225+
└─────────────────────────────────────────────────────────────────────┘
226+
```
227+
228+
### Socket Protocol (Non-Isolated Environments)
229+
230+
Non-isolated environments like `LocalREPL` communicate directly with the `LMHandler` via TCP sockets using a length-prefixed JSON protocol:
231+
232+
**Protocol Format**: `4-byte big-endian length prefix + UTF-8 JSON payload`
233+
234+
```python
235+
# Sending a message (from rlm/core/comms_utils.py)
236+
def socket_send(sock: socket.socket, data: dict) -> None:
237+
payload = json.dumps(data).encode("utf-8")
238+
sock.sendall(struct.pack(">I", len(payload)) + payload)
239+
```
240+
241+
**Request Flow**:
242+
1. Environment's `llm_query(prompt)` is called during code execution
243+
2. Creates `LMRequest` dataclass and calls `send_lm_request(address, request)`
244+
3. Opens TCP connection to `LMHandler` at `(host, port)`
245+
4. Sends length-prefixed JSON request
246+
5. `LMHandler` processes via `LMRequestHandler.handle()`
247+
6. Returns `LMResponse` with `RLMChatCompletion` or error
248+
249+
**Key Components**:
250+
- `LMHandler` (`rlm/core/lm_handler.py`): Multi-threaded TCP server wrapping LM clients
251+
- `LMRequest` / `LMResponse` (`rlm/core/comms_utils.py`): Typed request/response dataclasses
252+
- `send_lm_request()` / `send_lm_request_batched()`: Helper functions for socket communication
253+
254+
### HTTP Broker Pattern (Isolated Environments)
255+
256+
Isolated environments (Modal, Prime) cannot directly connect to the host's socket server. They use an HTTP broker pattern:
257+
258+
```
259+
┌─────────────────────────────────────────────────────────────────────────────┐
260+
│ Host Machine │
261+
│ ┌─────────┐ Socket ┌────────────┐ HTTP Poll ┌────────────────┐ │
262+
│ │ RLM │◄────────────►│ LMHandler │◄────────────────│ ModalREPL │ │
263+
│ └─────────┘ └────────────┘ │ (poller) │ │
264+
│ └────────────────┘ │
265+
│ │ │
266+
│ HTTP (tunnel) │
267+
│ │ │
268+
└──────────────────────────────────────────────────────────────────┼──────────┘
269+
│
270+
┌──────────────────────────────────────────────────────────────────┼──────────┐
271+
│ Cloud Sandbox (Modal/Prime) ▼ │
272+
│ ┌─────────────┐ HTTP (localhost) ┌─────────────────────────────┐ │
273+
│ │ Exec Script │◄────────────────────────►│ Broker Server (Flask) │ │
274+
│ │ (exec code) │ /enqueue, etc. │ - /enqueue (submit req) │ │
275+
│ └─────────────┘ │ - /pending (poll reqs) │ │
276+
│ │ - /respond (return resp) │ │
277+
│ └─────────────────────────────┘ │
278+
└─────────────────────────────────────────────────────────────────────────────┘
279+
```
280+
281+
**How It Works**:
282+
283+
1. **Sandbox Setup**: Environment creates a cloud sandbox with an HTTP broker server running inside
284+
2. **Tunnel Exposure**: Broker server is exposed via encrypted tunnel (e.g., Modal's `encrypted_ports`)
285+
3. **Code Execution**: When `llm_query()` is called inside sandbox, it POSTs to `http://localhost:8080/enqueue`
286+
4. **Request Queuing**: Broker queues the request and blocks waiting for response
287+
5. **Host Polling**: `ModalREPL` on host polls `{tunnel_url}/pending` for new requests
288+
6. **LM Forwarding**: Host forwards requests to `LMHandler` via socket, gets response
289+
7. **Response Delivery**: Host POSTs response to `{tunnel_url}/respond`
290+
8. **Unblocking**: Broker unblocks the original `/enqueue` call with the response
291+
292+
**Broker Endpoints**:
293+
| Endpoint | Method | Purpose |
294+
|----------|--------|---------|
295+
| `/enqueue` | POST | Submit LLM request from sandbox code (blocks until response) |
296+
| `/pending` | GET | Get list of pending requests (called by host poller) |
297+
| `/respond` | POST | Submit response for a request ID (called by host poller) |
298+
| `/health` | GET | Health check |
299+
300+
**Key Implementation Details**:
301+
- Broker runs as a Flask server inside the sandbox
302+
- Uses `threading.Event` for request/response synchronization
303+
- Poller thread on host runs in background with 100ms polling interval
304+
- State persistence via `dill` serialization to `/tmp/rlm_state.dill`
305+
306+
### Implementing a New Isolated Environment
307+
308+
When building a new isolated environment (e.g., for a new cloud provider):
309+
310+
1. **Create broker server** - Flask/HTTP server with `/enqueue`, `/pending`, `/respond` endpoints
311+
2. **Expose tunnel** - Use provider's tunnel/port forwarding to expose broker to host
312+
3. **Implement poller** - Background thread on host to poll and forward requests
313+
4. **Build exec script** - Script that runs inside sandbox with `llm_query()` calling broker
314+
5. **Handle state** - Serialize/deserialize execution state between code blocks
315+
316+
See `rlm/environments/modal_repl.py` as the canonical reference implementation.
317+

0 commit comments

Comments
 (0)