rasa-pro 3.12.0.dev2__py3-none-any.whl → 3.12.0.dev4__py3-none-any.whl

rasa/dialogue_understanding/generator/command_parser.py

@@ -0,0 +1,201 @@
+import re
+from functools import lru_cache
+from typing import Any, Callable, Dict, List, Optional, Type, Union
+
+import structlog
+
+from rasa.dialogue_understanding.commands import (
+    CancelFlowCommand,
+    ChitChatAnswerCommand,
+    ClarifyCommand,
+    Command,
+    HumanHandoffCommand,
+    KnowledgeAnswerCommand,
+    RepeatBotMessagesCommand,
+    SetSlotCommand,
+    SkipQuestionCommand,
+    StartFlowCommand,
+)
+from rasa.dialogue_understanding.commands.utils import start_flow_by_name
+from rasa.shared.core.flows import FlowsList
+
+structlogger = structlog.get_logger()
+
+
+DEFAULT_COMMANDS = [
+    SetSlotCommand,
+    StartFlowCommand,
+    CancelFlowCommand,
+    ChitChatAnswerCommand,
+    SkipQuestionCommand,
+    KnowledgeAnswerCommand,
+    HumanHandoffCommand,
+    ClarifyCommand,
+    RepeatBotMessagesCommand,
+]
+
+
+@lru_cache(maxsize=128)
+def _get_compiled_pattern(pattern_str: str) -> re.Pattern:
+    return re.compile(pattern_str)
+
+
+def _create_default_commands(
+    default_commands_to_remove: Union[List[Type[Command]], List[str]],
+) -> List[Type[Command]]:
+    """Return an updated list of default commands after removing the provided commands.
+
+    Args:
+        default_commands_to_remove: A list of commands to remove from the default
+            commands.
+
+    Returns:
+        A copy of the default commands with the provided commands removed.
+    """
+    commands_after_removal = []
+    for command in DEFAULT_COMMANDS:
+        if (
+            command not in default_commands_to_remove
+            and command.__name__ not in default_commands_to_remove
+        ):
+            commands_after_removal.append(command)
+
+    return commands_after_removal
+
+
+def _get_additional_parsing_logic(
+    command_clz: Type[Command],
+) -> Optional[Callable[[Command, FlowsList], Optional[Command]]]:
+    """Get additional parsing logic for a command."""
+    command_to_parsing_fn_mapper: Dict[
+        Type[Command], Callable[..., Optional[Command]]
+    ] = {
+        ClarifyCommand: _parse_clarify_command,
+        SetSlotCommand: _parse_set_slot_command,
+        StartFlowCommand: _parse_start_flow_command,
+    }
+    return command_to_parsing_fn_mapper.get(command_clz)
+
+
+def parse_commands(
+    actions: Optional[str],
+    flows: FlowsList,
+    is_handle_flows_prompt: bool = False,
+    additional_commands: Optional[List[Type[Command]]] = None,
+    default_commands_to_remove: Optional[Union[List[Type[Command]], List[str]]] = None,
+    **kwargs: Any,
+) -> List[Command]:
+    """Parse a list of action commands."""
+    if not actions:
+        return []
+
+    commands: List[Command] = []
+    default_commands = DEFAULT_COMMANDS
+    if default_commands_to_remove:
+        default_commands = _create_default_commands(default_commands_to_remove)
+
+    for action in actions.strip().splitlines():
+        if is_handle_flows_prompt:
+            # Needed for multistep command generator.
+            if (
+                len(commands) >= 2
+                or len(commands) == 1
+                and isinstance(commands[0], ClarifyCommand)
+            ):
+                break
+        commands.extend(
+            _parse_standard_commands(default_commands, action.strip(), flows, **kwargs)
+        )
+        if additional_commands:
+            commands.extend(
+                _parse_custom_commands(
+                    additional_commands, action.strip(), flows, **kwargs
+                )
+            )
+    return commands
+
+
+def _parse_standard_commands(
+    standard_commands: List[Type[Command]],
+    action: str,
+    flows: FlowsList,
+    **kwargs: Any,
+) -> List[Command]:
+    """Parse predefined standard commands."""
+    commands: List[Command] = []
+    for command_clz in standard_commands:
+        pattern = _get_compiled_pattern(command_clz.regex_pattern())
+        if match := pattern.search(action):
+            parsed_command = command_clz.from_dsl(match, **kwargs)
+            if _additional_parsing_fn := _get_additional_parsing_logic(command_clz):
+                parsed_command = _additional_parsing_fn(parsed_command, flows, **kwargs)
+            if parsed_command:
+                commands.append(parsed_command)
+    return commands
+
+
+def _parse_custom_commands(
+    custom_commands: List[Type[Command]],
+    action: str,
+    flows: FlowsList,
+    **kwargs: Any,
+) -> List[Command]:
+    """Parse custom commands."""
+    commands: List[Command] = []
+    for command_clz in custom_commands:
+        pattern = _get_compiled_pattern(command_clz.regex_pattern())
+        if match := pattern.search(action):
+            if parsed_command := command_clz.from_dsl(match, flows=flows, **kwargs):
+                commands.append(parsed_command)
+    return commands
+
+
+def _parse_set_slot_command(
+    parsed_command: Optional[SetSlotCommand], flows: FlowsList, **kwargs: Any
+) -> Optional[Command]:
+    """Additional parsing logic for the SetSlotCommand."""
+    if not parsed_command:
+        return None
+    # error case where the llm tries to start a flow using a slot set
+    if parsed_command.name == "flow_name":
+        return start_flow_by_name(parsed_command.value, flows)
+    return parsed_command
+
+
+def _parse_clarify_command(
+    parsed_command: Optional[ClarifyCommand], flows: FlowsList, **kwargs: Any
+) -> Optional[Command]:
+    """Additional parsing logic for the ClarifyCommand."""
+    if not parsed_command:
+        return None
+
+    # if no options are available
+    if not parsed_command.options:
+        # Return the command if options are optional; else, return an empty list
+        return (
+            ClarifyCommand([])
+            if kwargs.get("clarify_options_optional", False)
+            else None
+        )
+    valid_options = [
+        flow for flow in parsed_command.options if flow in flows.user_flow_ids
+    ]
+    unique_valid_options = list(set(valid_options))
+
+    # if there is only one valid option, start the flow
+    if len(unique_valid_options) == 1:
+        return start_flow_by_name(valid_options[0], flows)
+
+    # if there are multiple valid options, return a clarify command
+    if len(valid_options) > 1:
+        return ClarifyCommand(valid_options)
+    return None
+
+
+def _parse_start_flow_command(
+    parsed_command: Optional[StartFlowCommand], flows: FlowsList, **kwargs: Any
+) -> Optional[Command]:
+    """Additional parsing logic for the StartFlowCommand."""
+    if not parsed_command:
+        return None
+    return start_flow_by_name(parsed_command.flow, flows)

rasa/dialogue_understanding/generator/constants.py

@@ -1,7 +1,7 @@
 from rasa.shared.constants import (
-    PROVIDER_CONFIG_KEY,
-    OPENAI_PROVIDER,
     MODEL_CONFIG_KEY,
+    OPENAI_PROVIDER,
+    PROVIDER_CONFIG_KEY,
     TIMEOUT_CONFIG_KEY,
 )
 from rasa.shared.utils.llm import (

rasa/dialogue_understanding/generator/flow_retrieval.py

@@ -18,7 +18,7 @@ flows within a conversational context.
 """
 
 import importlib
-from typing import Dict, Text, Any, List, Optional
+from typing import Any, Dict, List, Optional, Text
 
 import structlog
 from jinja2 import Template
@@ -27,19 +27,19 @@ from langchain.schema.embeddings import Embeddings
 from langchain_community.vectorstores.faiss import FAISS
 from langchain_community.vectorstores.utils import DistanceStrategy
 
+import rasa.shared.utils.io
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
-import rasa.shared.utils.io
 from rasa.shared.constants import (
     EMBEDDINGS_CONFIG_KEY,
-    PROVIDER_CONFIG_KEY,
     OPENAI_PROVIDER,
+    PROVIDER_CONFIG_KEY,
 )
 from rasa.shared.core.domain import Domain
 from rasa.shared.core.flows import FlowsList
 from rasa.shared.core.trackers import DialogueStateTracker
 from rasa.shared.exceptions import ProviderClientAPIException
-from rasa.shared.nlu.constants import TEXT, FLOWS_FROM_SEMANTIC_SEARCH
+from rasa.shared.nlu.constants import FLOWS_FROM_SEMANTIC_SEARCH, TEXT
 from rasa.shared.nlu.training_data.message import Message
 from rasa.shared.providers.embedding._langchain_embedding_client_adapter import (
     _LangchainEmbeddingClientAdapter,
@@ -48,13 +48,13 @@ from rasa.shared.utils.health_check.embeddings_health_check_mixin import (
     EmbeddingsHealthCheckMixin,
 )
 from rasa.shared.utils.llm import (
-    tracker_as_readable_transcript,
-    embedder_factory,
     DEFAULT_OPENAI_EMBEDDING_MODEL_NAME,
     USER,
-    get_prompt_template,
     allowed_values_for_slot,
+    embedder_factory,
+    get_prompt_template,
     resolve_model_client_config,
+    tracker_as_readable_transcript,
 )
 
 DEFAULT_FLOW_DOCUMENT_TEMPLATE = importlib.resources.read_text(

rasa/dialogue_understanding/generator/llm_based_command_generator.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from functools import lru_cache
-from typing import Dict, Any, List, Optional, Tuple, Union, Text
+from typing import Any, Dict, List, Optional, Text, Tuple, Union
 
 import structlog
 from jinja2 import Template
@@ -8,30 +8,29 @@ from jinja2 import Template
 import rasa.shared.utils.io
 from rasa.dialogue_understanding.commands import (
     Command,
-    StartFlowCommand,
 )
 from rasa.dialogue_understanding.generator import CommandGenerator
 from rasa.dialogue_understanding.generator.constants import (
     DEFAULT_LLM_CONFIG,
-    LLM_CONFIG_KEY,
-    FLOW_RETRIEVAL_KEY,
     FLOW_RETRIEVAL_ACTIVE_KEY,
     FLOW_RETRIEVAL_FLOW_THRESHOLD,
+    FLOW_RETRIEVAL_KEY,
+    LLM_CONFIG_KEY,
 )
 from rasa.dialogue_understanding.generator.flow_retrieval import FlowRetrieval
-from rasa.engine.graph import GraphComponent, ExecutionContext
+from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
 from rasa.shared.core.domain import Domain
-from rasa.shared.core.flows import FlowStep, Flow, FlowsList
+from rasa.shared.core.flows import Flow, FlowsList, FlowStep
 from rasa.shared.core.flows.steps.collect import CollectInformationFlowStep
 from rasa.shared.core.trackers import DialogueStateTracker
-from rasa.shared.exceptions import FileIOException
-from rasa.shared.exceptions import ProviderClientAPIException
+from rasa.shared.exceptions import FileIOException, ProviderClientAPIException
 from rasa.shared.nlu.constants import FLOWS_IN_PROMPT
 from rasa.shared.nlu.training_data.message import Message
 from rasa.shared.nlu.training_data.training_data import TrainingData
+from rasa.shared.providers.llm.llm_response import LLMResponse
 from rasa.shared.utils.health_check.llm_health_check_mixin import LLMHealthCheckMixin
 from rasa.shared.utils.llm import (
     allowed_values_for_slot,
@@ -304,22 +303,21 @@ class LLMBasedCommandGenerator(
         )
         return filtered_flows
 
-    async def invoke_llm(self, prompt: Text) -> Optional[Text]:
+    async def invoke_llm(self, prompt: Text) -> Optional[LLMResponse]:
         """Use LLM to generate a response.
 
         Args:
             prompt: The prompt to send to the LLM.
 
         Returns:
-            The generated text.
+            An LLMResponse object.
 
         Raises:
-            ProviderClientAPIException if an error during API call.
+            ProviderClientAPIException: If an error occurs during the LLM API call.
         """
         llm = llm_factory(self.config.get(LLM_CONFIG_KEY), DEFAULT_LLM_CONFIG)
         try:
-            llm_response = await llm.acompletion(prompt)
-            return llm_response.choices[0]
+            return await llm.acompletion(prompt)
         except Exception as e:
             # unfortunately, langchain does not wrap LLM exceptions which means
             # we have to catch all exceptions here
@@ -328,50 +326,6 @@ class LLMBasedCommandGenerator(
                 message="LLM call exception", original_exception=e
             )
 
-    @staticmethod
-    def start_flow_by_name(flow_name: str, flows: FlowsList) -> List[Command]:
-        """Start a flow by name.
-
-        If the flow does not exist, no command is returned.
-        """
-        if flow_name in flows.user_flow_ids:
-            return [StartFlowCommand(flow=flow_name)]
-        else:
-            structlogger.debug(
-                "llm_command_generator.flow.start_invalid_flow_id", flow=flow_name
-            )
-            return []
-
-    @staticmethod
-    def is_none_value(value: str) -> bool:
-        """Check if the value is a none value."""
-        return value in {
-            "[missing information]",
-            "[missing]",
-            "None",
-            "undefined",
-            "null",
-        }
-
-    @staticmethod
-    def clean_extracted_value(value: str) -> str:
-        """Clean up the extracted value from the llm."""
-        # replace any combination of single quotes, double quotes, and spaces
-        # from the beginning and end of the string
-        return value.strip("'\" ")
-
-    @classmethod
-    def get_nullable_slot_value(cls, slot_value: str) -> Union[str, None]:
-        """Get the slot value or None if the value is a none value.
-
-        Args:
-            slot_value: the value to coerce
-
-        Returns:
-            The slot value or None if the value is a none value.
-        """
-        return slot_value if not cls.is_none_value(slot_value) else None
-
     def prepare_flows_for_template(
         self, flows: FlowsList, tracker: DialogueStateTracker
     ) -> List[Dict[str, Any]]:

rasa/dialogue_understanding/generator/llm_command_generator.py

@@ -10,6 +10,7 @@ from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
 from rasa.shared.exceptions import ProviderClientAPIException
+from rasa.shared.providers.llm.llm_response import LLMResponse
 from rasa.shared.utils.io import raise_deprecation_warning
 
 structlogger = structlog.get_logger()
@@ -53,7 +54,7 @@ class LLMCommandGenerator(SingleStepLLMCommandGenerator):
             **kwargs,
         )
 
-    async def invoke_llm(self, prompt: Text) -> Optional[Text]:
+    async def invoke_llm(self, prompt: Text) -> Optional[LLMResponse]:
         try:
             return await super().invoke_llm(prompt)
         except ProviderClientAPIException:

rasa/dialogue_understanding/generator/multi_step/multi_step_llm_command_generator.py

@@ -1,63 +1,63 @@
 import importlib.resources
-import re
-from typing import Dict, Any, List, Optional, Tuple, Union, Text
+from typing import Any, Dict, List, Optional, Text, Tuple, Union
 
 import structlog
 from jinja2 import Template
 
 import rasa.shared.utils.io
 from rasa.dialogue_understanding.commands import (
+    CannotHandleCommand,
+    ChangeFlowCommand,
     Command,
     ErrorCommand,
     SetSlotCommand,
-    CancelFlowCommand,
     StartFlowCommand,
-    HumanHandoffCommand,
-    ChitChatAnswerCommand,
-    SkipQuestionCommand,
-    KnowledgeAnswerCommand,
-    ClarifyCommand,
-    CannotHandleCommand,
 )
-from rasa.dialogue_understanding.commands.change_flow_command import ChangeFlowCommand
+from rasa.dialogue_understanding.commands.can_not_handle_command import (
+    DATA_KEY_CANNOT_HANDLE_REASON,
+)
+from rasa.dialogue_understanding.generator.command_parser import (
+    parse_commands as parse_commands_using_command_parsers,
+)
 from rasa.dialogue_understanding.generator.constants import (
+    DEFAULT_LLM_CONFIG,
+    FLOW_RETRIEVAL_KEY,
     LLM_CONFIG_KEY,
     USER_INPUT_CONFIG_KEY,
-    FLOW_RETRIEVAL_KEY,
-    DEFAULT_LLM_CONFIG,
 )
 from rasa.dialogue_understanding.generator.flow_retrieval import FlowRetrieval
 from rasa.dialogue_understanding.generator.llm_based_command_generator import (
     LLMBasedCommandGenerator,
 )
 from rasa.dialogue_understanding.stack.frames import UserFlowStackFrame
-from rasa.dialogue_understanding.stack.utils import (
-    top_flow_frame,
-    top_user_flow_frame,
-    user_flows_on_the_stack,
+from rasa.dialogue_understanding.stack.utils import top_flow_frame, top_user_flow_frame
+from rasa.dialogue_understanding.utils import (
+    add_commands_to_message_parse_data,
+    add_prompt_to_message_parse_data,
 )
 from rasa.engine.graph import ExecutionContext
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
 from rasa.shared.constants import (
-    RASA_PATTERN_CANNOT_HANDLE_NOT_SUPPORTED,
     EMBEDDINGS_CONFIG_KEY,
+    RASA_PATTERN_CANNOT_HANDLE_NOT_SUPPORTED,
+    ROUTE_TO_CALM_SLOT,
 )
-from rasa.shared.constants import ROUTE_TO_CALM_SLOT
-from rasa.shared.core.flows import FlowStep, Flow, FlowsList
+from rasa.shared.core.flows import Flow, FlowsList, FlowStep
 from rasa.shared.core.flows.steps.collect import CollectInformationFlowStep
 from rasa.shared.core.trackers import DialogueStateTracker
 from rasa.shared.exceptions import ProviderClientAPIException
 from rasa.shared.nlu.constants import TEXT
 from rasa.shared.nlu.training_data.message import Message
+from rasa.shared.providers.llm.llm_response import LLMResponse
 from rasa.shared.utils.io import deep_container_fingerprint
 from rasa.shared.utils.llm import (
-    get_prompt_template,
-    tracker_as_readable_transcript,
-    sanitize_message_for_prompt,
     allowed_values_for_slot,
+    get_prompt_template,
     resolve_model_client_config,
+    sanitize_message_for_prompt,
+    tracker_as_readable_transcript,
 )
 
 # multistep template keys
@@ -199,7 +199,7 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
                 message, flows, tracker
             )
             commands = self._clean_up_commands(commands)
-            self._add_commands_to_message_parse_data(
+            add_commands_to_message_parse_data(
                 message, MultiStepLLMCommandGenerator.__name__, commands
             )
         except ProviderClientAPIException:
@@ -242,86 +242,15 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
         Returns:
             The parsed commands.
         """
-        if not actions:
-            return []
-
-        commands: List[Command] = []
-
-        slot_set_re = re.compile(
-            r"""SetSlot\(['"]?([a-zA-Z_][a-zA-Z0-9_-]*)['"]?, ?['"]?(.*)['"]?\)"""
+        commands = parse_commands_using_command_parsers(
+            actions,
+            flows,
+            is_handle_flows_prompt=is_handle_flows_prompt,
+            additional_commands=[CannotHandleCommand, ChangeFlowCommand],
+            data={
+                DATA_KEY_CANNOT_HANDLE_REASON: RASA_PATTERN_CANNOT_HANDLE_NOT_SUPPORTED
+            },
         )
-        start_flow_re = re.compile(r"StartFlow\(['\"]?([a-zA-Z0-9_-]+)['\"]?\)")
-        change_flow_re = re.compile(r"ChangeFlow\(\)")
-        cancel_flow_re = re.compile(r"CancelFlow\(\)")
-        chitchat_re = re.compile(r"ChitChat\(\)")
-        skip_question_re = re.compile(r"SkipQuestion\(\)")
-        knowledge_re = re.compile(r"SearchAndReply\(\)")
-        humand_handoff_re = re.compile(r"HumanHandoff\(\)")
-        clarify_re = re.compile(r"Clarify\(([\"\'a-zA-Z0-9_, ]+)\)")
-        cannot_handle_re = re.compile(r"CannotHandle\(\)")
-
-        for action in actions.strip().splitlines():
-            if is_handle_flows_prompt:
-                if (
-                    len(commands) >= 2
-                    or len(commands) == 1
-                    and isinstance(commands[0], ClarifyCommand)
-                ):
-                    break
-
-            if cannot_handle_re.search(action):
-                commands.append(
-                    CannotHandleCommand(RASA_PATTERN_CANNOT_HANDLE_NOT_SUPPORTED)
-                )
-            if match := slot_set_re.search(action):
-                slot_name = cls.clean_extracted_value(match.group(1).strip())
-                slot_value = cls.clean_extracted_value(match.group(2))
-                # error case where the llm tries to start a flow using a slot set
-                if slot_name == "flow_name":
-                    commands.extend(cls.start_flow_by_name(slot_value, flows))
-                else:
-                    typed_slot_value = cls.get_nullable_slot_value(slot_value)
-                    commands.append(
-                        SetSlotCommand(name=slot_name, value=typed_slot_value)
-                    )
-            elif match := start_flow_re.search(action):
-                flow_name = match.group(1).strip()
-                commands.extend(cls.start_flow_by_name(flow_name, flows))
-            elif cancel_flow_re.search(action):
-                commands.append(CancelFlowCommand())
-            elif chitchat_re.search(action):
-                commands.append(ChitChatAnswerCommand())
-            elif skip_question_re.search(action):
-                commands.append(SkipQuestionCommand())
-            elif knowledge_re.search(action):
-                commands.append(KnowledgeAnswerCommand())
-            elif humand_handoff_re.search(action):
-                commands.append(HumanHandoffCommand())
-            elif match := clarify_re.search(action):
-                options = sorted([opt.strip() for opt in match.group(1).split(",")])
-                # Remove surrounding quotes if present
-                cleaned_options = []
-                for flow in options:
-                    if (flow.startswith('"') and flow.endswith('"')) or (
-                        flow.startswith("'") and flow.endswith("'")
-                    ):
-                        cleaned_options.append(flow[1:-1])
-                    else:
-                        cleaned_options.append(flow)
-                # check if flow is valid
-                valid_options = [
-                    flow
-                    for flow in cleaned_options
-                    if flow in flows.user_flow_ids
-                    and flow not in user_flows_on_the_stack(tracker.stack)
-                ]
-                if len(valid_options) == 1:
-                    commands.extend(cls.start_flow_by_name(valid_options[0], flows))
-                elif 1 < len(valid_options) <= 5:
-                    commands.append(ClarifyCommand(valid_options))
-            elif change_flow_re.search(action):
-                commands.append(ChangeFlowCommand())
-
         if not commands:
             structlogger.debug(
                 "multi_step_llm_command_generator.parse_commands",
@@ -535,7 +464,12 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
             prompt=prompt,
         )
 
-        actions = await self.invoke_llm(prompt)
+        response = await self.invoke_llm(prompt)
+        llm_response = LLMResponse.ensure_llm_response(response)
+        actions = None
+        if llm_response and llm_response.choices:
+            actions = llm_response.choices[0]
+
         structlogger.debug(
             "multi_step_llm_command_generator"
             ".predict_commands_for_active_flow"
@@ -546,11 +480,12 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
         commands = self.parse_commands(actions, tracker, available_flows)
 
         if commands:
-            self._add_prompt_to_message_parse_data(
-                message,
-                MultiStepLLMCommandGenerator.__name__,
-                "fill_slots_for_active_flow_prompt",
-                prompt,
+            add_prompt_to_message_parse_data(
+                message=message,
+                component_name=MultiStepLLMCommandGenerator.__name__,
+                prompt_name="fill_slots_for_active_flow_prompt",
+                user_prompt=prompt,
+                llm_response=llm_response,
             )
 
         return commands
@@ -584,7 +519,12 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
             prompt=prompt,
         )
 
-        actions = await self.invoke_llm(prompt)
+        response = await self.invoke_llm(prompt)
+        llm_response = LLMResponse.ensure_llm_response(response)
+        actions = None
+        if llm_response and llm_response.choices:
+            actions = llm_response.choices[0]
+
         structlogger.debug(
             "multi_step_llm_command_generator"
             ".predict_commands_for_handling_flows"
@@ -597,11 +537,12 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
         commands = self._filter_redundant_start_flow_commands(tracker, commands)
 
         if commands:
-            self._add_prompt_to_message_parse_data(
-                message,
-                MultiStepLLMCommandGenerator.__name__,
-                "handle_flows_prompt",
-                prompt,
+            add_prompt_to_message_parse_data(
+                message=message,
+                component_name=MultiStepLLMCommandGenerator.__name__,
+                prompt_name="handle_flows_prompt",
+                user_prompt=prompt,
+                llm_response=llm_response,
             )
 
         return commands
@@ -668,7 +609,12 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
             prompt=prompt,
         )
 
-        actions = await self.invoke_llm(prompt)
+        response = await self.invoke_llm(prompt)
+        llm_response = LLMResponse.ensure_llm_response(response)
+        actions = None
+        if llm_response and llm_response.choices:
+            actions = llm_response.choices[0]
+
         structlogger.debug(
             "multi_step_llm_command_generator"
             ".predict_commands_for_newly_started_flow"
@@ -694,11 +640,12 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
         )
 
         if commands:
-            self._add_prompt_to_message_parse_data(
-                message,
-                MultiStepLLMCommandGenerator.__name__,
-                "fill_slots_for_new_flow_prompt",
-                prompt,
+            add_prompt_to_message_parse_data(
+                message=message,
+                component_name=MultiStepLLMCommandGenerator.__name__,
+                prompt_name="fill_slots_for_new_flow_prompt",
+                user_prompt=prompt,
+                llm_response=llm_response,
             )
 
         return commands