refactor(backends): move generate_from_raw hook firing into Backend base class#1264
Merged
ajbozarth merged 2 commits intoJun 16, 2026
Conversation
…ase class Closes generative-computing#1183. Folds the additive half of generative-computing#1218 Part 1. - `generate_from_raw` becomes a `@final` wrapper on `Backend` that owns pre/post/error hook firing; backends implement `_generate_from_raw` returning `(results, usage)`. - All five backends drop the inline gen_id, latency timing, and three hook-fire blocks; wrapper catches `BaseException` (matches chat-path generative-computing#1181). - `generation_batch_pre_call` payload mutations now propagate to the backend impl (model_options/format/tool_calls), matching the chat path. Adds `model_options` field to `GenerationBatchPreCallPayload`. Closes a gap from generative-computing#1218 where the batch hook was wired for telemetry observation only. - Standardizes `self._model_id` and `self._provider` on every backend. Inlines model-id resolution into ollama/huggingface `__init__`s and deletes the `_get_*_model_id` helpers; renames `_hf_model_id` to `_model_id`. - Backends with per-MOT token counts (ollama, hf, watsonx) now populate `mot.generation.usage` per MOT; openai/litellm leave it `None` since their APIs only report whole-batch usage. `mot._meta["usage"]` writes preserved for generative-computing#1218's follow-up. - Adds `TestGenerationBatchHookCallSites` mirroring the chat-path tests; updates the custom-backend doc snippet. Assisted-by: Claude Code Signed-off-by: Alex Bozarth <ajbozart@us.ibm.com>
planetf1
reviewed
Jun 15, 2026
Without an entry in MELLEA_HOOK_PAYLOAD_POLICIES, cpex's default-DENY silently strips plugin mutations to model_options/format/tool_calls on the batch path, even though the new wrapper reads them back. The existing batch mutation test patches invoke_hook entirely, so it never exercises the cpex policy layer; the chat-path equivalent had the same gap. Adds real-plugin tests on both paths that go through register() + invoke_hook so the policy table stays honest. Assisted-by: Claude Code Signed-off-by: Alex Bozarth <ajbozart@us.ibm.com>
planetf1
approved these changes
Jun 16, 2026
planetf1
left a comment
Contributor
There was a problem hiding this comment.
LGTM. Fine leaving watsonx.
8 tasks
jakelorocco
approved these changes
Jun 16, 2026
Merged
via the queue into
generative-computing:main
with commit Jun 16, 2026
b703c46
9 checks passed
Merged
8 tasks
ajbozarth
added a commit
to ajbozarth/mellea
that referenced
this pull request
Jun 17, 2026
…meta["usage"] writes (generative-computing#1288) Migrate the only in-tree consumer (`think_budget_forcing`) off `mot._meta["usage"]` and onto `mot.generation.usage`, then drop the now-redundant `_meta["usage"]` writes from all five backends' `generate_from_raw` paths. The additive half (writing `mot.generation.usage` from raw paths) landed with generative-computing#1264. Budget forcing now raises a named `RuntimeError` when a sub-call returns `mot.generation.usage = None` instead of letting a confusing `TypeError` bubble out of arithmetic on `None`. Also replaces an implicit `IndexError` in the `think_max_tokens=0`/`answer_suffix=None` no-sub-calls edge case with a named error. Closes generative-computing#1218 Assisted-by: Claude Code Signed-off-by: Alex Bozarth <ajbozart@us.ibm.com>
This was referenced Jun 19, 2026
13 tasks
yelkurdi
added a commit
to yelkurdi/mellea
that referenced
this pull request
Jun 24, 2026
The Backend ABC's abstract method was renamed to _generate_from_raw (leading underscore) in generative-computing#1264; the public generate_from_raw is now a concrete hook-firing wrapper. The fakes in test_compactor.py overrode the wrapper but left the abstract method unimplemented, so each backend raised TypeError on instantiation. Rename generate_from_raw -> _generate_from_raw in FakeBackend, OtherBackend, BrokenBackend, BuggyBackend, and InterruptingBackend. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com>
AngeloDanducci
pushed a commit
to AngeloDanducci/mellea
that referenced
this pull request
Jul 2, 2026
…mputing#996) * feat: add context compaction strategies for react framework Adds CompactionStrategy abstraction and KeepLastN implementation to mellea/stdlib/compaction.py, wires an optional compaction parameter into the react() loop, and adds full test coverage in test/stdlib/test_compaction.py. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * refactor: express compaction threshold as token count Switches `CompactionStrategy.threshold` from a component-count trigger to a token-count trigger, read from the most recent `ModelOutputThunk.usage` populated by the backend. This aligns compaction with the real constraint (context size) and sidesteps per-backend tokenizer dependencies by using provider-reported usage; the trade-off is a one-turn lag since usage is recorded at the end of each model call. Also reorders the react loop so compaction runs after the final-answer check, skipping wasted work (and a wasted LLM call for LLMSummarize) on terminal turns. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * Fix mot.generation.usage Signed-off-by: ramon-astudillo <ramon@astudillo.com> * refactor: relocate compaction module into frameworks package Move the compaction strategies alongside the react framework they serve: - mellea/stdlib/compaction.py -> mellea/stdlib/frameworks/react_compaction.py - test/stdlib/test_compaction.py -> test/stdlib/frameworks/test_react_compaction.py Imports and module docstrings updated accordingly. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * docs: add Args/Returns sections to react_compaction compact overrides The docstring quality gate (tooling/docs-autogen/audit_coverage.py --quality --threshold 100) requires each documented symbol to have its own Args/Returns sections — inheritance from the abstract parent is not consulted. Six issues were reported against the compact() overrides on ClearAll, KeepLastN, and LLMSummarize. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * feat(compaction): per-turn Compactor protocol for ChatContext + ReACT Replaces the original async ``react_compaction`` strategies (ClearAll, KeepLastN, LLMSummarize) with a generic, sync ``Compactor`` protocol that operates on any ``Context``. ``ReACT`` and ``ChatContext`` are rewired around the new protocol; sample callers, tests, and docs are updated. Squash of 29 Mellea-side commits from context_compaction_for_react_2; the BCP eval harness commits in that branch are intentionally excluded. mellea/stdlib/context/ becomes a package - Compactor protocol: sync ``compact(ctx, *, backend=None) -> Context`` - WindowCompactor(size, pin_predicate) keep last-N body components; ``size=0`` clears the body and retains only the pinned prefix - ThresholdCompactor(inner, threshold) token-gated wrapper that reads cumulative context size from the most recent ModelOutputThunk's ``generation.usage`` and forwards to ``inner.compact`` only above the gate - LLMSummarizeCompactor(keep_n, pin_predicate, prompt_template) summarizes old body components via the backend; the (async) backend call is hidden behind a sync ``compact()`` via ``_run_coro_blocking`` so the protocol stays sync - PinPredicate API: ``pin_nothing``, ``pin_system``, ``pin_system_and_initial_user``; chat compactors compose freely mellea/stdlib/frameworks/react.py - ``react()`` gains a ``compactor: Compactor | None = None`` per-turn hook; invoked once after each tool observation - The old ``react_compaction`` module is removed mellea/stdlib/components/react.py - ``pin_react_initiator``: a PinPredicate that pins everything up to and including the first ``ReactInitiator`` - ``react_summary_prompt(goal=None, max_tokens_hint=None)``: factory that returns a research-flavoured summary prompt template (with the {conversation} placeholder LLMSummarizeCompactor expects). Optional ``GOAL: <goal>`` line and optional ``- Be at most ~N tokens`` bullet when callers want goal anchoring or length-cap hints. mellea/stdlib/context/chat.py - ``ChatContext()`` defaults to no compactor (full history); pass ``compactor=`` or ``window_size=`` for opt-in compaction. Matches upstream main's window_size=None unbounded semantics. Test coverage - test/stdlib/test_compactor.py (~500 LOC): protocol semantics; Window / Threshold / LLMSummarize behaviours; pin-predicate edge cases; ``size=0`` collapse; threshold gate edge cases - test/stdlib/frameworks/test_react_framework.py (~210 LOC): react() per-turn hook integration + react_summary_prompt (default, goal interpolation, brace escaping, max_tokens_hint bullet ordering, LLMSummarizeCompactor template-validation) - test/stdlib/test_base_context.py: pin-non-compacting ChatContext in the session-copy operations test (matches new opt-in default) Net diff: 17 files, +381 / -896 lines (drops the old react_compaction.py and its dedicated test file). Backwards-compatible default behaviour preserved: bare ``ChatContext()`` retains full history; ``react()`` without ``compactor=`` behaves identically to today; ``LLMSummarizeCompactor`` defaults to a generic conversation-summary prompt unless callers opt in to the research-flavoured variant via ``react_summary_prompt``. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * docs: fix stale references to ChatContext default and compaction examples - Correct comment in test_base_context.py: ChatContext defaults to compactor=None, not WindowCompactor(5). - Point compactor.py module docstring to docs/examples/context/ instead of the nonexistent docs/rewrite/. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * feat(compaction): InlineCompactor marker + required default_backend on LLMSummarize Restrict ChatContext(compactor=...) to compactors that inherit InlineCompactor. Wiring a backend-requiring compactor (e.g. LLMSummarizeCompactor) directly would invoke the backend on every add(); the new isinstance guard rejects that with a TypeError pointing at react(compactor=...), ThresholdCompactor, or manual compact() as alternatives. ThresholdCompactor remains accepted regardless of its inner -- it gates by token usage, so backend calls are sparse rather than per-add. LLMSummarizeCompactor now takes a required default_backend at construction and falls back to it when compact() is invoked without an explicit backend. A backend kwarg passed to compact() still overrides the default for that call. This makes ThresholdCompactor(LLMSummarizeCompactor(...)) work end-to-end when attached to ChatContext: at trip time, the inner uses its stored default_backend. InlineCompactor carries the compact() signature (raising NotImplementedError) so it's a usable static type without cast() workarounds. Specialized to ChatContext rather than parameterized over Context -- ThresholdCompactor's prior generic-T signature was unexercised, so the simpler shape applies. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * fix(session): track interaction_count out-of-band; doc compaction semantic shift Reviewer flagged that ChatContext(window_size=N) used to keep full history in as_list() and only window view_for_generation(); the per-turn compactor work made as_list() itself reflect the post-compaction state, which silently undercounted in MelleaSession.cleanup() (interaction_count = len(as_list())). - ChatContext docstring: add a Note describing the semantic shift and pointing callers at out-of-band turn tracking when full counts matter. - MelleaSession: turn ctx into a property; the setter increments _interaction_count, and reset() / __init__ bypass via _ctx so lifecycle events don't pollute the count. cleanup() now publishes self._interaction_count, stable under any compaction strategy. - SessionCleanupPayload: rewrite the interaction_count field doc to match the new semantics ("turns committed" rather than "items in context"). Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * fix(compaction): make LLMSummarizeCompactor backend errors non-fatal Reviewer flagged that an exception from the summarisation backend call (rate limit, network error, timeout) propagates through _run_coro_blocking and kills the entire react loop. For long-running research tasks that's quite painful, especially since compaction is best-effort by nature. LLMSummarizeCompactor.compact() now wraps _run_coro_blocking in a try/except Exception. On failure it logs a WARNING via MelleaLogger with the exception type and message, then returns ctx unchanged. The next compact() invocation retries; the conversation keeps growing in the meantime. BaseException (KeyboardInterrupt, SystemExit) still propagates so users can interrupt a stuck loop. Added a regression test with a backend that raises RuntimeError on every call: compact() returns the same ctx, original history is intact, and a warning naming the exception type is logged. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * docs: stronger warning on _run_coro_blocking event-loop blocking Reviewer flagged that _run_coro_blocking, when invoked from inside an async caller like react(), blocks the entire event loop for the full duration of the wrapped coroutine — not just the calling task. The previous "fine for a serial ReACT loop" wording undersold the implications. Beefed-up Warning: block now spells out: - The loop, not just the thread, is stalled — callbacks, telemetry, cancellation signals, other sessions sharing the loop, keepalives are all blocked. - Backends with per-loop resources (notably httpx.AsyncClient) may behave unexpectedly because the coroutine runs on a fresh loop in a worker thread; documents the typical failure signatures. - Long-term direction is an async variant on the Compactor protocol so callers can await natively. Docs-only; no behavior change. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * fix(compaction): make LLMSummarizeCompactor rendering loss-aware Reviewer flagged silent drops when rendering the slice fed to the summariser: - Message: append "[N image(s) attached]" / "[M document(s) attached]" markers; bytes not reproduced. - ModelOutputThunk: render "assistant called tools: name({args}), ..." for tool-call-only thunks (value=None) and "assistant: <empty>" for empty thunks. Eliminates "assistant: None". - Catch-all else: "<TypeName: content>" or "<TypeName>" instead of the default object repr (e.g. ReactInitiator when not pinned). Docstring gains a Note: that summaries are text-only and lossy for multimodal / heavy-tool sessions. Tests cover each branch. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * refactor(compaction): skip empty ModelOutputThunks instead of rendering "<empty>" A "<empty>" marker for thunks with neither value nor tool_calls tended to leak into the resulting summary verbatim. These turns carry no information worth summarising, so drop them from the rendered slice entirely. Test updated to assert no line is emitted for an empty thunk while neighbouring turns still come through. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * chore(compaction): set silence_context_type_warning on internal aact call Reviewer flagged that aact's context-type warning could be noisy under ThresholdCompactor-driven repeated compaction. Match react.py's pattern of setting silence_context_type_warning=True on internal framework calls so the warning stays quiet if the context argument is later changed to a non-SimpleContext, and to self-document the intent. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * feat(compaction): expose model_options on LLMSummarizeCompactor Reviewer flagged that the summary call uses the backend's default max_tokens (often 256-512 on local backends), silently truncating long summaries. react_summary_prompt(max_tokens_hint=N) is only a soft prompt- side nudge, not a real API parameter. Add model_options: dict | None to the constructor and forward it to mfuncs.aact so callers can set a hard token budget (or any other backend option). Default None preserves existing behaviour. Tests cover both forwarded and default paths. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * chore(docs): normalize source docstring backticks in context package Apply the same RST->Markdown backtick normalization upstream did in generative-computing#1155 (``symbol`` -> `symbol`) to the new mellea/stdlib/context/ package files that were authored on this branch in the old style. Mechanical sweep, no semantic changes. Files: - mellea/stdlib/context/__init__.py (4 lines) - mellea/stdlib/context/chat.py (36 lines) - mellea/stdlib/context/simple.py (4 lines) - mellea/stdlib/context/compactor.py (200 lines) Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * fix(docs): make custom_compactor.py example actually run TruncateOldest didn't inherit InlineCompactor, so the Pattern 1 demonstration crashed with TypeError when ChatContext's __init__ guard rejected it. The example pre-dated the InlineCompactor marker and was never updated. Inherit InlineCompactor on TruncateOldest, rewrite the module docstring to describe the two-pattern setup honestly (mentioning ThresholdCompactor as the wrap for backend-calling compactors), and move the "pure structural typing" demonstration to a fresh inline JustCompact class so it still proves the bare Compactor protocol works without inheritance. Also normalises backticks in this file (missed by the package-only sweep in 298ee27). Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * fix(compaction): re-raise programming errors instead of masking as backend failure The previous broad `except Exception` in LLMSummarizeCompactor.compact() caught both transient backend failures (the design intent) and programming errors like TypeError/AttributeError/AssertionError/LookupError. Genuine contract violations would be swallowed as a "warning logged, ctx unchanged" outcome and could go undetected for a long time. Add a narrow re-raise for the programming-error classes before the broad catch — the "best-effort retry" promise still holds for transient backend failures (RuntimeError, network, timeout, etc.), but real bugs surface immediately. LookupError covers KeyError and IndexError. Note section in the docstring is updated accordingly. Two regression tests added: - TypeError raised by the backend propagates out of compact() (existing RuntimeError test still passes — RuntimeError isn't in the re-raise set). - KeyboardInterrupt and other BaseException subclasses propagate unchanged across the _run_coro_blocking thread bridge — guards the docstring contract against future widening to `except BaseException`. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * docs(context): make ChatContext compaction docstring age-proof Drop "now" from the compaction notes in ChatContext — without an anchor it loses meaning over time. Also remove the "Earlier versions kept..." release-note framing; the undercount warning stands on its own. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * docs(context): drop caller list from _rebuild_chat_context docstring Reviewer pointed out that listing specific callers (WindowCompactor) rots — and indeed LLMSummarizeCompactor now uses it too. Drop the caller line; keep just what the function does and why (sidestep ChatContext.add to avoid compactor recursion). Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * refactor(compactor): hoist top-level imports out of function bodies Reviewer flagged the inline Message import in pin_system. Hoisted the imports that have no circular-dependency risk (Message, ToolMessage, MelleaLogger, _rebuild_chat_context, SimpleContext) to the top of the module. Left mfuncs lazy and added a comment explaining why: it imports SimpleContext via the context package init, which re-exports from this module — hoisting it would hit a partial-package ImportError. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * refactor(compactor): avoid wasted slice in LLMSummarizeCompactor.compact The body slice was only consumed by len() before being discarded (_async_compact recomputes its own body). Replace with body_len arithmetic, matching the existing pattern in WindowCompactor.compact. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * refactor(compactor): pass precomputed full/pin_end into _async_compact LLMSummarizeCompactor.compact() and _async_compact were independently calling ctx.as_list() and self.pin_predicate(full). Thread the values through instead so the async core skips the duplicate work. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * refactor(compactor): pythonic enumerate loop in pin_system Replace the while-with-index pattern in pin_system() with a cleaner enumerate-based early-return loop. Behavior unchanged. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * docs(react): note that compactor must return a ChatContext The Compactor protocol is generic in T bound to Context, so a custom compactor returning a non-chat Context (e.g. SimpleContext) type-checks at the react() call site but breaks the ChatContext assumptions the next aact call relies on. Document the constraint on the compactor= parameter so callers don't get surprised at runtime. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * docs(examples): note system-message eviction caveat in TruncateOldest Reviewer pointed out that items[1:] drops index 0 unconditionally, so a leading system message gets evicted on first compaction. Add a Note to the class docstring flagging the caveat and pointing readers at WindowCompactor + pin_system for the keep-the-system-prompt variant. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * docs(compactor): fix misleading "T = ChatContext" docstring framing Reviewer pointed out there's no class-level T = declaration — T is a per-call TypeVar on compact(). Reword to describe the override annotation directly: chat-only compactors override compact() with a ChatContext signature. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * docs(compactor): add Raises entry for InlineCompactor.compact Reviewer flagged the docstring gap: compact() raises NotImplementedError unconditionally on the marker base class, but the docstring had no Raises: section. Add one noting the always-raises behavior and that subclasses are expected to override. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * docs(compactor): add Args/Returns sections to satisfy doc quality gate The build-and-validate CI job's --fail-on-quality threshold flagged 8 issues across pin_nothing, pin_system, pin_system_and_initial_user, and InlineCompactor.compact (missing Args + Returns). Add the sections. Also collapse a multi-line _run_coro_blocking call back to one line so ruff format passes. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * test: override _generate_from_raw in compactor test fakes The Backend ABC's abstract method was renamed to _generate_from_raw (leading underscore) in generative-computing#1264; the public generate_from_raw is now a concrete hook-firing wrapper. The fakes in test_compactor.py overrode the wrapper but left the abstract method unimplemented, so each backend raised TypeError on instantiation. Rename generate_from_raw -> _generate_from_raw in FakeBackend, OtherBackend, BrokenBackend, BuggyBackend, and InterruptingBackend. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * test: assert compactor returns ChatContext in react loop Mirror the guard on the aact next_context returns: the LLMSummarizeCompactor docstring notes a non-ChatContext return would break the loop, but nothing enforced it. Add an isinstance assert after compact() to fail fast. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * fix(compaction): widen compactor types for upstream MOT as_list change as_list() now returns list[Component | CBlock | ModelOutputThunk]; align PinPredicate, pin_* helpers, and _async_compact. Update the react_compaction example backend to the new Backend API (widened _generate_from_context, override _generate_from_raw not the final generate_from_raw). Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> * docs(compaction): document window_size pin_system semantic change window_size=N maps to WindowCompactor (default pin_system), a deliberate change from the old raw last-N as_list(N) view: the system message is now pinned and size counts only the non-pinned body. Note the old drop-in (pin_nothing) in the docstring and add a CHANGELOG behavior-change entry covering both this and the as_list()-now-windowed move. Assisted-by: Claude Code Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> --------- Signed-off-by: Yousef El-Kurdi <yelkurdi@gmail.com> Signed-off-by: ramon-astudillo <ramon@astudillo.com> Co-authored-by: ramon-astudillo <ramon@astudillo.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Pull Request
Issue
Fixes #1183. Folds the additive half of #1218 Part 1.
Description
Brings the raw (batch) generation path in line with the chat path's
existing wrapper pattern:
Backend.generate_from_rawbecomes a@finalwrapper that owns hook firing, and backends implement only the new
_generate_from_rawabstract. The threegeneration_batch_*hooks nowfire from one place instead of being duplicated inline in all five
backends.
Reviewer call-outs:
Tuple return on
_generate_from_raw. Backends returntuple[list[ModelOutputThunk], dict | None]—(results, usage). Thewrapper unpacks the tuple, fires post_call with the aggregate, and
returns just
resultsto callers (public signature unchanged).Considered moving usage aggregation into the wrapper, but openai and
litellm only report whole-batch usage at the response root, so the
impl is the only place that knows the right shape. Matches the
existing
_generate_from_context -> tuple[MOT, Context]precedent.Standardized
self._model_idandself._providerinstead of newabstracts. The wrapper needs
modelandproviderfor hookpayloads (pre_call has no MOT yet; error may have no MOT). Considered
adding
_provider_name/_resolved_model_id()abstract members.Instead, finished the partial convention 3 of 5 backends already
used: every backend now sets
self._model_id: strandself._provider: strin__init__. Inlines model-id resolution intoollama and huggingface init bodies and deletes the now-unused
_get_*_model_idhelpers. RenamesLocalHFBackend._hf_model_idto_model_idfor consistency. ~10 chat-path setter sites switch toreading these attributes — no value changes, just one canonical
source per backend.
BaseExceptionon the raw-path error wrapper, matching chat. Thechat-path wrapper's
BaseExceptionwas added in refactor(telemetry)!: move backend tracing onto plugin/hook pattern #1181; the raw pathwas tentatively left at
except Exception. The same gap exists:synchronous
KeyboardInterrupt/asyncio.CancelledError/SystemExitinside the impl currently bypass the error hooksilently. This PR aligns to
BaseException. Behavior change: rawcancellation/interrupts now fire
generation_batch_errorbeforepropagating.
Batch pre_call payload mutations now propagate. When the batch
hooks were added in fix(backends): unify raw-path token usage on mot.generation.usage; guard eval_count=None #1218 they were wired for telemetry-only
observation: a plugin could mutate
model_options/format/tool_callson the pre_call payload, and the chat path would honorthose mutations, but the batch path silently dropped them. Same
plugin, same intent, different result depending on which API the
user reached. The wrapper now captures the post-hook payload and
reassigns the locals before calling
_generate_from_raw, identicalto the chat-path idiom. Adds
model_optionstoGenerationBatchPreCallPayload(it didn't exist on the batchpayload at all). Test inverted from
_are_not_propagatedto_propagates, mirroring the chat-path test.Folded the additive half of fix(backends): unify raw-path token usage on mot.generation.usage; guard eval_count=None #1218 Part 1. Backends with per-MOT
token counts (ollama, huggingface, watsonx) now also populate
mot.generation.usageper MOT. Backends with whole-batch-only usage(openai, litellm) leave per-MOT
mot.generation.usage = Noneandsurface the aggregate via the tuple return. Null-token policy:
all-or-nothing — if any of
prompt_tokens/completion_tokens/total_tokenscannot be determined for a MOT, that MOT'sgeneration.usagestaysNone(matchesdict | Nonetyping; ollamadocs document
Optional[int]as "not yet available" rather than"zero"). Watsonx's previous undocumented
or 0coerce switches tothe same policy.
The existing
mot._meta["usage"]writes are preserved everywhere —budget_forcing_alg.pystill reads from them. fix(backends): unify raw-path token usage on mot.generation.usage; guard eval_count=None #1218's remainingscope (consumer-side reads in
budget_forcing_alg.pyand themot._meta["usage"]deprecation path) stays open for follow-up.Originally folded in to support generic usage aggregation in the
wrapper; kept after the openai/litellm pivot to call-out 1's tuple
return because the per-MOT writes are still the right thing for the
3 backends whose APIs expose them.
Custom-backend doc updated.
docs/docs/community/building-extensions.mdwas teaching users tooverride the public
generate_from_context/generate_from_rawdirectly; those are
@finalnow. The example was updated toimplement
_generate_from_context/_generate_from_rawand set_model_id/_provider.Test coverage. Adds
TestGenerationBatchHookCallSitesintest/plugins/test_hook_call_sites.pymirroring the chat-pathfiring-site tests (9 tests).
_MockBackendextended to cover bothpaths with hardcoded behavior; error case via inline subclass per
the existing
RecordingBackend(_MockBackend)pattern. Mocks intest/stdlib/test_streaming.py,test/core/test_logger_plugin_hooks.py,and
test/stdlib/frameworks/test_react_framework.pyupdated to thenew abstract contract.
Testing
Attribution
Adding a new component, requirement, sampling strategy, or tool?
If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.
NOTE: Please ensure you have an issue that has been acknowledged by a core contributor and routed you to open a pull request against this repository. Otherwise, please open an issue before continuing with this pull request.