rlmness-cli 0.2.0__py3-none-any.whl

rlmagent/harness/engine.py

@@ -0,0 +1,489 @@
+"""Pure portable agent loop."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
+
+from rlmagent.harness.contracts.stream import (
+    AgentEvent,
+    MessageEndEvent,
+    MessageStartEvent,
+    MessageUpdateEvent,
+    RetryEvent,
+    RunEndEvent,
+    RunStartEvent,
+    ToolRunEndEvent,
+    ToolRunStartEvent,
+    ToolRunUpdateEvent,
+    TurnEndEvent,
+    TurnStartEvent,
+)
+from rlmagent.harness.contracts.tooling import CancelToken, ToolOutcome, ToolSpec
+from rlmagent.harness.contracts.transcript import (
+    CallBlock,
+    ModelEntry,
+    TextSegment,
+    ToolOutcomeEntry,
+    TranscriptEntry,
+)
+from rlmagent.harness.mending import mend_tool_history
+from rlmagent.harness.provider import wire
+from rlmagent.harness.provider.base import ModelProvider
+
+PreToolHook = Callable[[CallBlock], Awaitable[tuple[bool, str | None]]]
+PostToolHook = Callable[
+    [CallBlock, ToolOutcome, bool],
+    Awaitable[tuple[ToolOutcome, bool]],
+]
+
+
+# --- synthetic entries ------------------------------------------------------
+
+
+def _make_abort_result(text: str) -> ToolOutcome:
+    """Create an error tool outcome with the given message."""
+    return ToolOutcome(content=[TextSegment(text=text)])
+
+
+def _make_error_reply(model: str, reason: str) -> ModelEntry:
+    """Create a synthetic assistant entry representing a fault."""
+    return ModelEntry(
+        model=model,
+        content=[],
+        stop_reason="error",
+        error_message=reason,
+    )
+
+
+# --- provider streaming -----------------------------------------------------
+
+
+def _accepts_cache_key(provider: ModelProvider) -> bool:
+    """Whether ``provider.stream_response`` takes ``cache_key`` (or ``**kwargs``)."""
+    try:
+        params = inspect.signature(provider.stream_response).parameters.values()
+    except (TypeError, ValueError):
+        return False
+    return any(
+        p.name == "cache_key" or p.kind is inspect.Parameter.VAR_KEYWORD for p in params
+    )
+
+
+def _is_empty_failure(entry: TranscriptEntry) -> bool:
+    """A reply that failed before producing anything."""
+    return (
+        isinstance(entry, ModelEntry)
+        and entry.stop_reason in ("error", "aborted")
+        and not entry.content
+    )
+
+
+async def _stream_model_reply(
+    provider: ModelProvider,
+    model: str,
+    system: str,
+    transcript: list[TranscriptEntry],
+    tool_defs: list[ToolSpec],
+    token: CancelToken | None,
+    cache_key: str | None = None,
+) -> AsyncIterator[AgentEvent]:
+    """Relay wire events from the provider as agent-level message events."""
+    # Only forwarded to adapters that declare the parameter, so providers
+    # written before it existed keep working unchanged.
+    routing = {"cache_key": cache_key} if cache_key and _accepts_cache_key(provider) else {}
+    raw: AsyncIterator[wire.WireEvent] = provider.stream_response(
+        model=model,
+        system=system,
+        messages=transcript,
+        tools=tool_defs,
+        signal=token,
+        **routing,
+    )
+    saw_open = False
+    async for evt in raw:
+        if isinstance(evt, wire.StreamOpenEvent):
+            saw_open = True
+            yield MessageStartEvent(message=evt.partial)
+        elif isinstance(evt, wire.StreamCloseEvent):
+            if not saw_open:
+                yield MessageStartEvent(message=evt.message)
+            yield MessageEndEvent(message=evt.message)
+        elif isinstance(evt, wire.StreamFaultEvent):
+            if not saw_open:
+                yield MessageStartEvent(message=evt.error)
+            yield MessageEndEvent(message=evt.error)
+        elif isinstance(evt, wire.SourceRetryEvent):
+            yield RetryEvent(
+                attempt=evt.attempt,
+                max_attempts=evt.max_attempts,
+                delay_seconds=evt.delay_seconds,
+                message=evt.message,
+            )
+        else:
+            yield MessageUpdateEvent(
+                message=evt.partial,
+                assistant_message_event=evt.model_dump(),
+            )
+
+
+# --- single tool invocation -------------------------------------------------
+
+_JSON_TYPES = {
+    "string": lambda v: isinstance(v, str),
+    "boolean": lambda v: isinstance(v, bool),
+    "integer": lambda v: not isinstance(v, bool) and (
+        isinstance(v, int) or (isinstance(v, float) and v == int(v))),
+    "number": lambda v: isinstance(v, (int, float)) and not isinstance(v, bool),
+    "array": lambda v: isinstance(v, list),
+    "object": lambda v: isinstance(v, dict),
+    "null": lambda v: v is None,
+}
+
+
+def argument_problem(schema, arguments) -> str | None:
+    """What is wrong with a call's arguments: missing required ones or wrong top-level types."""
+    for name in schema.get("required") or ():
+        if name not in arguments:
+            return f"missing required argument {name!r}"
+    properties = schema.get("properties") or {}
+    for name, value in arguments.items():
+        declared = (properties.get(name) or {}).get("type") if isinstance(properties, dict) else None
+        allowed = declared if isinstance(declared, list) else [declared] if declared else []
+        known = [t for t in allowed if t in _JSON_TYPES]
+        if known and not any(_JSON_TYPES[t](value) for t in known):
+            return f"{name!r} must be {' or '.join(known)}, got {type(value).__name__}"
+    return None
+
+
+
+async def _run_single_tool(
+    spec: ToolSpec,
+    invocation: CallBlock,
+    token: CancelToken | None,
+) -> tuple[ToolOutcome, bool, list[ToolOutcome]]:
+    """Execute one tool, capturing progress snapshots."""
+    snapshots: list[ToolOutcome] = []
+    alive = True
+
+    def _on_partial(snapshot: ToolOutcome) -> None:
+        if alive:
+            snapshots.append(snapshot.model_copy(deep=True))
+
+    problem = argument_problem(spec.parameters, invocation.arguments)
+    if problem:
+        return _make_abort_result(f"Invalid arguments for {spec.name}: {problem}"), True, snapshots
+
+    try:
+        result = await spec.execute(invocation.id, invocation.arguments, token, _on_partial)
+        return result, False, snapshots
+    except asyncio.CancelledError:
+        raise
+    except Exception as exc:  # noqa: BLE001 - tools are an isolation boundary
+        return _make_abort_result(str(exc)), True, snapshots
+    finally:
+        alive = False
+
+
+# --- tool dispatch pipeline -------------------------------------------------
+
+
+async def _handle_invocation(
+    invocation: CallBlock,
+    registry: Mapping[str, ToolSpec],
+    token: CancelToken | None,
+    pre_hook: PreToolHook | None,
+    post_hook: PostToolHook | None,
+) -> AsyncIterator[AgentEvent]:
+    """Run one tool call through the full pre-hook / execute / post-hook pipeline."""
+    yield ToolRunStartEvent(
+        tool_call_id=invocation.id,
+        tool_name=invocation.name,
+        args=invocation.arguments,
+    )
+
+    rejected = False
+    rejection_msg: str | None = None
+    if pre_hook is not None:
+        rejected, rejection_msg = await pre_hook(invocation)
+
+    if rejected:
+        outcome = _make_abort_result(rejection_msg or "Tool execution was blocked")
+        errored = True
+    elif token is not None and token.is_cancelled():
+        outcome = _make_abort_result("Operation aborted")
+        errored = True
+    else:
+        spec = registry.get(invocation.name)
+        if spec is None:
+            outcome = _make_abort_result(f"Tool {invocation.name} not found")
+            errored = True
+        else:
+            outcome, errored, snapshots = await _run_single_tool(spec, invocation, token)
+            for snap in snapshots:
+                yield ToolRunUpdateEvent(
+                    tool_call_id=invocation.id,
+                    tool_name=invocation.name,
+                    args=invocation.arguments,
+                    partial_result=snap,
+                )
+
+    if post_hook is not None:
+        outcome, errored = await post_hook(invocation, outcome, errored)
+
+    yield ToolRunEndEvent(
+        tool_call_id=invocation.id,
+        tool_name=invocation.name,
+        result=outcome,
+        is_error=errored,
+    )
+    record = ToolOutcomeEntry(
+        tool_call_id=invocation.id,
+        tool_name=invocation.name,
+        content=outcome.content,
+        details=outcome.details,
+        added_tool_names=outcome.added_tool_names,
+        is_error=errored,
+    )
+    yield MessageStartEvent(message=record)
+    yield MessageEndEvent(message=record)
+
+
+# --- parallel / sequential strategy -----------------------------------------
+
+
+async def _gather_events_for(
+    invocation: CallBlock,
+    registry: Mapping[str, ToolSpec],
+    token: CancelToken | None,
+    pre_hook: PreToolHook | None,
+    post_hook: PostToolHook | None,
+) -> list[AgentEvent]:
+    """Buffer all events from a single tool invocation (used for gather)."""
+    buf: list[AgentEvent] = []
+    async for ev in _handle_invocation(invocation, registry, token, pre_hook, post_hook):
+        buf.append(ev)
+    return buf
+
+
+def _can_parallelize(
+    invocations: list[CallBlock],
+    registry: Mapping[str, ToolSpec],
+) -> bool:
+    """True when every invocation's tool permits concurrent execution."""
+    return all(
+        (spec := registry.get(c.name)) is None or spec.parallelism != "sequential"
+        for c in invocations
+    )
+
+
+def _cancelled_outcome_events(invocation: CallBlock) -> list[AgentEvent]:
+    """Build the event sequence for a call aborted before execution."""
+    abort = _make_abort_result("Operation aborted")
+    record = ToolOutcomeEntry(
+        tool_call_id=invocation.id,
+        tool_name=invocation.name,
+        content=abort.content,
+        is_error=True,
+    )
+    return [
+        ToolRunStartEvent(
+            tool_call_id=invocation.id,
+            tool_name=invocation.name,
+            args=invocation.arguments,
+        ),
+        ToolRunEndEvent(
+            tool_call_id=invocation.id,
+            tool_name=invocation.name,
+            result=abort,
+            is_error=True,
+        ),
+        MessageStartEvent(message=record),
+        MessageEndEvent(message=record),
+    ]
+
+
+async def _execute_calls(
+    invocations: list[CallBlock],
+    registry: Mapping[str, ToolSpec],
+    token: CancelToken | None,
+    pre_hook: PreToolHook | None,
+    post_hook: PostToolHook | None,
+) -> AsyncIterator[AgentEvent]:
+    """Execute tool calls, choosing parallel or sequential strategy."""
+    if len(invocations) > 1 and _can_parallelize(invocations, registry):
+        batches = await asyncio.gather(*(
+            _gather_events_for(inv, registry, token, pre_hook, post_hook)
+            for inv in invocations
+        ))
+        for batch in batches:
+            for ev in batch:
+                yield ev
+    else:
+        for inv in invocations:
+            if token is not None and token.is_cancelled():
+                for ev in _cancelled_outcome_events(inv):
+                    yield ev
+                continue
+            async for ev in _handle_invocation(inv, registry, token, pre_hook, post_hook):
+                yield ev
+
+
+def _pick_outcome_entries(events: list[AgentEvent]) -> list[ToolOutcomeEntry]:
+    """Extract ToolOutcomeEntry objects from MessageEndEvent events."""
+    return [
+        ev.message
+        for ev in events
+        if isinstance(ev, MessageEndEvent) and isinstance(ev.message, ToolOutcomeEntry)
+    ]
+
+
+# --- main loop --------------------------------------------------------------
+
+
+async def run_agent_loop(
+    *,
+    provider: ModelProvider,
+    model: str,
+    system: str,
+    messages: list[TranscriptEntry],
+    tools: list[ToolSpec],
+    prompts: Sequence[TranscriptEntry] = (),
+    max_turns: int | None = None,
+    signal: CancelToken | None = None,
+    get_steering_messages: Callable[[], Sequence[TranscriptEntry]] | None = None,
+    get_follow_up_messages: Callable[[], Sequence[TranscriptEntry]] | None = None,
+    before_tool_call: PreToolHook | None = None,
+    after_tool_call: PostToolHook | None = None,
+    cache_key: str | None = None,
+) -> AsyncIterator[AgentEvent]:
+    """Run the provider/tool loop and emit portable agent events."""
+    emitted: list[TranscriptEntry] = list(prompts)
+    if prompts:
+        messages.extend(prompts)
+
+    yield RunStartEvent()
+    yield TurnStartEvent()
+    for p in prompts:
+        yield MessageStartEvent(message=p)
+        yield MessageEndEvent(message=p)
+
+    if max_turns is not None and max_turns < 1:
+        err = _make_error_reply(model, "max_turns must be at least 1")
+        messages.append(err)
+        emitted.append(err)
+        yield MessageStartEvent(message=err)
+        yield MessageEndEvent(message=err)
+        yield TurnEndEvent(message=err)
+        yield RunEndEvent(messages=emitted)
+        return
+
+    lookup = {t.name: t for t in tools}
+    step = 1
+    first = True
+    pending_steering = tuple(get_steering_messages() if get_steering_messages else ())
+
+    while True:
+        has_calls = True
+        while has_calls or pending_steering:
+            # A cancelled run must not begin a *further* turn. The provider
+            # stops mid-stream on its own, but a partial reply often carries
+            # stop_reason="stop", which would otherwise look like a normal
+            # turn and send the loop round again.
+            #
+            # The first turn is always allowed to proceed: `_execute_calls`
+            # emits "aborted" outcomes for pending tool calls, and skipping it
+            # would leave dangling tool calls that providers reject.
+            if not first and signal is not None and signal.is_cancelled():
+                yield RunEndEvent(messages=emitted)
+                return
+
+            if not first:
+                yield TurnStartEvent()
+            first = False
+
+            for entry in pending_steering:
+                messages.append(entry)
+                emitted.append(entry)
+                yield MessageStartEvent(message=entry)
+                yield MessageEndEvent(message=entry)
+            pending_steering = ()
+
+            if max_turns is not None and step > max_turns:
+                err = _make_error_reply(
+                    model, f"Agent stopped after max_turns={max_turns}"
+                )
+                messages.append(err)
+                emitted.append(err)
+                yield MessageStartEvent(message=err)
+                yield MessageEndEvent(message=err)
+                yield TurnEndEvent(message=err)
+                yield RunEndEvent(messages=emitted)
+                return
+
+            assistant: ModelEntry | None = None
+            # Last safety net for callers that never mended stored history:
+            # the provider sees valid call/result pairing, the transcript the
+            # harness owns is left exactly as it is.
+            outbound = [
+                entry for entry in mend_tool_history(messages).entries
+                if not _is_empty_failure(entry)
+            ]
+            async for ev in _stream_model_reply(
+                provider, model, system, outbound, tools, signal, cache_key,
+            ):
+                yield ev
+                if isinstance(ev, MessageEndEvent) and isinstance(ev.message, ModelEntry):
+                    assistant = ev.message
+
+            if assistant is None:
+                assistant = _make_error_reply(
+                    model, "Provider produced no assistant message"
+                )
+                yield MessageStartEvent(message=assistant)
+                yield MessageEndEvent(message=assistant)
+
+            messages.append(assistant)
+            emitted.append(assistant)
+            if assistant.stop_reason in {"error", "aborted"}:
+                yield TurnEndEvent(message=assistant)
+                yield RunEndEvent(messages=emitted)
+                return
+
+            calls = list(assistant.tool_calls)
+            has_calls = bool(calls)
+
+            tool_events: list[AgentEvent] = []
+            async for ev in _execute_calls(
+                calls, lookup, signal, before_tool_call, after_tool_call,
+            ):
+                yield ev
+                tool_events.append(ev)
+
+            # Tools can ask for the turn to end here. Only an unanimous batch
+            # ends it: one tool cannot silence results the model still needs.
+            ends = [ev.result for ev in tool_events if isinstance(ev, ToolRunEndEvent)]
+            if ends and all(result.terminate is True for result in ends):
+                has_calls = False
+
+            outcomes = _pick_outcome_entries(tool_events)
+            for oc in outcomes:
+                messages.append(oc)
+                emitted.append(oc)
+
+            yield TurnEndEvent(message=assistant, tool_results=outcomes)
+            step += 1
+            pending_steering = tuple(
+                get_steering_messages() if get_steering_messages else ()
+            )
+
+        followups = tuple(
+            get_follow_up_messages() if get_follow_up_messages else ()
+        )
+        if followups:
+            pending_steering = followups
+            continue
+        break
+
+    yield RunEndEvent(messages=emitted)

rlmagent/harness/mending.py

@@ -0,0 +1,153 @@
+"""Deterministic repair of malformed tool-call history."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from rlmagent.harness.contracts.transcript import (
+    CallBlock,
+    ModelEntry,
+    TextSegment,
+    ToolOutcomeEntry,
+    TranscriptEntry,
+)
+
+INTERRUPTED_RESULT_TEXT = "Tool call interrupted by user"
+
+# A call is identified by (index of its assistant entry, 1-based position
+# among that entry's calls); its ideal result slot is entry index + position.
+type _CallSlot = tuple[int, int]
+type _Pick = tuple[int | None, ToolOutcomeEntry]
+
+
+@dataclass(frozen=True, slots=True)
+class HistoryMend:
+    """Repaired history plus counters describing what was changed."""
+
+    entries: tuple[TranscriptEntry, ...]
+    changed: bool = False
+    synthesized: int = 0
+    dropped_orphans: int = 0
+    dropped_duplicates: int = 0
+    reordered: int = 0
+
+    def counters(self) -> dict[str, int]:
+        """JSON-safe summary for a durable diagnostic record."""
+        return {
+            "synthesized": self.synthesized,
+            "dropped_orphans": self.dropped_orphans,
+            "dropped_duplicates": self.dropped_duplicates,
+            "reordered": self.reordered,
+        }
+
+
+def mend_tool_history(entries: Sequence[TranscriptEntry]) -> HistoryMend:
+    """Return history in which every tool call has exactly one adjacent result."""
+    history = tuple(entries)
+    calls = _call_slots(history)
+    results: dict[str, list[tuple[int, ToolOutcomeEntry]]] = defaultdict(list)
+    for position, entry in enumerate(history):
+        if isinstance(entry, ToolOutcomeEntry):
+            results[entry.tool_call_id].append((position, entry))
+
+    chosen: dict[_CallSlot, _Pick] = {}
+    taken: set[int] = set()
+
+    # Pairs that are already adjacent are claimed first, so a legitimately
+    # reused id stays with its own turn instead of an earlier call taking it.
+    for slot, call, ideal in calls:
+        if ideal < len(history) and ideal not in taken:
+            candidate = history[ideal]
+            if isinstance(candidate, ToolOutcomeEntry) and candidate.tool_call_id == call.id:
+                chosen[slot] = (ideal, candidate)
+                taken.add(ideal)
+
+    synthesized = 0
+    for slot, call, _ideal in calls:
+        if slot in chosen:
+            continue
+        free = [c for c in results.get(call.id, []) if c[0] not in taken]
+        if not free:
+            chosen[slot] = (None, _interruption(call))
+            synthesized += 1
+            continue
+        pool = [c for c in free if c[0] > slot[0]] or free
+        pick = next((c for c in pool if not _is_interruption(c[1])), pool[0])
+        chosen[slot] = pick
+        taken.add(pick[0])
+
+    # A call paired with a synthetic interruption switches to a real result
+    # for the same id if one is still unclaimed.
+    for slot, call, _ideal in calls:
+        claimed, result = chosen[slot]
+        if claimed is None or not _is_interruption(result):
+            continue
+        better = next(
+            (c for c in results.get(call.id, [])
+             if c[0] not in taken and not _is_interruption(c[1])),
+            None,
+        )
+        if better is not None:
+            taken.discard(claimed)
+            taken.add(better[0])
+            chosen[slot] = better
+
+    mended: list[TranscriptEntry] = []
+    reordered = 0
+    for position, entry in enumerate(history):
+        if isinstance(entry, ToolOutcomeEntry):
+            continue
+        mended.append(entry)
+        if not isinstance(entry, ModelEntry):
+            continue
+        for offset, _call in enumerate(entry.tool_calls, start=1):
+            source, result = chosen[(position, offset)]
+            mended.append(result)
+            if source is not None and source != position + offset:
+                reordered += 1
+
+    known_ids = {call.id for _slot, call, _ideal in calls}
+    leftovers = [
+        result
+        for group in results.values()
+        for position, result in group
+        if position not in taken
+    ]
+    changed = len(mended) != len(history) or any(
+        a is not b for a, b in zip(mended, history, strict=False)
+    )
+    return HistoryMend(
+        entries=tuple(mended),
+        changed=changed,
+        synthesized=synthesized,
+        dropped_orphans=sum(r.tool_call_id not in known_ids for r in leftovers),
+        dropped_duplicates=sum(r.tool_call_id in known_ids for r in leftovers),
+        reordered=reordered,
+    )
+
+
+def _call_slots(history: tuple[TranscriptEntry, ...]) -> list[tuple[_CallSlot, CallBlock, int]]:
+    slots: list[tuple[_CallSlot, CallBlock, int]] = []
+    for position, entry in enumerate(history):
+        if isinstance(entry, ModelEntry):
+            for offset, call in enumerate(entry.tool_calls, start=1):
+                slots.append(((position, offset), call, position + offset))
+    return slots
+
+
+def _interruption(call: CallBlock) -> ToolOutcomeEntry:
+    return ToolOutcomeEntry(
+        tool_call_id=call.id,
+        tool_name=call.name,
+        content=[TextSegment(text=INTERRUPTED_RESULT_TEXT)],
+        is_error=True,
+    )
+
+
+def _is_interruption(result: ToolOutcomeEntry) -> bool:
+    return result.is_error and result.text == INTERRUPTED_RESULT_TEXT
+
+
+__all__ = ["INTERRUPTED_RESULT_TEXT", "HistoryMend", "mend_tool_history"]

rlmagent/harness/provider/__init__.py

@@ -0,0 +1 @@
+"""rlmagent.harness provider seam: wire events + the ModelProvider protocol."""

rlmagent/harness/provider/base.py

@@ -0,0 +1,27 @@
+"""Structural protocol for model adapters consumed by the agent loop."""
+#provider
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Sequence
+from typing import Protocol
+
+from rlmagent.harness.contracts.tooling import CancelToken, ToolSpec
+from rlmagent.harness.contracts.transcript import TranscriptEntry
+from rlmagent.harness.provider.wire import WireEvent
+
+
+class ModelProvider(Protocol):
+    """Structural interface for anything that can produce model completions."""
+
+    def stream_response(
+        self,
+        *,
+        model: str,
+        system: str,
+        messages: Sequence[TranscriptEntry],
+        tools: Sequence[ToolSpec],
+        signal: CancelToken | None = None,
+        cache_key: str | None = None,
+    ) -> AsyncIterator[WireEvent]:
+        """Stream typed wire events for one model round-trip."""
+        ...