railtracks.prebuilt.middleware

15class ContextInjection(ModelMiddleware):
16    """Inject ``rt.context`` values into prompt placeholders before each model call.
17
18    Model-level middleware (``model_middleware=`` only). Fills ``{placeholder}``
19    templates in the message history from the active session's context::
20
21        import railtracks as rt
22        from railtracks.prebuilt import middleware
23
24        rt.agent_node(
25            "Agent",
26            llm=rt.llm.OpenAILLM(model_name="gpt-5.4-mini"),
27            system_message="You are helping {user_name}.",
28            model_middleware=[middleware.ContextInjection()],
29        )
30
31    List position matters: place it before (outside) any middleware that must see
32    the injected prompt, e.g. an input guard listed after this entry sees the
33    filled-in template.
34    """
35
36    def __init__(self):
37        super().__init__(self._middleware_fn)
38
39    async def _middleware_fn(
40        self,
41        call: Callable[
42            [MessageHistory, type[BaseModel] | None, list[Tool] | None],
43            Awaitable[Response],
44        ],
45        message_history: MessageHistory,
46        schema: type[BaseModel] | None,
47        tools: list[Tool] | None,
48    ):
49        inject_context(message_history)
50        return await call(message_history, schema, tools)

 34class ConversationMemory(Middleware):
 35    """Automatically cache and append conversation history across node invocations.
 36
 37    Node-level middleware (``middleware=``). Automatically preserves multi-turn
 38    conversations in session context and prepends prior turns to incoming inputs.
 39    Do not manually pass prior message history into invocations when this middleware
 40    is attached, as history is accumulated automatically.
 41
 42    See: https://docs.railtracks.org/documentation/agent_design/middleware/prebuilt/list/conversation_memory/
 43
 44    Args:
 45        context_key: Optional explicit session context key for sharing, querying,
 46            or seeding history. If None, an isolated per-instance key is generated,
 47            which no external caller can address.
 48        max_messages: Optional limit on the number of recent messages to retain.
 49            ``None`` and ``0`` both mean no limit.
 50
 51    Raises:
 52        ValueError: If ``context_key`` is blank or ``max_messages`` is negative.
 53    """
 54
 55    def __init__(
 56        self,
 57        context_key: str | None = None,
 58        *,
 59        max_messages: int | None = None,
 60    ):
 61        if context_key is not None and not context_key.strip():
 62            raise ValueError("context_key must be a non-empty string when provided.")
 63        if max_messages is not None and max_messages < 0:
 64            raise ValueError(f"max_messages must not be negative, got {max_messages}.")
 65
 66        self._instance_id = uuid.uuid4().hex[:8]
 67        self._context_key = context_key or f"conversation_history_{self._instance_id}"
 68        self._max_messages = max_messages or None
 69        self._state: dict[str, MessageHistory | None] = {"history": None}
 70        super().__init__(self._middleware_fn)
 71
 72    @property
 73    def context_key(self) -> str:
 74        """The session context key holding the conversation history."""
 75        return self._context_key
 76
 77    def get_history(self) -> MessageHistory | None:
 78        """Return the current conversation history from session context or instance."""
 79        return self._get_existing_history()
 80
 81    def clear(self) -> None:
 82        """Clear the stored conversation history.
 83
 84        The instance copy is always dropped. The session context entry is removed
 85        only when called during a run, since a finished run's context is no longer
 86        reachable from here: a ``FlowConnection`` held open on that run can still
 87        read the pre-clear value through ``conn.context``.
 88        """
 89        self._state["history"] = None
 90        if is_context_present():
 91            context = safe_get_runner_context()
 92            try:
 93                context.external_context.delete(self._context_key)
 94            except KeyError:
 95                pass
 96
 97    def __deepcopy__(self, memo: dict) -> ConversationMemory:
 98        # ``Node.extend_middleware()`` copies the middleware already attached to a
 99        # node when ``rt.couple()`` layers another one on. The history dict is shared
100        # by reference so the handle the caller holds keeps reading what the node
101        # actually spends, rather than diverging from it once the run ends.
102        cls = self.__class__
103        result = cls.__new__(cls)
104        memo[id(self)] = result
105        for k, v in self.__dict__.items():
106            if k == "_state":
107                setattr(result, k, v)
108            else:
109                setattr(result, k, deepcopy(v, memo))
110        return result
111
112    def _extract_user_input(
113        self, args: tuple[object, ...], kwargs: dict[str, object]
114    ) -> tuple[bool, bool, object | None]:
115        if args:
116            return True, True, args[0]
117        if "user_input" in kwargs:
118            return True, False, kwargs["user_input"]
119        return False, False, None
120
121    def _append_to_history(self, history: MessageHistory, user_input: object) -> None:
122        if isinstance(user_input, str):
123            history.append(UserMessage(user_input))
124        elif isinstance(user_input, UserMessage):
125            history.append(user_input)
126        elif isinstance(user_input, (list, MessageHistory)):
127            for msg in user_input:
128                if isinstance(msg, Message):
129                    history.append(msg)
130        else:
131            history.append(UserMessage(str(user_input)))
132
133    def _combine_history(
134        self, existing: MessageHistory, user_input: object
135    ) -> MessageHistory:
136        if isinstance(user_input, (list, MessageHistory)) and _continues(
137            existing, user_input
138        ):
139            # The caller handed back history this middleware already holds, so the
140            # input supersedes the stored copy rather than being appended to it.
141            combined = MessageHistory(deepcopy(list(user_input)))
142        else:
143            combined = deepcopy(existing)
144            self._append_to_history(combined, user_input)
145        if self._max_messages is not None and len(combined) > self._max_messages:
146            return MessageHistory(combined[-self._max_messages :])
147        return combined
148
149    def _get_existing_history(self) -> MessageHistory | None:
150        if is_context_present():
151            context = safe_get_runner_context()
152            try:
153                hist = context.external_context.get(self._context_key)
154                if hist is not None:
155                    return hist
156            except KeyError:
157                pass
158        return self._state.get("history")
159
160    def _save_result(self, result: object) -> None:
161        if isinstance(result, LLMResponse):
162            history = result.message_history
163            if self._max_messages is not None and len(history) > self._max_messages:
164                history = MessageHistory(history[-self._max_messages :])
165            saved_copy = deepcopy(history)
166            self._state["history"] = saved_copy
167            if is_context_present():
168                context = safe_get_runner_context()
169                context.external_context.put(self._context_key, saved_copy)
170
171    async def _middleware_fn(self, call, *args, **kwargs):
172        existing = self._get_existing_history()
173        has_input, is_pos, user_input = self._extract_user_input(args, kwargs)
174
175        if has_input and existing:
176            combined = self._combine_history(existing, user_input)
177            if is_pos:
178                args = (combined, *args[1:])
179            else:
180                kwargs = {**kwargs, "user_input": combined}
181
182        result = await call(*args, **kwargs)
183        self._save_result(result)
184        return result

50def pre_verifier(
51    approve_fn: _ApproveFn[_P] | None = None,
52    /,
53    *,
54    timeout: float | None = None,
55    name: str | None = None,
56) -> (
57    Middleware[_P, Any, _MiddlewareSignature[_P, Never]]
58    | Callable[[_ApproveFn[_P]], Middleware[_P, Any, _MiddlewareSignature[_P, Never]]]
59):
60    """Build a node-verification middleware around ``approve_fn`` that gates a call
61    BEFORE it runs.
62
63    ``approve_fn`` is called with the exact ``*args, **kwargs`` the wrapped
64    node was called with — sync or async, both supported — and must return a
65    `Verdict`. On decline, `VerifierRejectedError` is raised and the node's own
66    body never runs. On accept, the call is forwarded onward, using the
67    verdict's ``args``/``kwargs`` if it supplied overrides, otherwise the
68    original ones unchanged.
69
70    If ``timeout`` is set and ``approve_fn`` doesn't respond in time, the call
71    is treated as declined with reason ``"timeout"``.
72
73    See also :func:`~railtracks.prebuilt.middleware.post_verifier.post_verifier`,
74    which gates a call's output AFTER it has already run. For the full
75    picture (composing with other middleware, custom approval backends,
76    guided walkthroughs), see the Verifiers docs.
77    """
78
79    if approve_fn is None:
80        return lambda fn: pre_verifier(fn, timeout=timeout, name=name)
81
82    return wrap_node(_wrapper(approve_fn, timeout), name=name)

 68def post_verifier(
 69    approve_fn: _ApproveFn[_R, _P] | None = None,
 70    /,
 71    *,
 72    timeout: float | None = None,
 73    name: str | None = None,
 74) -> Middleware[_P, _R] | Callable[[_ApproveFn[_R, _P]], Middleware[_P, _R]]:
 75    """Build a node-verification middleware around ``approve_fn`` that gates a
 76    call's OUTPUT AFTER it has already run.
 77
 78    The wrapped node always runs first. ``approve_fn`` is then called with the
 79    produced value as its first positional parameter, followed by the node's
 80    own ``*args, **kwargs`` — sync or async, both supported — and must return
 81    a `Verdict`. This shape is validated eagerly, at ``post_verifier(...)``
 82    call time: an ``approve_fn`` that doesn't take ``result`` first raises
 83    `TypeError` immediately, naming what was found instead.
 84
 85    Decline can't undo the call (it already happened) but still raises
 86    `VerifierRejectedError`, stopping the result from propagating onward. On
 87    accept, the result propagates using the verdict's ``result`` if it
 88    supplied an override, otherwise the original result unchanged.
 89
 90    If ``timeout`` is set and ``approve_fn`` doesn't respond in time, the call
 91    is treated as declined with reason ``"timeout"``.
 92
 93    See also :func:`~railtracks.prebuilt.middleware.pre_verifier.pre_verifier`,
 94    which gates whether a call happens at all, BEFORE it runs. For the full
 95    picture (composing with other middleware, custom approval backends,
 96    guided walkthroughs), see the Verifiers docs.
 97    """
 98
 99    if approve_fn is None:
100        return lambda fn: post_verifier(fn, timeout=timeout, name=name)
101
102    _require_result_first(approve_fn)
103    return wrap_node(_wrapper(approve_fn, timeout), name=name)

12class Lock(Middleware[Any, Any, Never]):
13    """Serialize concurrent invocations of the wrapped call.
14
15    Reuse one instance across nodes that must not execute concurrently.
16    """
17
18    def __init__(self):
19        self._lock = asyncio.Lock()
20        super().__init__(self._middleware_fn)
21
22    async def _middleware_fn(self, call, *args, **kwargs):
23        async with self._lock:
24            return await call(*args, **kwargs)

 18class MaxCalls(Middleware[Any, Any, Never]):
 19    """Fail the wrapped call once it has been invoked ``max_calls`` times.
 20
 21    Slot-agnostic: works both as node middleware (``middleware=``) and as model
 22    middleware (``model_middleware=``) — it only counts invocations of ``call``
 23    and never inspects the arguments::
 24
 25        import railtracks as rt
 26        from railtracks.prebuilt import middleware
 27
 28        rt.agent_node(
 29            "Agent",
 30            llm=rt.llm.OpenAILLM(model_name="gpt-5.4-mini"),
 31            middleware=[middleware.MaxCalls(5)],  # cap calls to the whole node
 32            model_middleware=[middleware.MaxCalls(5)],  # cap raw model calls
 33        )
 34
 35    The count is tracked per session, resetting to zero when a new one begins. A
 36    single ``MaxCalls`` instance shared across nodes within the same session
 37    enforces a combined budget, while a fresh instance per node gives each its
 38    own limit. A bare top-level ``rt.call`` is its own run, so hold several calls
 39    under one budget by running them inside a ``Flow``.
 40
 41    Finished sessions' counters stay readable via :attr:`call_count`, capped at
 42    the 64 most recent.
 43
 44    Args:
 45        max_calls: Number of calls allowed before the limit is enforced.
 46        custom_message: Message to raise once the limit is exceeded. Defaults
 47            to ``"Maximum number of calls exceeded"``.
 48
 49    Raises:
 50        MaxCallsExceededError: Once ``call`` has already been invoked
 51            ``max_calls`` times.
 52    """
 53
 54    def __init__(
 55        self,
 56        max_calls: int,
 57        custom_message: str | None = None,
 58    ):
 59        self._max_calls = max_calls
 60        self._custom_message = custom_message
 61        # Keyed by session_id, or None outside a run; least-recently-used first.
 62        self._session_counts: OrderedDict[str | None, int] = OrderedDict()
 63        self._last_session_id: str | None = None
 64        super().__init__(self._middleware_fn)
 65
 66    @property
 67    def max_calls(self) -> int:
 68        """The maximum number of calls allowed before raising."""
 69        return self._max_calls
 70
 71    def _current_session_id(self) -> str | None:
 72        """Return the active session ID, or ``None`` if not inside a run."""
 73        if is_context_present():
 74            try:
 75                return get_session_identity().session_id
 76            except ContextError:
 77                return None
 78        return None
 79
 80    def _inspection_key(self) -> str | None:
 81        """Return the live session, or the last counted one when outside a run."""
 82        if is_context_present():
 83            return self._current_session_id()
 84        return self._last_session_id
 85
 86    @property
 87    def call_count(self) -> int:
 88        """Calls in the active session, or the run that just finished.
 89
 90        Returns 0 for a session evicted by the 64-session retention cap.
 91        """
 92        return self._session_counts.get(self._inspection_key(), 0)
 93
 94    def reset(self) -> None:
 95        """Reset the counter that :attr:`call_count` reads."""
 96        self._session_counts.pop(self._inspection_key(), None)
 97
 98    def reset_all(self) -> None:
 99        """Reset call counters for all sessions."""
100        self._session_counts.clear()
101        self._last_session_id = None
102
103    async def _middleware_fn(self, call, *args, **kwargs):
104        # Live session only: a call outside a run must not spend a finished
105        # run's budget, so unlike _inspection_key this never falls back.
106        sess_id = self._current_session_id()
107        # Read-modify-write is safe under one event loop: no await between them.
108        current_count = self._session_counts.get(sess_id, 0)
109        if current_count >= self._max_calls:
110            if self._custom_message:
111                raise MaxCallsExceededError(self._custom_message)
112            raise MaxCallsExceededError("Maximum number of calls exceeded")
113        self._session_counts[sess_id] = current_count + 1
114        self._session_counts.move_to_end(sess_id)
115        self._last_session_id = sess_id
116        while len(self._session_counts) > _MAX_TRACKED_SESSIONS:
117            self._session_counts.popitem(last=False)
118        return await call(*args, **kwargs)

121class MaxCallsExceededError(Exception):
122    """Raised when a :class:`MaxCalls`-wrapped call is invoked past its limit."""

12class Retry(Middleware[Any, Any, Never]):
13    """Retry the wrapped call when it raises a transient error.
14
15    Slot-agnostic: works both as node middleware (``middleware=``) and as model
16    middleware (``model_middleware=``) — it only re-invokes ``call`` and never
17    inspects the arguments::
18
19        import railtracks as rt
20        from railtracks.prebuilt import middleware
21
22        rt.agent_node(
23            "Agent",
24            llm=rt.llm.OpenAILLM(model_name="gpt-5.4-mini"),
25            middleware=[middleware.Retry(3)],  # retry the whole node
26            model_middleware=[middleware.Retry(3)],  # retry each raw model call
27        )
28
29    The backoff schedule is delegated to a :class:`~railtracks.llm.retries.RetryApproach`
30    (the same strategies used by ``retry_approach`` on model objects).
31
32    Args:
33        max_tries: Total attempts including the first call. Ignored when
34            ``approach`` is given.
35        approach: Backoff strategy; defaults to
36            :class:`~railtracks.llm.retries.ExponentialRetry` with ``max_tries``.
37        retry_on: Exception types worth retrying. Defaults to the transient LLM
38            provider errors (rate limits, timeouts, connection failures) — pass
39            your own tuple (e.g. ``(Exception,)``) for node-level use. Anything
40            not in the tuple propagates immediately.
41    """
42
43    def __init__(
44        self,
45        max_tries: int = 3,
46        *,
47        approach: RetryApproach | None = None,
48        retry_on: tuple[type[Exception], ...] | None = None,
49    ):
50        self._approach = (
51            approach if approach is not None else ExponentialRetry(max_tries=max_tries)
52        )
53        self._retry_on = retry_on
54        super().__init__(self._middleware_fn)
55
56    async def _middleware_fn(self, call, *args, **kwargs):
57        return await self._approach.acall_with_retry(
58            lambda: call(*args, **kwargs), retry_on=self._retry_on
59        )

12class Timeout(Middleware[Any, Any, Never]):
13    """Fail the wrapped call when it runs longer than ``seconds``.
14
15    The timeout applies to the complete wrapped call. When the deadline expires,
16    the call is cancelled and :class:`TimeoutError` is raised.
17
18    Args:
19        seconds: Maximum number of seconds to wait for the wrapped call.
20    """
21
22    def __init__(self, seconds: float):
23        self._seconds = seconds
24        super().__init__(self._middleware_fn)
25
26    async def _middleware_fn(self, call, *args, **kwargs):
27        return await asyncio.wait_for(call(*args, **kwargs), timeout=self._seconds)