@hicaru/pi-rlm 0.2.1 → 0.3.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +28 -47
- package/README.ru.md +18 -23
- package/README.zh-CN.md +17 -28
- package/package.json +22 -19
- package/src/bridge/add-context.ts +322 -0
- package/src/bridge/subcall-handlers.ts +63 -17
- package/src/commands/rlm-config.ts +47 -18
- package/src/commands/rlm.ts +3 -152
- package/src/config/defaults.ts +8 -18
- package/src/config/settings.ts +13 -34
- package/src/context/anydoc.ts +67 -0
- package/src/context/listing.ts +70 -0
- package/src/context/md-cache.ts +112 -0
- package/src/context/merge.ts +97 -0
- package/src/context/namespace.ts +180 -0
- package/src/context/resolve.ts +122 -0
- package/src/context/source-dir.ts +166 -0
- package/src/context/source-doc.ts +71 -0
- package/src/context/source-git.ts +51 -0
- package/src/context/source-text.ts +45 -0
- package/src/context/types.ts +88 -0
- package/src/context/walk.ts +250 -0
- package/src/core/engine.ts +61 -345
- package/src/core/history.ts +1 -1
- package/src/core/limits.ts +5 -12
- package/src/core/resource-limits.ts +0 -2
- package/src/core/types.ts +10 -38
- package/src/index.ts +92 -54
- package/src/mode/llm-model.ts +54 -0
- package/src/mode/rlm-mode.ts +28 -58
- package/src/prompts/glossary.ts +290 -0
- package/src/prompts/native.ts +127 -0
- package/src/prompts/system.ts +15 -408
- package/src/sandbox/context-file.ts +154 -0
- package/src/sandbox/interrupts.ts +160 -0
- package/src/sandbox/protocol.ts +20 -75
- package/src/sandbox/py/__pycache__/guards.cpython-314.pyc +0 -0
- package/src/sandbox/py/__pycache__/retrieval.cpython-314.pyc +0 -0
- package/src/sandbox/py/__pycache__/tasks.cpython-314.pyc +0 -0
- package/src/sandbox/py/guards.py +150 -0
- package/src/sandbox/py/retrieval.py +265 -0
- package/src/sandbox/py/tasks.py +129 -0
- package/src/sandbox/py/worker.py +856 -0
- package/src/sandbox/sandbox-manager.ts +24 -9
- package/src/sandbox/sandbox.ts +99 -193
- package/src/text/tokens.ts +31 -5
- package/src/tool/repl-details.ts +2 -2
- package/src/tool/repl-render.ts +58 -0
- package/src/tool/repl-result.ts +70 -0
- package/src/tool/repl-tool.ts +60 -170
- package/src/tool/rlm-aggregator.ts +2 -10
- package/src/tool/rlm-details.ts +0 -2
- package/src/tool/rlm-events.ts +0 -14
- package/src/tool/rlm-tool.ts +2 -13
- package/src/ui/config-panel.ts +12 -20
- package/src/ui/intro.ts +1 -2
- package/src/ui/model-picker.ts +34 -10
- package/src/ui/status.ts +3 -7
- package/src/util/concurrency.ts +9 -5
- package/src/bridge/fallback-todo.ts +0 -148
- package/src/bridge/interactive.ts +0 -65
- package/src/bridge/library.ts +0 -155
- package/src/bridge/pi-interactive.ts +0 -41
- package/src/context/library-context.ts +0 -266
- package/src/context/repomix-context.ts +0 -204
- package/src/core/artifacts.ts +0 -89
- package/src/core/critique.ts +0 -92
- package/src/core/gates.ts +0 -301
- package/src/core/pipeline-handlers.ts +0 -319
- package/src/core/pipeline.ts +0 -268
- package/src/prompts/phases.ts +0 -104
- package/src/sandbox/worker.py +0 -1456
- package/src/state/index.ts +0 -24
- package/src/state/internal.ts +0 -46
- package/src/state/paths.ts +0 -44
- package/src/state/reads.ts +0 -133
- package/src/state/resume.ts +0 -173
- package/src/state/rows.ts +0 -123
- package/src/state/writes.ts +0 -58
package/src/sandbox/worker.py
DELETED
|
@@ -1,1456 +0,0 @@
|
|
|
1
|
-
"""RLM sandbox worker: a persistent Python REPL driven over a JSONL stdio protocol.
|
|
2
|
-
|
|
3
|
-
Executes model-authored Python with secrets stripped from the environment.
|
|
4
|
-
This is NOT a security sandbox: __import__ and open are available, so code can import networking modules
|
|
5
|
-
(socket, urllib, subprocess) and read/write local files. Trust the root model's code.
|
|
6
|
-
|
|
7
|
-
Protocol (parent -> worker): {"id","type":"exec"|"load_context"|"shutdown", ...}
|
|
8
|
-
Protocol (worker -> parent): {"id","ok",...result} # response to a request
|
|
9
|
-
{"type":"llm_query"|"llm_query_batched"|"rlm_query"|...
|
|
10
|
-
"advance_phase"|"save_artifact"|"ask_user_question"|"todo","rid",...}
|
|
11
|
-
# mid-exec helper request
|
|
12
|
-
When sandbox code calls llm_query/rlm_query/advance_phase/save_artifact/ask_user_question/todo, the worker writes a request line
|
|
13
|
-
and BLOCKS reading stdin until the matching {"type":"llm_reply","rid",...} arrives. The parent
|
|
14
|
-
services the request in-process (it holds API keys).
|
|
15
|
-
|
|
16
|
-
Requests and replies are decoupled: `_post` writes a request and returns its rid without
|
|
17
|
-
waiting, and replies are parked in `_inbox` keyed by rid until something asks for them. That
|
|
18
|
-
is what makes `spawn()` / `rlm_await()` / `rlm_await_all()` possible — many requests can be
|
|
19
|
-
in flight at once (the parent already services interrupts concurrently), and a task may be
|
|
20
|
-
awaited in a LATER exec than the one that started it.
|
|
21
|
-
"""
|
|
22
|
-
|
|
23
|
-
from __future__ import annotations
|
|
24
|
-
|
|
25
|
-
import argparse
|
|
26
|
-
import fnmatch
|
|
27
|
-
import heapq
|
|
28
|
-
import io
|
|
29
|
-
import json
|
|
30
|
-
import math
|
|
31
|
-
import os
|
|
32
|
-
import pickle
|
|
33
|
-
import re
|
|
34
|
-
import signal
|
|
35
|
-
import sys
|
|
36
|
-
import time
|
|
37
|
-
import traceback
|
|
38
|
-
from contextlib import contextmanager
|
|
39
|
-
from typing import Any
|
|
40
|
-
|
|
41
|
-
# Capture the REAL stdio before exec() redirects sys.stdout/sys.stderr into buffers.
|
|
42
|
-
# All protocol writes must go to the real stdout even while user code's prints are captured.
|
|
43
|
-
_REAL_STDOUT = sys.stdout
|
|
44
|
-
_REAL_STDIN = sys.stdin
|
|
45
|
-
_REAL_STDERR = sys.stderr
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
def _builtin(name: str):
|
|
49
|
-
return __builtins__[name] if isinstance(__builtins__, dict) else getattr(__builtins__, name, None)
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
# Restricted builtins: enough for real data work, minus the dangerous reflection escapes.
|
|
53
|
-
_SAFE_BUILTINS = {
|
|
54
|
-
name: _builtin(name)
|
|
55
|
-
for name in (
|
|
56
|
-
"abs", "all", "any", "ascii", "bin", "bool", "bytearray", "bytes", "callable",
|
|
57
|
-
"chr", "classmethod", "complex", "dict", "dir", "divmod", "enumerate", "filter",
|
|
58
|
-
"float", "format", "frozenset", "getattr", "hasattr", "hash", "hex", "id", "int",
|
|
59
|
-
"isinstance", "issubclass", "iter", "len", "list", "map", "max", "min", "next",
|
|
60
|
-
"object", "oct", "ord", "pow", "print", "property", "range", "repr", "reversed",
|
|
61
|
-
"round", "set", "setattr", "slice", "sorted", "staticmethod", "str", "sum", "super",
|
|
62
|
-
"tuple", "type", "vars", "zip", "delattr", "memoryview", "__import__", "__build_class__",
|
|
63
|
-
"Exception", "BaseException", "ValueError", "TypeError", "KeyError", "IndexError",
|
|
64
|
-
"AttributeError", "FileNotFoundError", "OSError", "IOError", "RuntimeError",
|
|
65
|
-
"NameError", "ImportError", "StopIteration", "AssertionError", "NotImplementedError",
|
|
66
|
-
"ArithmeticError", "ZeroDivisionError", "LookupError", "Warning", "True", "False", "None",
|
|
67
|
-
)
|
|
68
|
-
}
|
|
69
|
-
# `open` is allowed for data work; eval/exec/compile/input/globals/locals are not.
|
|
70
|
-
# When read_only=True (pipeline runs), write modes raise PermissionError via
|
|
71
|
-
# builtins.open, io.open (pathlib), and os.open. Steering, not a security sandbox.
|
|
72
|
-
_WRITE_MODE_CHARS = frozenset("wax+")
|
|
73
|
-
_OS_WRITE_FLAGS = os.O_WRONLY | os.O_RDWR | os.O_CREAT | os.O_APPEND | os.O_TRUNC
|
|
74
|
-
|
|
75
|
-
_REAL_IO_OPEN = io.open
|
|
76
|
-
_REAL_OS_OPEN = os.open
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
def _install_read_only_guards():
|
|
80
|
-
"""Route every common file-open path through the read-only check.
|
|
81
|
-
|
|
82
|
-
Steering, not a sandbox: closes builtins.open, io.open (hence pathlib), and
|
|
83
|
-
os.open. A determined model can still reach the filesystem via ctypes or a
|
|
84
|
-
subprocess — the goal is that ACCIDENTAL writes cannot pass silently.
|
|
85
|
-
Worker-internal I/O keeps using _REAL_IO_OPEN / _REAL_OS_OPEN.
|
|
86
|
-
"""
|
|
87
|
-
def guarded_io_open(file, mode="r", *args, **kwargs):
|
|
88
|
-
if _WRITE_MODE_CHARS & set(str(mode)):
|
|
89
|
-
raise PermissionError(
|
|
90
|
-
f"read-only RLM run: refusing to open {file!r} with mode {mode!r}. "
|
|
91
|
-
"This pipeline produces a plan; file changes go through the host edit tool."
|
|
92
|
-
)
|
|
93
|
-
return _REAL_IO_OPEN(file, mode, *args, **kwargs)
|
|
94
|
-
|
|
95
|
-
def guarded_os_open(path, flags, *args, **kwargs):
|
|
96
|
-
if flags & _OS_WRITE_FLAGS:
|
|
97
|
-
raise PermissionError(
|
|
98
|
-
f"read-only RLM run: refusing os.open({path!r}) with write flags."
|
|
99
|
-
)
|
|
100
|
-
return _REAL_OS_OPEN(path, flags, *args, **kwargs)
|
|
101
|
-
|
|
102
|
-
io.open = guarded_io_open
|
|
103
|
-
os.open = guarded_os_open
|
|
104
|
-
return guarded_io_open
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
# _builtin()'s getattr(..., None) fallback would silently inject None for a name this
|
|
108
|
-
# interpreter lacks, surfacing much later as "'NoneType' object is not callable" inside model
|
|
109
|
-
# code. Fail at startup instead. Note "None" is legitimately None, and the block-list below is
|
|
110
|
-
# deliberate — which is why this check runs BEFORE it.
|
|
111
|
-
_MISSING = sorted(name for name, value in _SAFE_BUILTINS.items() if value is None and name != "None")
|
|
112
|
-
if _MISSING:
|
|
113
|
-
raise RuntimeError(f"unsupported Python interpreter: missing builtins {_MISSING}")
|
|
114
|
-
|
|
115
|
-
def _blocked_builtin(name: str):
|
|
116
|
-
"""Bind a disabled builtin to a callable that explains itself.
|
|
117
|
-
|
|
118
|
-
Binding these to None made `eval(...)` fail with a bare "'NoneType' object is not callable",
|
|
119
|
-
which reads as a broken sandbox rather than a deliberate block: an audit session spent six
|
|
120
|
-
execs on it and filed a phantom "namespace corruption" bug. Saying so at the point of failure
|
|
121
|
-
fixes it for every model without spending native-prompt budget on a rule most runs never hit.
|
|
122
|
-
"""
|
|
123
|
-
def blocked(*_args, **_kwargs):
|
|
124
|
-
raise PermissionError(
|
|
125
|
-
f"{name}() is disabled in the RLM sandbox by design — it is not missing and the "
|
|
126
|
-
"namespace is not corrupt. Names are already bound, so reference them directly; "
|
|
127
|
-
"inspect `context` with search() / grep_context() / outline()."
|
|
128
|
-
)
|
|
129
|
-
blocked.__name__ = name
|
|
130
|
-
return blocked
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
# Blocked on purpose (NOT missing) — see _blocked_builtin. The _MISSING check above runs first,
|
|
134
|
-
# so a genuinely absent builtin is still a startup failure rather than a silent None.
|
|
135
|
-
for _blocked in ("eval", "exec", "compile", "input", "globals", "locals"):
|
|
136
|
-
_SAFE_BUILTINS[_blocked] = _blocked_builtin(_blocked)
|
|
137
|
-
|
|
138
|
-
RESERVED = frozenset(
|
|
139
|
-
{
|
|
140
|
-
"llm_query", "llm_query_batched", "llm_query_chunked",
|
|
141
|
-
"rlm_query", "rlm_query_batched",
|
|
142
|
-
"spawn", "rlm_await", "rlm_await_all",
|
|
143
|
-
"map_files", "llm_map_reduce",
|
|
144
|
-
"search", "grep_context", "outline",
|
|
145
|
-
"advance_phase", "save_artifact",
|
|
146
|
-
"ask_user_question", "todo",
|
|
147
|
-
"load_library",
|
|
148
|
-
"SHOW_VARS", "answer", "context",
|
|
149
|
-
}
|
|
150
|
-
)
|
|
151
|
-
# NOTE: `answers` and `plan` are deliberately NOT reserved. They are seeded by the scaffold but
|
|
152
|
-
# owned by the model, so they must appear in SHOW_VARS and be captured by snapshots — losing a
|
|
153
|
-
# memoized answer across a resume is exactly the failure the memo exists to prevent.
|
|
154
|
-
# Only the single name `context` is the packed world. Legacy context_N names are filtered out.
|
|
155
|
-
_CONTEXT_NAME = re.compile(r"context(_\d+)?\Z")
|
|
156
|
-
|
|
157
|
-
# Sizing for llm_query_chunked: leave room for the instruction and the chunk header.
|
|
158
|
-
_CHUNK_HEADER_OVERHEAD = 64
|
|
159
|
-
_MAX_CHUNK_BATCH = 20 # fan-out per llm_query_batched call (matches prompt guidance)
|
|
160
|
-
_MAX_CHUNKS = 500 # ceiling: above this, force pre-filtering in Python
|
|
161
|
-
_NUDGE_CHARS = 500_000 # str/bytes vars above this trigger a one-time stdout hint
|
|
162
|
-
|
|
163
|
-
|
|
164
|
-
def _chunk_text(text: str, chunk_chars: int) -> list[str]:
|
|
165
|
-
"""Split text into <=chunk_chars pieces, preferring newline boundaries."""
|
|
166
|
-
chunks: list[str] = []
|
|
167
|
-
n = len(text)
|
|
168
|
-
start = 0
|
|
169
|
-
while start < n:
|
|
170
|
-
end = min(start + chunk_chars, n)
|
|
171
|
-
if end < n:
|
|
172
|
-
nl = text.rfind("\n", start, end)
|
|
173
|
-
if nl > start:
|
|
174
|
-
end = nl + 1
|
|
175
|
-
chunks.append(text[start:end])
|
|
176
|
-
start = end
|
|
177
|
-
return chunks
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
# ---- deterministic retrieval over `context` -----------------------------------------------
|
|
181
|
-
#
|
|
182
|
-
# The RLM paper's trajectories retrieve by having the root model hand-write regex over the
|
|
183
|
-
# context (App. E.1). Frontier models do that well; small/fast models guess keywords badly and
|
|
184
|
-
# the first decomposition attempt disproportionately decides the outcome (paper §5, Fig. 4a).
|
|
185
|
-
# These primitives make retrieval deterministic and token-free: no sub-LLM call, no root tokens
|
|
186
|
-
# spent on printed file bodies — the model gets ranked pointers and decides what to delegate.
|
|
187
|
-
|
|
188
|
-
_INDEX_WINDOW_LINES = 40 # a window is the retrieval unit: big enough to carry meaning
|
|
189
|
-
_INDEX_MAX_WINDOWS = 20_000 # ceiling so a huge load_library() cannot exhaust worker memory
|
|
190
|
-
_SNIPPET_CHARS = 400
|
|
191
|
-
_GREP_HARD_CAP = 200 # absolute ceiling on returned grep hits, whatever k asks for
|
|
192
|
-
_BM25_K1 = 1.2
|
|
193
|
-
_BM25_B = 0.75
|
|
194
|
-
|
|
195
|
-
_TOKEN_SPLIT = re.compile(r"[^0-9A-Za-z]+") # also splits snake_case and paths
|
|
196
|
-
_CAMEL_SPLIT = re.compile(r"(?<=[a-z0-9])(?=[A-Z])")
|
|
197
|
-
|
|
198
|
-
# Definition-ish lines across the languages this plugin is likely to meet. Deliberately
|
|
199
|
-
# lexical: an outline is an orientation aid, not a parse tree.
|
|
200
|
-
_OUTLINE_LINE = re.compile(
|
|
201
|
-
r"^\s*(?:"
|
|
202
|
-
r"(?:export\s+)?(?:default\s+)?(?:async\s+)?(?:function|class|interface|type|enum|struct|impl|trait|namespace)\s+\w+"
|
|
203
|
-
r"|(?:export\s+)?(?:const|let|var)\s+\w+\s*[:=]\s*(?:async\s*)?(?:function|\(|<)"
|
|
204
|
-
r"|(?:pub\s+)?(?:async\s+)?fn\s+\w+"
|
|
205
|
-
r"|def\s+\w+|class\s+\w+"
|
|
206
|
-
r"|func\s+\w+"
|
|
207
|
-
r"|#{1,4}\s+\S"
|
|
208
|
-
r")"
|
|
209
|
-
)
|
|
210
|
-
|
|
211
|
-
|
|
212
|
-
def _tokenize(text: str) -> list[str]:
|
|
213
|
-
"""Lowercased alphanumeric runs, plus camelCase parts so `resolveModelId` matches `model id`."""
|
|
214
|
-
out: list[str] = []
|
|
215
|
-
for raw in _TOKEN_SPLIT.split(text):
|
|
216
|
-
if not raw:
|
|
217
|
-
continue
|
|
218
|
-
lowered = raw.lower()
|
|
219
|
-
out.append(lowered)
|
|
220
|
-
if len(raw) > 3:
|
|
221
|
-
parts = _CAMEL_SPLIT.split(raw)
|
|
222
|
-
if len(parts) > 1:
|
|
223
|
-
for part in parts:
|
|
224
|
-
piece = part.lower()
|
|
225
|
-
if piece and piece != lowered:
|
|
226
|
-
out.append(piece)
|
|
227
|
-
return out
|
|
228
|
-
|
|
229
|
-
|
|
230
|
-
def _context_entries(context: Any) -> list[tuple[str, str]]:
|
|
231
|
-
"""(path, content) pairs for either context shape: list[dict] bundles or a raw string."""
|
|
232
|
-
if isinstance(context, str):
|
|
233
|
-
return [("<context>", context)]
|
|
234
|
-
if not isinstance(context, list):
|
|
235
|
-
return []
|
|
236
|
-
out: list[tuple[str, str]] = []
|
|
237
|
-
for i, item in enumerate(context):
|
|
238
|
-
if isinstance(item, dict):
|
|
239
|
-
content = item.get("content", "")
|
|
240
|
-
out.append((
|
|
241
|
-
str(item.get("path", f"<context[{i}]>")),
|
|
242
|
-
content if isinstance(content, str) else str(content),
|
|
243
|
-
))
|
|
244
|
-
elif isinstance(item, str):
|
|
245
|
-
out.append((f"<context[{i}]>", item))
|
|
246
|
-
return out
|
|
247
|
-
|
|
248
|
-
|
|
249
|
-
class _Bm25Index:
|
|
250
|
-
"""Okapi BM25 over fixed-line windows of `context`. Built lazily, discarded on change."""
|
|
251
|
-
|
|
252
|
-
__slots__ = ("paths", "starts", "texts", "postings", "doc_len", "avg_len", "truncated")
|
|
253
|
-
|
|
254
|
-
def __init__(self, entries: list[tuple[str, str]]) -> None:
|
|
255
|
-
self.paths: list[str] = []
|
|
256
|
-
self.starts: list[int] = []
|
|
257
|
-
self.texts: list[str] = []
|
|
258
|
-
self.postings: dict[str, list[tuple[int, int]]] = {}
|
|
259
|
-
self.doc_len: list[int] = []
|
|
260
|
-
self.truncated = False
|
|
261
|
-
|
|
262
|
-
for path, content in entries:
|
|
263
|
-
if not content:
|
|
264
|
-
continue
|
|
265
|
-
lines = content.split("\n")
|
|
266
|
-
for start in range(0, len(lines), _INDEX_WINDOW_LINES):
|
|
267
|
-
if len(self.texts) >= _INDEX_MAX_WINDOWS:
|
|
268
|
-
self.truncated = True
|
|
269
|
-
break
|
|
270
|
-
window = "\n".join(lines[start:start + _INDEX_WINDOW_LINES])
|
|
271
|
-
idx = len(self.texts)
|
|
272
|
-
self.paths.append(path)
|
|
273
|
-
self.starts.append(start + 1)
|
|
274
|
-
self.texts.append(window)
|
|
275
|
-
terms = _tokenize(window)
|
|
276
|
-
self.doc_len.append(len(terms))
|
|
277
|
-
freq: dict[str, int] = {}
|
|
278
|
-
for term in terms:
|
|
279
|
-
freq[term] = freq.get(term, 0) + 1
|
|
280
|
-
for term, tf in freq.items():
|
|
281
|
-
self.postings.setdefault(term, []).append((idx, tf))
|
|
282
|
-
if self.truncated:
|
|
283
|
-
break
|
|
284
|
-
|
|
285
|
-
total = len(self.doc_len)
|
|
286
|
-
self.avg_len = (sum(self.doc_len) / total) if total else 1.0
|
|
287
|
-
|
|
288
|
-
def query(self, terms: list[str], k: int, path_glob: str | None) -> list[dict[str, Any]]:
|
|
289
|
-
total = len(self.texts)
|
|
290
|
-
if total == 0:
|
|
291
|
-
return []
|
|
292
|
-
scores: dict[int, float] = {}
|
|
293
|
-
for term in set(terms):
|
|
294
|
-
posting = self.postings.get(term)
|
|
295
|
-
if not posting:
|
|
296
|
-
continue
|
|
297
|
-
df = len(posting)
|
|
298
|
-
idf = math.log(1.0 + (total - df + 0.5) / (df + 0.5))
|
|
299
|
-
for idx, tf in posting:
|
|
300
|
-
norm = _BM25_K1 * (1.0 - _BM25_B + _BM25_B * self.doc_len[idx] / self.avg_len)
|
|
301
|
-
scores[idx] = scores.get(idx, 0.0) + idf * (tf * (_BM25_K1 + 1.0)) / (tf + norm)
|
|
302
|
-
if path_glob:
|
|
303
|
-
scores = {i: s for i, s in scores.items() if fnmatch.fnmatch(self.paths[i], path_glob)}
|
|
304
|
-
if not scores:
|
|
305
|
-
return []
|
|
306
|
-
top = heapq.nlargest(k, scores.items(), key=lambda kv: kv[1])
|
|
307
|
-
out: list[dict[str, Any]] = [None] * len(top) # type: ignore[list-item]
|
|
308
|
-
for i, (idx, score) in enumerate(top):
|
|
309
|
-
text = self.texts[idx]
|
|
310
|
-
out[i] = {
|
|
311
|
-
"path": self.paths[idx],
|
|
312
|
-
"line": self.starts[idx],
|
|
313
|
-
"score": round(score, 3),
|
|
314
|
-
"snippet": text[:_SNIPPET_CHARS],
|
|
315
|
-
}
|
|
316
|
-
return out
|
|
317
|
-
|
|
318
|
-
|
|
319
|
-
class _AnswerDict(dict):
|
|
320
|
-
"""`answer` dict; flipping `ready` True captures the final answer for the parent."""
|
|
321
|
-
|
|
322
|
-
def __init__(self, on_ready):
|
|
323
|
-
super().__init__()
|
|
324
|
-
super().__setitem__("content", "")
|
|
325
|
-
super().__setitem__("ready", False)
|
|
326
|
-
self._on_ready = on_ready
|
|
327
|
-
|
|
328
|
-
def __setitem__(self, key, value):
|
|
329
|
-
super().__setitem__(key, value)
|
|
330
|
-
if key == "ready" and value:
|
|
331
|
-
self._on_ready(self.get("content", ""))
|
|
332
|
-
|
|
333
|
-
|
|
334
|
-
def _send(obj: dict[str, Any]) -> None:
|
|
335
|
-
_REAL_STDOUT.write(json.dumps(obj, ensure_ascii=False) + "\n")
|
|
336
|
-
_REAL_STDOUT.flush()
|
|
337
|
-
|
|
338
|
-
|
|
339
|
-
class _StallTimeout(Exception):
|
|
340
|
-
"""No frame from the host while a sub-call was pending."""
|
|
341
|
-
|
|
342
|
-
|
|
343
|
-
@contextmanager
|
|
344
|
-
def _stall_alarm(exec_timeout_s: float, stall_timeout_s: float):
|
|
345
|
-
"""Swap the per-cell alarm for a stall alarm while blocked on the parent.
|
|
346
|
-
|
|
347
|
-
Sub-LLM latency is network time, not cell compute time, so it must not count against the
|
|
348
|
-
```repl``` block timeout — but an unbounded wait is exactly how a lost reply turns into a
|
|
349
|
-
dead session. The yielded `rearm()` restarts the stall clock on every frame, so a healthy
|
|
350
|
-
long-running child never trips it.
|
|
351
|
-
"""
|
|
352
|
-
use = hasattr(signal, "SIGALRM")
|
|
353
|
-
remaining = signal.getitimer(signal.ITIMER_REAL)[0] if (use and exec_timeout_s > 0) else 0.0
|
|
354
|
-
|
|
355
|
-
def _fire(signum, frame): # noqa: ARG001
|
|
356
|
-
raise _StallTimeout(
|
|
357
|
-
f"sub-call stalled — no reply from the host for {stall_timeout_s:g}s "
|
|
358
|
-
"(the task may still be running; rlm_await it again in a later block)"
|
|
359
|
-
)
|
|
360
|
-
|
|
361
|
-
old = signal.signal(signal.SIGALRM, _fire) if use else None
|
|
362
|
-
|
|
363
|
-
def rearm() -> None:
|
|
364
|
-
if use and stall_timeout_s > 0:
|
|
365
|
-
signal.setitimer(signal.ITIMER_REAL, stall_timeout_s)
|
|
366
|
-
|
|
367
|
-
rearm()
|
|
368
|
-
try:
|
|
369
|
-
yield rearm
|
|
370
|
-
finally:
|
|
371
|
-
if use:
|
|
372
|
-
signal.setitimer(signal.ITIMER_REAL, 0)
|
|
373
|
-
if old is not None:
|
|
374
|
-
signal.signal(signal.SIGALRM, old)
|
|
375
|
-
if remaining > 0: # restore the cell's remaining budget
|
|
376
|
-
signal.setitimer(signal.ITIMER_REAL, remaining)
|
|
377
|
-
|
|
378
|
-
|
|
379
|
-
def _surfaced_error(message: str) -> str:
|
|
380
|
-
"""The "Error: …" contract value, ALSO written to the cell's stderr.
|
|
381
|
-
|
|
382
|
-
A spawn/await misuse whose only trace is the returned value reads to the model as a random
|
|
383
|
-
string much later — which is exactly how `tasks.items()` blew up on a str.
|
|
384
|
-
"""
|
|
385
|
-
print(f"[rlm] {message}", file=sys.stderr)
|
|
386
|
-
return f"Error: {message}"
|
|
387
|
-
|
|
388
|
-
|
|
389
|
-
# ---- reply reducers: raw parent replies -> the value the scaffold fn returns ----------------
|
|
390
|
-
# One reducer per result shape, shared by the sync helpers and their spawned equivalents.
|
|
391
|
-
|
|
392
|
-
|
|
393
|
-
def _reduce_one(replies: list[dict[str, Any]]) -> str:
|
|
394
|
-
r = replies[0]
|
|
395
|
-
return f"Error: {r['error']}" if r.get("error") else r.get("response", "")
|
|
396
|
-
|
|
397
|
-
|
|
398
|
-
def _reduce_batch(n: int):
|
|
399
|
-
"""Reducer for a single *_query_batched reply of n prompts."""
|
|
400
|
-
def reduce(replies: list[dict[str, Any]]) -> list[str]:
|
|
401
|
-
r = replies[0]
|
|
402
|
-
if r.get("error"):
|
|
403
|
-
return [f"Error: {r['error']}"] * n
|
|
404
|
-
out = r.get("responses")
|
|
405
|
-
if not isinstance(out, list) or len(out) != n:
|
|
406
|
-
return ["Error: malformed batched response"] * n
|
|
407
|
-
return [s if isinstance(s, str) else f"Error: {s}" for s in out]
|
|
408
|
-
return reduce
|
|
409
|
-
|
|
410
|
-
|
|
411
|
-
def _reduce_chunked(sizes: list[int]):
|
|
412
|
-
"""Concatenate several llm_query_batched replies back into one flat chunk list."""
|
|
413
|
-
per = [_reduce_batch(n) for n in sizes]
|
|
414
|
-
|
|
415
|
-
def reduce(replies: list[dict[str, Any]]) -> list[str]:
|
|
416
|
-
out: list[str] = []
|
|
417
|
-
for red, rep in zip(per, replies):
|
|
418
|
-
out.extend(red([rep]))
|
|
419
|
-
return out
|
|
420
|
-
return reduce
|
|
421
|
-
|
|
422
|
-
|
|
423
|
-
def _reduce_map_files(sizes: list[int], spans: list[tuple[str, int]]):
|
|
424
|
-
"""Flatten the batch replies (same as chunked), then regroup them per path.
|
|
425
|
-
|
|
426
|
-
A file larger than the per-prompt budget contributed several requests; its answers rejoin
|
|
427
|
-
in order, which is what makes map_files a {path: answer} dict rather than a flat list.
|
|
428
|
-
"""
|
|
429
|
-
flatten = _reduce_chunked(sizes)
|
|
430
|
-
|
|
431
|
-
def reduce(replies: list[dict[str, Any]]) -> dict[str, str]:
|
|
432
|
-
responses = flatten(replies)
|
|
433
|
-
out: dict[str, str] = {}
|
|
434
|
-
cursor = 0
|
|
435
|
-
for path, count in spans:
|
|
436
|
-
part = responses[cursor:cursor + count]
|
|
437
|
-
cursor += count
|
|
438
|
-
out[path] = part[0] if count == 1 and part else "\n\n".join(part)
|
|
439
|
-
return out
|
|
440
|
-
return reduce
|
|
441
|
-
|
|
442
|
-
|
|
443
|
-
def _spawnable(name: str):
|
|
444
|
-
"""Tag a sync scaffold fn with the request kind spawn() should route it to."""
|
|
445
|
-
def mark(fn):
|
|
446
|
-
fn._rlm_name = name
|
|
447
|
-
return fn
|
|
448
|
-
return mark
|
|
449
|
-
|
|
450
|
-
|
|
451
|
-
class Task:
|
|
452
|
-
"""Handle for parent-side work in flight, returned by spawn().
|
|
453
|
-
|
|
454
|
-
Opaque to model code apart from `done` and repr. A Task may be awaited in a later
|
|
455
|
-
```repl``` block than the one that created it.
|
|
456
|
-
"""
|
|
457
|
-
|
|
458
|
-
__slots__ = ("kind", "label", "_worker", "_rids", "_reduce", "_value", "_settled")
|
|
459
|
-
|
|
460
|
-
def __init__(self, worker: "Worker", kind: str, rids, reduce, label: str = ""):
|
|
461
|
-
self.kind = kind
|
|
462
|
-
self.label = label
|
|
463
|
-
self._worker = worker
|
|
464
|
-
self._rids = tuple(rids)
|
|
465
|
-
self._reduce = reduce
|
|
466
|
-
self._value: Any = None
|
|
467
|
-
self._settled = False
|
|
468
|
-
|
|
469
|
-
@staticmethod
|
|
470
|
-
def resolved(worker: "Worker", kind: str, value: Any, label: str = "") -> "Task":
|
|
471
|
-
"""A Task that never hit the wire — validation errors and empty inputs."""
|
|
472
|
-
task = Task(worker, kind, (), lambda _replies: value, label)
|
|
473
|
-
task._value = value
|
|
474
|
-
task._settled = True
|
|
475
|
-
return task
|
|
476
|
-
|
|
477
|
-
@property
|
|
478
|
-
def done(self) -> bool:
|
|
479
|
-
"""True once every reply has landed — awaiting will not block."""
|
|
480
|
-
return self._settled or all(r in self._worker.inbox for r in self._rids)
|
|
481
|
-
|
|
482
|
-
def __repr__(self) -> str:
|
|
483
|
-
return f"<Task {self.kind} {'done' if self.done else 'running'} {self.label}>"
|
|
484
|
-
|
|
485
|
-
|
|
486
|
-
class Worker:
|
|
487
|
-
def __init__(
|
|
488
|
-
self,
|
|
489
|
-
depth: int,
|
|
490
|
-
exec_timeout_s: float,
|
|
491
|
-
max_prompt_chars: int,
|
|
492
|
-
read_only: bool = False,
|
|
493
|
-
await_timeout_s: float = 600.0,
|
|
494
|
-
):
|
|
495
|
-
self.depth = depth
|
|
496
|
-
self.exec_timeout_s = exec_timeout_s
|
|
497
|
-
self.max_prompt_chars = max_prompt_chars
|
|
498
|
-
self.read_only = read_only
|
|
499
|
-
self.await_timeout_s = await_timeout_s
|
|
500
|
-
self._rid = 0
|
|
501
|
-
self._final_answer: str | None = None
|
|
502
|
-
# Replies parked by rid until something awaits them. Unbounded by design: a task
|
|
503
|
-
# the model spawns and never awaits keeps its entry for the life of the process.
|
|
504
|
-
# Bounded in practice by session length; evicting would silently hang a later
|
|
505
|
-
# rlm_await, which is strictly worse than the memory.
|
|
506
|
-
self.inbox: dict[str, dict[str, Any]] = {}
|
|
507
|
-
self._inflight: set[str] = set()
|
|
508
|
-
# Requests (exec/snapshot/shutdown) that arrived mid-exec; main() replays them.
|
|
509
|
-
self._deferred: list[Any] = []
|
|
510
|
-
# True only while spawn() runs a builder — marks requests that may outlive this exec.
|
|
511
|
-
self._detached = False
|
|
512
|
-
self.ns: dict[str, Any] = {}
|
|
513
|
-
self._setup()
|
|
514
|
-
|
|
515
|
-
def _setup(self) -> None:
|
|
516
|
-
builtins = _SAFE_BUILTINS.copy()
|
|
517
|
-
if self.read_only:
|
|
518
|
-
builtins["open"] = _install_read_only_guards()
|
|
519
|
-
else:
|
|
520
|
-
builtins["open"] = open
|
|
521
|
-
self.ns = {"__builtins__": builtins, "__name__": "__main__"}
|
|
522
|
-
self._context_payload: Any | None = None # pristine restore for the single `context` var
|
|
523
|
-
self._nudged: set[str] = set()
|
|
524
|
-
self._index: _Bm25Index | None = None
|
|
525
|
-
self._index_stamp: tuple[int, int] | None = None # (id(context), len(context))
|
|
526
|
-
self._restore_scaffold()
|
|
527
|
-
|
|
528
|
-
def _capture_answer(self, content: Any) -> None:
|
|
529
|
-
self._final_answer = str(content)
|
|
530
|
-
|
|
531
|
-
def _restore_scaffold(self) -> None:
|
|
532
|
-
# Re-inject any scaffolding the user code clobbered.
|
|
533
|
-
ns = self.ns
|
|
534
|
-
ns["llm_query"] = self._llm_query
|
|
535
|
-
ns["llm_query_batched"] = self._llm_query_batched
|
|
536
|
-
ns["llm_query_chunked"] = self._llm_query_chunked
|
|
537
|
-
ns["rlm_query"] = self._rlm_query
|
|
538
|
-
ns["rlm_query_batched"] = self._rlm_query_batched
|
|
539
|
-
ns["spawn"] = self._spawn
|
|
540
|
-
ns["rlm_await"] = self._await_task
|
|
541
|
-
ns["rlm_await_all"] = self._await_all
|
|
542
|
-
ns["map_files"] = self._map_files
|
|
543
|
-
ns["llm_map_reduce"] = self._llm_map_reduce
|
|
544
|
-
ns["search"] = self._search
|
|
545
|
-
ns["grep_context"] = self._grep_context
|
|
546
|
-
ns["outline"] = self._outline
|
|
547
|
-
# env_tips memo (paper App. C.3): "If a value isn't in `answers`, it doesn't exist."
|
|
548
|
-
# Re-created only when deleted — contents must survive every turn.
|
|
549
|
-
if not isinstance(ns.get("answers"), dict):
|
|
550
|
-
ns["answers"] = {}
|
|
551
|
-
if not isinstance(ns.get("plan"), dict):
|
|
552
|
-
ns["plan"] = {}
|
|
553
|
-
ns["advance_phase"] = self._advance_phase
|
|
554
|
-
ns["save_artifact"] = self._save_artifact
|
|
555
|
-
ns["ask_user_question"] = self._ask_user_question
|
|
556
|
-
ns["todo"] = self._todo
|
|
557
|
-
ns["load_library"] = self._load_library
|
|
558
|
-
ns["SHOW_VARS"] = self._show_vars
|
|
559
|
-
if not isinstance(ns.get("answer"), _AnswerDict):
|
|
560
|
-
cur = ns.get("answer")
|
|
561
|
-
ans = _AnswerDict(self._capture_answer)
|
|
562
|
-
if isinstance(cur, dict):
|
|
563
|
-
for k, v in cur.items():
|
|
564
|
-
dict.__setitem__(ans, k, v)
|
|
565
|
-
if cur.get("ready") and self._final_answer is None:
|
|
566
|
-
self._final_answer = str(cur.get("content", ""))
|
|
567
|
-
ns["answer"] = ans
|
|
568
|
-
# Single context variable (RLM paper: the context lives in the environment and
|
|
569
|
-
# the model may transform it in place). Re-inject only if the model deleted the
|
|
570
|
-
# name entirely; mutations and re-binds persist within the run.
|
|
571
|
-
if self._context_payload is not None:
|
|
572
|
-
ns.setdefault("context", self._context_payload)
|
|
573
|
-
# Scrub any legacy context_N names so the model never sees multi-slot APIs.
|
|
574
|
-
for k in list(ns.keys()):
|
|
575
|
-
if k != "context" and _CONTEXT_NAME.match(k):
|
|
576
|
-
del ns[k]
|
|
577
|
-
|
|
578
|
-
def _user_var_names(self) -> list[str]:
|
|
579
|
-
"""User-created variable names — filters builtins, scaffold, and `context`.
|
|
580
|
-
|
|
581
|
-
Shared by SHOW_VARS() and the exec result so both expose the same namespace view.
|
|
582
|
-
This is the cheap orientation hint that goes into history instead of full stdout.
|
|
583
|
-
"""
|
|
584
|
-
return [
|
|
585
|
-
k for k in self.ns
|
|
586
|
-
if not k.startswith("_")
|
|
587
|
-
and not _CONTEXT_NAME.match(k)
|
|
588
|
-
and k not in RESERVED
|
|
589
|
-
]
|
|
590
|
-
|
|
591
|
-
def _show_vars(self) -> str:
|
|
592
|
-
avail = {k: type(self.ns[k]).__name__ for k in self._user_var_names()}
|
|
593
|
-
return f"Available variables: {avail}" if avail else "No variables created yet."
|
|
594
|
-
|
|
595
|
-
# ---- sub-LLM bridge over stdio --------------------------------------------------------
|
|
596
|
-
|
|
597
|
-
def _post(self, kind: str, payload: dict[str, Any]) -> str:
|
|
598
|
-
"""Write one parent request and return its rid WITHOUT waiting for the reply."""
|
|
599
|
-
self._rid += 1
|
|
600
|
-
rid = f"q{self._rid}"
|
|
601
|
-
# Register only after the write succeeds — a broken pipe must not leave an
|
|
602
|
-
# _inflight entry that nothing will ever settle.
|
|
603
|
-
_send({"type": kind, "rid": rid, "depth": self.depth,
|
|
604
|
-
"detached": self._detached, **payload})
|
|
605
|
-
self._inflight.add(rid)
|
|
606
|
-
return rid
|
|
607
|
-
|
|
608
|
-
def park_reply(self, msg: Any) -> bool:
|
|
609
|
-
"""File an llm_reply against its rid. True when the frame was a reply.
|
|
610
|
-
|
|
611
|
-
Public because main() needs it too: a spawned task can settle while the worker
|
|
612
|
-
sits idle between execs, and that reply must not fall through to "unknown type".
|
|
613
|
-
"""
|
|
614
|
-
if not isinstance(msg, dict) or msg.get("type") != "llm_reply":
|
|
615
|
-
return False
|
|
616
|
-
rid = msg.get("rid")
|
|
617
|
-
if isinstance(rid, str) and rid in self._inflight:
|
|
618
|
-
self._inflight.discard(rid)
|
|
619
|
-
self.inbox[rid] = msg
|
|
620
|
-
else:
|
|
621
|
-
# Late reply to an abandoned rid (e.g. a request from a discarded sandbox).
|
|
622
|
-
print(f"[rlm-sandbox] dropping reply for unknown rid: {rid!r}", file=_REAL_STDERR)
|
|
623
|
-
return True
|
|
624
|
-
|
|
625
|
-
def take_deferred(self) -> Any | None:
|
|
626
|
-
"""Pop a request that arrived mid-exec, for main() to replay. None when empty."""
|
|
627
|
-
return self._deferred.pop(0) if self._deferred else None
|
|
628
|
-
|
|
629
|
-
def _pump(self) -> bool:
|
|
630
|
-
"""Read one frame from the parent into the inbox. False when the pipe closed."""
|
|
631
|
-
line = _REAL_STDIN.readline()
|
|
632
|
-
if not line:
|
|
633
|
-
return False
|
|
634
|
-
try:
|
|
635
|
-
msg = json.loads(line)
|
|
636
|
-
except ValueError:
|
|
637
|
-
print(f"[rlm-sandbox] skipping non-JSON parent frame: {line[:200]}", file=_REAL_STDERR)
|
|
638
|
-
return True
|
|
639
|
-
if self.park_reply(msg):
|
|
640
|
-
return True
|
|
641
|
-
# A request (exec/snapshot/shutdown) arriving mid-exec: main() replays it.
|
|
642
|
-
self._deferred.append(msg)
|
|
643
|
-
return True
|
|
644
|
-
|
|
645
|
-
def _drain_until(self, rids) -> None:
|
|
646
|
-
"""Block until every rid in `rids` has its reply parked in the inbox.
|
|
647
|
-
|
|
648
|
-
Bounded: a host that goes silent raises inside the ```repl``` block instead of hanging
|
|
649
|
-
the session forever.
|
|
650
|
-
"""
|
|
651
|
-
if all(r in self.inbox for r in rids):
|
|
652
|
-
return
|
|
653
|
-
with _stall_alarm(self.exec_timeout_s, self.await_timeout_s) as rearm:
|
|
654
|
-
while not all(r in self.inbox for r in rids):
|
|
655
|
-
if not self._pump():
|
|
656
|
-
raise RuntimeError("parent closed the pipe during a sub-LLM request")
|
|
657
|
-
rearm()
|
|
658
|
-
|
|
659
|
-
def _take(self, rids) -> list[dict[str, Any]]:
|
|
660
|
-
return [self.inbox.pop(r) for r in rids]
|
|
661
|
-
|
|
662
|
-
def _rpc(self, kind: str, payload: dict[str, Any]) -> dict[str, Any]:
|
|
663
|
-
"""Post one request and block for its reply — the synchronous single-shot path."""
|
|
664
|
-
rid = self._post(kind, payload)
|
|
665
|
-
self._drain_until((rid,))
|
|
666
|
-
return self._take((rid,))[0]
|
|
667
|
-
|
|
668
|
-
# ---- spawn / await ---------------------------------------------------------------------
|
|
669
|
-
|
|
670
|
-
def _start_prompt(self, kind: str, prompt, model) -> Task:
|
|
671
|
-
text = str(prompt)
|
|
672
|
-
# A sub-LLM asked nothing answers something: the confabulation then sits in `answers`
|
|
673
|
-
# looking exactly like data. Refuse instead of spending a call on it.
|
|
674
|
-
if not text.strip():
|
|
675
|
-
return Task.resolved(self, kind, _surfaced_error(
|
|
676
|
-
f"{kind}() got an empty prompt — a sub-LLM would confabulate an answer to nothing"))
|
|
677
|
-
rid = self._post(kind, {"prompt": text, "model": model})
|
|
678
|
-
return Task(self, kind, (rid,), _reduce_one, text[:40])
|
|
679
|
-
|
|
680
|
-
def _start_prompts(self, kind: str, prompts, model) -> Task:
|
|
681
|
-
prompts = [str(p) for p in prompts]
|
|
682
|
-
if not prompts:
|
|
683
|
-
return Task.resolved(self, kind, [])
|
|
684
|
-
# Only the all-blank case: one blank prompt among twenty is the caller's business.
|
|
685
|
-
if not any(p.strip() for p in prompts):
|
|
686
|
-
return Task.resolved(self, kind, [
|
|
687
|
-
_surfaced_error(f"{kind}() got only empty prompts")
|
|
688
|
-
] * len(prompts))
|
|
689
|
-
rid = self._post(kind, {"prompts": prompts, "model": model})
|
|
690
|
-
return Task(self, kind, (rid,), _reduce_batch(len(prompts)), f"×{len(prompts)}")
|
|
691
|
-
|
|
692
|
-
def _start_llm_query(self, prompt, model: str | None = None) -> Task:
|
|
693
|
-
return self._start_prompt("llm_query", prompt, model)
|
|
694
|
-
|
|
695
|
-
def _start_rlm_query(self, prompt, model: str | None = None) -> Task:
|
|
696
|
-
return self._start_prompt("rlm_query", prompt, model)
|
|
697
|
-
|
|
698
|
-
def _start_llm_query_batched(self, prompts, model: str | None = None) -> Task:
|
|
699
|
-
return self._start_prompts("llm_query_batched", prompts, model)
|
|
700
|
-
|
|
701
|
-
def _start_rlm_query_batched(self, prompts, model: str | None = None) -> Task:
|
|
702
|
-
return self._start_prompts("rlm_query_batched", prompts, model)
|
|
703
|
-
|
|
704
|
-
def _start_llm_query_chunked(self, text, prompt: str, model: str | None = None) -> Task:
|
|
705
|
-
"""Split oversized text into cap-sized chunks and post EVERY batch at once.
|
|
706
|
-
|
|
707
|
-
One answer per chunk, order preserved. No exceptions escape: errors come back as
|
|
708
|
-
"Error: ..." strings per chunk (same contract as llm_query_batched). Because all
|
|
709
|
-
batches go on the wire together, a large input costs one round-trip of latency
|
|
710
|
-
rather than one per 20 chunks.
|
|
711
|
-
|
|
712
|
-
NOTE: budget uses Python code-point length (len) while the parent-side cap check counts
|
|
713
|
-
UTF-16 units (JS string.length); astral/emoji-heavy text may be marginally larger on the
|
|
714
|
-
parent and get per-chunk rejected. Acceptable trade-off for typical code/log/profile text.
|
|
715
|
-
"""
|
|
716
|
-
text, prompt = str(text), str(prompt)
|
|
717
|
-
if not text:
|
|
718
|
-
return Task.resolved(self, "llm_query_chunked", [])
|
|
719
|
-
budget = self.max_prompt_chars - len(prompt) - _CHUNK_HEADER_OVERHEAD
|
|
720
|
-
if budget < 1_000:
|
|
721
|
-
return Task.resolved(self, "llm_query_chunked", [
|
|
722
|
-
f"Error: prompt leaves under 1,000 chars per chunk (cap {self.max_prompt_chars:,}) — shorten the instruction"
|
|
723
|
-
])
|
|
724
|
-
chunks = _chunk_text(text, budget)
|
|
725
|
-
total = len(chunks)
|
|
726
|
-
if total > _MAX_CHUNKS:
|
|
727
|
-
return Task.resolved(self, "llm_query_chunked", [
|
|
728
|
-
f"Error: {total} chunks would be needed — filter/slice the text in Python first"
|
|
729
|
-
])
|
|
730
|
-
rids: list[str] = []
|
|
731
|
-
sizes: list[int] = []
|
|
732
|
-
for i in range(0, total, _MAX_CHUNK_BATCH):
|
|
733
|
-
batch = [
|
|
734
|
-
f"{prompt}\n\n[chunk {i + j + 1}/{total} of the input]\n{c}"
|
|
735
|
-
for j, c in enumerate(chunks[i:i + _MAX_CHUNK_BATCH])
|
|
736
|
-
]
|
|
737
|
-
rids.append(self._post("llm_query_batched", {"prompts": batch, "model": model}))
|
|
738
|
-
sizes.append(len(batch))
|
|
739
|
-
return Task(self, "llm_query_chunked", tuple(rids), _reduce_chunked(sizes), f"{total} chunks")
|
|
740
|
-
|
|
741
|
-
def _builder_for(self, name: str):
|
|
742
|
-
# llm_map_reduce is deliberately absent: its reduce step is a SECOND sub-LLM call that
|
|
743
|
-
# depends on its own map results, so it cannot be one (rids, pure reduce) Task.
|
|
744
|
-
return {
|
|
745
|
-
"llm_query": self._start_llm_query,
|
|
746
|
-
"llm_query_batched": self._start_llm_query_batched,
|
|
747
|
-
"llm_query_chunked": self._start_llm_query_chunked,
|
|
748
|
-
"map_files": self._start_map_files,
|
|
749
|
-
"rlm_query": self._start_rlm_query,
|
|
750
|
-
"rlm_query_batched": self._start_rlm_query_batched,
|
|
751
|
-
}.get(name)
|
|
752
|
-
|
|
753
|
-
def _spawn(self, fn, *args, **kwargs) -> Task:
|
|
754
|
-
"""Start a sub-call without waiting for it. `fn` is the scaffold function itself.
|
|
755
|
-
|
|
756
|
-
Returns a Task for rlm_await / rlm_await_all, possibly in a later ```repl``` block.
|
|
757
|
-
Misuse returns an already-resolved error Task rather than raising, matching the
|
|
758
|
-
"Error: ..." contract of the synchronous helpers.
|
|
759
|
-
"""
|
|
760
|
-
name = getattr(fn, "_rlm_name", None)
|
|
761
|
-
builder = self._builder_for(name) if isinstance(name, str) else None
|
|
762
|
-
if builder is None:
|
|
763
|
-
return Task.resolved(self, "spawn", _surfaced_error(
|
|
764
|
-
"spawn() takes llm_query, llm_query_batched, llm_query_chunked, map_files, "
|
|
765
|
-
"rlm_query or rlm_query_batched — not llm_map_reduce, whose reduce step depends "
|
|
766
|
-
"on its own map results and so cannot be a single Task"))
|
|
767
|
-
# Mark every request this builder posts as detached: the parent routes them to its
|
|
768
|
-
# session-scoped registry, since they may outlive the exec that started them.
|
|
769
|
-
self._detached = True
|
|
770
|
-
try:
|
|
771
|
-
return builder(*args, **kwargs)
|
|
772
|
-
except TypeError as e:
|
|
773
|
-
return Task.resolved(self, "spawn", _surfaced_error(f"bad spawn arguments — {e}"))
|
|
774
|
-
finally:
|
|
775
|
-
self._detached = False
|
|
776
|
-
|
|
777
|
-
def _await_task(self, task) -> Any:
|
|
778
|
-
"""Block until `task` has its result. Idempotent — the value is memoized."""
|
|
779
|
-
if not isinstance(task, Task):
|
|
780
|
-
return _surfaced_error(
|
|
781
|
-
f"rlm_await expects a Task from spawn(), got {type(task).__name__}"
|
|
782
|
-
)
|
|
783
|
-
if not task._settled:
|
|
784
|
-
self._drain_until(task._rids)
|
|
785
|
-
task._value = task._reduce(self._take(task._rids))
|
|
786
|
-
task._settled = True
|
|
787
|
-
return task._value
|
|
788
|
-
|
|
789
|
-
def _await_all(self, tasks) -> list:
|
|
790
|
-
"""Block until every task has its result. Order matches the input."""
|
|
791
|
-
tasks = list(tasks)
|
|
792
|
-
# One union drain so the tasks overlap instead of settling one after another.
|
|
793
|
-
union: list[str] = []
|
|
794
|
-
seen: set[str] = set()
|
|
795
|
-
for t in tasks:
|
|
796
|
-
if not isinstance(t, Task) or t._settled:
|
|
797
|
-
continue
|
|
798
|
-
for rid in t._rids:
|
|
799
|
-
if rid not in seen:
|
|
800
|
-
seen.add(rid)
|
|
801
|
-
union.append(rid)
|
|
802
|
-
if union:
|
|
803
|
-
self._drain_until(union)
|
|
804
|
-
return [self._await_task(t) for t in tasks]
|
|
805
|
-
|
|
806
|
-
# ---- deterministic retrieval (no sub-LLM calls, no root tokens) -----------------------
|
|
807
|
-
|
|
808
|
-
def _entries(self) -> list[tuple[str, str]]:
|
|
809
|
-
return _context_entries(self.ns.get("context"))
|
|
810
|
-
|
|
811
|
-
def _get_index(self) -> _Bm25Index:
|
|
812
|
-
"""Build the BM25 index on first use; rebuild when `context` was replaced or resized.
|
|
813
|
-
|
|
814
|
-
Identity+length is a cheap stamp that catches the two ways context actually changes:
|
|
815
|
-
load_library() extending the list, and the model re-binding the name. In-place edits
|
|
816
|
-
that preserve length are not detected — documented, and rare in practice.
|
|
817
|
-
"""
|
|
818
|
-
ctx = self.ns.get("context")
|
|
819
|
-
stamp = (id(ctx), len(ctx) if isinstance(ctx, (list, str)) else 0)
|
|
820
|
-
if self._index is None or self._index_stamp != stamp:
|
|
821
|
-
self._index = _Bm25Index(self._entries())
|
|
822
|
-
self._index_stamp = stamp
|
|
823
|
-
return self._index
|
|
824
|
-
|
|
825
|
-
def _search(self, query: str, k: int = 10, path_glob: str | None = None) -> list[dict[str, Any]]:
|
|
826
|
-
"""Rank `context` windows against a natural-language query (BM25).
|
|
827
|
-
|
|
828
|
-
Returns [{path, line, score, snippet}] — pointers, not bodies. Follow up by slicing the
|
|
829
|
-
named files out of `context` and delegating them to llm_query / map_files.
|
|
830
|
-
"""
|
|
831
|
-
terms = _tokenize(str(query))
|
|
832
|
-
if not terms:
|
|
833
|
-
return []
|
|
834
|
-
try:
|
|
835
|
-
limit = max(1, min(int(k), 100))
|
|
836
|
-
except (TypeError, ValueError):
|
|
837
|
-
limit = 10
|
|
838
|
-
return self._get_index().query(terms, limit, path_glob)
|
|
839
|
-
|
|
840
|
-
def _grep_context(
|
|
841
|
-
self,
|
|
842
|
-
pattern: str,
|
|
843
|
-
k: int = 50,
|
|
844
|
-
path_glob: str | None = None,
|
|
845
|
-
before: int = 0,
|
|
846
|
-
after: int = 0,
|
|
847
|
-
) -> dict[str, Any]:
|
|
848
|
-
"""Regex over `context`, capped and shaped.
|
|
849
|
-
|
|
850
|
-
Returns {"hits": [{path, line, text}], "counts": {path: n}, "total": n, "truncated": bool}.
|
|
851
|
-
`counts` is complete even when `hits` is capped, so a wide pattern reports its shape
|
|
852
|
-
instead of flooding stdout.
|
|
853
|
-
"""
|
|
854
|
-
try:
|
|
855
|
-
rx = re.compile(pattern)
|
|
856
|
-
except re.error as e:
|
|
857
|
-
return {"hits": [], "counts": {}, "total": 0, "truncated": False, "error": f"bad regex: {e}"}
|
|
858
|
-
try:
|
|
859
|
-
limit = max(1, min(int(k), _GREP_HARD_CAP))
|
|
860
|
-
except (TypeError, ValueError):
|
|
861
|
-
limit = 50
|
|
862
|
-
pad_before = max(0, min(int(before or 0), 10))
|
|
863
|
-
pad_after = max(0, min(int(after or 0), 10))
|
|
864
|
-
|
|
865
|
-
hits: list[dict[str, Any]] = []
|
|
866
|
-
counts: dict[str, int] = {}
|
|
867
|
-
total = 0
|
|
868
|
-
for path, content in self._entries():
|
|
869
|
-
if path_glob and not fnmatch.fnmatch(path, path_glob):
|
|
870
|
-
continue
|
|
871
|
-
if not rx.search(content):
|
|
872
|
-
continue
|
|
873
|
-
lines = content.split("\n")
|
|
874
|
-
for i, line in enumerate(lines):
|
|
875
|
-
if not rx.search(line):
|
|
876
|
-
continue
|
|
877
|
-
total += 1
|
|
878
|
-
counts[path] = counts.get(path, 0) + 1
|
|
879
|
-
if len(hits) >= limit:
|
|
880
|
-
continue
|
|
881
|
-
lo = max(0, i - pad_before)
|
|
882
|
-
hi = min(len(lines), i + pad_after + 1)
|
|
883
|
-
hits.append({"path": path, "line": i + 1, "text": "\n".join(lines[lo:hi])[:_SNIPPET_CHARS]})
|
|
884
|
-
return {"hits": hits, "counts": counts, "total": total, "truncated": total > len(hits)}
|
|
885
|
-
|
|
886
|
-
def _outline(self, path: str) -> str:
|
|
887
|
-
"""Definition/heading skeleton of one context file — orient in ~200 chars, not 20K.
|
|
888
|
-
|
|
889
|
-
`path` matches exactly, then by suffix, then as a glob.
|
|
890
|
-
"""
|
|
891
|
-
target = str(path)
|
|
892
|
-
entries = self._entries()
|
|
893
|
-
content: str | None = None
|
|
894
|
-
for p, c in entries:
|
|
895
|
-
if p == target:
|
|
896
|
-
content = c
|
|
897
|
-
break
|
|
898
|
-
if content is None:
|
|
899
|
-
for p, c in entries:
|
|
900
|
-
if p.endswith(target) or fnmatch.fnmatch(p, target):
|
|
901
|
-
content = c
|
|
902
|
-
target = p
|
|
903
|
-
break
|
|
904
|
-
if content is None:
|
|
905
|
-
return f"Error: no context file matching {path!r} — use search() or list paths from context"
|
|
906
|
-
out: list[str] = [f"# {target}"]
|
|
907
|
-
for i, line in enumerate(content.split("\n")):
|
|
908
|
-
if _OUTLINE_LINE.match(line):
|
|
909
|
-
out.append(f"{i + 1}: {line.strip()[:160]}")
|
|
910
|
-
if len(out) == 1:
|
|
911
|
-
return f"# {target}\n(no definition-like lines found)"
|
|
912
|
-
return "\n".join(out)
|
|
913
|
-
|
|
914
|
-
# ---- one-line delegation (structural: orchestrating must be easier than solving) -------
|
|
915
|
-
|
|
916
|
-
def _start_map_files(self, files: Any, prompt: str, model: str | None = None) -> Task:
|
|
917
|
-
"""Post every batch map_files needs, WITHOUT waiting. Contract: see _map_files.
|
|
918
|
-
|
|
919
|
-
All batches go on the wire together, so a 100-file map costs one round-trip of latency
|
|
920
|
-
rather than one per 20 files.
|
|
921
|
-
"""
|
|
922
|
-
prompt = str(prompt)
|
|
923
|
-
by_path: list[tuple[str, str]] = []
|
|
924
|
-
lookup: dict[str, str] | None = None
|
|
925
|
-
for item in files if isinstance(files, (list, tuple)) else [files]:
|
|
926
|
-
if isinstance(item, dict):
|
|
927
|
-
content = item.get("content", "")
|
|
928
|
-
by_path.append((str(item.get("path", "?")), content if isinstance(content, str) else str(content)))
|
|
929
|
-
elif isinstance(item, str):
|
|
930
|
-
if lookup is None:
|
|
931
|
-
lookup = {p: c for p, c in self._entries()}
|
|
932
|
-
if item in lookup:
|
|
933
|
-
by_path.append((item, lookup[item]))
|
|
934
|
-
else:
|
|
935
|
-
by_path.append((item, ""))
|
|
936
|
-
if not by_path:
|
|
937
|
-
return Task.resolved(self, "map_files", {})
|
|
938
|
-
|
|
939
|
-
# Per-file prompt budget; anything larger is chunked and its answers concatenated.
|
|
940
|
-
budget = self.max_prompt_chars - len(prompt) - _CHUNK_HEADER_OVERHEAD - 256
|
|
941
|
-
if budget < 1_000:
|
|
942
|
-
return Task.resolved(self, "map_files", {
|
|
943
|
-
p: "Error: prompt too long to leave room for file content" for p, _ in by_path
|
|
944
|
-
})
|
|
945
|
-
|
|
946
|
-
requests: list[str] = []
|
|
947
|
-
spans: list[tuple[str, int]] = [] # (path, number of chunks contributed)
|
|
948
|
-
for path, content in by_path:
|
|
949
|
-
chunks = _chunk_text(content, budget) if len(content) > budget else [content]
|
|
950
|
-
spans.append((path, len(chunks)))
|
|
951
|
-
for j, chunk in enumerate(chunks):
|
|
952
|
-
header = f"[file {path}" + (f", part {j + 1}/{len(chunks)}]" if len(chunks) > 1 else "]")
|
|
953
|
-
requests.append(f"{prompt}\n\n{header}\n{chunk}")
|
|
954
|
-
|
|
955
|
-
rids: list[str] = []
|
|
956
|
-
sizes: list[int] = []
|
|
957
|
-
for i in range(0, len(requests), _MAX_CHUNK_BATCH):
|
|
958
|
-
batch = requests[i:i + _MAX_CHUNK_BATCH]
|
|
959
|
-
rids.append(self._post("llm_query_batched", {"prompts": batch, "model": model}))
|
|
960
|
-
sizes.append(len(batch))
|
|
961
|
-
return Task(self, "map_files", tuple(rids),
|
|
962
|
-
_reduce_map_files(sizes, spans), f"{len(by_path)} files")
|
|
963
|
-
|
|
964
|
-
@_spawnable("map_files")
|
|
965
|
-
def _map_files(self, files: Any, prompt: str, model: str | None = None) -> dict[str, str]:
|
|
966
|
-
"""Ask `prompt` of every given file, batched, and return {path: answer}.
|
|
967
|
-
|
|
968
|
-
`files` accepts context entries (dicts), paths (strings), or a mix — the whole
|
|
969
|
-
chunk/batch/collect loop the system prompt used to spell out, as one call.
|
|
970
|
-
Oversized files are split and their per-chunk answers joined.
|
|
971
|
-
"""
|
|
972
|
-
return self._await_task(self._start_map_files(files, prompt, model))
|
|
973
|
-
|
|
974
|
-
def _llm_map_reduce(
|
|
975
|
-
self,
|
|
976
|
-
items: Any,
|
|
977
|
-
map_prompt: str,
|
|
978
|
-
reduce_prompt: str,
|
|
979
|
-
model: str | None = None,
|
|
980
|
-
) -> str:
|
|
981
|
-
"""Map `map_prompt` over `items` in one batch, then reduce the answers with one call.
|
|
982
|
-
|
|
983
|
-
The paper's canonical strategy ("query an LLM per chunk ... then query an LLM with all
|
|
984
|
-
the buffers") as a single call, so the root never hand-rolls the loop.
|
|
985
|
-
"""
|
|
986
|
-
map_prompt, reduce_prompt = str(map_prompt), str(reduce_prompt)
|
|
987
|
-
seq = list(items) if isinstance(items, (list, tuple)) else [items]
|
|
988
|
-
if not seq:
|
|
989
|
-
return "Error: llm_map_reduce got no items"
|
|
990
|
-
texts = [
|
|
991
|
-
(str(it.get("content", "")) if isinstance(it, dict) else str(it))
|
|
992
|
-
for it in seq
|
|
993
|
-
]
|
|
994
|
-
labels = [
|
|
995
|
-
(str(it.get("path", f"item {i + 1}")) if isinstance(it, dict) else f"item {i + 1}")
|
|
996
|
-
for i, it in enumerate(seq)
|
|
997
|
-
]
|
|
998
|
-
mapped: list[str] = []
|
|
999
|
-
for i in range(0, len(texts), _MAX_CHUNK_BATCH):
|
|
1000
|
-
batch = [
|
|
1001
|
-
f"{map_prompt}\n\n[{labels[i + j]}]\n{t}"
|
|
1002
|
-
for j, t in enumerate(texts[i:i + _MAX_CHUNK_BATCH])
|
|
1003
|
-
]
|
|
1004
|
-
mapped.extend(self._llm_query_batched(batch, model))
|
|
1005
|
-
joined = "\n\n".join(f"[{labels[i]}]\n{a}" for i, a in enumerate(mapped))
|
|
1006
|
-
return self._llm_query(f"{reduce_prompt}\n\nPartial answers:\n{joined}", model)
|
|
1007
|
-
|
|
1008
|
-
# ---- sync helpers: await(start(...)), so there is exactly one code path -----------------
|
|
1009
|
-
|
|
1010
|
-
@_spawnable("llm_query")
|
|
1011
|
-
def _llm_query(self, prompt: str, model: str | None = None) -> str:
|
|
1012
|
-
return self._await_task(self._start_llm_query(prompt, model))
|
|
1013
|
-
|
|
1014
|
-
@_spawnable("llm_query_batched")
|
|
1015
|
-
def _llm_query_batched(self, prompts, model: str | None = None) -> list[str]:
|
|
1016
|
-
return self._await_task(self._start_llm_query_batched(prompts, model))
|
|
1017
|
-
|
|
1018
|
-
@_spawnable("llm_query_chunked")
|
|
1019
|
-
def _llm_query_chunked(self, text, prompt: str, model: str | None = None) -> list[str]:
|
|
1020
|
-
return self._await_task(self._start_llm_query_chunked(text, prompt, model))
|
|
1021
|
-
|
|
1022
|
-
@_spawnable("rlm_query")
|
|
1023
|
-
def _rlm_query(self, prompt: str, model: str | None = None) -> str:
|
|
1024
|
-
return self._await_task(self._start_rlm_query(prompt, model))
|
|
1025
|
-
|
|
1026
|
-
def _ask_user_question(self, questions: list[dict]) -> list[dict]:
|
|
1027
|
-
"""Present structured questions to the user; blocks until answered.
|
|
1028
|
-
|
|
1029
|
-
Returns a list of {question, selected, custom?} dicts.
|
|
1030
|
-
Each dict has: question (str), selected (list[str]), custom (str|None).
|
|
1031
|
-
Only valid at root depth; sub-RLM calls return an error answer.
|
|
1032
|
-
"""
|
|
1033
|
-
if self.depth > 0:
|
|
1034
|
-
qlist = questions if isinstance(questions, list) else []
|
|
1035
|
-
return [
|
|
1036
|
-
{"question": str(q.get("question", "")) if isinstance(q, dict) else "",
|
|
1037
|
-
"selected": [],
|
|
1038
|
-
"custom": "Error: ask_user_question not available inside rlm_query sub-calls"}
|
|
1039
|
-
for q in qlist
|
|
1040
|
-
] or [{"question": "", "selected": [],
|
|
1041
|
-
"custom": "Error: ask_user_question not available inside rlm_query sub-calls"}]
|
|
1042
|
-
if not isinstance(questions, list) or not questions:
|
|
1043
|
-
return [{"question": "", "selected": [], "custom": "Error: questions must be a non-empty list"}]
|
|
1044
|
-
cleaned = []
|
|
1045
|
-
for q in questions:
|
|
1046
|
-
if not isinstance(q, dict) or "question" not in q or "options" not in q:
|
|
1047
|
-
return [{"question": "", "selected": [], "custom": "Error: each question needs 'question', 'header', 'options'"}]
|
|
1048
|
-
opts = q.get("options")
|
|
1049
|
-
if not isinstance(opts, list):
|
|
1050
|
-
return [{"question": str(q.get("question", "")), "selected": [], "custom": "Error: options must be a list"}]
|
|
1051
|
-
cleaned_opts = []
|
|
1052
|
-
for o in opts:
|
|
1053
|
-
if not isinstance(o, dict) or "label" not in o:
|
|
1054
|
-
return [{"question": str(q.get("question", "")), "selected": [], "custom": "Error: each option needs 'label'"}]
|
|
1055
|
-
item = {"label": str(o["label"]), "description": str(o.get("description", ""))}
|
|
1056
|
-
if "preview" in o:
|
|
1057
|
-
item["preview"] = str(o["preview"])
|
|
1058
|
-
cleaned_opts.append(item)
|
|
1059
|
-
cleaned.append({
|
|
1060
|
-
"question": str(q["question"]),
|
|
1061
|
-
"header": str(q.get("header", "Q")),
|
|
1062
|
-
"multiSelect": bool(q.get("multiSelect", False)),
|
|
1063
|
-
"options": cleaned_opts,
|
|
1064
|
-
})
|
|
1065
|
-
r = self._rpc("ask_user_question", {"questions": cleaned})
|
|
1066
|
-
if r.get("error"):
|
|
1067
|
-
return [{"question": q["question"], "selected": [], "custom": f"Error: {r['error']}"} for q in cleaned]
|
|
1068
|
-
answers = r.get("answers", [])
|
|
1069
|
-
return answers if isinstance(answers, list) else []
|
|
1070
|
-
|
|
1071
|
-
def _todo(self, action: str, **kwargs) -> str:
|
|
1072
|
-
"""Manage the run's task list.
|
|
1073
|
-
|
|
1074
|
-
action: "create" | "update" | "list" | "get" | "delete" | "clear"
|
|
1075
|
-
kwargs: id, subject, description, status, activeForm, blockedBy, owner, filterStatus
|
|
1076
|
-
Returns a human-readable string result.
|
|
1077
|
-
"""
|
|
1078
|
-
params = {k: v for k, v in kwargs.items() if v is not None}
|
|
1079
|
-
r = self._rpc("todo", {"action": str(action), **params})
|
|
1080
|
-
if r.get("error"):
|
|
1081
|
-
return f"Error: {r['error']}"
|
|
1082
|
-
return str(r.get("response", "ok"))
|
|
1083
|
-
|
|
1084
|
-
def _load_library(self, source: str) -> dict[str, Any] | str:
|
|
1085
|
-
"""Pack an external dir/file/git-URL on the host and append it into `context`.
|
|
1086
|
-
|
|
1087
|
-
Paths are namespaced under lib/<source_id>/ (host). Content is always in the
|
|
1088
|
-
single `context` list — never a new context_N variable.
|
|
1089
|
-
Host-side idempotency may return already_loaded without a payload path.
|
|
1090
|
-
"""
|
|
1091
|
-
r = self._rpc("load_library", {"source": str(source)})
|
|
1092
|
-
if r.get("error"):
|
|
1093
|
-
return f"Error: {r['error']}"
|
|
1094
|
-
if r.get("already_loaded"):
|
|
1095
|
-
source_id = r.get("source_id") if isinstance(r.get("source_id"), str) else "lib"
|
|
1096
|
-
path_prefix = r.get("path_prefix") if isinstance(r.get("path_prefix"), str) else f"lib/{source_id}/"
|
|
1097
|
-
ctx = self.ns.get("context")
|
|
1098
|
-
ctx_len = len(ctx) if isinstance(ctx, list) else 0
|
|
1099
|
-
print(
|
|
1100
|
-
f"[rlm] load_library: already loaded {source_id} "
|
|
1101
|
-
f"(paths under {path_prefix}, context len={ctx_len})"
|
|
1102
|
-
)
|
|
1103
|
-
return {
|
|
1104
|
-
"source": str(source),
|
|
1105
|
-
"source_id": source_id,
|
|
1106
|
-
"path_prefix": path_prefix,
|
|
1107
|
-
"files": 0,
|
|
1108
|
-
"chars": r.get("chars"),
|
|
1109
|
-
"context_len": ctx_len,
|
|
1110
|
-
"already_loaded": True,
|
|
1111
|
-
}
|
|
1112
|
-
path = r.get("path")
|
|
1113
|
-
if not isinstance(path, str):
|
|
1114
|
-
return "Error: malformed load_library reply (no path)"
|
|
1115
|
-
try:
|
|
1116
|
-
# Worker-internal read — use real io.open so read-only guards never block us.
|
|
1117
|
-
with _REAL_IO_OPEN(path, "r") as f:
|
|
1118
|
-
payload = json.load(f) if r.get("json") else f.read()
|
|
1119
|
-
finally:
|
|
1120
|
-
try:
|
|
1121
|
-
os.remove(path) # worker owns temp-file cleanup (host does NOT unlink)
|
|
1122
|
-
except OSError:
|
|
1123
|
-
pass
|
|
1124
|
-
return self._append_library(str(source), payload, r)
|
|
1125
|
-
|
|
1126
|
-
def _append_library(self, source: str, payload: Any, meta: dict[str, Any]) -> dict[str, Any] | str:
|
|
1127
|
-
"""Append host-packed library files into `context` (idempotent by path prefix)."""
|
|
1128
|
-
ctx = self.ns.get("context")
|
|
1129
|
-
if not isinstance(ctx, list):
|
|
1130
|
-
kind = type(ctx).__name__ if ctx is not None else "None"
|
|
1131
|
-
return f"Error: load_library requires list context (file bundle); got {kind}"
|
|
1132
|
-
|
|
1133
|
-
source_id = meta.get("source_id")
|
|
1134
|
-
if not isinstance(source_id, str) or not source_id:
|
|
1135
|
-
source_id = "lib"
|
|
1136
|
-
path_prefix = meta.get("path_prefix")
|
|
1137
|
-
if not isinstance(path_prefix, str) or not path_prefix:
|
|
1138
|
-
path_prefix = f"lib/{source_id}/"
|
|
1139
|
-
|
|
1140
|
-
# Idempotent: already present if any path uses this library prefix.
|
|
1141
|
-
for item in ctx:
|
|
1142
|
-
if isinstance(item, dict) and str(item.get("path", "")).startswith(path_prefix):
|
|
1143
|
-
print(
|
|
1144
|
-
f"[rlm] load_library: already loaded {source_id} "
|
|
1145
|
-
f"(paths under {path_prefix}, context len={len(ctx)})"
|
|
1146
|
-
)
|
|
1147
|
-
return {
|
|
1148
|
-
"source": source,
|
|
1149
|
-
"source_id": source_id,
|
|
1150
|
-
"path_prefix": path_prefix,
|
|
1151
|
-
"files": 0,
|
|
1152
|
-
"chars": meta.get("chars"),
|
|
1153
|
-
"context_len": len(ctx),
|
|
1154
|
-
"already_loaded": True,
|
|
1155
|
-
}
|
|
1156
|
-
|
|
1157
|
-
files = self._library_file_entries(payload, path_prefix)
|
|
1158
|
-
if not files:
|
|
1159
|
-
return "Error: load_library produced no files"
|
|
1160
|
-
|
|
1161
|
-
ctx.extend(files)
|
|
1162
|
-
# Keep restore payload in sync with the live list.
|
|
1163
|
-
self._context_payload = ctx
|
|
1164
|
-
self.ns["context"] = ctx
|
|
1165
|
-
|
|
1166
|
-
print(
|
|
1167
|
-
f"[rlm] load_library: +{len(files)} files into context "
|
|
1168
|
-
f"(len={len(ctx)}); paths under {path_prefix}"
|
|
1169
|
-
)
|
|
1170
|
-
return {
|
|
1171
|
-
"source": source,
|
|
1172
|
-
"source_id": source_id,
|
|
1173
|
-
"path_prefix": path_prefix,
|
|
1174
|
-
"files": len(files),
|
|
1175
|
-
"chars": meta.get("chars"),
|
|
1176
|
-
"context_len": len(ctx),
|
|
1177
|
-
"already_loaded": False,
|
|
1178
|
-
}
|
|
1179
|
-
|
|
1180
|
-
@staticmethod
|
|
1181
|
-
def _library_file_entries(payload: Any, path_prefix: str) -> list[dict[str, Any]]:
|
|
1182
|
-
"""Normalize host payload to list[dict]. Host already namespaces; string is fallback."""
|
|
1183
|
-
if isinstance(payload, str):
|
|
1184
|
-
return [{
|
|
1185
|
-
"path": f"{path_prefix}content",
|
|
1186
|
-
"content": payload,
|
|
1187
|
-
"tokens": max(1, (len(payload) + 3) // 4),
|
|
1188
|
-
}]
|
|
1189
|
-
if not isinstance(payload, list):
|
|
1190
|
-
return []
|
|
1191
|
-
out: list[dict[str, Any]] = []
|
|
1192
|
-
for item in payload:
|
|
1193
|
-
if isinstance(item, dict) and "path" in item and "content" in item:
|
|
1194
|
-
out.append(item)
|
|
1195
|
-
return out
|
|
1196
|
-
|
|
1197
|
-
def _advance_phase(self, phase: str, summary: str | None = None) -> str:
|
|
1198
|
-
"""Transition the root RLM pipeline to a new phase.
|
|
1199
|
-
|
|
1200
|
-
Only callable at depth 0. The parent handler validates the transition
|
|
1201
|
-
against the phase state machine (research → blueprint → validate)
|
|
1202
|
-
and runs deterministic artifact gates before accepting the transition.
|
|
1203
|
-
Returns a short confirmation, or an `Error: …` string the model can act on.
|
|
1204
|
-
"""
|
|
1205
|
-
if self.depth > 0:
|
|
1206
|
-
return "Error: advance_phase is only available at the root RLM depth"
|
|
1207
|
-
r = self._rpc("advance_phase", {"phase": str(phase), "summary": summary})
|
|
1208
|
-
if r.get("error"):
|
|
1209
|
-
return f"Error: {r['error']}"
|
|
1210
|
-
response = r.get("response", "ok")
|
|
1211
|
-
if isinstance(response, str) and response.startswith("Error:"):
|
|
1212
|
-
return response
|
|
1213
|
-
return response if isinstance(response, str) else "ok"
|
|
1214
|
-
|
|
1215
|
-
def _save_artifact(self, kind: str, content: str) -> str:
|
|
1216
|
-
"""Persist a stage artifact (research/plan/validation) under .rlm/artifacts/.
|
|
1217
|
-
|
|
1218
|
-
Only callable at depth 0. The engine gates advance_phase against the latest
|
|
1219
|
-
saved artifact for the current stage.
|
|
1220
|
-
"""
|
|
1221
|
-
if self.depth > 0:
|
|
1222
|
-
return "Error: save_artifact is only available at the root RLM depth"
|
|
1223
|
-
r = self._rpc("save_artifact", {"artifactKind": str(kind), "content": str(content)})
|
|
1224
|
-
if r.get("error"):
|
|
1225
|
-
return f"Error: {r['error']}"
|
|
1226
|
-
response = r.get("response", "ok")
|
|
1227
|
-
if isinstance(response, str) and response.startswith("Error:"):
|
|
1228
|
-
return response
|
|
1229
|
-
return response if isinstance(response, str) else "ok"
|
|
1230
|
-
|
|
1231
|
-
@_spawnable("rlm_query_batched")
|
|
1232
|
-
def _rlm_query_batched(self, prompts, model: str | None = None) -> list[str]:
|
|
1233
|
-
return self._await_task(self._start_rlm_query_batched(prompts, model))
|
|
1234
|
-
|
|
1235
|
-
# ---- context + execution --------------------------------------------------------------
|
|
1236
|
-
|
|
1237
|
-
def load_context(self, path: str, index: int | None = None, is_json: bool = False) -> int:
|
|
1238
|
-
"""Load the packed world into the single REPL variable `context`.
|
|
1239
|
-
|
|
1240
|
-
`index` is accepted for protocol compatibility but ignored — there is only
|
|
1241
|
-
one context slot. Libraries are merged on the host (or via load_library).
|
|
1242
|
-
"""
|
|
1243
|
-
with open(path, "r") as f:
|
|
1244
|
-
payload = json.load(f) if is_json else f.read()
|
|
1245
|
-
self._context_payload = payload
|
|
1246
|
-
self.ns["context"] = payload
|
|
1247
|
-
# Drop legacy multi-slot names if present.
|
|
1248
|
-
for k in list(self.ns.keys()):
|
|
1249
|
-
if k != "context" and _CONTEXT_NAME.match(k):
|
|
1250
|
-
del self.ns[k]
|
|
1251
|
-
return 0
|
|
1252
|
-
|
|
1253
|
-
@contextmanager
|
|
1254
|
-
def _capture(self):
|
|
1255
|
-
out, err = io.StringIO(), io.StringIO()
|
|
1256
|
-
old_out, old_err = sys.stdout, sys.stderr
|
|
1257
|
-
sys.stdout, sys.stderr = out, err
|
|
1258
|
-
try:
|
|
1259
|
-
yield out, err
|
|
1260
|
-
finally:
|
|
1261
|
-
sys.stdout, sys.stderr = old_out, old_err
|
|
1262
|
-
|
|
1263
|
-
def _exec(self, code: str, ns: dict[str, Any]) -> None:
|
|
1264
|
-
t = self.exec_timeout_s
|
|
1265
|
-
if t <= 0 or not hasattr(signal, "SIGALRM"):
|
|
1266
|
-
exec(compile(code, "<repl>", "exec"), ns, ns) # noqa: S102
|
|
1267
|
-
return
|
|
1268
|
-
|
|
1269
|
-
def _alarm(signum, frame): # noqa: ARG001
|
|
1270
|
-
raise TimeoutError(f"```repl``` block exceeded {t:g}s timeout")
|
|
1271
|
-
|
|
1272
|
-
old = signal.signal(signal.SIGALRM, _alarm)
|
|
1273
|
-
signal.setitimer(signal.ITIMER_REAL, t)
|
|
1274
|
-
try:
|
|
1275
|
-
exec(compile(code, "<repl>", "exec"), ns, ns) # noqa: S102
|
|
1276
|
-
finally:
|
|
1277
|
-
signal.setitimer(signal.ITIMER_REAL, 0)
|
|
1278
|
-
signal.signal(signal.SIGALRM, old)
|
|
1279
|
-
|
|
1280
|
-
def _nudge_lines(self) -> list[str]:
|
|
1281
|
-
"""One-time hint for newly created huge raw-text variables (single line).
|
|
1282
|
-
|
|
1283
|
-
Collapses to one line so it survives headless stdout elision (head 200 + tail 200).
|
|
1284
|
-
"""
|
|
1285
|
-
names: list[str] = []
|
|
1286
|
-
for k in self._user_var_names():
|
|
1287
|
-
v = self.ns.get(k)
|
|
1288
|
-
if isinstance(v, (str, bytes)) and len(v) > _NUDGE_CHARS and k not in self._nudged:
|
|
1289
|
-
self._nudged.add(k)
|
|
1290
|
-
names.append(f"{k} ({len(v):,} chars)")
|
|
1291
|
-
if not names:
|
|
1292
|
-
return []
|
|
1293
|
-
return [
|
|
1294
|
-
f"[rlm] huge raw-text variable(s): {', '.join(names)} — do NOT analyze them yourself; "
|
|
1295
|
-
'delegate with llm_query_chunked(name, "your question") or slice + llm_query_batched.'
|
|
1296
|
-
]
|
|
1297
|
-
|
|
1298
|
-
def execute(self, code: str) -> dict[str, Any]:
|
|
1299
|
-
start = time.perf_counter()
|
|
1300
|
-
raised = False
|
|
1301
|
-
with self._capture() as (out, err):
|
|
1302
|
-
try:
|
|
1303
|
-
self._restore_scaffold()
|
|
1304
|
-
self._exec(code, self.ns)
|
|
1305
|
-
self._restore_scaffold()
|
|
1306
|
-
stdout, stderr = out.getvalue(), err.getvalue()
|
|
1307
|
-
except BaseException as e: # noqa: BLE001
|
|
1308
|
-
raised = True
|
|
1309
|
-
self._restore_scaffold()
|
|
1310
|
-
stdout = out.getvalue()
|
|
1311
|
-
stderr = err.getvalue() + f"\n{type(e).__name__}: {e}\n" + traceback.format_exc()
|
|
1312
|
-
final, self._final_answer = self._final_answer, None
|
|
1313
|
-
answer = self.ns.get("answer")
|
|
1314
|
-
answer_content = answer.get("content", "") if isinstance(answer, dict) else ""
|
|
1315
|
-
# ready may have been flipped with empty content before content was assigned later
|
|
1316
|
-
# in the same block; the dict's current content is the real submission.
|
|
1317
|
-
if final is not None and not final.strip() and str(answer_content).strip():
|
|
1318
|
-
final = str(answer_content)
|
|
1319
|
-
nudges = self._nudge_lines()
|
|
1320
|
-
if nudges:
|
|
1321
|
-
parts = [stdout] if stdout else []
|
|
1322
|
-
parts.extend(nudges)
|
|
1323
|
-
stdout = "\n".join(parts) + "\n"
|
|
1324
|
-
return {
|
|
1325
|
-
"stdout": stdout,
|
|
1326
|
-
"stderr": stderr,
|
|
1327
|
-
"final_answer": final,
|
|
1328
|
-
"answer_content": str(answer_content),
|
|
1329
|
-
"raised": raised,
|
|
1330
|
-
"execution_time": time.perf_counter() - start,
|
|
1331
|
-
"var_names": self._user_var_names(),
|
|
1332
|
-
}
|
|
1333
|
-
|
|
1334
|
-
def _serializer(self):
|
|
1335
|
-
try:
|
|
1336
|
-
import dill as s
|
|
1337
|
-
return s
|
|
1338
|
-
except ImportError:
|
|
1339
|
-
return pickle
|
|
1340
|
-
|
|
1341
|
-
def snapshot(self, path: str, nonce: str) -> dict:
|
|
1342
|
-
"""Pickle user variables atomically to path. Stores session nonce for restore verification.
|
|
1343
|
-
|
|
1344
|
-
Writes to path.tmp then os.rename — atomic on POSIX, so no .tmp leak and no
|
|
1345
|
-
TypeScript-side finalize step needed. On resume (fresh session = different nonce),
|
|
1346
|
-
restore fails — caller falls back to history-only replay.
|
|
1347
|
-
"""
|
|
1348
|
-
s = self._serializer()
|
|
1349
|
-
out, skipped = {}, []
|
|
1350
|
-
MAX_VAR_BYTES = 50 * 1024 * 1024
|
|
1351
|
-
for k, v in self.ns.items():
|
|
1352
|
-
if k.startswith("_") or _CONTEXT_NAME.match(k) or k in RESERVED or k == "__builtins__":
|
|
1353
|
-
continue
|
|
1354
|
-
# A Task holds a back-reference to this Worker, so dill would happily pickle the
|
|
1355
|
-
# whole process. Top-level guard only: a Task nested inside a list/dict still
|
|
1356
|
-
# falls to the generic `except` below and skips the variable — which is why this
|
|
1357
|
-
# guard is explicit rather than left to that fallback.
|
|
1358
|
-
if isinstance(v, Task):
|
|
1359
|
-
skipped.append(k)
|
|
1360
|
-
continue
|
|
1361
|
-
try:
|
|
1362
|
-
blob = s.dumps(v)
|
|
1363
|
-
if len(blob) > MAX_VAR_BYTES:
|
|
1364
|
-
skipped.append(k)
|
|
1365
|
-
continue
|
|
1366
|
-
out[k] = v
|
|
1367
|
-
except Exception:
|
|
1368
|
-
skipped.append(k)
|
|
1369
|
-
if skipped:
|
|
1370
|
-
print(f"[rlm-sandbox] snapshot skipped {len(skipped)} unpicklable/oversized vars: {skipped}", file=_REAL_STDERR)
|
|
1371
|
-
tmp = path + ".tmp"
|
|
1372
|
-
with _REAL_IO_OPEN(tmp, "wb") as f:
|
|
1373
|
-
s.dump({"nonce": nonce, "vars": out}, f)
|
|
1374
|
-
os.rename(tmp, path) # atomic rename
|
|
1375
|
-
return {"skipped": skipped}
|
|
1376
|
-
|
|
1377
|
-
def restore(self, path: str, nonce: str) -> dict:
|
|
1378
|
-
"""Restore user variables from a pickle file. Verifies session nonce before deserializing.
|
|
1379
|
-
|
|
1380
|
-
SECURITY: pickle.load executes arbitrary code. The session nonce check ensures the
|
|
1381
|
-
.pkl was written by THIS engine session. Cross-session resume falls back to
|
|
1382
|
-
history-only replay (caller skips restore when sessionNonce is undefined).
|
|
1383
|
-
"""
|
|
1384
|
-
s = self._serializer()
|
|
1385
|
-
with _REAL_IO_OPEN(path, "rb") as f:
|
|
1386
|
-
data = s.load(f)
|
|
1387
|
-
if not isinstance(data, dict) or data.get("nonce") != nonce:
|
|
1388
|
-
raise ValueError("snapshot nonce mismatch — not from this session")
|
|
1389
|
-
self.ns.update(data.get("vars", {}))
|
|
1390
|
-
self._restore_scaffold()
|
|
1391
|
-
return {"restored": list(data.get("vars", {}).keys())}
|
|
1392
|
-
|
|
1393
|
-
|
|
1394
|
-
def main() -> None:
|
|
1395
|
-
ap = argparse.ArgumentParser()
|
|
1396
|
-
ap.add_argument("--depth", type=int, default=int(os.environ.get("RLM_DEPTH", "1")))
|
|
1397
|
-
ap.add_argument("--timeout", type=float, default=float(os.environ.get("RLM_EXEC_TIMEOUT_S", "600")))
|
|
1398
|
-
ap.add_argument("--await-timeout", type=float,
|
|
1399
|
-
default=float(os.environ.get("RLM_AWAIT_TIMEOUT_S", "600")))
|
|
1400
|
-
ap.add_argument("--max-prompt-chars", type=int,
|
|
1401
|
-
default=int(os.environ.get("RLM_MAX_PROMPT_CHARS", "400000")))
|
|
1402
|
-
ap.add_argument("--read-only", action="store_true",
|
|
1403
|
-
default=os.environ.get("RLM_READ_ONLY", "").lower() in ("1", "true", "yes"),
|
|
1404
|
-
help="Reject open() write modes (pipeline runs)")
|
|
1405
|
-
args = ap.parse_args()
|
|
1406
|
-
|
|
1407
|
-
worker = Worker(depth=args.depth, exec_timeout_s=args.timeout,
|
|
1408
|
-
max_prompt_chars=args.max_prompt_chars, read_only=args.read_only,
|
|
1409
|
-
await_timeout_s=args.await_timeout)
|
|
1410
|
-
_send({"id": "_init", "ok": True})
|
|
1411
|
-
|
|
1412
|
-
while True:
|
|
1413
|
-
# Requests that arrived mid-exec were parked by _pump; replay them before reading.
|
|
1414
|
-
req = worker.take_deferred()
|
|
1415
|
-
if req is None:
|
|
1416
|
-
line = _REAL_STDIN.readline()
|
|
1417
|
-
if not line:
|
|
1418
|
-
return
|
|
1419
|
-
raw = line.strip()
|
|
1420
|
-
if not raw:
|
|
1421
|
-
continue
|
|
1422
|
-
try:
|
|
1423
|
-
req = json.loads(raw)
|
|
1424
|
-
except json.JSONDecodeError as e:
|
|
1425
|
-
_send({"id": "?", "ok": False, "error": f"bad json: {e}"})
|
|
1426
|
-
continue
|
|
1427
|
-
# A task spawned in an earlier exec settling while the worker is idle. Park it for
|
|
1428
|
-
# a later rlm_await; without this it would fall through to "unknown type" and the
|
|
1429
|
-
# result would be lost.
|
|
1430
|
-
if worker.park_reply(req):
|
|
1431
|
-
continue
|
|
1432
|
-
if not isinstance(req, dict):
|
|
1433
|
-
_send({"id": "?", "ok": False, "error": f"expected an object, got {type(req).__name__}"})
|
|
1434
|
-
continue
|
|
1435
|
-
rid, kind = req.get("id", "?"), req.get("type")
|
|
1436
|
-
try:
|
|
1437
|
-
if kind == "exec":
|
|
1438
|
-
_send({"id": rid, "ok": True, **worker.execute(req.get("code", ""))})
|
|
1439
|
-
elif kind == "load_context":
|
|
1440
|
-
idx = worker.load_context(req.get("path"), req.get("index"), req.get("json"))
|
|
1441
|
-
_send({"id": rid, "ok": True, "index": idx})
|
|
1442
|
-
elif kind == "shutdown":
|
|
1443
|
-
_send({"id": rid, "ok": True})
|
|
1444
|
-
return
|
|
1445
|
-
elif kind == "snapshot":
|
|
1446
|
-
_send({"id": rid, "ok": True, **worker.snapshot(req.get("path", ""), req.get("nonce", ""))})
|
|
1447
|
-
elif kind == "restore":
|
|
1448
|
-
_send({"id": rid, "ok": True, **worker.restore(req.get("path", ""), req.get("nonce", ""))})
|
|
1449
|
-
else:
|
|
1450
|
-
_send({"id": rid, "ok": False, "error": f"unknown type: {kind!r}"})
|
|
1451
|
-
except BaseException as e: # noqa: BLE001
|
|
1452
|
-
_send({"id": rid, "ok": False, "error": f"{type(e).__name__}: {e}\n{traceback.format_exc()}"})
|
|
1453
|
-
|
|
1454
|
-
|
|
1455
|
-
if __name__ == "__main__":
|
|
1456
|
-
main()
|