refactor(generation): decompose LLMGenerationActions into single-responsibility helpers#2168
Open
Pouyanpi wants to merge 19 commits into
Open
refactor(generation): decompose LLMGenerationActions into single-responsibility helpers#2168Pouyanpi wants to merge 19 commits into
Pouyanpi wants to merge 19 commits into
Conversation
github-actions
Bot
added
size: XL
status: needs triage
1 task
Phase 0 of the generation.py decomposition: a four-channel characterization gate that pins the observable behavior of generate_user_intent, generate_next_steps, generate_bot_message, generate_value and generate_intent_steps_message before any refactor. Channels: emitted internal-event sequence and per-call LLMCallInfo.task (from GenerationOptions log), the stop argument and stream mode (RecordingFakeLLM, the only channel that distinguishes the four Task.GENERAL call sites), and the streamed chunk sequence (stream_async). Covers dialog multi-call, embeddings -only, single-call (incl. the <<STREAMING[uid]>> handoff), passthrough/general, context-var bot intent, reasoning-trace packaging, and generate_value. Two findings from building this: safe_eval's error path is effectively unreachable, and the general-path streaming push emits no duplicate visible chunk.
Phase 1 of the LLMGenerationActions decomposition. The four Task.GENERAL LLM-call sites (the general user-intent fallback, the passthrough completion, the single-call general branch, and the passthrough bot-message branch) duplicated the same LLMCallInfo-set + llm_call + parse core. Extract it into a private _generate_general_response parameterized by prompt, stop, stream_during_call, llm_call_task and parse_task so each site keeps its exact behavior (including the passthrough bot-message branch reporting GENERATE_BOT_MESSAGE while parsing as GENERAL). Prompt rendering, chunk retrieval and output stripping stay at the call sites. Also extract the duplicated passthrough-fn return unpacking into _unpack_passthrough_output. Behavior preserved: the Phase 0 equivalence suite and full test suite pass unchanged.
Phase 2 of the LLMGenerationActions decomposition. generate_user_intent returned a UserIntent with user messages but a BotMessage/BotToolCalls turn without them; that polymorphism was inlined in one large if/else. Split it into two private helpers so the action reads as a dispatcher: - _detect_user_intent: the canonical-form path (embeddings-only lookup plus LLM-based intent generation). - _emit_general_bot_turn: the no-user-messages path (passthrough with/without a passthrough fn and the non-passthrough general path), including the push/reasoning-trace/tool-calls/BotMessage packaging and passthrough_output context event. generate_intent_steps_message's else-branch is intentionally NOT folded into _emit_general_bot_turn: its current behavior is a reduced subset (it does not emit BotThinking, BotToolCalls, or handle passthrough), so sharing the full helper would change behavior. Converging the two is left as a follow-up. Removed a redundant streaming_handler re-fetch that the parameterization made dead. Behavior preserved: the Phase 0 equivalence suite and full suite pass.
Phase 3 of the LLMGenerationActions decomposition. The single-call path passed its precomputed bot intent/message events between actions via free-form event["additional_info"]["bot_intent_event" | "bot_message_event"] dict access, an implicit cross-action contract. Introduce build_single_call_payload (producer, in generate_intent_steps_message) and read_single_call_payload returning a SingleCallPayload dataclass (consumers, in generate_next_steps and generate_bot_message) so the contract is explicit and typed. The wire format is unchanged: build still emits the same plain nested additional_info dict, so the event serializes byte-identically; the dataclass is only a transient read view. KeyError on a missing payload is preserved (guarding it is a follow-up). Behavior preserved: the Phase 0 equivalence suite and full suite pass.
Phase 4 of the LLMGenerationActions decomposition. The single-call streaming path coordinated with the bot-message phase through a module-global dict (local_streaming_handlers) and a hand-built '<<STREAMING[uid]>>' marker parsed with magic offsets (text[26:-4]). Wrap both behind _StreamingHandoffRegistry (register -> marker text, parse_marker -> uid|None, get -> handler), so the marker format and the handler store live in one place. Handlers are still kept for the module lifetime; the get accessor does not evict, preserving the existing handler leak (cleanup is a follow-up). Unify the two divergent set_pattern rules into _streaming_pattern_for: the include_bot_message_parser flag preserves generate_bot_message treating both verbose_v1 and bot_message as the verbose pattern while generate_intent_steps_message treats only verbose_v1 that way. A cast (no runtime effect) resolves a pre-existing Optional-handler type gap that the registry's typed accessor surfaced. Behavior preserved: the Phase 0 equivalence suite, streaming suites, and full suite pass.
…tion
Follow-up cleanup. Two unreachable branches:
- In _detect_user_intent, the `if user_intent is None` check after the
preceding `user_intent = "unknown message"` fallback could never be true.
- In generate_next_steps, next_step is only ever {"bot": ...}, so the
`next_step.get("execute")` branch (and the next_step dict itself) was dead.
Behavior-preserving: the equivalence suite and full suite pass unchanged.
generate_intent_steps_message passed event["text"] straight to the user-message index search and KB lookup. For multimodal input event["text"] is a list of content parts, not a string, so the single-call path mishandled it while generate_user_intent already normalized it. Join the text parts to a string up front, mirroring generate_user_intent, and use it for both the index search and the KB lookup. Adds an equivalence scenario for multimodal (list) content in single-call mode.
The single-call streaming handoff stored each inner StreamingHandler in a module-level registry but never removed it, so handlers accumulated for the process lifetime (a slow leak under sustained single-call streaming). Rename the registry accessor get -> take and pop the handler on retrieval; it is consumed exactly once by generate_bot_message, so eviction is safe. A unit test covers the marker round-trip and that a handler is gone after being taken.
…l path When single_call is enabled but there are no user messages, generate_intent_steps_message took a reduced general-response path: it rendered the GENERAL prompt with no context, called the LLM with no stop sequence, and emitted only a BotMessage. It dropped reasoning traces (BotThinking), tool calls (BotToolCalls), and passthrough handling that generate_user_intent's general path produces. Thread context into generate_intent_steps_message and delegate its else-branch to the shared _emit_general_bot_turn, so the no-user-messages general turn is identical regardless of single_call. Observable changes in this mode: the LLM call now uses stop=["User:"], and reasoning traces / tool calls / passthrough output are emitted. The rendered prompt is unchanged when relevant_chunks is empty. Updates the equivalence suite: the single-call general scenario now asserts the converged stop, and a new test covers reasoning-trace emission in this mode.
generate_bot_message repeated the same "optional BotThinking from the reasoning trace, then the final event" packaging in four return paths (the two single-call cache branches and the two final bot-message branches). Extract it into a module-level _bot_turn_output_events helper. The context_updates parameter preserves the existing difference: the generated bot-message path records bot_thinking on the context, the single-call cached path does not. Pure refactor; the full test suite passes unchanged.
_extract_user_message_example and _extract_bot_message_example repeated the same `asdict(spec_op.spec) if isinstance(spec_op.spec, Spec) else cast(Dict, ...)` normalization (the latter carried a TODO about the duplication). Extract it into a module-level _spec_op_as_dict helper. Pure refactor; full suite passes.
…message Move the single-call cache-consumption block (the streaming handoff and the direct cached-message paths) into _bot_message_from_single_call_cache, returning None to signal "fall back to regeneration". The `event` rebinding to the last user-intent event stays in generate_bot_message (with the None check that narrows it), preserving the existing fall-through behavior where the non-passthrough branch reads that rebound event. Pure refactor; the equivalence suite and full suite pass.
generate_intent_steps_message inlined ~80 lines of deeply nested logic to build the few-shot examples (search the user-message index, pair each candidate intent with a flow and its bot message). Extract it into _build_intent_steps_examples, returning (examples, potential_user_intents). Pure refactor; the equivalence suite and full suite pass.
Use explicit is-not-None guards instead of the (opts and opts.llm_params) or {}
idiom, which ty types as GenerationOptions | dict and rejects against the
_generate_general_response llm_params: Optional[dict] parameter. Behavior is
unchanged and matches the existing guard used in the generate_value path.
Pouyanpi
force-pushed
the
pouyanpi/improve-generate-colang
branch
from
July 13, 2026 17:38
b31fe12 to
5f4c248
Compare
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Pouyanpi
marked this pull request as ready for review
July 14, 2026 06:02
Pouyanpi
added
status: triaged
Contributor
Greptile SummaryThis PR decomposes the Colang 1.0 dialog generation actions into smaller helpers. The main changes are:
|
| Filename | Overview |
|---|---|
| nemoguardrails/actions/llm/generation.py | Refactors dialog generation into helpers and updates single-call cache fallback behavior. |
| tests/test_generation_equivalence.py | Adds behavioral coverage for generation events, call metadata, stop arguments, and streaming. |
Reviews (5): Last reviewed commit: "test(recorded): add dialog bot-message g..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Contributor
📝 WalkthroughWalkthroughLLM generation actions now use shared response, intent, bot-turn, payload, and streaming-handoff helpers. Single-call and general paths are covered by a new behavioral equivalence test harness spanning dialog, embeddings, streaming, passthrough, reasoning, and value generation. ChangesLLM generation refactor
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant IntentSteps as generate_intent_steps_message
participant Registry as _StreamingHandoffRegistry
participant BotMessage as generate_bot_message
IntentSteps->>Registry: register inner streaming handler
IntentSteps->>BotMessage: cache encoded streaming marker
BotMessage->>Registry: parse and take handler
Registry->>BotMessage: pipe inner stream to main handler
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@nemoguardrails/actions/llm/generation.py`:
- Around line 990-998: Update the non-passthrough cache-miss fallthrough in
generate_bot_message to search bot_message_index using the preserved bot_intent
value instead of the rebound event["intent"]. Keep the UserIntent event binding
for _bot_message_from_single_call_cache unchanged.
- Around line 886-937: The cached bot intent mismatch path in
_bot_message_from_single_call_cache must consume any embedded streaming handoff
before falling back. Parse bot_message_event["text"] and call
_streaming_handoff.take() when a marker exists before returning None, while
preserving the existing match validation and normal streaming behavior.
In `@tests/test_generation_equivalence.py`:
- Around line 408-440: Add a non-text multimodal part, specifically an image_url
item, alongside the existing text item in the messages passed to
chat.app.generate within test_single_call_multimodal_input. Keep the text
content and existing assertions unchanged so the test verifies image parts are
ignored while text normalization still produces the same intent and response
sequence.
- Around line 121-122: Add a single-call general-path case to the generation
equivalence tests and assert the complete emitted BotToolCalls payload,
including each tool call’s ID, function name, and arguments, rather than only
the count produced by the event_type == "BotToolCalls" handling. Keep
passthrough coverage intact and verify the public response shape exactly.
- Around line 500-516: Extend
test_streaming_handoff_registry_round_trip_and_eviction to cover malformed
streaming markers, including a marker with the expected prefix but missing or
trailing suffix, and assert _StreamingHandoffRegistry.parse_marker() returns
None for each. Preserve the existing valid round-trip, non-marker, and eviction
assertions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: bbda0e00-200f-4c74-b33d-d3e70a427b27
📒 Files selected for processing (2)
nemoguardrails/actions/llm/generation.pytests/test_generation_equivalence.py
…treaming handoff - generate_bot_message: on a single-call cache miss, search bot messages with the bot intent instead of the rebound user intent. - _bot_message_from_single_call_cache: evict an embedded streaming handoff on the cache-miss fall-through so the registered handler is not leaked. - _StreamingHandoffRegistry.parse_marker: require the closing suffix so malformed markers are rejected instead of returning a bad uid. Tests: reorganize the equivalence suite into classes and cover the single-call cache miss (search intent + handoff eviction), multimodal non-text parts, the full BotToolCalls payload, and malformed streaming markers.
generate_bot_message only released a single-call streaming handoff on the cache path; the predefined-message and $context_var branches returned first, leaking the registry entry (and stranding a background llm_call). Add _discard_single_call_handoff on both bypass branches so the handler is released regardless of which branch phase 3 takes, restoring the registry invariant.
- test the streaming-handoff release on the predefined-message bypass path - add call_prompts assertions to the dialog and single-call tests so a dropped render-context key is no longer hidden by the prompt-independent fake - add direct contract tests for _streaming_pattern_for, _bot_turn_output_events and _build_intent_steps_examples (prompt-only, unobservable end-to-end) - cover generate_next_steps single-step cleanup, the passthrough raw_llm_request list branch, the _generate_general_response task-vs-parse divergence, and the non-single-call multimodal path
…anches - TestSingleCall: the single-call cache returns None (falls back) when the last user-intent event is a UserMessage rather than a UserIntent - TestGeneralAndPassthrough: user_messages + passthrough drives generate_bot_message's passthrough branch, which reports the generate_bot_message task while parsing GENERAL The multi-step parse-error reduce loop is left uncovered by design: Colang 1.0's parser is lenient enough that the parse failure almost never fires.
A dialog_generation recorded config whose intent maps (via a flow) to a bot intent with no predefined message, so generate_bot_message calls the LLM with few-shot examples + retrieved kb chunks. VCR matches on recorded_body, so the cassette replays only if all three dialog prompts (intent, next-step, bot-message) are byte-identical -- a prompt-fidelity equivalence check the prompt-independent FakeLLMModel oracle cannot make.
Pouyanpi
force-pushed
the
pouyanpi/improve-generate-colang
branch
from
July 14, 2026 13:57
7b7d706 to
28f0440
Compare
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
1 participant
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.
Description
Decompose
LLMGenerationActions(the Colang 1.0 dialog generation actions innemoguardrails/actions/llm/generation.py).Today
generate_user_intent,generate_bot_messageandgenerate_intent_steps_messageeach mix several responsibilities, and the sameTask.GENERAL"general response" logic is duplicated across four call sites withsubtle divergences. The result is that you cannot tell, from a config alone,
which code path a request will take (dialog vs single-call vs passthrough vs
general), and the passthrough/tool-call path is easy to break silently.
This PR pins the observable behavior first, then decomposes behind that safety
net:
tests/test_generation_equivalence.py).It captures four channels per scenario that an event-only check would miss:
the emitted internal-event sequence, the per-call
LLMCallInfo.task, thestopargument / stream mode (via a recording fake LLM), and the streamedchunk sequence. This is the gate every later commit keeps green.
_generate_general_responsehelper for the fourTask.GENERALsites, a splitof intent detection from the general bot turn (
_emit_general_bot_turn), atyped single-call cache seam, and a
_StreamingHandoffRegistrythatencapsulates the
<<STREAMING[uid]>>handoff.refactor): normalize multimodal user input in the single-call path; evict
streaming-handoff handlers when consumed (fixes a module-lifetime leak);
converge the single-call no-user-messages general turn onto the multi-call
path (it now emits
BotThinking/ tool calls / passthrough output and usesstop=["User:"]); and remove genuinely dead branches.This also unblocks per-request OpenAI
toolssupport for non-passthrough configs(#2057): the extracted
_generate_general_response/_emit_general_bot_turnseams are exactly what that feature builds on.
The change is a curated, gradual commit series (harness first, then one
transformation per commit) intended to be reviewed commit-by-commit.
Related Issue(s)
toolsfor non-passthrough configs)AI Assistance
Checklist
Summary by CodeRabbit
Bug Fixes
Tests