diff --git a/nemoguardrails/actions/llm/generation.py b/nemoguardrails/actions/llm/generation.py index 7e5d1a8484..58a7dcaa2b 100644 --- a/nemoguardrails/actions/llm/generation.py +++ b/nemoguardrails/actions/llm/generation.py @@ -20,10 +20,10 @@ import random import re import sys -from dataclasses import asdict +from dataclasses import asdict, dataclass from functools import lru_cache from time import time -from typing import Any, Awaitable, Callable, Dict, List, Optional, cast +from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, cast from jinja2 import meta from jinja2.sandbox import SandboxedEnvironment @@ -70,7 +70,91 @@ log = logging.getLogger(__name__) -local_streaming_handlers = {} +class _StreamingHandoffRegistry: + """Encodes the ``<>`` handoff used by the single-call path. + + The single-call streaming path generates the bot message on an inner + ``StreamingHandler``, registers it here and leaves a marker on the cached bot + message. The later bot-message phase parses the marker, takes the handler and + pipes it to the main handler. Each handler is removed when taken, so handlers + do not accumulate for the module lifetime. + """ + + _MARKER_PREFIX = 'Bot message: "< str: + """Register an inner handler and return the marker text for it.""" + self._handlers[handler.uid] = handler + return f"{self._MARKER_PREFIX}{handler.uid}{self._MARKER_SUFFIX}" + + def parse_marker(self, text: str) -> Optional[str]: + """Return the handler uid encoded in ``text``, or None if not a marker.""" + if text.startswith(self._MARKER_PREFIX) and text.endswith(self._MARKER_SUFFIX): + return text[len(self._MARKER_PREFIX) : -len(self._MARKER_SUFFIX)] + return None + + def take(self, uid: str) -> "StreamingHandler": + """Return the registered handler for ``uid`` and remove it from the registry.""" + return self._handlers.pop(uid) + + +_streaming_handoff = _StreamingHandoffRegistry() + + +def _streaming_pattern_for(output_parser, *, include_bot_message_parser: bool): + """Return the ``(prefix, suffix)`` streaming pattern for the bot message. + + The two call sites differ: ``generate_bot_message`` treats both ``verbose_v1`` + and ``bot_message`` as the verbose pattern, while + ``generate_intent_steps_message`` treats only ``verbose_v1`` that way; + ``include_bot_message_parser`` selects which rule applies. + """ + verbose = output_parser == "verbose_v1" or (include_bot_message_parser and output_parser == "bot_message") + if verbose: + return 'Bot message: "', '"' + return ' "', '"' + + +@dataclass +class SingleCallPayload: + """The bot intent/message events computed by the single-call generation. + + Carried on the ``additional_info`` of the ``UserIntent`` event so that the + later ``generate_next_steps`` and ``generate_bot_message`` phases can reuse + the results of the single LLM call instead of calling the LLM again. + """ + + bot_intent_event: dict + bot_message_event: dict + + +def build_single_call_payload(bot_intent_event: dict, bot_message_event: dict) -> dict: + """Build the single-call cache carried on the ``UserIntent`` event. + + Returns the plain nested ``additional_info`` dict (not a dataclass) so the + event serializes byte-identically as it passes between actions. + """ + return { + "bot_intent_event": bot_intent_event, + "bot_message_event": bot_message_event, + } + + +def read_single_call_payload(event: dict) -> SingleCallPayload: + """Read the single-call cache back from a ``UserIntent`` event. + + Mirrors the existing direct dict access (raising ``KeyError`` when the + payload is absent); guarding that is a separate follow-up. + """ + additional_info = event["additional_info"] + return SingleCallPayload( + bot_intent_event=additional_info["bot_intent_event"], + bot_message_event=additional_info["bot_message_event"], + ) class LLMGenerationActions: @@ -138,11 +222,8 @@ def _extract_user_message_example(self, flow: Flow) -> None: self.user_messages[flow.name] = [message] elif spec_op.op == "await": - # The SpecOp.spec type is Union[Spec, dict]. Need to convert to Dict to have `elements` field - # which isn't in the Spec definition - await_spec_dict: Dict[str, Any] = ( - asdict(spec_op.spec) if isinstance(spec_op.spec, Spec) else cast(Dict, spec_op.spec) - ) + # Convert to Dict to have the `elements` field, which isn't in the Spec definition. + await_spec_dict: Dict[str, Any] = _spec_op_as_dict(spec_op) if isinstance(await_spec_dict, dict) and await_spec_dict.get("_type") == "spec_or": specs = await_spec_dict.get("elements", None) @@ -171,11 +252,7 @@ def _extract_bot_message_example(self, flow: Flow): return spec_op: SpecOp = el - spec: Dict[str, Any] = ( - asdict(spec_op.spec) # TODO! Refactor this function as it's duplicated in many places - if isinstance(spec_op.spec, Spec) - else cast(Dict, spec_op.spec) - ) + spec: Dict[str, Any] = _spec_op_as_dict(spec_op) if spec.get("_type") in ("spec_or", "spec_and"): return @@ -347,235 +424,303 @@ def _get_sample_conversation_two_turns(self): return sample_conversation - @action(is_system_action=True) - async def generate_user_intent( + async def _generate_general_response( self, - events: List[dict], - context: dict, - config: RailsConfig, - llm: Optional[LLMModel] = None, - kb: Optional[KnowledgeBase] = None, - ): - """Generate the canonical form for what the user said i.e. user intent.""" - # If using a single LLM call, use the specific action defined for this task. - if self.config.rails.dialog.single_call.enabled: - return await self.generate_intent_steps_message(events=events, llm=llm, kb=kb) - # The last event should be the "StartInternalSystemAction" and the one before it the "UtteranceUserActionFinished". - event = get_last_user_utterance_event(events) - if not event: - raise ValueError("No user message found in event stream. Unable to generate user intent.") - if event["type"] != "UserMessage": - raise ValueError( - f"Expected UserMessage event, but found {event['type']}. " - "Cannot generate user intent from this event type." + *, + generation_llm: Optional[LLMModel], + prompt, + streaming_handler: Optional[StreamingHandler] = None, + stream_during_call: bool = False, + stop: Optional[List[str]] = None, + llm_call_task: Task = Task.GENERAL, + parse_task: Task = Task.GENERAL, + llm_params: Optional[dict] = None, + ) -> str: + """Make a single general-response LLM call and parse its output. + + This is the duplicated core of the four ``Task.GENERAL`` call sites: the + general user-intent fallback, the passthrough completion, the single-call + general branch and the passthrough bot-message branch. Prompt rendering, + chunk retrieval and any output stripping stay at the call sites; only the + ``LLMCallInfo`` set, the ``llm_call`` and the parse are shared here. The + ``stop``, ``stream_during_call``, ``llm_call_task`` and ``parse_task`` + parameters preserve the per-site differences (e.g. the passthrough + bot-message branch reports ``GENERATE_BOT_MESSAGE`` but parses as + ``GENERAL``). + """ + llm_call_info_var.set(LLMCallInfo(task=llm_call_task.value)) + + result = ( + await llm_call( + generation_llm, + prompt, + streaming_handler=streaming_handler if stream_during_call else None, + stop=stop, + llm_params=llm_params, ) + ).content - # Use action specific llm if registered else fallback to main llm - # This can be None as some code-paths use embedding lookups rather than LLM generation - generation_llm: Optional[LLMModel] = llm if llm else self.llm + return self.llm_task_manager.parse_task_output(parse_task, output=result) - streaming_handler = streaming_handler_var.get() + async def _detect_user_intent( + self, + events: List[dict], + event: dict, + config: RailsConfig, + generation_llm: Optional[LLMModel], + ) -> ActionResult: + """Detect the canonical form (user intent) for the given user message. - await self._ensure_user_message_index() + Covers both the embeddings-only lookup and the LLM-based canonical-form + generation. Always returns a single ``UserIntent`` event. + """ + # TODO: based on the config we can use a specific canonical forms model + # or use the LLM to detect the canonical form. The below implementation + # is for the latter. - # TODO: check for an explicit way of enabling the canonical form detection + log.info("Phase 1 :: Generating user intent") - if self.user_messages: - # TODO: based on the config we can use a specific canonical forms model - # or use the LLM to detect the canonical form. The below implementation - # is for the latter. + # We search for the most relevant similar user utterance + examples = "" + potential_user_intents = [] + if isinstance(event["text"], list): + text = " ".join([item["text"] for item in event["text"] if item["type"] == "text"]) + else: + text = event["text"] - log.info("Phase 1 :: Generating user intent") + if self.user_message_index is not None: + if config.rails.dialog.user_messages and config.rails.dialog.user_messages.embeddings_only: + threshold = config.rails.dialog.user_messages.embeddings_only_similarity_threshold + results = await self.user_message_index.search(text=text, max_results=5, threshold=threshold) - # We search for the most relevant similar user utterance - examples = "" - potential_user_intents = [] - if isinstance(event["text"], list): - text = " ".join([item["text"] for item in event["text"] if item["type"] == "text"]) - else: - text = event["text"] + if results: + intent = results[0].meta["intent"] + return ActionResult(events=[new_event_dict("UserIntent", intent=intent)]) + elif config.rails.dialog.user_messages.embeddings_only_fallback_intent: + intent = config.rails.dialog.user_messages.embeddings_only_fallback_intent + return ActionResult(events=[new_event_dict("UserIntent", intent=intent)]) - if self.user_message_index is not None: - if config.rails.dialog.user_messages and config.rails.dialog.user_messages.embeddings_only: - threshold = config.rails.dialog.user_messages.embeddings_only_similarity_threshold - results = await self.user_message_index.search(text=text, max_results=5, threshold=threshold) + results = await self.user_message_index.search(text=text, max_results=5, threshold=None) + # We add these in reverse order so the most relevant is towards the end. + for result in reversed(results): + examples += f'user "{result.text}"\n {result.meta["intent"]}\n\n' + if result.meta["intent"] not in potential_user_intents: + potential_user_intents.append(result.meta["intent"]) - if results: - intent = results[0].meta["intent"] - return ActionResult(events=[new_event_dict("UserIntent", intent=intent)]) - elif config.rails.dialog.user_messages.embeddings_only_fallback_intent: - intent = config.rails.dialog.user_messages.embeddings_only_fallback_intent - return ActionResult(events=[new_event_dict("UserIntent", intent=intent)]) + prompt = self.llm_task_manager.render_task_prompt( + task=Task.GENERATE_USER_INTENT, + events=events, + context={ + "examples": examples, + "potential_user_intents": ", ".join(potential_user_intents), + }, + ) - results = await self.user_message_index.search(text=text, max_results=5, threshold=None) - # We add these in reverse order so the most relevant is towards the end. - for result in reversed(results): - examples += f'user "{result.text}"\n {result.meta["intent"]}\n\n' - if result.meta["intent"] not in potential_user_intents: - potential_user_intents.append(result.meta["intent"]) + # Initialize the LLMCallInfo object + llm_call_info_var.set(LLMCallInfo(task=Task.GENERATE_USER_INTENT.value)) - prompt = self.llm_task_manager.render_task_prompt( - task=Task.GENERATE_USER_INTENT, - events=events, - context={ - "examples": examples, - "potential_user_intents": ", ".join(potential_user_intents), - }, + # We make this call with temperature 0 to have it as deterministic as possible. + result = ( + await llm_call( + generation_llm, + prompt, + llm_params={"temperature": self.config.lowest_temperature}, ) + ).content - # Initialize the LLMCallInfo object - llm_call_info_var.set(LLMCallInfo(task=Task.GENERATE_USER_INTENT.value)) - - # We make this call with temperature 0 to have it as deterministic as possible. - result = ( - await llm_call( - generation_llm, - prompt, - llm_params={"temperature": self.config.lowest_temperature}, - ) - ).content + # Parse the output using the associated parser + result = self.llm_task_manager.parse_task_output(Task.GENERATE_USER_INTENT, output=result) - # Parse the output using the associated parser - result = self.llm_task_manager.parse_task_output(Task.GENERATE_USER_INTENT, output=result) + user_intent = get_first_nonempty_line(result) + if user_intent is None: + user_intent = "unknown message" - user_intent = get_first_nonempty_line(result) - if user_intent is None: - user_intent = "unknown message" + if user_intent.startswith("user "): + user_intent = user_intent[5:] - if user_intent and user_intent.startswith("user "): - user_intent = user_intent[5:] + log.info("Canonical form for user intent: " + user_intent) - log.info("Canonical form for user intent: " + (user_intent if user_intent else "None")) + return ActionResult(events=[new_event_dict("UserIntent", intent=user_intent)]) - if user_intent is None: - return ActionResult(events=[new_event_dict("UserIntent", intent="unknown message")]) - else: - return ActionResult(events=[new_event_dict("UserIntent", intent=user_intent)]) - else: - output_events = [] - context_updates = {} + async def _emit_general_bot_turn( + self, + events: List[dict], + context: dict, + event: dict, + generation_llm: Optional[LLMModel], + streaming_handler: Optional[StreamingHandler], + kb: Optional[KnowledgeBase], + ) -> ActionResult: + """Generate and package a general bot turn when there are no user messages. + + Handles the passthrough (with or without a passthrough fn) and the + non-passthrough general paths, then packages the result into the + BotMessage/BotToolCalls/BotThinking events and context updates that + ``generate_user_intent`` returns in this mode. + """ + output_events = [] + context_updates = {} - # If we are in passthrough mode, we just use the input for prompting - if self.config.passthrough: - # We check if we have a raw request. If the guardrails API is using - # the `generate_events` API, this will not be set. - raw_prompt = raw_llm_request.get() + # If we are in passthrough mode, we just use the input for prompting + if self.config.passthrough: + # We check if we have a raw request. If the guardrails API is using + # the `generate_events` API, this will not be set. + raw_prompt = raw_llm_request.get() - if raw_prompt is None: + if raw_prompt is None: + prompt = event["text"] + else: + if isinstance(raw_prompt, str): + # If we're in completion mode, we use directly the last $user_message + # as it may have been altered by the input rails. prompt = event["text"] + elif isinstance(raw_prompt, list): + prompt = raw_prompt.copy() + + # In this case, if the last message is from the user, we replace the text + # just in case the input rails may have altered it. + if prompt[-1]["role"] == "user": + raw_prompt[-1]["content"] = event["text"] else: - if isinstance(raw_prompt, str): - # If we're in completion mode, we use directly the last $user_message - # as it may have been altered by the input rails. - prompt = event["text"] - elif isinstance(raw_prompt, list): - prompt = raw_prompt.copy() + raise ValueError(f"Unsupported type for raw prompt: {type(raw_prompt)}") - # In this case, if the last message is from the user, we replace the text - # just in case the input rails may have altered it. - if prompt[-1]["role"] == "user": - raw_prompt[-1]["content"] = event["text"] - else: - raise ValueError(f"Unsupported type for raw prompt: {type(raw_prompt)}") + if self._passthrough_fn: + raw_output = await self._passthrough_fn(context=context, events=events) + text, passthrough_output = _unpack_passthrough_output(raw_output) - if self._passthrough_fn: - raw_output = await self._passthrough_fn(context=context, events=events) + # We record the passthrough output in the context + output_events.append( + new_event_dict( + "ContextUpdate", + data={"passthrough_output": passthrough_output}, + ) + ) + else: + gen_options: Optional[GenerationOptions] = generation_options_var.get() - # If the passthrough action returns a single value, we consider that - # to be the text output - if isinstance(raw_output, tuple) or isinstance(raw_output, list): - text, passthrough_output = raw_output[0], raw_output[1] - else: - text = raw_output - passthrough_output = None + llm_params = ( + gen_options.llm_params if gen_options is not None and gen_options.llm_params is not None else {} + ) - # We record the passthrough output in the context - output_events.append( - new_event_dict( - "ContextUpdate", - data={"passthrough_output": passthrough_output}, - ) - ) - else: - # Initialize the LLMCallInfo object - llm_call_info_var.set(LLMCallInfo(task=Task.GENERAL.value)) + text = await self._generate_general_response( + generation_llm=generation_llm, + prompt=prompt, + streaming_handler=streaming_handler, + stream_during_call=True, + llm_params=llm_params, + ) - gen_options: Optional[GenerationOptions] = generation_options_var.get() + else: + if kb: + chunks = await kb.search_relevant_chunks(event["text"]) + relevant_chunks = "\n".join([chunk["body"] for chunk in chunks]) + else: + # in case there is no user flow (user message) then we need the context update to work for relevant_chunks + relevant_chunks = get_retrieved_relevant_chunks(events, skip_user_message=True) - llm_params = (gen_options and gen_options.llm_params) or {} + # Otherwise, we still create an altered prompt. + prompt = self.llm_task_manager.render_task_prompt( + task=Task.GENERAL, + events=events, + context={"relevant_chunks": relevant_chunks}, + ) - streaming_handler: Optional[StreamingHandler] = streaming_handler_var.get() + generation_options: Optional[GenerationOptions] = generation_options_var.get() + llm_params = ( + generation_options.llm_params + if generation_options is not None and generation_options.llm_params is not None + else {} + ) - text = ( - await llm_call( - generation_llm, - prompt, - streaming_handler=streaming_handler, - llm_params=llm_params, - ) - ).content - text = self.llm_task_manager.parse_task_output(Task.GENERAL, output=text) + text = await self._generate_general_response( + generation_llm=generation_llm, + prompt=prompt, + streaming_handler=streaming_handler, + stream_during_call=True, + stop=["User:"], + llm_params=llm_params, + ) + text = text.strip() + if text.startswith('"'): + text = text[1:-1] - else: - # Initialize the LLMCallInfo object - llm_call_info_var.set(LLMCallInfo(task=Task.GENERAL.value)) + # In streaming mode, we also push this. + if streaming_handler: + await streaming_handler.push_chunk(text) - if kb: - chunks = await kb.search_relevant_chunks(event["text"]) - relevant_chunks = "\n".join([chunk["body"] for chunk in chunks]) - else: - # in case there is no user flow (user message) then we need the context update to work for relevant_chunks - relevant_chunks = get_retrieved_relevant_chunks(events, skip_user_message=True) + reasoning_trace = get_and_clear_reasoning_trace_contextvar() + if reasoning_trace: + context_updates["bot_thinking"] = reasoning_trace + output_events.append(new_event_dict("BotThinking", content=reasoning_trace)) - # Otherwise, we still create an altered prompt. - prompt = self.llm_task_manager.render_task_prompt( - task=Task.GENERAL, - events=events, - context={"relevant_chunks": relevant_chunks}, - ) + if self.config.passthrough: + from nemoguardrails.actions.llm.utils import ( + get_and_clear_tool_calls_contextvar, + ) - generation_options: Optional[GenerationOptions] = generation_options_var.get() - llm_params = (generation_options and generation_options.llm_params) or {} + tool_calls = get_and_clear_tool_calls_contextvar() - result = ( - await llm_call( - generation_llm, - prompt, - streaming_handler=streaming_handler, - stop=["User:"], - llm_params=llm_params, - ) - ).content + if tool_calls: + output_events.append(new_event_dict("BotToolCalls", tool_calls=tool_calls)) + else: + output_events.append(new_event_dict("BotMessage", text=text)) + else: + output_events.append(new_event_dict("BotMessage", text=text)) - text = self.llm_task_manager.parse_task_output(Task.GENERAL, output=result) - text = text.strip() - if text.startswith('"'): - text = text[1:-1] + return ActionResult(events=output_events, context_updates=context_updates) - # In streaming mode, we also push this. - if streaming_handler: - await streaming_handler.push_chunk(text) + @action(is_system_action=True) + async def generate_user_intent( + self, + events: List[dict], + context: dict, + config: RailsConfig, + llm: Optional[LLMModel] = None, + kb: Optional[KnowledgeBase] = None, + ): + """Generate the canonical form for what the user said i.e. user intent.""" + # If using a single LLM call, use the specific action defined for this task. + if self.config.rails.dialog.single_call.enabled: + return await self.generate_intent_steps_message(events=events, context=context, llm=llm, kb=kb) + # The last event should be the "StartInternalSystemAction" and the one before it the "UtteranceUserActionFinished". + event = get_last_user_utterance_event(events) + if not event: + raise ValueError("No user message found in event stream. Unable to generate user intent.") + if event["type"] != "UserMessage": + raise ValueError( + f"Expected UserMessage event, but found {event['type']}. " + "Cannot generate user intent from this event type." + ) - reasoning_trace = get_and_clear_reasoning_trace_contextvar() - if reasoning_trace: - context_updates["bot_thinking"] = reasoning_trace - output_events.append(new_event_dict("BotThinking", content=reasoning_trace)) + # Use action specific llm if registered else fallback to main llm + # This can be None as some code-paths use embedding lookups rather than LLM generation + generation_llm: Optional[LLMModel] = llm if llm else self.llm - if self.config.passthrough: - from nemoguardrails.actions.llm.utils import ( - get_and_clear_tool_calls_contextvar, - ) + streaming_handler = streaming_handler_var.get() - tool_calls = get_and_clear_tool_calls_contextvar() + await self._ensure_user_message_index() - if tool_calls: - output_events.append(new_event_dict("BotToolCalls", tool_calls=tool_calls)) - else: - output_events.append(new_event_dict("BotMessage", text=text)) - else: - output_events.append(new_event_dict("BotMessage", text=text)) + # TODO: check for an explicit way of enabling the canonical form detection - return ActionResult(events=output_events, context_updates=context_updates) + # With user messages we detect the canonical form (user intent); without + # them we fall back to generating a general bot turn directly. The two + # paths are split into helpers so the polymorphic return is legible. + if self.user_messages: + return await self._detect_user_intent( + events=events, + event=event, + config=config, + generation_llm=generation_llm, + ) + else: + return await self._emit_general_bot_turn( + events=events, + context=context, + event=event, + generation_llm=generation_llm, + streaming_handler=streaming_handler, + kb=kb, + ) async def _search_flows_index(self, text, max_results): """Search the index of flows.""" @@ -616,7 +761,7 @@ async def generate_next_steps(self, events: List[dict], llm: Optional[LLMModel] if event["type"] == "UserIntent": # If using a single LLM call, use the results computed in the first call. if self.config.rails.dialog.single_call.enabled: - bot_intent_event = event["additional_info"]["bot_intent_event"] + bot_intent_event = read_single_call_payload(event).bot_intent_event return ActionResult(events=[bot_intent_event]) user_intent = event["intent"] @@ -668,25 +813,10 @@ async def generate_next_steps(self, events: List[dict], llm: Optional[LLMModel] # Also, sometimes, there's a comma and more content if "," in bot_intent: bot_intent = bot_intent.split(",")[0].strip() - - next_step = {"bot": bot_intent} else: - next_step = {"bot": "general response"} - - # If we have to execute an action, we return the event to start it - if next_step.get("execute"): - return ActionResult( - events=[ - new_event_dict( - "StartInternalSystemAction", - action_name=next_step["execute"], - ) - ] - ) - else: - bot_intent = next_step.get("bot") + bot_intent = "general response" - return ActionResult(events=[new_event_dict("BotIntent", intent=bot_intent)]) + return ActionResult(events=[new_event_dict("BotIntent", intent=bot_intent)]) else: # Otherwise, we parse the output as a single flow. # If we have a parsing error, we try to reduce size of the flow, potentially @@ -753,6 +883,82 @@ def _render_string( return template.render(render_context) + async def _bot_message_from_single_call_cache( + self, + user_intent_event: dict, + events: List[dict], + streaming_handler: Optional[StreamingHandler], + ) -> Optional[ActionResult]: + """Return the bot message cached by the single LLM call, or None to fall back. + + Returns None when the cache cannot be used -- the last user-intent event + is not a ``UserIntent``, or its cached bot intent does not match the bot + intent now being generated -- in which case the caller regenerates. + """ + if user_intent_event["type"] != "UserIntent": + return None + + payload = read_single_call_payload(user_intent_event) + bot_message_event = payload.bot_message_event + + # We only need to use the bot message if it corresponds to the + # generate bot intent as well. + last_bot_intent = get_last_bot_intent_event(events) + if not last_bot_intent: + raise RuntimeError("No last bot intent found to generate bot message") + if last_bot_intent["intent"] != payload.bot_intent_event["intent"]: + # If the cached message embedded a streaming handoff, evict it so the + # registered handler does not leak when we fall back to regeneration. + streaming_handler_uid = _streaming_handoff.parse_marker(bot_message_event["text"]) + if streaming_handler_uid is not None: + _streaming_handoff.take(streaming_handler_uid) + return None + + text = bot_message_event["text"] + # If the bot message is being generated in streaming mode + streaming_handler_uid = _streaming_handoff.parse_marker(text) + if streaming_handler_uid is not None: + _streaming_handler = _streaming_handoff.take(streaming_handler_uid) + + # We pipe the content from this handler to the main one. + # The marker is only present when generation streamed, + # so the main handler is set here. + _streaming_handler.set_pipe_to(cast(StreamingHandler, streaming_handler)) + await _streaming_handler.disable_buffering() + + # And wait for it to finish. + # We stop after the closing double quotes for the bot message. + _streaming_handler.stop = [ + '"\n', + ] + text = await _streaming_handler.wait() + + return ActionResult(events=_bot_turn_output_events(new_event_dict("BotMessage", text=text))) + + if streaming_handler: + await streaming_handler.push_chunk(bot_message_event["text"]) + + return ActionResult(events=_bot_turn_output_events(bot_message_event)) + + def _discard_single_call_handoff(self, events: List[dict]) -> None: + """Evict a pending single-call streaming handoff for a bypassed bot message. + + When ``generate_bot_message`` short-circuits to a predefined message or a + ``$context_var`` it never consumes the single-call cache, so a handoff + registered while streaming the (now unused) cached bot message would leak. + This releases it. Safe to call when there is no pending handoff. + """ + user_intent_event = get_last_user_intent_event(events) + if not user_intent_event or user_intent_event.get("type") != "UserIntent": + return + additional_info = user_intent_event.get("additional_info") or {} + bot_message_event = additional_info.get("bot_message_event") + if not bot_message_event: + return + streaming_handler_uid = _streaming_handoff.parse_marker(bot_message_event.get("text", "")) + if streaming_handler_uid is not None: + _streaming_handoff.take(streaming_handler_uid) + @action(is_system_action=True) async def generate_bot_message(self, events: List[dict], context: dict, llm: Optional[LLMModel] = None): """Generate a bot message based on the desired bot intent.""" @@ -796,66 +1002,29 @@ async def generate_bot_message(self, events: List[dict], context: dict, llm: Opt # We skip output rails for predefined messages. context_updates["skip_output_rails"] = True + if self.config.rails.dialog.single_call.enabled: + self._discard_single_call_handoff(events) + # Check if the output is supposed to be the content of a context variable elif bot_intent and bot_intent[0] == "$" and bot_intent[1:] in context: bot_utterance = context[bot_intent[1:]] + if self.config.rails.dialog.single_call.enabled: + self._discard_single_call_handoff(events) + else: # Generate the bot message using an LLM call # If using a single LLM call, use the results computed in the first call. if self.config.rails.dialog.single_call.enabled: - event = get_last_user_intent_event(events) - - if not event: + user_intent_event = get_last_user_intent_event(events) + if not user_intent_event: raise RuntimeError("No last user intent found to generate bot message") - if event["type"] == "UserIntent": - bot_message_event = event["additional_info"]["bot_message_event"] - - # We only need to use the bot message if it corresponds to the - # generate bot intent as well. - last_bot_intent = get_last_bot_intent_event(events) - if not last_bot_intent: - raise RuntimeError("No last bot intent found to generate bot message") - - if last_bot_intent["intent"] == event["additional_info"]["bot_intent_event"]["intent"]: - text = bot_message_event["text"] - # If the bot message is being generated in streaming mode - if text.startswith('Bot message: "<>"` - # Extract the streaming handler uid and get a reference. - streaming_handler_uid = text[26:-4] - _streaming_handler = local_streaming_handlers[streaming_handler_uid] - - # We pipe the content from this handler to the main one. - _streaming_handler.set_pipe_to(streaming_handler) - await _streaming_handler.disable_buffering() - - # And wait for it to finish. - # We stop after the closing double quotes for the bot message. - _streaming_handler.stop = [ - '"\n', - ] - text = await _streaming_handler.wait() - - output_events = [] - reasoning_trace = get_and_clear_reasoning_trace_contextvar() - if reasoning_trace: - output_events.append(new_event_dict("BotThinking", content=reasoning_trace)) - output_events.append(new_event_dict("BotMessage", text=text)) - - return ActionResult(events=output_events) - else: - if streaming_handler: - await streaming_handler.push_chunk(bot_message_event["text"]) - - output_events = [] - reasoning_trace = get_and_clear_reasoning_trace_contextvar() - if reasoning_trace: - output_events.append(new_event_dict("BotThinking", content=reasoning_trace)) - output_events.append(bot_message_event) - - return ActionResult(events=output_events) + cached_result = await self._bot_message_from_single_call_cache( + user_intent_event, events, streaming_handler + ) + if cached_result is not None: + return cached_result # If we are in passthrough mode, we just use the input for prompting if self.config.passthrough: @@ -863,14 +1032,7 @@ async def generate_bot_message(self, events: List[dict], context: dict, llm: Opt if self._passthrough_fn: prompt = None raw_output = await self._passthrough_fn(context=context, events=events) - - # If the passthrough action returns a single value, we consider that - # to be the text output - if isinstance(raw_output, tuple) or isinstance(raw_output, list): - result, passthrough_output = raw_output[0], raw_output[1] - else: - result = raw_output - passthrough_output = None + result, passthrough_output = _unpack_passthrough_output(raw_output) # We record the passthrough output in the context context_updates["passthrough_output"] = passthrough_output @@ -879,9 +1041,6 @@ async def generate_bot_message(self, events: List[dict], context: dict, llm: Opt t0 = time() - # Initialize the LLMCallInfo object - llm_call_info_var.set(LLMCallInfo(task=Task.GENERATE_BOT_MESSAGE.value)) - # In passthrough mode, we should use the full conversation history # instead of just the last user message to preserve tool message context raw_prompt = raw_llm_request.get() @@ -901,20 +1060,21 @@ async def generate_bot_message(self, events: List[dict], context: dict, llm: Opt prompt = context.get("user_message") gen_options: Optional[GenerationOptions] = generation_options_var.get() - llm_params = (gen_options and gen_options.llm_params) or {} + llm_params = ( + gen_options.llm_params if gen_options is not None and gen_options.llm_params is not None else {} + ) if not prompt: raise RuntimeError("No prompt found to generate bot message") - result = ( - await llm_call( - generation_llm, - prompt, - streaming_handler=streaming_handler, - llm_params=llm_params, - ) - ).content - - result = self.llm_task_manager.parse_task_output(Task.GENERAL, output=result) + result = await self._generate_general_response( + generation_llm=generation_llm, + prompt=prompt, + streaming_handler=streaming_handler, + stream_during_call=True, + llm_call_task=Task.GENERATE_BOT_MESSAGE, + parse_task=Task.GENERAL, + llm_params=llm_params, + ) log.info( "--- :: LLM Bot Message Generation passthrough call took %.2f seconds", @@ -950,10 +1110,10 @@ async def generate_bot_message(self, events: List[dict], context: dict, llm: Opt if streaming_handler: # TODO: Figure out a more generic way to deal with this - if prompt_config.output_parser in ["verbose_v1", "bot_message"]: - streaming_handler.set_pattern(prefix='Bot message: "', suffix='"') - else: - streaming_handler.set_pattern(prefix=' "', suffix='"') + prefix, suffix = _streaming_pattern_for( + prompt_config.output_parser, include_bot_message_parser=True + ) + streaming_handler.set_pattern(prefix=prefix, suffix=suffix) # Initialize the LLMCallInfo object llm_call_info_var.set(LLMCallInfo(task=Task.GENERATE_BOT_MESSAGE.value)) @@ -996,15 +1156,8 @@ async def generate_bot_message(self, events: List[dict], context: dict, llm: Opt if streaming_handler: await streaming_handler.push_chunk(bot_utterance) - output_events = [] - reasoning_trace = get_and_clear_reasoning_trace_contextvar() - if reasoning_trace: - context_updates["bot_thinking"] = reasoning_trace - output_events.append(new_event_dict("BotThinking", content=reasoning_trace)) - output_events.append(new_event_dict("BotMessage", text=bot_utterance)) - return ActionResult( - events=output_events, + events=_bot_turn_output_events(new_event_dict("BotMessage", text=bot_utterance), context_updates), context_updates=context_updates, ) else: @@ -1013,15 +1166,8 @@ async def generate_bot_message(self, events: List[dict], context: dict, llm: Opt if streaming_handler: await streaming_handler.push_chunk(bot_utterance) - output_events = [] - reasoning_trace = get_and_clear_reasoning_trace_contextvar() - if reasoning_trace: - context_updates["bot_thinking"] = reasoning_trace - output_events.append(new_event_dict("BotThinking", content=reasoning_trace)) - output_events.append(new_event_dict("BotMessage", text=bot_utterance)) - return ActionResult( - events=output_events, + events=_bot_turn_output_events(new_event_dict("BotMessage", text=bot_utterance), context_updates), context_updates=context_updates, ) @@ -1105,10 +1251,103 @@ async def generate_value( log.error(f"Error evaluating value: {value}. Error: {str(e)}") raise ValueError(f"Invalid LLM response: `{value}`") + async def _build_intent_steps_examples(self, text: str) -> Tuple[List[str], List[str]]: + """Build the few-shot examples and candidate intents for the single call. + + Searches the user-message index for utterances similar to ``text``, pairs + each candidate intent with a flow (and its bot message, if any) from the + flows / bot-message indexes, and returns up to five formatted examples + plus the list of candidate user intents. + """ + examples: List[str] = [] + potential_user_intents: List[str] = [] + intent_results = [] + flow_results = {} + + if self.user_message_index: + # Get the top 10 intents even if we use less in the selected examples. + # Some of these intents might not have an associated flow and will be + # skipped from the few-shot examples. + intent_results = await self.user_message_index.search(text=text, max_results=10, threshold=None) + + # We fill in the list of potential user intents + for result in intent_results: + if result.meta["intent"] not in potential_user_intents: + potential_user_intents.append(result.meta["intent"]) + + if self.flows_index: + for intent in potential_user_intents: + flow_results_intent = await self._search_flows_index(text=intent, max_results=2) + flow_results[intent] = flow_results_intent + + # We add the intent to the examples in reverse order + # so the most relevant is towards the end. + for result in intent_results: + # Stop after the first 5 flow examples, in case more than 5 intents + # have been selected from the index. + if len(examples) >= 5: + break + + intent = result.meta["intent"] + example = f'user "{result.text}"\n {intent}\n' + + flow_results_intent = flow_results.get(intent, []) + found_flow_for_intent = False + for result_flow in flow_results_intent: + # Assumption: each flow should contain at least two lines, the first is the user intent. + # Just in case there are some flows with only one line + if "\n" not in result_flow.text: + continue + (flow_user_intent, flow_continuation) = result_flow.text.split("\n", 1) + flow_user_intent = flow_user_intent[5:] + if flow_user_intent == intent: + found_flow_for_intent = True + example += f"{flow_continuation}\n" + + # Also add the bot message if the last line in the flow is a bot canonical form + last_flow_line = flow_continuation + if "\n" in flow_continuation: + (_, last_flow_line) = flow_continuation.rsplit("\n", 1) + if last_flow_line.startswith("bot "): + bot_canonical_form = last_flow_line[4:] + + found_bot_message = False + if self.bot_message_index: + bot_messages_results = await self.bot_message_index.search( + text=bot_canonical_form, + max_results=1, + threshold=None, + ) + + for bot_message_result in bot_messages_results: + if bot_message_result.text == bot_canonical_form: + found_bot_message = True + example += f' "{bot_message_result.meta["text"]}"\n' + # Only use the first bot message for now + break + + if not found_bot_message: + # This is for canonical forms that do not have an associated message. + # Create a simple message for the bot canonical form. + # In a later version we could generate a message with the LLM at app initialization. + example += f" # On the next line generate a bot message related to {bot_canonical_form}\n" + + # For now, only use the first flow for each intent. + break + if not found_flow_for_intent: + # Skip intents that do not have an associated flow. + continue + + example += "\n" + examples.append(example) + + return examples, potential_user_intents + @action(is_system_action=True) async def generate_intent_steps_message( self, events: List[dict], + context: dict, llm: Optional[LLMModel] = None, kb: Optional[KnowledgeBase] = None, ): @@ -1142,95 +1381,17 @@ async def generate_intent_steps_message( log.info("Generate all three phases in one LLM call...") - # We search for the most relevant similar user utterance - examples = [] - potential_user_intents = [] - intent_results = [] - flow_results = {} - - if self.user_message_index: - # Get the top 10 intents even if we use less in the selected examples. - # Some of these intents might not have an associated flow and will be - # skipped from the few-shot examples. - intent_results = await self.user_message_index.search( - text=event["text"], max_results=10, threshold=None - ) - - # We fill in the list of potential user intents - for result in intent_results: - if result.meta["intent"] not in potential_user_intents: - potential_user_intents.append(result.meta["intent"]) - - if self.flows_index: - for intent in potential_user_intents: - flow_results_intent = await self._search_flows_index(text=intent, max_results=2) - flow_results[intent] = flow_results_intent - - # We add the intent to the examples in reverse order - # so the most relevant is towards the end. - for result in intent_results: - # Stop after the first 5 flow examples, in case more than 5 intents - # have been selected from the index. - if len(examples) >= 5: - break - - intent = result.meta["intent"] - example = f'user "{result.text}"\n {intent}\n' - - flow_results_intent = flow_results.get(intent, []) - found_flow_for_intent = False - for result_flow in flow_results_intent: - # Assumption: each flow should contain at least two lines, the first is the user intent. - # Just in case there are some flows with only one line - if "\n" not in result_flow.text: - continue - (flow_user_intent, flow_continuation) = result_flow.text.split("\n", 1) - flow_user_intent = flow_user_intent[5:] - if flow_user_intent == intent: - found_flow_for_intent = True - example += f"{flow_continuation}\n" - - # Also add the bot message if the last line in the flow is a bot canonical form - last_flow_line = flow_continuation - if "\n" in flow_continuation: - (_, last_flow_line) = flow_continuation.rsplit("\n", 1) - if last_flow_line.startswith("bot "): - bot_canonical_form = last_flow_line[4:] - - found_bot_message = False - if self.bot_message_index: - bot_messages_results = await self.bot_message_index.search( - text=bot_canonical_form, - max_results=1, - threshold=None, - ) - - for bot_message_result in bot_messages_results: - if bot_message_result.text == bot_canonical_form: - found_bot_message = True - example += f' "{bot_message_result.meta["text"]}"\n' - # Only use the first bot message for now - break - - if not found_bot_message: - # This is for canonical forms that do not have an associated message. - # Create a simple message for the bot canonical form. - # In a later version we could generate a message with the LLM at app initialization. - example += ( - f" # On the next line generate a bot message related to {bot_canonical_form}\n" - ) - - # For now, only use the first flow for each intent. - break - if not found_flow_for_intent: - # Skip intents that do not have an associated flow. - continue + # Normalize multimodal input to text, mirroring generate_user_intent. + if isinstance(event["text"], list): + text = " ".join([item["text"] for item in event["text"] if item["type"] == "text"]) + else: + text = event["text"] - example += "\n" - examples.append(example) + # We search for the most relevant similar user utterance + examples, potential_user_intents = await self._build_intent_steps_examples(text) if kb: - chunks = await kb.search_relevant_chunks(event["text"]) + chunks = await kb.search_relevant_chunks(text) relevant_chunks = "\n".join([chunk["body"] for chunk in chunks]) else: relevant_chunks = "" @@ -1254,7 +1415,6 @@ async def generate_intent_steps_message( if streaming_handler: # Create a new "inner" streaming handler and save the reference _streaming_handler = StreamingHandler() - local_streaming_handlers[_streaming_handler.uid] = _streaming_handler # We buffer the content, so we can get a chance to look at the # first k lines. @@ -1272,15 +1432,13 @@ async def generate_intent_steps_message( # We also mark that the message is still being generated # by a streaming handler. - result += f'\nBot message: "<>"' + result += f"\n{_streaming_handoff.register(_streaming_handler)}" # Moving forward we need to set the expected pattern to correctly # parse the message. # TODO: Figure out a more generic way to deal with this. - if prompt_config.output_parser == "verbose_v1": - _streaming_handler.set_pattern(prefix='Bot message: "', suffix='"') - else: - _streaming_handler.set_pattern(prefix=' "', suffix='"') + prefix, suffix = _streaming_pattern_for(prompt_config.output_parser, include_bot_message_parser=False) + _streaming_handler.set_pattern(prefix=prefix, suffix=suffix) else: # Initialize the LLMCallInfo object llm_call_info_var.set(LLMCallInfo(task=Task.GENERATE_INTENT_STEPS_MESSAGE.value)) @@ -1343,37 +1501,63 @@ async def generate_intent_steps_message( log.info("Canonical form for bot intent: " + (bot_intent if bot_intent else "None")) log.info("Generated bot message: " + (bot_message if bot_message else "None")) - additional_info = { - "bot_intent_event": new_event_dict("BotIntent", intent=bot_intent), - "bot_message_event": new_event_dict("BotMessage", text=bot_message), - } + additional_info = build_single_call_payload( + bot_intent_event=new_event_dict("BotIntent", intent=bot_intent), + bot_message_event=new_event_dict("BotMessage", text=bot_message), + ) events = [new_event_dict("UserIntent", intent=user_intent, additional_info=additional_info)] return ActionResult(events=events) else: - prompt = self.llm_task_manager.render_task_prompt(task=Task.GENERAL, events=events) + # No user messages: this is a general bot turn, identical to the + # one generate_user_intent emits in this mode. + return await self._emit_general_bot_turn( + events=events, + context=context, + event=event, + generation_llm=generation_llm, + streaming_handler=streaming_handler, + kb=kb, + ) - # Initialize the LLMCallInfo object - llm_call_info_var.set(LLMCallInfo(task=Task.GENERAL.value)) - # We make this call with temperature 0 to have it as deterministic as possible. - gen_options: Optional[GenerationOptions] = generation_options_var.get() - llm_params = (gen_options and gen_options.llm_params) or {} - result = (await llm_call(generation_llm, prompt, llm_params=llm_params)).content +def _spec_op_as_dict(spec_op: SpecOp) -> Dict[str, Any]: + """Return a ``SpecOp``'s spec as a dict. - result = self.llm_task_manager.parse_task_output(Task.GENERAL, output=result) - text = result.strip() - if text.startswith('"'): - text = text[1:-1] + ``SpecOp.spec`` is ``Union[Spec, dict]``; callers that need dict-only fields + (such as ``elements``) use this to normalize it. + """ + return asdict(spec_op.spec) if isinstance(spec_op.spec, Spec) else cast(Dict, spec_op.spec) - # In streaming mode, we also push this. - if streaming_handler: - await streaming_handler.push_chunk(text) - return ActionResult( - events=[new_event_dict("BotMessage", text=text)], - ) +def _unpack_passthrough_output(raw_output): + """Split a passthrough fn result into ``(text, passthrough_output)``. + + A tuple/list result carries an explicit passthrough output as its second + element; any other return value is treated as the text output alone. + """ + if isinstance(raw_output, (tuple, list)): + return raw_output[0], raw_output[1] + return raw_output, None + + +def _bot_turn_output_events(final_event: dict, context_updates: Optional[dict] = None) -> List[dict]: + """Build the output events for a bot turn: an optional BotThinking then the final event. + + Pops any reasoning trace into a ``BotThinking`` event placed before + ``final_event``. When ``context_updates`` is provided, the reasoning trace is + also recorded as ``bot_thinking`` (the generated bot-message path does this; + the single-call cached path does not). + """ + output_events = [] + reasoning_trace = get_and_clear_reasoning_trace_contextvar() + if reasoning_trace: + if context_updates is not None: + context_updates["bot_thinking"] = reasoning_trace + output_events.append(new_event_dict("BotThinking", content=reasoning_trace)) + output_events.append(final_event) + return output_events def clean_utterance_content(utterance: str) -> str: diff --git a/tests/recorded/rails/public_api/cassettes/test_dialog/test_dialog_generation_bot_message_public_contract.yaml b/tests/recorded/rails/public_api/cassettes/test_dialog/test_dialog_generation_bot_message_public_contract.yaml new file mode 100644 index 0000000000..f8a4f56fbc --- /dev/null +++ b/tests/recorded/rails/public_api/cassettes/test_dialog/test_dialog_generation_bot_message_public_contract.yaml @@ -0,0 +1,384 @@ +version: 1 +interactions: +- request: + method: POST + uri: https://api.openai.com/v1/chat/completions + headers: + host: + - api.openai.com + accept: + - '*/*' + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + user-agent: + - python-httpx/0.28.1 + content-type: + - application/json + parsed_body: + model: gpt-5.4-nano + messages: + - role: user + content: | + """ + Below is a conversation between a helpful AI assistant and a user. The bot is designed to generate human-like text based on the input that it receives. The bot is talkative and provides lots of specific details. If the bot does not know the answer to a question, it truthfully says it does not know. + """ + + # This is how a conversation between a user and the bot can go: + User message: "Hello there!" + User intent: express greeting + Bot intent: express greeting + Bot message: "Hello! How can I assist you today?" + User message: "What can you do for me?" + User intent: ask about capabilities + Bot intent: respond about capabilities + Bot message: "As an AI assistant, I can help you with a wide range of tasks. This includes question answering on various topics, generating text for various purposes and providing suggestions based on your preferences." + User message: "Tell me a bit about the history of NVIDIA." + User intent: ask general question + Bot intent: response for general question + Bot message: "NVIDIA is a technology company that specializes in designing and manufacturing graphics processing units (GPUs) and other computer hardware. The company was founded in 1993 by Jen-Hsun Huang, Chris Malachowsky, and Curtis Priem." + User message: "tell me more" + User intent: request more information + Bot intent: provide more information + Bot message: "Initially, the company focused on developing 3D graphics processing technology for the PC gaming market. In 1999, NVIDIA released the GeForce 256, the world's first GPU, which was a major breakthrough for the gaming industry. The company continued to innovate in the GPU space, releasing new products and expanding into other markets such as professional graphics, mobile devices, and artificial intelligence." + User message: "thanks" + User intent: express appreciation + Bot intent: express appreciation and offer additional help + Bot message: "You're welcome. If you have any more questions or if there's anything else I can help you with, please don't hesitate to ask." + + + # This is how the user talks: + User message: "tell me about your product" + User intent: ask about the product + + + + # This is the current conversation between the user and the bot: + # Choose intent from this list: ask about the product + User message: "Hello there!" + User intent: express greeting + Bot intent: express greeting + Bot message: "Hello! How can I assist you today?" + User message: "What can you do for me?" + User intent: ask about capabilities + Bot intent: respond about capabilities + Bot message: "As an AI assistant, I can help you with a wide range of tasks. This includes question answering on various topics, generating text for various purposes and providing suggestions based on your preferences." + User message: "tell me about your product" + response: + status: + code: 200 + message: OK + headers: + content-type: + - application/json + transfer-encoding: + - chunked + connection: + - keep-alive + server: + - cloudflare + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + access-control-expose-headers: + - X-Request-ID + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + body: + parsed_body: + id: '[RECORDED_RESPONSE_ID]' + object: chat.completion + created: 0 + model: gpt-5.4-nano-2026-03-17 + choices: + - index: 0 + message: + role: assistant + content: |- + Hello! I’m an AI assistant designed to help with a wide range of tasks through conversation. Here are a few things I can do that often feel “product-like”: + + - **Answer questions**: Explain topics, help with general knowledge, and walk through concepts. + - **Generate text**: Draft emails, messages, summaries, stories, posts, or structured documents. + - **Help with writing & editing**: Improve clarity, tone, grammar, and rewrite content for different audiences. + - **Assist with problem-solving**: Break down steps for coding, math, study plans, brainstorming, and troubleshooting (as long as the details you provide are enough). + - **Support planning & ideation**: Generate ideas, outlines, checklists, and options for decisions. + + If you tell me what you want to accomplish (e.g., writing something, learning a topic, coding help, planning), I can tailor the assistance to that goal. What kind of task are you thinking about? + refusal: null + annotations: [] + finish_reason: stop + usage: + prompt_tokens: 580 + completion_tokens: 200 + total_tokens: 780 + prompt_tokens_details: + cached_tokens: 0 + audio_tokens: 0 + completion_tokens_details: + reasoning_tokens: 0 + audio_tokens: 0 + accepted_prediction_tokens: 0 + rejected_prediction_tokens: 0 + service_tier: default + system_fingerprint: null +- request: + method: POST + uri: https://api.openai.com/v1/chat/completions + headers: + host: + - api.openai.com + accept: + - '*/*' + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + user-agent: + - python-httpx/0.28.1 + content-type: + - application/json + parsed_body: + model: gpt-5.4-nano + messages: + - role: user + content: | + """ + Below is a conversation between a helpful AI assistant and a user. The bot is designed to generate human-like text based on the input that it receives. The bot is talkative and provides lots of specific details. If the bot does not know the answer to a question, it truthfully says it does not know. + """ + + # This is how a conversation between a user and the bot can go: + User intent: express greeting + Bot intent: express greeting + User intent: ask about capabilities + Bot intent: respond about capabilities + User intent: ask general question + Bot intent: response for general question + User intent: request more information + Bot intent: provide more information + User intent: express appreciation + Bot intent: express appreciation and offer additional help + + + # This is how the bot thinks: + User intent: ask about the product + Bot intent: describe the product + + User intent: express greeting + Bot intent: express greeting + + + + # This is the current conversation between the user and the bot: + User intent: express greeting + Bot intent: express greeting + User intent: ask about capabilities + Bot intent: respond about capabilities + User intent: Hello! I’m an AI assistant designed to help with a wide range of tasks through conversation. Here are a few things I can do that often feel “product-like”: + response: + status: + code: 200 + message: OK + headers: + content-type: + - application/json + transfer-encoding: + - chunked + connection: + - keep-alive + server: + - cloudflare + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + access-control-expose-headers: + - X-Request-ID + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + body: + parsed_body: + id: '[RECORDED_RESPONSE_ID]' + object: chat.completion + created: 0 + model: gpt-5.4-nano-2026-03-17 + choices: + - index: 0 + message: + role: assistant + content: |- + Hello! I’m here to help with a wide range of tasks through conversation. A few things I can do that often feel “product-like” include: + + - **Answer questions and explain concepts** (from quick facts to deeper walkthroughs) + - **Help you write and edit** (emails, essays, resumes, cover letters, posts, etc.) + - **Brainstorm ideas** (projects, names, outlines, strategies, content themes) + - **Summarize and extract key points** from text you provide + - **Assist with planning** (study plans, travel itineraries, schedules, checklists) + - **Draft and refine messages** in a particular tone (friendly, formal, persuasive, etc.) + - **Provide coding help** (debugging, writing snippets, explaining errors, algorithm ideas) + - **Role-play or simulate conversations** (interviews, customer support, practice dialogues) + + If you tell me what you’re trying to do today, I can jump in right away. What would you like help with? + refusal: null + annotations: [] + finish_reason: stop + usage: + prompt_tokens: 269 + completion_tokens: 208 + total_tokens: 477 + prompt_tokens_details: + cached_tokens: 0 + audio_tokens: 0 + completion_tokens_details: + reasoning_tokens: 0 + audio_tokens: 0 + accepted_prediction_tokens: 0 + rejected_prediction_tokens: 0 + service_tier: default + system_fingerprint: null +- request: + method: POST + uri: https://api.openai.com/v1/chat/completions + headers: + host: + - api.openai.com + accept: + - '*/*' + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + user-agent: + - python-httpx/0.28.1 + content-type: + - application/json + parsed_body: + model: gpt-5.4-nano + messages: + - role: user + content: | + """ + Below is a conversation between a helpful AI assistant and a user. The bot is designed to generate human-like text based on the input that it receives. The bot is talkative and provides lots of specific details. If the bot does not know the answer to a question, it truthfully says it does not know. + """ + + # This is how a conversation between a user and the bot can go: + User message: "Hello there!" + User intent: express greeting + Bot intent: express greeting + Bot message: "Hello! How can I assist you today?" + User message: "What can you do for me?" + User intent: ask about capabilities + Bot intent: respond about capabilities + Bot message: "As an AI assistant, I can help you with a wide range of tasks. This includes question answering on various topics, generating text for various purposes and providing suggestions based on your preferences." + User message: "Tell me a bit about the history of NVIDIA." + User intent: ask general question + Bot intent: response for general question + Bot message: "NVIDIA is a technology company that specializes in designing and manufacturing graphics processing units (GPUs) and other computer hardware. The company was founded in 1993 by Jen-Hsun Huang, Chris Malachowsky, and Curtis Priem." + User message: "tell me more" + User intent: request more information + Bot intent: provide more information + Bot message: "Initially, the company focused on developing 3D graphics processing technology for the PC gaming market. In 1999, NVIDIA released the GeForce 256, the world's first GPU, which was a major breakthrough for the gaming industry. The company continued to innovate in the GPU space, releasing new products and expanding into other markets such as professional graphics, mobile devices, and artificial intelligence." + User message: "thanks" + User intent: express appreciation + Bot intent: express appreciation and offer additional help + Bot message: "You're welcome. If you have any more questions or if there's anything else I can help you with, please don't hesitate to ask." + + + + # This is some additional context: + ```markdown + The product is a guardrails toolkit for large language models. It lets + developers add programmable rails that check user input, steer the dialog, and + validate bot output before it reaches the user. + ``` + + + # This is how the bot talks: + Bot intent: inform cannot engage with inappropriate content + Bot message: "I will not engage with inappropriate content." + + Bot intent: inform cannot engage in self harm behavior + Bot message: "I will not engage in any self harm behavior." + + Bot intent: inform cannot engage in abusive or harmful behavior + Bot message: "I will not engage in any abusive or harmful behavior." + + Bot intent: offer help + Bot message: "I'm here to help with anything you need." + + Bot intent: express greeting + Bot message: "Hey there!" + + + + # This is the current conversation between the user and the bot: + User message: "Hello there!" + User intent: express greeting + Bot intent: express greeting + Bot message: "Hello! How can I assist you today?" + User message: "What can you do for me?" + User intent: ask about capabilities + Bot intent: respond about capabilities + Bot message: "As an AI assistant, I can help you with a wide range of tasks. This includes question answering on various topics, generating text for various purposes and providing suggestions based on your preferences." + User message: "tell me about your product" + User intent: Hello! I’m an AI assistant designed to help with a wide range of tasks through conversation. Here are a few things I can do that often feel “product-like”: + Bot intent: general response + response: + status: + code: 200 + message: OK + headers: + content-type: + - application/json + transfer-encoding: + - chunked + connection: + - keep-alive + server: + - cloudflare + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + access-control-expose-headers: + - X-Request-ID + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + body: + parsed_body: + id: '[RECORDED_RESPONSE_ID]' + object: chat.completion + created: 0 + model: gpt-5.4-nano-2026-03-17 + choices: + - index: 0 + message: + role: assistant + content: |- + Sure! If you mean the **guardrails toolkit for large language models** you mentioned in the context, here’s what it is and what it helps you do: + + - **Programmable “rails” for LLMs:** Lets developers define rules and logic that run during the conversation. + - **Check user input:** You can validate or filter what users send (e.g., detect disallowed content, enforce formatting, constrain requests). + - **Steer the dialog:** You can route or modify the assistant’s behavior based on the input or conversation state (e.g., ask clarifying questions, refuse safely, switch modes). + - **Validate bot output:** Before the response is shown to the user, rails can verify it meets requirements (e.g., must include/avoid certain content, follow a policy, adhere to a schema). + - **Reduce unsafe or off-policy responses:** The overall goal is to make LLM behavior more controlled, consistent, and safer. + + If you tell me what kind of use case you’re building (customer support, tutoring, legal/health Q&A, internal tools, etc.), I can suggest what kinds of rails you’d typically implement. + refusal: null + annotations: [] + finish_reason: stop + usage: + prompt_tokens: 753 + completion_tokens: 235 + total_tokens: 988 + prompt_tokens_details: + cached_tokens: 0 + audio_tokens: 0 + completion_tokens_details: + reasoning_tokens: 0 + audio_tokens: 0 + accepted_prediction_tokens: 0 + rejected_prediction_tokens: 0 + service_tier: default + system_fingerprint: null diff --git a/tests/recorded/rails/public_api/configs.py b/tests/recorded/rails/public_api/configs.py index b126bc6ae2..b29abbf8df 100644 --- a/tests/recorded/rails/public_api/configs.py +++ b/tests/recorded/rails/public_api/configs.py @@ -53,6 +53,7 @@ NIM_BASELINE_CONFIG = RailsConfigSource.from_path(CONFIGS_DIR, "nim_baseline") NEMOGUARDS_FULL_CONFIG = RailsConfigSource.from_path(CONFIGS_DIR, "nemoguards_full") DIALOG_CONFIG = RailsConfigSource.from_path(CONFIGS_DIR, "dialog") +DIALOG_GENERATION_CONFIG = RailsConfigSource.from_path(CONFIGS_DIR, "dialog_generation") SINGLE_CALL_CONFIG = RailsConfigSource.from_path(CONFIGS_DIR, "single_call") TASK_MODELS_CONFIG = RailsConfigSource.from_path(CONFIGS_DIR, "task_models") STREAMING_OUTPUT_RAILS_CONFIG = RailsConfigSource.from_path(CONFIGS_DIR, "streaming_output_rails") diff --git a/tests/recorded/rails/public_api/configs/dialog_generation/config.py b/tests/recorded/rails/public_api/configs/dialog_generation/config.py new file mode 100644 index 0000000000..08664f1a21 --- /dev/null +++ b/tests/recorded/rails/public_api/configs/dialog_generation/config.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tests.recorded.rails.public_api.simple_embedding_provider import SimpleEmbeddingSearchProvider + + +def init(app): + app.register_embedding_search_provider("simple", SimpleEmbeddingSearchProvider) diff --git a/tests/recorded/rails/public_api/configs/dialog_generation/config.yml b/tests/recorded/rails/public_api/configs/dialog_generation/config.yml new file mode 100644 index 0000000000..38a21e8346 --- /dev/null +++ b/tests/recorded/rails/public_api/configs/dialog_generation/config.yml @@ -0,0 +1,12 @@ +models: + - type: main + engine: openai + model: gpt-5.4-nano + parameters: + max_retries: 0 +core: + embedding_search_provider: + name: simple +knowledge_base: + embedding_search_provider: + name: simple diff --git a/tests/recorded/rails/public_api/configs/dialog_generation/kb/product.md b/tests/recorded/rails/public_api/configs/dialog_generation/kb/product.md new file mode 100644 index 0000000000..b0fba19957 --- /dev/null +++ b/tests/recorded/rails/public_api/configs/dialog_generation/kb/product.md @@ -0,0 +1,5 @@ +# Product overview + +The product is a guardrails toolkit for large language models. It lets +developers add programmable rails that check user input, steer the dialog, and +validate bot output before it reaches the user. diff --git a/tests/recorded/rails/public_api/configs/dialog_generation/rails.co b/tests/recorded/rails/public_api/configs/dialog_generation/rails.co new file mode 100644 index 0000000000..a676042583 --- /dev/null +++ b/tests/recorded/rails/public_api/configs/dialog_generation/rails.co @@ -0,0 +1,21 @@ +define user express greeting + "hello" + "hi" + +define user ask about the product + "tell me about your product" + "what does the product do" + +define bot express greeting + "Hey there!" + +define bot offer help + "I'm here to help with anything you need." + +define flow greeting + user express greeting + bot express greeting + +define flow product + user ask about the product + bot describe the product diff --git a/tests/recorded/rails/public_api/test_dialog.py b/tests/recorded/rails/public_api/test_dialog.py index 120dd7cd74..931e5a1407 100644 --- a/tests/recorded/rails/public_api/test_dialog.py +++ b/tests/recorded/rails/public_api/test_dialog.py @@ -21,7 +21,7 @@ from nemoguardrails.rails.llm.options import RailStatus from tests.recorded.assertions import assert_generated_message, assert_rails_result from tests.recorded.normalization import normalize_rails_result -from tests.recorded.rails.public_api.configs import DIALOG_CONFIG, SINGLE_CALL_CONFIG +from tests.recorded.rails.public_api.configs import DIALOG_CONFIG, DIALOG_GENERATION_CONFIG, SINGLE_CALL_CONFIG from tests.recorded.rails_config import load_config from tests.recorded.snapshots import snapshot @@ -43,6 +43,41 @@ async def test_dialog_generate_async_public_contract(openai_api_key): ) +@pytest.mark.vcr +async def test_dialog_generation_bot_message_public_contract(openai_api_key): + """Bot-message GENERATION prompt carries few-shot examples + retrieved kb chunks. + + "tell me about your product" maps (via a flow) to the bot intent + "describe the product", which has no predefined message, so + ``generate_bot_message`` calls the LLM with the bot-message index examples and + the retrieved kb chunks. VCR matches on ``recorded_body``, so this cassette + pins all three dialog prompts (intent, next-step, bot-message) by their exact + request bodies. It was verified to replay byte-identically against both the + pre-refactor ``generation.py`` on ``develop`` and the refactored branch, so + the prompts are unchanged pre/post -- a prompt-fidelity equivalence check the + prompt-independent ``FakeLLMModel`` oracle cannot make. + """ + rails = LLMRails(load_config(DIALOG_GENERATION_CONFIG), verbose=False) + + result = await rails.generate_async(messages=[{"role": "user", "content": "tell me about your product"}]) + + assert_generated_message(result) + assert result == snapshot( + { + "role": "assistant", + "content": """\ +Sure! If you mean the **guardrails toolkit for large language models** you mentioned in the context, here’s what it is and what it helps you do: +- **Programmable “rails” for LLMs:** Lets developers define rules and logic that run during the conversation. +- **Check user input:** You can validate or filter what users send (e.g., detect disallowed content, enforce formatting, constrain requests). +- **Steer the dialog:** You can route or modify the assistant’s behavior based on the input or conversation state (e.g., ask clarifying questions, refuse safely, switch modes). +- **Validate bot output:** Before the response is shown to the user, rails can verify it meets requirements (e.g., must include/avoid certain content, follow a policy, adhere to a schema). +- **Reduce unsafe or off-policy responses:** The overall goal is to make LLM behavior more controlled, consistent, and safer. +If you tell me what kind of use case you’re building (customer support, tutoring, legal/health Q&A, internal tools, etc.), I can suggest what kinds of rails you’d typically implement.\ +""", + } + ) + + @pytest.mark.vcr async def test_single_call_generate_async_public_contract(openai_api_key): rails = LLMRails(load_config(SINGLE_CALL_CONFIG), verbose=False) diff --git a/tests/test_generation_equivalence.py b/tests/test_generation_equivalence.py new file mode 100644 index 0000000000..4e535d5754 --- /dev/null +++ b/tests/test_generation_equivalence.py @@ -0,0 +1,1198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Behavioral regression suite for ``LLMGenerationActions``. + +Pins the observable behavior of ``generate_user_intent``, ``generate_next_steps``, +``generate_bot_message``, ``generate_value`` and ``generate_intent_steps_message`` +across the decomposition of that module -- including the behavior changes the +fix commits intentionally introduce (multimodal normalization, the converged +single-call general turn, and streaming-handoff eviction). + +An event-only oracle is insufficient because some observable behavior is not in +the event stream, so each scenario is pinned across four channels: + +1. Emitted event list -- the ordered, generation-relevant projection of + ``GenerationLog.internal_events`` (``event_sequence``). The random ``flow_id`` + of ``start_flow`` is normalized away by keeping only the flow body; the + ``passthrough_output`` ``ContextUpdate`` event and the ``skip_output_rails`` / + ``_last_bot_prompt`` context updates appear here in order, so the "event form + vs dict form" distinction is captured by position. +2. Per-call task -- ``GenerationLog.llm_calls[i].task`` (``llm_tasks``), e.g. the + passthrough bot-message branch sets a ``generate_bot_message`` LLMCallInfo but + parses as ``general``, which shows up as a ``general`` task in that path. +3. The ``stop`` argument and whether the call streamed -- captured by + ``RecordingFakeLLM`` because ``LLMCallInfo`` has no ``stop`` field and the base + fake discards it, so a dropped ``stop=["User:"]`` is invisible to every other + channel. This is the only channel that distinguishes the ``Task.GENERAL`` call + sites (rendered-template vs raw-prompt). +4. Streaming chunk sequence -- collected from ``stream_async`` for the streaming + scenarios. +""" + +import asyncio +import os +from typing import cast +from unittest.mock import MagicMock + +import pytest + +from nemoguardrails import RailsConfig +from nemoguardrails.actions.llm.generation import LLMGenerationActions, build_single_call_payload +from nemoguardrails.llm.taskmanager import LLMTaskManager +from nemoguardrails.rails.llm.options import GenerationResponse +from nemoguardrails.testing.fake_model import FakeLLMModel +from nemoguardrails.types import LLMResponse, ToolCall, ToolCallFunction +from nemoguardrails.utils import new_event_dict +from tests.utils import TestChat + +LOG_OPTS = {"log": {"llm_calls": True, "internal_events": True}} + +# Context-update keys emitted by the generation actions that the oracle tracks. +_GENERATION_CONTEXT_KEYS = [ + "skip_output_rails", + "passthrough_output", + "bot_thinking", + "_last_bot_prompt", +] + + +class RecordingFakeLLM(FakeLLMModel): + """A ``FakeLLMModel`` that records the ``stop`` argument and call mode. + + ``LLMCallInfo`` does not carry ``stop`` and the base fake discards it, so this + is the only place a dropped or added ``stop`` is observable. ``calls`` is a + list of ``(mode, stop)`` where ``mode`` is ``"generate"`` or ``"stream"`` + (i.e. whether a streaming handler was passed into ``llm_call``). + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.calls = [] + + async def generate_async(self, prompt, *, stop=None, **kwargs): + self.calls.append(("generate", tuple(stop) if stop else None)) + return await super().generate_async(prompt, stop=stop, **kwargs) + + async def stream_async(self, prompt, *, stop=None, **kwargs): + self.calls.append(("stream", tuple(stop) if stop else None)) + async for chunk in super().stream_async(prompt, stop=stop, **kwargs): + yield chunk + + +def event_sequence(response): + """Project ``internal_events`` to the ordered, generation-relevant events. + + Volatile fields (uids, timestamps) are dropped; the ``start_flow`` flow_id + (a random uuid) is normalized away by keeping only the flow body. The large + ``_last_bot_prompt`` value is reduced to ```` -- its content is + asserted through ``call_prompts`` where it matters. + """ + sequence = [] + for event in response.log.internal_events or []: + event_type = event.get("type") + if event_type == "UserIntent": + tag = f"UserIntent:{event['intent']}" + if event.get("additional_info"): + cache_keys = ",".join(sorted(event["additional_info"].keys())) + tag += f"|cache:{cache_keys}" + sequence.append(tag) + elif event_type == "BotIntent": + sequence.append(f"BotIntent:{event['intent']}") + elif event_type == "BotMessage": + sequence.append(f"BotMessage:{event['text']}") + elif event_type == "BotThinking": + sequence.append(f"BotThinking:{event['content']}") + elif event_type == "BotToolCalls": + sequence.append(f"BotToolCalls:{len(event.get('tool_calls') or [])}") + elif event_type == "start_flow": + sequence.append(f"start_flow:{event.get('flow_body')!r}") + elif event_type == "ContextUpdate": + data = event.get("data") or {} + for key in _GENERATION_CONTEXT_KEYS: + if key in data: + value = "" if key == "_last_bot_prompt" else data[key] + sequence.append(f"ctx:{key}={value}") + return sequence + + +def llm_tasks(response): + """The ordered list of ``LLMCallInfo.task`` values for the generation.""" + return [call.task for call in (response.log.llm_calls or [])] + + +def call_prompts(response): + """The prompt string passed to each LLM call, in order. + + Used to distinguish the ``Task.GENERAL`` call sites: the templated path + renders the GENERAL prompt (starts with the general instruction), while the + raw passthrough path sends the conversation verbatim with no template. + """ + return [call.prompt or "" for call in (response.log.llm_calls or [])] + + +def generate(config, completions, message): + """Run a single-turn generation, returning ``(response, recording_llm)``.""" + recording_llm = RecordingFakeLLM(responses=completions) + chat = TestChat(config, llm=recording_llm) + response = cast(GenerationResponse, chat.app.generate(message, options=LOG_OPTS)) + return response, recording_llm + + +class TestDialogMode: + """Multi-call dialog: the classic three-phase pipeline.""" + + def test_three_phase(self): + """User intent without a flow -> next-step prediction -> bot-message LLM.""" + config = RailsConfig.from_content( + """ + define user ask booking + "I want to book a flight" + """ + ) + response, llm = generate( + config, + [" ask booking", "bot respond booking", ' "Sure, I can help with booking."'], + "I want to book a flight", + ) + + assert response.response == "Sure, I can help with booking." + assert event_sequence(response) == [ + "UserIntent:ask booking", + "BotIntent:respond booking", + "ctx:_last_bot_prompt=", + "BotMessage:Sure, I can help with booking.", + ] + assert llm_tasks(response) == [ + "generate_user_intent", + "generate_next_steps", + "generate_bot_message", + ] + assert llm.calls == [("generate", None), ("generate", None), ("generate", None)] + # Each rendered prompt carries the dynamic context (guards a dropped + # render-context key, which the prompt-independent fake would otherwise hide). + prompts = call_prompts(response) + assert 'user "I want to book a flight"' in prompts[0] + assert "user ask booking" in prompts[1] + assert "bot respond booking" in prompts[2] + + def test_predefined_bot_message(self): + """A predefined bot message short-circuits the phase-3 LLM call.""" + config = RailsConfig.from_content( + """ + define user express greeting + "hello" + define flow + user express greeting + bot express greeting + define bot express greeting + "Hello there!" + """ + ) + response, llm = generate(config, [" express greeting"], "hello there!") + + assert response.response == "Hello there!" + assert event_sequence(response) == [ + "UserIntent:express greeting", + "BotIntent:express greeting", + "ctx:skip_output_rails=True", + "BotMessage:Hello there!", + "ctx:skip_output_rails=False", + ] + assert llm_tasks(response) == ["generate_user_intent"] + assert llm.calls == [("generate", None)] + + def test_context_var_bot_intent(self): + """A ``bot $var`` intent renders the bot message from a context variable. + + Exercises the ``generate_bot_message`` branch where the bot intent starts + with ``$`` and names a context variable; no phase-3 LLM call is made. + """ + config = RailsConfig.from_content( + """ + define user give name + "my name is X" + define flow + user give name + $name = "World" + bot $name + """ + ) + response, llm = generate(config, [" give name"], "my name is World") + + assert response.response == "World" + assert event_sequence(response) == [ + "UserIntent:give name", + "BotIntent:$name", + "BotMessage:World", + ] + assert llm_tasks(response) == ["generate_user_intent"] + assert llm.calls == [("generate", None)] + + def test_multi_step_generation(self): + """``enable_multi_step_generation`` emits a ``start_flow`` with a parsed flow. + + The phase-2 LLM returns a multi-line flow; the runtime starts it (random + ``flow_id`` normalized away) and runs two bot intents/messages from it. + """ + config = RailsConfig.from_path(os.path.join(os.path.dirname(__file__), "test_configs", "multi_step_generation")) + llm = RecordingFakeLLM( + responses=[ + " express greeting", + " request appointment", + ' "What\'s your name?"', + " provide date", + "bot acknowledge the date\nbot confirm appointment", + ' "Ok, an appointment for tomorrow."', + ' "Your appointment is now confirmed."', + ] + ) + chat = TestChat(config, llm=llm) + # The completions are consumed across turns, so drive the first two turns + # to advance state before capturing the multi-step turn (the third). + chat.app.generate("hi") + chat.app.generate( + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "Hey there!"}, + {"role": "user", "content": "i need to make an appointment"}, + ] + ) + response = cast( + GenerationResponse, + chat.app.generate( + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "Hey there!"}, + {"role": "user", "content": "i need an appointment"}, + {"role": "assistant", "content": "What's your name?"}, + {"role": "user", "content": "I want to come tomorrow"}, + ], + options=LOG_OPTS, + ), + ) + + assert event_sequence(response) == [ + "UserIntent:provide date", + "start_flow:'bot acknowledge the date\\nbot confirm appointment'", + "BotIntent:acknowledge the date", + "ctx:_last_bot_prompt=", + "BotMessage:Ok, an appointment for tomorrow.", + "BotIntent:confirm appointment", + "ctx:_last_bot_prompt=", + "BotMessage:Your appointment is now confirmed.", + ] + assert llm_tasks(response) == [ + "generate_user_intent", + "generate_next_steps", + "generate_bot_message", + "generate_bot_message", + ] + + +_EMBEDDINGS_COLANG = """ +define user express greeting + "hi" +define bot express greeting + "Hello!" +define flow + user express greeting + bot express greeting +""" + + +def _embeddings_config(threshold, fallback_intent="express greeting"): + return RailsConfig.from_content( + _EMBEDDINGS_COLANG, + f""" + rails: + dialog: + user_messages: + embeddings_only: True + embeddings_only_similarity_threshold: {threshold} + embeddings_only_fallback_intent: {fallback_intent!r} + """, + ) + + +class TestEmbeddingsOnly: + """Dialog mode with embeddings-only intent detection.""" + + def test_index_hit(self): + """A similarity hit returns the matched intent with zero LLM calls. + + Threshold 0.0 guarantees the index search returns a hit. + """ + config = _embeddings_config(threshold=0.0) + response, llm = generate(config, [], "hi") + + assert response.response == "Hello!" + assert "UserIntent:express greeting" in event_sequence(response) + assert llm_tasks(response) == [] + assert llm.calls == [] + + def test_threshold_miss_uses_fallback_intent(self): + """A below-threshold search falls back to the configured intent, no LLM. + + Threshold 1.0 guarantees a miss, exercising the fallback-intent branch. + """ + config = _embeddings_config(threshold=1.0) + response, llm = generate(config, [], "totally unrelated request") + + assert response.response == "Hello!" + assert "UserIntent:express greeting" in event_sequence(response) + assert llm_tasks(response) == [] + assert llm.calls == [] + + def test_threshold_miss_falls_through_to_llm(self): + """A miss with no fallback intent falls through to LLM intent detection. + + Canonical-form detection, not the general path, so there is no ``stop``. + """ + config = _embeddings_config(threshold=1.0) + config.rails.dialog.user_messages.embeddings_only_fallback_intent = None + response, llm = generate(config, [" express greeting"], "totally unrelated request") + + assert "UserIntent:express greeting" in event_sequence(response) + assert llm_tasks(response) == ["generate_user_intent"] + assert llm.calls == [("generate", None)] + + +class TestSingleCall: + """Single-call mode via ``generate_intent_steps_message``.""" + + def test_intent_steps_message(self): + """One LLM call produces UserIntent + cached bot intent/message events.""" + config = RailsConfig.from_content( + """ + define user express greeting + "hello" + define flow + user express greeting + bot express greeting + """ + ) + config.rails.dialog.single_call.enabled = True + response, llm = generate( + config, + [' express greeting\nbot express greeting\n "Hello, there!"'], + "hello there!", + ) + + assert response.response == "Hello, there!" + assert event_sequence(response) == [ + "UserIntent:express greeting|cache:bot_intent_event,bot_message_event", + "BotIntent:express greeting", + "BotMessage:Hello, there!", + ] + # A single LLM call computes all three phases; 2 and 3 unpack the cache. + assert llm_tasks(response) == ["generate_intent_steps_message"] + assert llm.calls == [("generate", None)] + assert 'user "hello there!"' in call_prompts(response)[0] + + def test_multimodal_input(self): + """Multimodal (list) user content is normalized before intent detection. + + ``generate_intent_steps_message`` previously passed ``event["text"]`` (a + list) to the index search unchanged; it now joins the text parts like + ``generate_user_intent``. Non-text parts (e.g. ``image_url``) are ignored, + so text normalization still yields the same intent and response. + """ + config = RailsConfig.from_content( + """ + define user express greeting + "hello" + define flow + user express greeting + bot express greeting + """ + ) + config.rails.dialog.single_call.enabled = True + llm = RecordingFakeLLM(responses=[' express greeting\nbot express greeting\n "Hello, there!"']) + chat = TestChat(config, llm=llm) + response = cast( + GenerationResponse, + chat.app.generate( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello there!"}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}, + ], + } + ], + options=LOG_OPTS, + ), + ) + + assert event_sequence(response) == [ + "UserIntent:express greeting|cache:bot_intent_event,bot_message_event", + "BotIntent:express greeting", + "BotMessage:Hello, there!", + ] + assert llm_tasks(response) == ["generate_intent_steps_message"] + + def test_streaming_handoff(self): + """The ``<>`` handoff streams the bot message via stop tokens. + + Streaming uses ``stream_async`` (no ``GenerationLog``), so only the + model-call channel (mode + ``stop``) and the streamed chunk sequence are + asserted. + """ + + async def run(): + config = RailsConfig.from_content( + """ + define user express greeting + "hello" + define flow + user express greeting + bot express greeting + """, + "streaming: True\n", + ) + config.rails.dialog.single_call.enabled = True + llm = RecordingFakeLLM(responses=[' express greeting\nbot express greeting\n "Hello, there!"']) + chat = TestChat(config, llm=llm, streaming=True) + chunks = [] + async for chunk in chat.app.stream_async(messages=[{"role": "user", "content": "hello there!"}]): + chunks.append(chunk) + return llm, chunks + + llm, chunks = asyncio.run(run()) + + # The single intent-steps call streams with the intent stop tokens. + assert llm.calls == [("stream", ("\nuser ", "\nUser "))] + assert chunks == ["Hello, ", "there!"] + + @pytest.mark.asyncio + async def test_cache_miss_searches_with_bot_intent(self): + """On a single-call cache miss, the bot-message search uses the bot intent. + + Regression: the single-call branch previously rebound ``event`` to the + user-intent event, so when the cached bot intent did not match the + ``BotIntent`` being generated, the fall-through ``bot_message_index`` + search ran with the user intent instead of the bot intent. + """ + config = RailsConfig.from_content( + """ + define user express greeting + "hello" + """, + yaml_content="rails:\n dialog:\n single_call:\n enabled: True\n", + ) + actions = LLMGenerationActions( + config=config, + llm=FakeLLMModel(responses=[' "regenerated bot message"']), + llm_task_manager=LLMTaskManager(config), + get_embedding_search_provider_instance=MagicMock(return_value=None), + ) + + # Record the text the bot-message index is searched with. + searched_with = [] + + async def recording_search(text, max_results, threshold): + searched_with.append(text) + return [] + + actions.bot_message_index = MagicMock() + actions.bot_message_index.search = recording_search + + # A single-call cache whose bot intent does NOT match the BotIntent being + # generated, so the cache is a miss and generation falls through. + events = [ + new_event_dict( + "UserIntent", + intent="express greeting", + additional_info=build_single_call_payload( + bot_intent_event=new_event_dict("BotIntent", intent="cached bot intent"), + bot_message_event=new_event_dict("BotMessage", text="cached message"), + ), + ), + new_event_dict("BotIntent", intent="respond kindly"), + ] + + await actions.generate_bot_message(events=events, context={}) + + assert searched_with == ["respond kindly"] + + @pytest.mark.asyncio + async def test_cache_miss_evicts_streaming_handoff(self): + """A single-call cache miss evicts an embedded streaming handoff (no leak). + + If the cached bot message carried a ``<>`` marker and the + bot intent no longer matches, the registered handler must be evicted + rather than lingering for the module lifetime. + """ + from nemoguardrails.actions.llm.generation import _streaming_handoff + from nemoguardrails.streaming import StreamingHandler + + config = RailsConfig.from_content( + """ + define user express greeting + "hello" + """, + yaml_content="rails:\n dialog:\n single_call:\n enabled: True\n", + ) + actions = LLMGenerationActions( + config=config, + llm=FakeLLMModel(responses=[' "regenerated bot message"']), + llm_task_manager=LLMTaskManager(config), + get_embedding_search_provider_instance=MagicMock(return_value=None), + ) + + async def _empty_search(text, max_results, threshold): + return [] + + actions.bot_message_index = MagicMock() + actions.bot_message_index.search = _empty_search + + handler = StreamingHandler() + marker = _streaming_handoff.register(handler) + + events = [ + new_event_dict( + "UserIntent", + intent="express greeting", + additional_info=build_single_call_payload( + bot_intent_event=new_event_dict("BotIntent", intent="cached bot intent"), + bot_message_event=new_event_dict("BotMessage", text=marker), + ), + ), + new_event_dict("BotIntent", intent="respond kindly"), + ] + + await actions.generate_bot_message(events=events, context={}) + + # The handoff was consumed on the cache-miss fall-through, not leaked. + with pytest.raises(KeyError): + _streaming_handoff.take(handler.uid) + + @pytest.mark.asyncio + async def test_predefined_message_evicts_pending_handoff(self): + """A predefined-message bot turn still evicts a pending single-call handoff. + + If phase 1 (streaming) registered a handoff but phase 3 resolves the bot + intent to a predefined ``define bot`` message, the handler must not leak. + """ + from nemoguardrails.actions.llm.generation import _streaming_handoff + from nemoguardrails.streaming import StreamingHandler + + config = RailsConfig.from_content( + """ + define user express greeting + "hello" + define bot express greeting + "Hello!" + """, + yaml_content="rails:\n dialog:\n single_call:\n enabled: True\n", + ) + actions = LLMGenerationActions( + config=config, + llm=FakeLLMModel(responses=[]), + llm_task_manager=LLMTaskManager(config), + get_embedding_search_provider_instance=MagicMock(return_value=None), + ) + + handler = StreamingHandler() + marker = _streaming_handoff.register(handler) + + events = [ + new_event_dict( + "UserIntent", + intent="express greeting", + additional_info=build_single_call_payload( + bot_intent_event=new_event_dict("BotIntent", intent="express greeting"), + bot_message_event=new_event_dict("BotMessage", text=marker), + ), + ), + new_event_dict("BotIntent", intent="express greeting"), + ] + + await actions.generate_bot_message(events=events, context={}) + + # The predefined-message branch bypassed the cache but still released the handoff. + with pytest.raises(KeyError): + _streaming_handoff.take(handler.uid) + + @pytest.mark.asyncio + async def test_cache_returns_none_when_last_user_event_not_user_intent(self): + """The cache falls back when the last user-intent event is a ``UserMessage`` + rather than a ``UserIntent`` (current silent fall-through).""" + config = RailsConfig.from_content( + 'define user express greeting\n "hello"\n', + yaml_content="rails:\n dialog:\n single_call:\n enabled: True\n", + ) + actions = LLMGenerationActions( + config=config, + llm=FakeLLMModel(responses=[]), + llm_task_manager=LLMTaskManager(config), + get_embedding_search_provider_instance=MagicMock(return_value=None), + ) + result = await actions._bot_message_from_single_call_cache( + {"type": "UserMessage", "text": "hi"}, + events=[new_event_dict("BotIntent", intent="respond")], + streaming_handler=None, + ) + assert result is None + + +class TestStreaming: + """Streaming surfaces and the handoff registry.""" + + def test_general_streaming_push(self): + """The no-user-messages general path streams during the call. + + The call passes the streaming handler into ``llm_call`` (mode ``stream``) + with ``stop=["User:"]``; the post-call ``push_chunk`` does not add a + duplicate visible chunk. + """ + + async def run(): + config = RailsConfig.from_content(yaml_content="models: []\nstreaming: True\n") + llm = RecordingFakeLLM(responses=["hello world foo"]) + chat = TestChat(config, llm=llm, streaming=True) + chunks = [] + async for chunk in chat.app.stream_async(messages=[{"role": "user", "content": "hi"}]): + chunks.append(chunk) + return llm, chunks + + llm, chunks = asyncio.run(run()) + + assert llm.calls == [("stream", ("User:",))] + assert chunks == ["hello ", "world ", "foo"] + + def test_handoff_registry_round_trip_and_eviction(self): + """The handoff registry round-trips the marker and evicts on take (no leak).""" + from nemoguardrails.actions.llm.generation import _StreamingHandoffRegistry + from nemoguardrails.streaming import StreamingHandler + + registry = _StreamingHandoffRegistry() + handler = StreamingHandler() + + marker = registry.register(handler) + assert marker == f'Bot message: "<>"' + assert registry.parse_marker(marker) == handler.uid + assert registry.parse_marker("not a marker") is None + + # Malformed markers (prefix present but suffix missing/trailing) are not + # markers -- parse_marker must reject them rather than return a bad uid. + assert registry.parse_marker(f'Bot message: "<> trailing') is None + + assert registry.take(handler.uid) is handler + # Taken once, the handler is gone -- it does not linger for the module lifetime. + with pytest.raises(KeyError): + registry.take(handler.uid) + + +class TestGeneralAndPassthrough: + """No-user-messages general path and passthrough mode.""" + + def test_general_no_user_messages(self): + """No user messages, not passthrough: the rendered-GENERAL path. + + Renders the GENERAL prompt with ``relevant_chunks`` and calls the LLM with + ``stop=["User:"]``; the prompt and ``stop`` guard the helper extraction. + """ + config = RailsConfig.from_content(yaml_content="models: []\n") + response, llm = generate(config, ["I am a general answer."], "tell me something") + + assert response.response == "I am a general answer." + assert event_sequence(response) == ["BotMessage:I am a general answer."] + assert llm_tasks(response) == ["general"] + assert llm.calls == [("generate", ("User:",))] + prompt = call_prompts(response)[0] + assert prompt.startswith("Below is a conversation") + assert prompt.endswith("User: tell me something\nAssistant:") + + def test_general_reasoning_trace_emits_bot_thinking(self): + """A reasoning trace produces ``bot_thinking`` context + a ``BotThinking`` event. + + The context update precedes the ``BotThinking`` event, which precedes the + ``BotMessage``. + """ + config = RailsConfig.from_content(yaml_content="models: []\n") + llm = RecordingFakeLLM(llm_responses=[LLMResponse(content="the answer", reasoning="let me think")]) + chat = TestChat(config, llm=llm) + response = cast(GenerationResponse, chat.app.generate("question?", options=LOG_OPTS)) + + assert response.response == "the answer" + assert event_sequence(response) == [ + "ctx:bot_thinking=let me think", + "BotThinking:let me think", + "BotMessage:the answer", + ] + assert llm_tasks(response) == ["general"] + + def test_passthrough_no_fn(self): + """Passthrough without a passthrough fn: the raw-prompt path. + + The prompt is the raw conversation (no rendered GENERAL template) and the + LLM is called with no ``stop`` -- the divergence from the rendered path + that only the model-call channel can see. + """ + config = RailsConfig.from_content(yaml_content="passthrough: true\n") + response, llm = generate(config, ["passthrough llm answer"], "hello") + + assert response.response == "passthrough llm answer" + assert event_sequence(response) == ["BotMessage:passthrough llm answer"] + assert llm_tasks(response) == ["general"] + assert llm.calls == [("generate", None)] + prompt = call_prompts(response)[0] + assert "Below is a conversation" not in prompt + assert "hello" in prompt + + def test_passthrough_with_fn(self): + """A passthrough fn supplies the output and a ``passthrough_output`` event. + + The ``passthrough_output`` rides as a ``ContextUpdate`` event emitted + before the ``BotMessage`` (distinct from ``generate_bot_message``, which + returns it as a ``context_updates`` dict entry); the ordering pins that + form. + """ + config = RailsConfig.from_content(yaml_content="passthrough: true\n") + llm = RecordingFakeLLM(responses=["unused"]) + chat = TestChat(config, llm=llm) + + async def passthrough_fn(context, events): + return "fn output text", {"extra": 1} + + chat.app.passthrough_fn = passthrough_fn + response = cast(GenerationResponse, chat.app.generate("hello", options=LOG_OPTS)) + + assert response.response == "fn output text" + assert event_sequence(response) == [ + "ctx:passthrough_output={'extra': 1}", + "BotMessage:fn output text", + ] + assert llm_tasks(response) == [] + assert llm.calls == [] + + def test_passthrough_tool_calls(self): + """Tool calls in passthrough mode emit a ``BotToolCalls`` event.""" + config = RailsConfig.from_content( + yaml_content=""" + passthrough: true + models: + - type: main + engine: openai + model: gpt-4 + """ + ) + llm = RecordingFakeLLM( + llm_responses=[ + LLMResponse( + content="", + tool_calls=[ + ToolCall( + id="call_1", + type="function", + function=ToolCallFunction(name="t", arguments={"p": "v"}), + ) + ], + ) + ] + ) + chat = TestChat(config, llm=llm) + response = cast(GenerationResponse, chat.app.generate("call the tool", options=LOG_OPTS)) + + assert event_sequence(response) == ["BotToolCalls:1"] + assert llm_tasks(response) == ["general"] + # The full tool-call payload (id, function name, arguments) is surfaced on + # the public response, not just the count. + assert response.response == "" + assert response.tool_calls == [ + {"id": "call_1", "type": "function", "function": {"name": "t", "arguments": {"p": "v"}}} + ] + + def test_single_call_general_no_user_messages(self): + """Single-call enabled but no user messages: the converged general bot turn. + + This branch delegates to ``_emit_general_bot_turn``, so it behaves like + ``generate_user_intent``'s general path: a GENERAL prompt rendered with + ``relevant_chunks`` and an LLM call with ``stop=["User:"]`` (previously it + rendered with no context and no stop). Rendering GENERAL with empty + ``relevant_chunks`` is identical to rendering it with no context. + """ + config = RailsConfig.from_content(yaml_content="passthrough: false\n") + config.rails.dialog.single_call.enabled = True + response, llm = generate(config, ["a general single-call answer"], "hi") + + assert response.response == "a general single-call answer" + assert event_sequence(response) == ["BotMessage:a general single-call answer"] + assert llm_tasks(response) == ["general"] + assert llm.calls == [("generate", ("User:",))] + assert call_prompts(response)[0].startswith("Below is a conversation") + + def test_single_call_general_reasoning_trace(self): + """The single-call general path emits ``BotThinking`` like the multi-call one. + + ``generate_intent_steps_message``'s else-branch previously dropped + reasoning traces; converging it onto ``_emit_general_bot_turn`` now + packages them. + """ + config = RailsConfig.from_content(yaml_content="passthrough: false\n") + config.rails.dialog.single_call.enabled = True + llm = RecordingFakeLLM(llm_responses=[LLMResponse(content="the answer", reasoning="let me think")]) + chat = TestChat(config, llm=llm) + response = cast(GenerationResponse, chat.app.generate("hi", options=LOG_OPTS)) + + assert response.response == "the answer" + assert event_sequence(response) == [ + "ctx:bot_thinking=let me think", + "BotThinking:let me think", + "BotMessage:the answer", + ] + assert llm_tasks(response) == ["general"] + + def test_dialog_with_passthrough_uses_bot_message_branch(self): + """user_messages + passthrough: generate_bot_message takes the passthrough + branch, which reports the ``generate_bot_message`` task but parses GENERAL.""" + config = RailsConfig.from_content( + """ + define user express greeting + "hello" + define flow + user express greeting + bot respond kindly + """, + yaml_content="passthrough: true\n", + ) + response, llm = generate(config, [" express greeting", "passthrough bot answer"], "hello") + + assert response.response == "passthrough bot answer" + assert llm_tasks(response) == ["generate_user_intent", "generate_bot_message"] + assert event_sequence(response)[-1] == "BotMessage:passthrough bot answer" + + +class TestGenerateValue: + """``generate_value`` action ($var = ...).""" + + def test_three_phase(self): + """``$var = ...`` triggers a ``generate_value`` LLM call between phases.""" + config = RailsConfig.from_path(os.path.join(os.path.dirname(__file__), "test_configs", "generate_value")) + llm = RecordingFakeLLM( + responses=[ + " ask math question", + '"What is the largest prime factor for 1024?"', + ' "The largest prime factor for 1024 is 2."', + ] + ) + chat = TestChat(config, llm=llm) + + async def mock_wolfram_alpha_request_action(query): + return "2" + + chat.app.register_action(mock_wolfram_alpha_request_action, "wolfram alpha request") + response = cast( + GenerationResponse, + chat.app.generate("What is the largest prime factor for 1024", options=LOG_OPTS), + ) + + assert response.response == "The largest prime factor for 1024 is 2." + assert llm_tasks(response) == [ + "generate_user_intent", + "generate_value", + "generate_bot_message", + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "completion,expected", + [ + ("42", "42"), + ("42;", "42"), # trailing ``;`` is stripped before safe_eval + ('"hello"', "hello"), + ], + ) + async def test_parsing(self, completion, expected): + """The value-parsing tail (first line, ``;`` strip, safe_eval) is preserved. + + ``safe_eval`` is intentionally lenient (it wraps unparseable input as a + string rather than raising), so the action's ``except -> ValueError`` + branch is effectively unreachable in practice and is not exercised here. + """ + config = RailsConfig.from_content(yaml_content="models: []\n") + actions = LLMGenerationActions( + config=config, + llm=FakeLLMModel(responses=[completion]), + llm_task_manager=LLMTaskManager(config), + get_embedding_search_provider_instance=MagicMock(return_value=None), + ) + events = [ + {"type": "UserMessage", "text": "give me a number"}, + {"type": "StartInternalSystemAction", "action_name": "generate_value", "action_result_key": "x"}, + ] + value = await actions.generate_value(instructions="give a number", events=events, var_name="x") + assert value == expected + + +class TestStreamingPatternFor: + """Contract of ``_streaming_pattern_for(output_parser, *, include_bot_message_parser)``.""" + + def test_pattern_matrix(self): + from nemoguardrails.actions.llm.generation import _streaming_pattern_for + + verbose = ('Bot message: "', '"') + plain = (' "', '"') + + # verbose_v1 is always verbose, regardless of the flag. + assert _streaming_pattern_for("verbose_v1", include_bot_message_parser=False) == verbose + assert _streaming_pattern_for("verbose_v1", include_bot_message_parser=True) == verbose + # bot_message is verbose only when the flag is set (the generate_bot_message site). + assert _streaming_pattern_for("bot_message", include_bot_message_parser=True) == verbose + assert _streaming_pattern_for("bot_message", include_bot_message_parser=False) == plain + # Anything else is the plain pattern. + assert _streaming_pattern_for("something_else", include_bot_message_parser=True) == plain + assert _streaming_pattern_for(None, include_bot_message_parser=False) == plain + + +class TestBotTurnOutputEvents: + """Contract of the module-level ``_bot_turn_output_events`` helper.""" + + def test_reasoning_trace_prepends_bot_thinking_and_records_context(self): + from nemoguardrails.actions.llm.generation import _bot_turn_output_events + from nemoguardrails.context import reasoning_trace_var + + final = new_event_dict("BotMessage", text="hi") + context_updates = {} + reasoning_trace_var.set("let me think") + try: + events = _bot_turn_output_events(final, context_updates) + finally: + reasoning_trace_var.set(None) + + assert [event["type"] for event in events] == ["BotThinking", "BotMessage"] + assert events[0]["content"] == "let me think" + assert events[1] is final + assert context_updates["bot_thinking"] == "let me think" + + def test_reasoning_trace_without_context_updates_does_not_record(self): + from nemoguardrails.actions.llm.generation import _bot_turn_output_events + from nemoguardrails.context import reasoning_trace_var + + reasoning_trace_var.set("let me think") + try: + events = _bot_turn_output_events(new_event_dict("BotMessage", text="hi")) + finally: + reasoning_trace_var.set(None) + + # BotThinking is still emitted; there is just no dict to record into. + assert [event["type"] for event in events] == ["BotThinking", "BotMessage"] + + def test_no_reasoning_trace_returns_only_final(self): + from nemoguardrails.actions.llm.generation import _bot_turn_output_events + from nemoguardrails.context import reasoning_trace_var + + reasoning_trace_var.set(None) + final = new_event_dict("BotMessage", text="hi") + events = _bot_turn_output_events(final, {}) + assert events == [final] + + +class TestBuildIntentStepsExamples: + """Contract of the single-call few-shot example builder (prompt-only, so the + equivalence scenarios cannot observe it end-to-end).""" + + def _actions(self): + config = RailsConfig.from_content('define user express greeting\n "hi"\n') + return LLMGenerationActions( + config=config, + llm=FakeLLMModel(responses=[]), + llm_task_manager=LLMTaskManager(config), + get_embedding_search_provider_instance=MagicMock(return_value=None), + ) + + @pytest.mark.asyncio + async def test_pairs_intent_with_flow_and_bot_message(self): + from types import SimpleNamespace + + actions = self._actions() + + async def user_search(text, max_results, threshold): + return [SimpleNamespace(text="hello", meta={"intent": "express greeting"})] + + async def flows_search(text, max_results, threshold): + flow = "user express greeting\nbot express greeting" + return [SimpleNamespace(text=flow, meta={"flow": flow})] + + async def bot_search(text, max_results, threshold): + return [SimpleNamespace(text="express greeting", meta={"text": "Hello!"})] + + actions.user_message_index = MagicMock() + actions.user_message_index.search = user_search + actions.flows_index = MagicMock() + actions.flows_index.search = flows_search + actions.bot_message_index = MagicMock() + actions.bot_message_index.search = bot_search + + examples, intents = await actions._build_intent_steps_examples("hi") + + assert intents == ["express greeting"] + assert len(examples) == 1 + assert 'user "hello"' in examples[0] + assert " express greeting" in examples[0] + assert "bot express greeting" in examples[0] + assert '"Hello!"' in examples[0] + + @pytest.mark.asyncio + async def test_intent_without_flow_is_skipped(self): + from types import SimpleNamespace + + actions = self._actions() + + async def user_search(text, max_results, threshold): + return [SimpleNamespace(text="hello", meta={"intent": "express greeting"})] + + async def empty_flows_search(text, max_results, threshold): + return [] + + actions.user_message_index = MagicMock() + actions.user_message_index.search = user_search + actions.flows_index = MagicMock() + actions.flows_index.search = empty_flows_search + actions.bot_message_index = None + + examples, intents = await actions._build_intent_steps_examples("hi") + + # The intent is still a candidate, but with no flow it yields no example. + assert intents == ["express greeting"] + assert examples == [] + + +class TestNextStepsBranches: + """Single-step branches of ``generate_next_steps`` (non-multi-step).""" + + def _actions(self, response): + config = RailsConfig.from_content('define user express greeting\n "hi"\n') + return LLMGenerationActions( + config=config, + llm=FakeLLMModel(responses=[response]), + llm_task_manager=LLMTaskManager(config), + get_embedding_search_provider_instance=MagicMock(return_value=None), + ) + + @pytest.mark.asyncio + async def test_bot_intent_comma_cleanup(self): + actions = self._actions("bot respond politely, and more") + events = [new_event_dict("UserIntent", intent="express greeting")] + result = await actions.generate_next_steps(events=events) + assert result.events[0]["type"] == "BotIntent" + assert result.events[0]["intent"] == "respond politely" + + @pytest.mark.asyncio + async def test_non_bot_line_yields_general_response(self): + actions = self._actions("this is not a bot line") + events = [new_event_dict("UserIntent", intent="express greeting")] + result = await actions.generate_next_steps(events=events) + assert result.events[0]["type"] == "BotIntent" + assert result.events[0]["intent"] == "general response" + + +class TestPassthroughRawPromptList: + """The passthrough ``raw_llm_request`` is-a-list branch of ``_emit_general_bot_turn``.""" + + @pytest.mark.asyncio + async def test_last_user_content_replaced_in_prompt(self): + from nemoguardrails.context import raw_llm_request + + captured = {} + + class CapturingFake(FakeLLMModel): + async def generate_async(self, prompt, *, stop=None, **kwargs): + captured["prompt"] = prompt + return await super().generate_async(prompt, stop=stop, **kwargs) + + config = RailsConfig.from_content(yaml_content="passthrough: true\n") + actions = LLMGenerationActions( + config=config, + llm=CapturingFake(responses=["passthrough answer"]), + llm_task_manager=LLMTaskManager(config), + get_embedding_search_provider_instance=MagicMock(return_value=None), + ) + + raw_list = [{"role": "user", "content": "original"}] + token = raw_llm_request.set(raw_list) + try: + events = [new_event_dict("UserMessage", text="altered by input rails")] + await actions.generate_user_intent(events=events, context={}, config=config) + finally: + raw_llm_request.reset(token) + + # The prompt sent to the LLM reflects the input-rail-altered user content + # (the shallow copy shares the element dict, so the mutation is visible); + # the original content is gone. + assert isinstance(captured["prompt"], list) + prompt_text = str(captured["prompt"]) + assert "altered by input rails" in prompt_text + assert "original" not in prompt_text + + +class TestGenerateGeneralResponse: + """Contract of the shared ``_generate_general_response`` helper.""" + + @pytest.mark.asyncio + async def test_reports_llm_call_task_distinct_from_parse_task(self): + from nemoguardrails.context import llm_call_info_var + from nemoguardrails.llm.types import Task + + config = RailsConfig.from_content(yaml_content="models: []\n") + actions = LLMGenerationActions( + config=config, + llm=FakeLLMModel(responses=["output"]), + llm_task_manager=LLMTaskManager(config), + get_embedding_search_provider_instance=MagicMock(return_value=None), + ) + + llm_call_info_var.set(None) + await actions._generate_general_response( + generation_llm=actions.llm, + prompt="hello", + stream_during_call=False, + llm_call_task=Task.GENERATE_BOT_MESSAGE, + parse_task=Task.GENERAL, + ) + + info = llm_call_info_var.get() + assert info is not None + # The call is reported under llm_call_task even though it parses as parse_task. + assert info.task == Task.GENERATE_BOT_MESSAGE.value + + +class TestMultiCallMultimodal: + """Non-single-call multimodal normalization in ``generate_user_intent``.""" + + def test_non_text_parts_ignored(self): + config = RailsConfig.from_content('define user express greeting\n "hello"\n') + llm = RecordingFakeLLM(responses=[" express greeting", "bot respond", ' "Hi there!"']) + chat = TestChat(config, llm=llm) + response = cast( + GenerationResponse, + chat.app.generate( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello there"}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}, + ], + } + ], + options=LOG_OPTS, + ), + ) + assert "UserIntent:express greeting" in event_sequence(response)