monkeybot 2.1.1__py3-none-any.whl

monkeybot/core/hooks/__init__.py

@@ -0,0 +1,261 @@
+"""In-process lifecycle hooks for the agent loop.
+
+Hooks are async callbacks registered against well-known points in the agent
+loop (turn start/end, tool call before/after, provider request/response, etc.).
+They are the extensibility point used by the memory subsystem, but the manager
+itself is memory-agnostic.
+
+Design rules (mirrored from rohitg00/agentmemory's hook contract):
+
+* **Bounded.** Every hook call is wrapped in ``asyncio.wait_for`` (or scheduled
+  as fire-and-forget when ``timeout_s == 0``). Hooks cannot stall the loop.
+* **Silent.** Hook errors and timeouts are logged at WARNING and never
+  propagate. A broken hook never breaks a turn.
+* **Re-entrancy safe.** Hooks running inside ``fire()`` cannot recursively
+  trigger hooks; the inner call returns the payload unchanged.
+* **Mutable payload.** Hooks may set ``inject_text`` / ``inject_memory_lines``
+  on the shared :class:`HookPayload`; later hooks for the same event observe
+  and may extend those fields. Provider/tool hooks may also replace
+  ``provider_messages`` / ``tools``.
+
+This module has no dependency on :mod:`monkeybot.core.memory` or the agent
+loop; both consume it.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextvars
+import logging
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from monkeybot.core.context import TurnContext
+    from monkeybot.core.llm.provider import Message
+    from monkeybot.core.types.types_tools import ToolDef
+
+_log = logging.getLogger(__name__)
+
+
+class HookEvent(StrEnum):
+    """Lifecycle points at which hooks may fire.
+
+    Read-side events (``USER_MESSAGE``, ``PRE_TURN``, ``PRE_TOOL``,
+    ``TOOL_DEFINITION``, ``BEFORE_PROVIDER_REQUEST``) run synchronously with a
+    short timeout and may mutate the payload (injection fields, tools, or
+    provider messages). Write-side events (``POST_TOOL``, ``POST_TURN``,
+    ``SESSION_END``, ``AFTER_PROVIDER_RESPONSE``) are fire-and-forget; their
+    return values are ignored.
+    """
+
+    USER_MESSAGE = "user_message"
+    PRE_TURN = "pre_turn"
+    PRE_TOOL = "pre_tool"
+    POST_TOOL = "post_tool"
+    POST_TURN = "post_turn"
+    SESSION_END = "session_end"
+    TOOL_DEFINITION = "tool.definition"
+    BEFORE_PROVIDER_REQUEST = "before_provider_request"
+    AFTER_PROVIDER_RESPONSE = "after_provider_response"
+
+
+@dataclass
+class HookPayload:
+    """Per-fire payload passed to every hook for one event.
+
+    The payload is **mutable**: hooks may assign ``inject_text`` or extend
+    ``inject_memory_lines`` to surface context to the loop. The loop reads
+    those fields after ``fire()`` returns.
+
+    Event-specific fields (``user_message``, ``tool_*``, ``provider_messages``,
+    ``tools``, response fields) are populated by the caller based on which
+    event is firing; hooks for an unrelated event should ignore them.
+
+    ``TOOL_DEFINITION`` / ``BEFORE_PROVIDER_REQUEST`` may replace ``tools``
+    and/or ``provider_messages`` (``Message`` is frozen — replace list entries,
+    do not assign to message fields).
+    """
+
+    event: HookEvent
+    thread_id: str
+    request_id: str
+    ctx: "TurnContext"
+
+    user_message: str | None = None
+    tool_name: str | None = None
+    tool_args: dict[str, Any] | None = None
+    tool_result: str | None = None
+    tool_error: str | None = None
+
+    # TOOL_DEFINITION / BEFORE_PROVIDER_REQUEST (mutable; replace list to filter)
+    tools: list["ToolDef"] | None = None
+    provider_messages: list["Message"] | None = None
+    inner_turn: int | None = None
+
+    # AFTER_PROVIDER_RESPONSE (observational)
+    assistant_text: str | None = None
+    thinking_text: str | None = None
+    tool_requests: list[dict[str, Any]] | None = None
+    usage: dict[str, int] | None = None
+    provider_error: str | None = None
+
+    inject_text: str | None = None
+    inject_memory_lines: list[str] = field(default_factory=list)
+
+
+HookFn = Callable[[HookPayload], Awaitable[None]]
+
+
+_in_hook: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    "monkeybot_in_hook", default=False
+)
+
+
+_DEFAULT_TIMEOUT_S = 2.0
+
+
+class HookManager:
+    """Register and dispatch hooks across :class:`HookEvent` points.
+
+    One ``HookManager`` is constructed per agent (gateway-owned), shared by
+    the loop and any subscribers (e.g. the memory hook). Subagents receive a
+    no-op manager so duplicate writes do not race the parent's hooks.
+
+    Fire-and-forget handlers (``timeout_s == 0``) are tracked so
+    :meth:`drain_settlement` can wait for them before ``TurnComplete`` or the
+    next provider call — without requiring SSE client ACKs.
+    """
+
+    def __init__(self) -> None:
+        self._handlers: dict[HookEvent, list[HookFn]] = {}
+        self._pending: set[asyncio.Task[None]] = set()
+
+    def register(self, event: HookEvent, fn: HookFn) -> None:
+        """Append ``fn`` to the handler list for ``event``.
+
+        Handlers fire in registration order. Duplicates are allowed; callers
+        must not register the same function twice unless duplicate firing is
+        intended.
+        """
+        self._handlers.setdefault(event, []).append(fn)
+
+    def clear(self, event: HookEvent | None = None) -> None:
+        """Remove all handlers for ``event`` (or every event when ``None``)."""
+        if event is None:
+            self._handlers.clear()
+        else:
+            self._handlers.pop(event, None)
+
+    async def fire(
+        self,
+        payload: HookPayload,
+        *,
+        timeout_s: float = _DEFAULT_TIMEOUT_S,
+    ) -> HookPayload:
+        """Run every handler registered for ``payload.event`` with bounded time.
+
+        ``timeout_s``:
+            * ``> 0``: each handler is awaited up to that many seconds; on
+              timeout the handler is cancelled and a warning logged.
+            * ``== 0``: the handler is scheduled as a background task; this
+              method returns immediately without awaiting completion. Use for
+              write-side events where the agent should not wait on memory.
+              Detached tasks may not finish before the process exits; in
+              short-lived handlers (Lambda, Cloud Functions) pass
+              ``hook_manager=None`` or use ``timeout_s > 0`` so hooks complete
+              before returning. Call :meth:`drain_settlement` before idle
+              boundaries when side effects must land.
+
+        Hooks may not recursively trigger this method; nested calls return
+        ``payload`` unchanged (after a single debug log line).
+        """
+        if _in_hook.get():
+            _log.debug("hook re-entrancy blocked for event=%s", payload.event)
+            return payload
+
+        handlers = self._handlers.get(payload.event)
+        if not handlers:
+            return payload
+
+        if timeout_s == 0:
+            for fn in handlers:
+                self._schedule_background(fn, payload)
+            return payload
+
+        token = _in_hook.set(True)
+        try:
+            for fn in handlers:
+                await self._run_one(fn, payload, timeout_s)
+        finally:
+            _in_hook.reset(token)
+        return payload
+
+    async def drain_settlement(self, *, timeout_s: float = _DEFAULT_TIMEOUT_S) -> None:
+        """Await outstanding fire-and-forget hook tasks up to ``timeout_s``.
+
+        Does not cancel slow hooks (they may still finish after the barrier).
+        Never raises for hook failures — only logs a warning on timeout.
+        Safe to call when no tasks are pending.
+        """
+        pending = [t for t in self._pending if not t.done()]
+        if not pending:
+            return
+        _done, still = await asyncio.wait(pending, timeout=max(0.0, timeout_s))
+        if still:
+            _log.warning(
+                "hook settlement timed out pending=%d timeout_s=%.2f",
+                len(still),
+                timeout_s,
+            )
+
+    @staticmethod
+    async def _run_one(fn: HookFn, payload: HookPayload, timeout_s: float) -> None:
+        try:
+            await asyncio.wait_for(fn(payload), timeout=timeout_s)
+        except asyncio.TimeoutError:
+            _log.warning(
+                "hook timed out event=%s fn=%s timeout_s=%.2f",
+                payload.event,
+                getattr(fn, "__qualname__", repr(fn)),
+                timeout_s,
+            )
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            _log.warning(
+                "hook error event=%s fn=%s err=%r",
+                payload.event,
+                getattr(fn, "__qualname__", repr(fn)),
+                exc,
+            )
+
+    def _schedule_background(self, fn: HookFn, payload: HookPayload) -> None:
+        async def _wrap() -> None:
+            token = _in_hook.set(True)
+            try:
+                await fn(payload)
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                _log.warning(
+                    "background hook error event=%s fn=%s err=%r",
+                    payload.event,
+                    getattr(fn, "__qualname__", repr(fn)),
+                    exc,
+                )
+            finally:
+                _in_hook.reset(token)
+
+        task = asyncio.create_task(_wrap())
+        self._pending.add(task)
+
+        def _done(t: asyncio.Task[None]) -> None:
+            self._pending.discard(t)
+
+        task.add_done_callback(_done)
+
+
+__all__ = ["HookEvent", "HookPayload", "HookFn", "HookManager"]

monkeybot/core/llm/__init__.py

@@ -0,0 +1,4 @@
+"""LLM streaming protocol (:mod:`monkeybot.core.llm.provider`) and turn usage (:mod:`monkeybot.core.llm.usage`).
+
+Concrete vendor backends live in :mod:`monkeybot.providers`.
+"""

monkeybot/core/llm/provider.py

@@ -0,0 +1,296 @@
+"""Thin provider streaming contract (LLM adapter boundary)."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any, Literal, Protocol, TypeAlias, cast
+
+from monkeybot.core.types.content_blocks import ContentBlock, Text, ToolRequest, ToolResponse
+from monkeybot.core.types.types_tools import ToolDef
+
+Role: TypeAlias = Literal["user", "assistant", "system"]
+
+
+@dataclass(frozen=True, kw_only=True)
+class Message:
+    """Single turn for tool-calling chat models (typed content blocks)."""
+
+    role: Role
+    content: list[ContentBlock]
+
+    def __post_init__(self) -> None:
+        if self.role not in ("user", "assistant", "system"):
+            raise ValueError(f"invalid role: {self.role!r}")
+        for i, block in enumerate(self.content):
+            if not isinstance(block, ContentBlock):
+                raise ValueError(
+                    f"content[{i}] must be ContentBlock, got {type(block).__name__}"
+                )
+
+    def to_dict(self) -> dict[str, object]:
+        return {"role": self.role, "content": [b.to_dict() for b in self.content]}
+
+    @classmethod
+    def from_dict(cls, d: Mapping[str, object]) -> Message:
+        role = d.get("role")
+        if role not in ("user", "assistant", "system"):
+            raise ValueError(f"invalid role: {role!r}")
+        raw_blocks = d.get("content", [])
+        if raw_blocks is None:
+            raw_blocks = []
+        if not isinstance(raw_blocks, list):
+            raise ValueError("content must be a list")
+        blocks: list[ContentBlock] = []
+        for item in raw_blocks:
+            if not isinstance(item, dict):
+                raise ValueError("content list elements must be JSON objects")
+            blocks.append(ContentBlock.from_dict(item))
+        return cls(role=role, content=blocks)
+
+    @classmethod
+    def text(cls, role: Role, text: str) -> Message:
+        """Build a message whose content is a single :class:`Text` block."""
+        return cls(role=role, content=[Text(text=text)])
+
+
+@dataclass(frozen=True, kw_only=True)
+class TextDelta:
+    kind: Literal["text_delta"] = "text_delta"
+    text: str
+
+
+@dataclass(frozen=True, kw_only=True)
+class ThinkingDelta:
+    """Incremental thinking/reasoning chunk from a streaming provider."""
+
+    kind: Literal["thinking_delta"] = "thinking_delta"
+    text: str
+    signature: str | None = None
+
+
+@dataclass(frozen=True, kw_only=True)
+class ToolCall:
+    kind: Literal["tool_call"] = "tool_call"
+    call_id: str
+    name: str
+    args: dict[str, object]
+    parse_error: str | None = None
+    metadata: dict[str, object] | None = None
+
+
+@dataclass(frozen=True, kw_only=True)
+class ToolInputDelta:
+    """Incremental tool-argument JSON fragment while args are still streaming.
+
+    ``delta`` is a raw, opaque fragment of a single streaming JSON document keyed
+    by ``call_id`` (e.g. Anthropic ``input_json_delta``). It is not valid JSON on
+    its own — consumers must concatenate all fragments for a given ``call_id``
+    (in arrival order) before attempting to parse the result. The final,
+    validated arguments are delivered separately on :class:`ToolCall`.
+    """
+
+    kind: Literal["tool_input_delta"] = "tool_input_delta"
+    call_id: str
+    name: str
+    delta: str
+
+
+@dataclass(frozen=True, kw_only=True)
+class GroundingEvent:
+    """Provider-native web-search grounding metadata (e.g. Gemini ``google_search``).
+
+    Additive to the harness's pluggable ``web_search`` custom tool — this carries
+    citations/search-suggestion data from a provider-hosted search tool invoked
+    server-side, not a tool call the harness dispatched itself.
+    """
+
+    kind: Literal["grounding"] = "grounding"
+    sources: list[dict[str, str]]
+    search_queries: list[str]
+
+
+@dataclass(frozen=True, kw_only=True)
+class ProviderCallHints:
+    """Optional transport metadata for prompt-cache / session affinity.
+
+    Content strategy (stable/volatile split + epoch) lives in prompts; these
+    hints are provider-specific request options layered on top.
+    """
+
+    session_id: str | None = None
+    cache_retention: Literal["none", "short", "long"] = "short"
+
+
+def cache_retention_from_env() -> Literal["none", "short", "long"]:
+    """Resolve ``MODEL_CACHE_RETENTION`` (default ``short``)."""
+    import os
+
+    raw = os.environ.get("MODEL_CACHE_RETENTION", "short").strip().lower()
+    if raw in ("none", "short", "long"):
+        return raw  # type: ignore[return-value]
+    return "short"
+
+
+def gemini_extra_kwargs(provider: Provider, *, vertex_google_search: bool) -> dict[str, bool]:
+    if vertex_google_search and provider.name == "gemini":
+        return {"vertex_google_search": True}
+    return {}
+
+
+def provider_call_hints_kwargs(
+    provider: Provider,
+    hints: ProviderCallHints | None,
+) -> dict[str, Any]:
+    """Return kwargs accepted by providers that opt into ``ProviderCallHints``."""
+    if hints is None:
+        return {}
+    # Anthropic-family + OpenAI accept hints; others ignore via Protocol default.
+    if provider.name in ("claude", "vertex-claude", "bedrock", "openai"):
+        return {"hints": hints}
+    return {}
+
+
+async def provider_count_input_tokens(
+    provider: Provider,
+    messages: Sequence[Message],
+    tools: Sequence[ToolDef],
+    *,
+    model: str,
+    thinking_budget: int | None = None,
+    vertex_google_search: bool = False,
+    hints: ProviderCallHints | None = None,
+) -> int:
+    kwargs: dict[str, Any] = {"model": model, "thinking_budget": thinking_budget}
+    kwargs.update(gemini_extra_kwargs(provider, vertex_google_search=vertex_google_search))
+    kwargs.update(provider_call_hints_kwargs(provider, hints))
+    if kwargs.keys() - {"model", "thinking_budget"}:
+        return int(
+            await cast(Any, provider).count_input_tokens(messages, tools, **kwargs)
+        )
+    return await provider.count_input_tokens(messages, tools, **kwargs)
+
+
+def provider_stream(
+    provider: Provider,
+    messages: Sequence[Message],
+    tools: Sequence[ToolDef],
+    *,
+    model: str,
+    thinking_budget: int | None = None,
+    vertex_google_search: bool = False,
+    hints: ProviderCallHints | None = None,
+) -> AsyncIterator[ProviderEvent]:
+    kwargs: dict[str, Any] = {"model": model, "thinking_budget": thinking_budget}
+    kwargs.update(gemini_extra_kwargs(provider, vertex_google_search=vertex_google_search))
+    kwargs.update(provider_call_hints_kwargs(provider, hints))
+    if kwargs.keys() - {"model", "thinking_budget"}:
+        return cast(
+            AsyncIterator[ProviderEvent],
+            cast(Any, provider).stream(messages, tools, **kwargs),
+        )
+    return provider.stream(messages, tools, **kwargs)
+
+
+@dataclass(frozen=True, kw_only=True)
+class UsageEvent:
+    kind: Literal["usage"] = "usage"
+    input_tokens: int
+    output_tokens: int
+    cached_tokens: int = 0
+    cache_read_tokens: int = 0
+    cache_creation_tokens: int = 0
+
+
+@dataclass(frozen=True)
+class Done:
+    """End of a provider stream. ``truncated`` means the vendor hit an output length limit."""
+
+    kind: Literal["done"] = "done"
+    truncated: bool = False
+
+
+ProviderEvent: TypeAlias = (
+    TextDelta
+    | ThinkingDelta
+    | ToolCall
+    | ToolInputDelta
+    | GroundingEvent
+    | UsageEvent
+    | Done
+)
+
+
+class Provider(Protocol):
+    """Streams model output as :class:`ProviderEvent` values.
+
+    Exactly one consumer should iterate a given ``stream`` at a time; concurrent
+    overlapping calls on the same instance are intentionally undefined.
+
+    ``stream`` is annotated as a synchronous method returning :class:`AsyncIterator`
+    so async-generator implementations match under strict mypy (``async def``
+    with ``yield`` is *not* typed as returning ``Coroutine[..., AsyncIterator]``).
+    """
+
+    @property
+    def name(self) -> str:
+        """Stable provider id (e.g. ``\"gemini\"``)."""
+
+    @property
+    def supports_streaming(self) -> bool:
+        """Whether partial output is exposed as incremental deltas."""
+
+    def stream(
+        self,
+        messages: Sequence[Message],
+        tools: Sequence[ToolDef],
+        *,
+        model: str,
+        thinking_budget: int | None = None,
+    ) -> AsyncIterator[ProviderEvent]:
+        """Yield provider events for one model request.
+
+        ``thinking_budget`` overrides the configured reasoning budget for this call
+        when the provider supports it (Gemini, Claude). ``None`` uses the default.
+        """
+
+    async def count_input_tokens(
+        self,
+        messages: Sequence[Message],
+        tools: Sequence[ToolDef],
+        *,
+        model: str,
+        thinking_budget: int | None = None,
+    ) -> int:
+        """Return the provider-aligned input (prompt) token count for one outbound request.
+
+        Must reflect the same payload shape as :meth:`stream` (messages, tools, model),
+        typically via the vendor's tokenizer or count API — not post-hoc usage from a
+        prior response.
+        ``thinking_budget`` mirrors :meth:`stream` for providers whose token count
+        changes with reasoning configuration.
+        """
+
+
+__all__ = [
+    "Done",
+    "GroundingEvent",
+    "Message",
+    "Provider",
+    "ProviderCallHints",
+    "ProviderEvent",
+    "Role",
+    "Text",
+    "TextDelta",
+    "ThinkingDelta",
+    "ToolCall",
+    "ToolDef",
+    "ToolRequest",
+    "ToolResponse",
+    "UsageEvent",
+    "cache_retention_from_env",
+    "gemini_extra_kwargs",
+    "provider_call_hints_kwargs",
+    "provider_count_input_tokens",
+    "provider_stream",
+]