rlmness-cli 0.2.0__py3-none-any.whl

rlmagent/agents/budget.py

@@ -0,0 +1,212 @@
+"""Every model call in the tree passes through one shared allowance."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import aclosing
+from dataclasses import dataclass
+
+from rlmness import Allowance, AllowanceSpent, Spend
+
+from rlmagent.harness.contracts.transcript import ModelEntry, TextSegment
+from rlmagent.harness.contracts.transcript.diagnostics import UsageStats
+from rlmagent.harness.provider.wire import StreamCloseEvent, StreamFaultEvent
+from rlmagent.model.token_bounds import TokenBoundError, TokenReservation
+
+
+def to_spend(usage: UsageStats) -> Spend:
+    cost = usage.cost.total
+    return Spend(
+        prompt_tokens=usage.input,
+        completion_tokens=usage.output,
+        total_tokens=usage.total_tokens or usage.input + usage.output,
+        # A provider that reports no price is not free; say "unknown".
+        cost=cost if cost > 0 else None,
+        cached_tokens=usage.cache_read or None,
+        reasoning_tokens=usage.reasoning,
+    )
+
+
+REFUSED = "budget reached"
+
+#: A heuristic for advisory sizing. Token density varies by model and input.
+CHARS_PER_TOKEN = 3
+#: Reply tokens assumed when the provider's own reply limit cannot be read.
+DEFAULT_REPLY_TOKENS = 8192
+
+
+def _text_of(value) -> str:
+    if isinstance(value, str):
+        return value
+    text = getattr(value, "text", None)
+    if isinstance(text, str):
+        return text
+    content = getattr(value, "content", None)
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        return " ".join(_text_of(part) for part in content)
+    return str(value) if value is not None else ""
+
+
+def estimate_prompt_tokens(system, messages) -> int:
+    """Estimate the prompt for advisory sizing, treating it as uncached."""
+    chars = len(_text_of(system)) + sum(len(_text_of(m)) for m in messages or ())
+    return chars // CHARS_PER_TOKEN + 1
+
+
+def reply_limit(provider) -> int:
+    """The most tokens the provider will write in one reply, as far as it says."""
+    seen = set()
+    while provider is not None and id(provider) not in seen:
+        seen.add(id(provider))
+        profile = getattr(provider, "_profile", None)
+        limit = getattr(profile, "max_tokens", None)
+        if isinstance(limit, int) and limit > 0:
+            return limit
+        provider = getattr(provider, "inner", None) or getattr(provider, "_inner", None)
+    return DEFAULT_REPLY_TOKENS
+
+
+class NewTokenBudget:
+    """Reserve the unmeasured balance for one provider request at a time."""
+
+    def __init__(self, limit: int, *, strict: bool = False) -> None:
+        self.limit = limit
+        self.strict = strict
+        self.used = 0
+        self.unmeasured = False
+        self.lock = asyncio.Lock()
+
+    def refusal(self, upcoming: int = 0) -> str | None:
+        """Why the next request may not go out; `upcoming` is its estimated size."""
+        if self.unmeasured:
+            return "new tokens could not be measured for the preceding request"
+        if self.used >= self.limit:
+            return f"{self.used} new tokens reaches the limit of {self.limit}"
+        if upcoming and self.used + upcoming > self.limit:
+            return (f"the next request could use about {upcoming} new tokens, and "
+                    f"{self.used} of {self.limit} are already used")
+        return None
+
+    def settle(self, usage: UsageStats) -> None:
+        tokens = max(usage.input + usage.cache_write + usage.output,
+                     usage.total_tokens - usage.cache_read)
+        if not (tokens or usage.cache_read or usage.total_tokens):
+            self.unmeasured = True
+        self.used += tokens
+
+
+@dataclass
+class ContractAllowance(Allowance):
+    new_token_budget: NewTokenBudget | None = None
+
+
+class BudgetedProvider:
+    def __init__(self, inner, allowance: Allowance) -> None:
+        self.inner = inner
+        self.allowance = allowance
+        self.refusal: str | None = None
+
+    def __getattr__(self, name):
+        return getattr(self.inner, name)
+
+    def stream_response(self, *, model, **kwargs) -> AsyncIterator:
+        return self.model_stream(self.inner, model=model, reserve=True, **kwargs)
+
+    def model_stream(self, provider, *, model, reserve=False, **kwargs) -> AsyncIterator:
+        budget = getattr(self.allowance, "new_token_budget", None)
+
+        def request():
+            return provider.stream_response(model=model, **kwargs)
+
+        def start(factory=request):
+            return self._reserved_stream(model, factory) if reserve else factory()
+
+        if budget is None:
+            return start()
+        return self.measured_stream(provider, model, kwargs, start)
+
+    def request_size(self, system, messages) -> int:
+        """Estimated new tokens of a request: its whole prompt plus the longest reply."""
+        return estimate_prompt_tokens(system, messages) + reply_limit(self.inner)
+
+    def _refuse(self, model, message, usage=None):
+        self.refusal = message
+        entry = ModelEntry(model=model, content=[TextSegment(text="")], stop_reason="error",
+                           error_message=f"{REFUSED}: {message}", usage=usage)
+        return StreamFaultEvent(reason="error", error=entry)
+
+    async def _token_plan(self, provider, model, kwargs, budget):
+        from rlmagent.model.external import is_external
+
+        prepare = getattr(provider, "prepare_token_request", None)
+        if is_external(model) or prepare is None:
+            raise TokenBoundError("this route has no verified input counter and provider-enforced reply cap")
+        plan = await prepare(model=model, remaining=budget.limit - budget.used, **kwargs)
+        if (not isinstance(plan, TokenReservation) or type(plan.input_tokens) is not int
+                or type(plan.output_tokens) is not int or plan.input_tokens < 0
+                or plan.output_tokens <= 0 or not callable(plan.start)
+                or plan.input_tokens + plan.output_tokens > budget.limit - budget.used):
+            raise TokenBoundError("provider did not establish a reservation within the remaining token budget")
+        return plan
+
+    async def measured_stream(self, provider, model, kwargs, start) -> AsyncIterator:
+        budget = self.allowance.new_token_budget
+        terminal = None
+        async with budget.lock:
+            plan = None
+            upcoming = 0 if budget.strict else self.request_size(kwargs.get("system"), kwargs.get("messages"))
+            reason = budget.refusal(upcoming)
+            if reason is None and budget.strict:
+                try:
+                    plan = await self._token_plan(provider, model, kwargs, budget)
+                except TokenBoundError as exc:
+                    reason = str(exc)
+            if reason is not None:
+                terminal = self._refuse(model, reason)
+            else:
+                settled = False
+                try:
+                    async with aclosing(start(plan.start) if plan else start()) as stream:
+                        async for event in stream:
+                            if isinstance(event, StreamCloseEvent):
+                                budget.settle(event.message.usage)
+                                settled = True
+                                terminal = event
+                                usage = event.message.usage
+                                if plan and (budget.unmeasured or usage.output > plan.output_tokens
+                                        or max(usage.input + usage.cache_read + usage.cache_write,
+                                               usage.total_tokens - usage.output) > plan.input_tokens):
+                                    budget.unmeasured = True
+                                    terminal = self._refuse(model, "provider usage did not match the strict reservation", usage)
+                                elif budget.used > budget.limit:
+                                    terminal = self._refuse(model, budget.refusal(), usage)
+                                break
+                            if isinstance(event, StreamFaultEvent):
+                                if not (event.error.error_message or "").startswith(REFUSED):
+                                    budget.settle(event.error.usage)
+                                    if plan:
+                                        budget.unmeasured = True
+                                settled = True
+                                terminal = event
+                                break
+                            yield event
+                finally:
+                    if not settled:
+                        budget.unmeasured = True
+        if terminal is not None:
+            yield terminal
+
+    async def _reserved_stream(self, model, start) -> AsyncIterator:
+        try:
+            self.allowance.reserve()
+        except AllowanceSpent as exc:
+            # No allowance.cancel(): later calls would say "abandoned", not the limit.
+            yield self._refuse(model, str(exc))
+            return
+        async for event in start():
+            if isinstance(event, StreamCloseEvent):
+                self.allowance.settle(to_spend(event.message.usage))
+            yield event

rlmagent/agents/child_handles.py

@@ -0,0 +1,46 @@
+"""Names and guest helpers for parent-owned child handles."""
+
+import re
+
+
+def checked_name(name: str) -> str:
+    if not isinstance(name, str) or not re.fullmatch(r"[A-Za-z][A-Za-z0-9_-]{0,63}", name):
+        raise ValueError("child name must start with a letter and contain up to 64 letters, digits, underscores or hyphens")
+    return name
+
+
+HELPERS = '''
+async def name_child(handle, name):
+    """Assign or change a name under the current parent and return the same handle."""
+    await _child_operation("name", handle.id, name)
+    return handle
+
+async def find_child(name):
+    """Look up one of your children by name."""
+    return _child_handle((await _child_operation("find", None, name))["id"])
+
+async def child_info(handle):
+    """Inspect a child's status, identity and recorded usage without consuming its output."""
+    return await _child_operation("info", handle.id)
+
+async def child_messages(handle, new_only=True):
+    """List metadata for child notes without marking them read."""
+    return await _child_operation("messages", handle.id, new_only=new_only)
+
+async def read_child_message(handle, message_id):
+    """Read a specific note and keep its receipt available for later lookup."""
+    return await _child_operation("read_message", handle.id, identity=message_id)
+
+async def steer_child(handle, text, message_id=None):
+    """Deliver an instruction at an active child's next model request."""
+    return await _child_operation("steer", handle.id, text=text, message_id=message_id)
+'''
+
+
+def compatibility_helpers(call: str) -> str:
+    return f'''
+_child_handle = _SubAgent
+
+async def _child_operation(action, key=None, name=None, **arguments):
+    return await {call}("child_handle", {{"action": action, "id": key, "name": name, **arguments}})
+''' + HELPERS

rlmagent/agents/child_mailbox.py

@@ -0,0 +1,30 @@
+"""Read-state bookkeeping for parent-owned child notes."""
+
+from time import time
+from uuid import uuid4
+
+
+def record_notes(notes, previous, current):
+    """Fold a mailbox snapshot. Each runtime serializes its appends and full drains."""
+    if previous == current:
+        return notes
+    if not current:
+        pending = len(previous)
+        return notes[:-pending] + [{**note, "status": "read"} for note in notes[-pending:]]
+    return notes + [{"id": uuid4().hex, "text": text, "status": "unread", "time": time()}
+                    for text in current[len(previous):]]
+
+
+def notes_metadata(notes, new_only):
+    return [{k: v for k, v in note.items() if k != "text"}
+            for note in notes if not new_only or note["status"] == "unread"]
+
+
+def read_note(notes, identity):
+    for index, note in enumerate(notes):
+        if note["id"] == identity:
+            if note["status"] == "read":
+                return notes, dict(note)
+            read = {**note, "status": "read"}
+            return notes[:index] + [read] + notes[index + 1:], dict(read)
+    return notes, None

rlmagent/agents/kernel_api.py

@@ -0,0 +1,156 @@
+"""Code run in every kernel before the model's first cell."""
+
+from textwrap import indent
+
+from rlmagent.agents.child_handles import compatibility_helpers
+
+KERNEL_API = r'''
+import asyncio as _rlmagent_asyncio
+import json as _rlmagent_json
+import os as _rlmagent_os
+import socket as _rlmagent_socket
+
+_RLMAGENT = {
+    "port": int(_rlmagent_os.environ["RLM_AGENT_BRIDGE_PORT"]),
+    "token": _rlmagent_os.environ["RLM_AGENT_BRIDGE_TOKEN"],
+}
+
+
+def _rlmagent_payload(op, args):
+    return (_rlmagent_json.dumps({"token": _RLMAGENT["token"], "op": op, "args": args})
+            + "\n").encode()
+
+
+def _rlmagent_reply(line):
+    reply = _rlmagent_json.loads(line)
+    if "error" in reply:
+        kind = TimeoutError if reply.get("type") == "TimeoutError" else RuntimeError
+        raise kind(reply["error"])
+    return reply["ok"]
+
+
+def _rlmagent_call_sync(op, args):
+    with _rlmagent_socket.create_connection(("127.0.0.1", _RLMAGENT["port"])) as sock:
+        sock.sendall(_rlmagent_payload(op, args))
+        chunks = []
+        while True:
+            chunk = sock.recv(1 << 16)
+            if not chunk:
+                break
+            chunks.append(chunk)
+            if chunk.endswith(b"\n"):
+                break
+    return _rlmagent_reply(b"".join(chunks))
+
+
+async def _rlmagent_call(op, args):
+    reader, writer = await _rlmagent_asyncio.open_connection(
+        "127.0.0.1", _RLMAGENT["port"], limit=2**26)
+    writer.write(_rlmagent_payload(op, args))
+    await writer.drain()
+    line = await reader.readline()
+    writer.close()
+    return _rlmagent_reply(line)
+
+
+if _rlmagent_os.environ.get("RLM_AGENT_CAN_DELEGATE") == "1":
+    async def rlm(task, context=None, role=None, model=None):
+        """Hand `task` to a sub-agent; `context` arrives in its kernel as CONTEXT. Returns its FINAL value."""
+        return await _rlmagent_call("rlm", {"task": str(task), "context": context, "role": role,
+                                            "model": model})
+
+    async def gather_rlm(jobs):
+        """Run several sub-agents at once. `jobs` is a list of (task, context) pairs or tasks. Results keep order."""
+        pairs = [list(j) if isinstance(j, (tuple, list)) else [j, None] for j in jobs]
+        return await _rlmagent_call("gather", {"jobs": [[str(t), c] for t, c in pairs]})
+
+class _SubAgent:
+    """A sub-agent started with spawn(). Losing this object does not stop it."""
+
+    def __init__(self, id):
+        self.id = id
+
+    def __repr__(self):
+        return f"<sub-agent {self.id}>"
+
+    async def status(self):
+        """'starting', 'running', 'idle', 'done', 'failed' or 'cancelled'."""
+        return await _rlmagent_call("status", {"id": self.id})
+
+    async def result(self, timeout=None):
+        """Its FINAL value; waits for it. TimeoutError if still running after `timeout` seconds."""
+        return await _rlmagent_call("result", {"id": self.id, "timeout": timeout})
+
+    async def cancel(self):
+        """Stop it and every sub-agent it started."""
+        return await _rlmagent_call("cancel", {"id": self.id})
+
+    async def send(self, text):
+        """Give it another instruction; wakes it if it is idle."""
+        return await _rlmagent_call("send", {"id": self.id, "text": str(text)})
+
+    async def messages(self):
+        """Notes it sent with tell_parent, oldest first; each is returned once."""
+        return await _rlmagent_call("messages", {"id": self.id})
+
+
+if _rlmagent_os.environ.get("RLM_AGENT_CAN_DELEGATE") == "1":
+    async def spawn(task, data=None, keep=False, role=None, model=None):
+        """Start a sub-agent and return at once. keep=True: it stays idle after FINAL, ready for send()."""
+        return _SubAgent(await _rlmagent_call(
+            "spawn", {"task": str(task), "context": data, "keep": bool(keep), "role": role,
+                      "model": model}))
+
+    async def children():
+        """Every sub-agent you started, in order."""
+        return [_SubAgent(i) for i in await _rlmagent_call("children", {})]
+
+async def improve(instructions=None, shared=False):
+    """Ask for the playbook to be updated from this conversation when the turn ends."""
+    return await _rlmagent_call("improve", {"instructions": instructions, "shared": bool(shared)})
+
+
+async def capabilities():
+    """The limits and approval setting in force, read from the harness."""
+    return await _rlmagent_call("capabilities", {})
+
+
+async def guidance(id=None):
+    return await _rlmagent_call("guidance", {"id": id})
+
+
+async def playbook(id=None):
+    """One playbook entry in full, or a list of all of them."""
+    return await _rlmagent_call("playbook", {"id": id})
+
+async def skill(id):
+    """Load a playbook skill and return its function."""
+    entry = await playbook(id)
+    if entry["kind"] != "skill":
+        raise ValueError(f"{id} is a {entry['kind']}, not a skill")
+    namespace = {}
+    exec(compile(entry["content"], f"<skill {id}>", "exec"), namespace)
+    return namespace[entry["callable"]]
+
+
+if _rlmagent_os.environ.get("RLM_AGENT_IS_CHILD") == "1":
+    async def tell_parent(text):
+        """Leave a note the agent that started you can read with messages()."""
+        await _rlmagent_call("tell_parent", {"text": str(text)})
+
+    def FINAL(value):
+        """Return `value` to the agent that started you. Call once, when done."""
+        try:
+            _rlmagent_json.dumps(value)
+            args = {"value": value, "note": ""}
+        except (TypeError, ValueError):
+            args = {"value": repr(value), "note": "value was not JSON-serialisable; sent as repr"}
+        _rlmagent_call_sync("final", args)
+        print("FINAL recorded.")
+
+    CONTEXT = _rlmagent_call_sync("context", {})
+'''
+
+KERNEL_API += '\nif _rlmagent_os.environ.get("RLM_AGENT_CAN_DELEGATE") == "1":\n' + indent(
+    compatibility_helpers("_rlmagent_call"), "    ",
+)

rlmagent/agents/output.py

@@ -0,0 +1,34 @@
+"""Bound retained command output while draining the entire pipe."""
+
+BYTE_LIMIT = 50 * 1024
+LINE_LIMIT = 2000
+
+
+class OutputTail:
+    def __init__(self):
+        self.data = bytearray()
+        self.total = 0
+        self.truncated = False
+
+    def feed(self, chunk):
+        self.total += len(chunk)
+        self.data.extend(chunk)
+        # Leave room for the truncation notice in the returned output.
+        capacity = BYTE_LIMIT - 128
+        if len(self.data) > capacity:
+            del self.data[:-capacity]
+            self.truncated = True
+        lines = self.data.count(b"\n") + bool(self.data and self.data[-1] != 10)
+        limit = LINE_LIMIT - 1 if self.truncated else LINE_LIMIT
+        while lines > limit:
+            del self.data[: self.data.index(10) + 1]
+            lines -= 1
+            self.truncated = True
+            limit = LINE_LIMIT - 1
+
+    def text(self):
+        # Ignore an incomplete leading UTF-8 character cut by the byte boundary.
+        text = bytes(self.data).decode("utf-8", errors="ignore")
+        if self.truncated:
+            return f"[output truncated; {self.total} bytes produced]\n" + text
+        return text

rlmagent/agents/patch.py

@@ -0,0 +1,107 @@
+"""Apply a unified diff to several files: everything is checked before anything is written."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+
+_HUNK = re.compile(r"@@ -(\d+)(?:,\d+)? \+\d+(?:,\d+)? @@")
+
+
+class PatchError(ValueError):
+    pass
+
+
+@dataclass
+class Hunk:
+    start: int  # 1-based line in the old file; 0 when it has no old lines
+    old: list[str] = field(default_factory=list)
+    new: list[str] = field(default_factory=list)
+
+
+@dataclass
+class FilePatch:
+    old_path: str | None  # None: the file is created
+    new_path: str | None  # None: the file is deleted
+    hunks: list[Hunk] = field(default_factory=list)
+
+    @property
+    def path(self) -> str:
+        return self.new_path or self.old_path
+
+
+def _name(raw: str) -> str | None:
+    name = raw.split("\t")[0].strip()
+    if name == "/dev/null":
+        return None
+    return name[2:] if name.startswith(("a/", "b/")) else name
+
+
+def parse(diff: str) -> list[FilePatch]:
+    files: list[FilePatch] = []
+    lines = diff.splitlines()
+    i = 0
+    while i < len(lines):
+        if not lines[i].startswith("--- "):
+            i += 1
+            continue
+        if i + 1 >= len(lines) or not lines[i + 1].startswith("+++ "):
+            raise PatchError(f"line {i + 1}: '---' without a following '+++'")
+        current = FilePatch(_name(lines[i][4:]), _name(lines[i + 1][4:]))
+        if current.old_path is None and current.new_path is None:
+            raise PatchError(f"line {i + 1}: both sides are /dev/null")
+        files.append(current)
+        i += 2
+        while i < len(lines) and lines[i].startswith("@@"):
+            found = _HUNK.match(lines[i])
+            if found is None:
+                raise PatchError(f"line {i + 1}: bad hunk header")
+            hunk = Hunk(int(found.group(1)))
+            current.hunks.append(hunk)
+            i += 1
+            while i < len(lines) and not lines[i].startswith(("@@", "--- ")):
+                line = lines[i]
+                if line.startswith("\\"):
+                    pass  # "No newline at end of file": the file keeps its own ending
+                elif line.startswith("-"):
+                    hunk.old.append(line[1:])
+                elif line.startswith("+"):
+                    hunk.new.append(line[1:])
+                elif line.startswith(" ") or not line:
+                    hunk.old.append(line[1:])
+                    hunk.new.append(line[1:])
+                else:
+                    break  # the text after the last hunk, such as a "diff --git" line
+                i += 1
+    if not files:
+        raise PatchError("no file changes found; expected '---' and '+++' lines")
+    return files
+
+
+def _locate(lines: list[str], hunk: Hunk, floor: int, path: str) -> int:
+    size = len(hunk.old)
+    stated = max(hunk.start - 1, 0)
+    if stated >= floor and lines[stated:stated + size] == hunk.old:
+        return stated
+    places = [at for at in range(floor, len(lines) - size + 1) if lines[at:at + size] == hunk.old]
+    if len(places) == 1:
+        return places[0]
+    what = "not found" if not places else f"found {len(places)} times"
+    raise PatchError(f"{path}: the hunk at line {hunk.start} was {what}; read the file again")
+
+
+def patched(text: str, patch: FilePatch) -> str:
+    """The file's new text; raises PatchError when a hunk does not fit."""
+    newline = "\r\n" if "\r\n" in text else "\n"
+    ends_open = bool(text) and not text.endswith("\n")
+    lines = text.splitlines()
+    result: list[str] = []
+    position = 0
+    for hunk in patch.hunks:
+        at = _locate(lines, hunk, position, patch.path)
+        result += lines[position:at] + hunk.new
+        position = at + len(hunk.old)
+    result += lines[position:]
+    if not result:
+        return ""
+    return newline.join(result) + ("" if ends_open else newline)