rasa-pro 3.11.5__py3-none-any.whl → 3.12.0__py3-none-any.whl

rasa/dialogue_understanding/generator/command_parser.py

@@ -0,0 +1,245 @@
+import re
+import sys
+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.prompt_command import PromptCommand
+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[PromptCommand]], List[str]],
+) -> List[Type[PromptCommand]]:
+    """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[PromptCommand],
+) -> Optional[Callable[[PromptCommand, FlowsList], Optional[PromptCommand]]]:
+    """Get additional parsing logic for a command."""
+    command_to_parsing_fn_mapper: Dict[
+        Type[PromptCommand], Callable[..., Optional[PromptCommand]]
+    ] = {
+        ClarifyCommand: _parse_clarify_command,
+        SetSlotCommand: _parse_set_slot_command,
+        StartFlowCommand: _parse_start_flow_command,
+    }
+    return command_to_parsing_fn_mapper.get(command_clz)
+
+
+def validate_custom_commands(command_classes: List[Type[PromptCommand]]) -> None:
+    clz_not_inheriting_from_command_clz = [
+        command_clz
+        for command_clz in command_classes
+        if not issubclass(command_clz, Command)
+    ]
+
+    if clz_not_inheriting_from_command_clz:
+        structlogger.error(
+            "command_parser.validate_custom_commands.invalid_command",
+            invalid_commands=clz_not_inheriting_from_command_clz,
+            event_info=(
+                "The additional command classes must be a subclass of the 'Command' "
+                "class. Please refer to the class in "
+                "`rasa.dialogue_understanding.commands.command.Command`"
+            ),
+        )
+        sys.exit(1)
+
+    clz_not_adhering_to_prompt_command_protocol = [
+        command_clz
+        for command_clz in command_classes
+        if not isinstance(command_clz, PromptCommand)
+    ]
+
+    if clz_not_adhering_to_prompt_command_protocol:
+        structlogger.error(
+            "command_parser.validate_custom_commands.invalid_command",
+            invalid_commands=clz_not_adhering_to_prompt_command_protocol,
+            event_info=(
+                "The additional command classes must adhere to the 'PromptCommand' "
+                "protocol. Please refer to the protocol in "
+                "`rasa.dialogue_understanding.commands.prompt_command.PromptCommand`"
+            ),
+        )
+        sys.exit(1)
+
+
+def parse_commands(
+    actions: Optional[str],
+    flows: FlowsList,
+    is_handle_flows_prompt: bool = False,
+    additional_commands: Optional[List[Type[PromptCommand]]] = None,
+    default_commands_to_remove: Optional[
+        Union[List[Type[PromptCommand]], List[str]]
+    ] = None,
+    **kwargs: Any,
+) -> List[Command]:
+    """Parse a list of action commands."""
+    if not actions:
+        return []
+
+    commands: List[Command] = []
+    validate_custom_commands(additional_commands or [])
+
+    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[PromptCommand]],
+    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[PromptCommand]],
+    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[PromptCommand]:
+    """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[PromptCommand]:
+    """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[PromptCommand]:
+    """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,9 @@
 from rasa.shared.constants import (
-    PROVIDER_CONFIG_KEY,
-    OPENAI_PROVIDER,
+    MAX_TOKENS_CONFIG_KEY,
     MODEL_CONFIG_KEY,
+    OPENAI_PROVIDER,
+    PROVIDER_CONFIG_KEY,
+    TEMPERATURE_CONFIG_KEY,
     TIMEOUT_CONFIG_KEY,
 )
 from rasa.shared.utils.llm import (
@@ -12,8 +14,8 @@ from rasa.shared.utils.llm import (
 DEFAULT_LLM_CONFIG = {
     PROVIDER_CONFIG_KEY: OPENAI_PROVIDER,
     MODEL_CONFIG_KEY: DEFAULT_OPENAI_CHAT_MODEL_NAME_ADVANCED,
-    "temperature": 0.0,
-    "max_tokens": DEFAULT_OPENAI_MAX_GENERATED_TOKENS,
+    TEMPERATURE_CONFIG_KEY: 0.0,
+    MAX_TOKENS_CONFIG_KEY: DEFAULT_OPENAI_MAX_GENERATED_TOKENS,
     TIMEOUT_CONFIG_KEY: 7,
 }
 
@@ -25,3 +27,9 @@ FLOW_RETRIEVAL_ACTIVE_KEY = "active"
 FLOW_RETRIEVAL_EMBEDDINGS_CONFIG_KEY = "embeddings"
 
 FLOW_RETRIEVAL_FLOW_THRESHOLD = 20
+
+COMMAND_PROMPT_FILE_NAME = "command_prompt.jinja2"
+LLM_BASED_COMMAND_GENERATOR_CONFIG_FILE = "config.json"
+
+MODEL_NAME_GPT_4O_2024_11_20 = "gpt-4o-2024-11-20"
+MODEL_NAME_CLAUDE_3_5_SONNET_20240620 = "claude-3-5-sonnet-20240620"

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,37 +1,43 @@
 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, Set, Text, Tuple, Union
 
 import structlog
 from jinja2 import Template
 
+import rasa.dialogue_understanding.generator.utils
 import rasa.shared.utils.io
 from rasa.dialogue_understanding.commands import (
     Command,
+    SetSlotCommand,
     StartFlowCommand,
 )
+from rasa.dialogue_understanding.constants import KEY_MINIMIZE_NUM_CALLS
 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.dialogue_understanding.stack.utils import top_flow_frame
+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.constants import SetSlotExtractor
 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.slot_mappings import SlotFillingManager
 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, measure_llm_latency
 from rasa.shared.utils.health_check.llm_health_check_mixin import LLMHealthCheckMixin
 from rasa.shared.utils.llm import (
     allowed_values_for_slot,
@@ -169,7 +175,7 @@ class LLMBasedCommandGenerator(
         """
         self.perform_llm_health_check(
             self.config.get(LLM_CONFIG_KEY),
-            DEFAULT_LLM_CONFIG,
+            self.get_default_llm_config(),
             "llm_based_command_generator.train",
             LLMBasedCommandGenerator.__name__,
         )
@@ -304,22 +310,33 @@ class LLMBasedCommandGenerator(
         )
         return filtered_flows
 
-    async def invoke_llm(self, prompt: Text) -> Optional[Text]:
+    @measure_llm_latency
+    async def invoke_llm(
+        self, prompt: Union[List[dict], List[str], str]
+    ) -> Optional[LLMResponse]:
         """Use LLM to generate a response.
 
         Args:
-            prompt: The prompt to send to the LLM.
+            prompt: The prompt can be,
+                - a list of preformatted messages. Each message should be a dictionary
+                    with the following keys:
+                    - content: The message content.
+                    - role: The role of the message (e.g. user or system).
+                - a list of messages. Each message is a string and will be formatted
+                    as a user message.
+                - a single message as a string which will be formatted as user message.
 
         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)
+        llm = llm_factory(
+            self.config.get(LLM_CONFIG_KEY), self.get_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 +345,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]]:
@@ -498,3 +471,158 @@ class LLMBasedCommandGenerator(
             if isinstance(current_step, CollectInformationFlowStep)
             else (None, None)
         )
+
+    @staticmethod
+    def _prior_commands_contain_start_flow(prior_commands: List[Command]) -> bool:
+        return any(isinstance(command, StartFlowCommand) for command in prior_commands)
+
+    @staticmethod
+    def _prior_commands_contain_set_slot_for_active_collect_step(
+        prior_commands: List[Command],
+        flows: FlowsList,
+        tracker: DialogueStateTracker,
+    ) -> bool:
+        latest_user_frame = top_flow_frame(tracker.stack, ignore_call_frames=False)
+
+        if latest_user_frame is None:
+            return False
+
+        active_flow = latest_user_frame.flow(flows)
+        active_step = active_flow.step_by_id(latest_user_frame.step_id)
+
+        if not isinstance(active_step, CollectInformationFlowStep):
+            return False
+
+        return any(
+            command.name == active_step.collect
+            for command in prior_commands
+            if isinstance(command, SetSlotCommand)
+        )
+
+    def _should_skip_llm_call(
+        self,
+        prior_commands: List[Command],
+        flows: FlowsList,
+        tracker: DialogueStateTracker,
+    ) -> bool:
+        """Skip invoking the LLM.
+
+        This returns True if the bot builder sets the property
+        KEY_MINIMIZE_NUM_CALLS to True and the prior commands
+        either contain a StartFlowCommand or a SetSlot command
+        for the current collect step.
+        """
+        return self.config.get(KEY_MINIMIZE_NUM_CALLS, False) and (
+            self._prior_commands_contain_start_flow(prior_commands)
+            or self._prior_commands_contain_set_slot_for_active_collect_step(
+                prior_commands, flows, tracker
+            )
+        )
+
+    @staticmethod
+    def _check_commands_against_slot_mappings(
+        commands: List[Command],
+        tracker: DialogueStateTracker,
+        domain: Optional[Domain] = None,
+    ) -> List[Command]:
+        """Check if the LLM-issued slot commands are fillable.
+
+        The LLM-issued slot commands are fillable if the slot
+        mappings are satisfied (in particular the mapping conditions).
+        """
+        if not domain:
+            return commands
+
+        llm_fillable_slots = [
+            tracker.slots.get(command.name)
+            for command in commands
+            if isinstance(command, SetSlotCommand)
+            and command.extractor == SetSlotExtractor.LLM.value
+            and tracker.slots.get(command.name) is not None
+        ]
+
+        if not llm_fillable_slots:
+            return commands
+
+        slot_filling_manager = SlotFillingManager(domain, tracker)
+        slots_to_be_removed = []
+
+        structlogger.debug(
+            "command_processor.check_commands_against_slot_mappings.active_flow",
+            active_flow=tracker.active_flow,
+        )
+
+        for slot in llm_fillable_slots:
+            should_fill_slot = False
+            for mapping in slot.mappings:  # type: ignore[union-attr]
+                should_fill_slot = slot_filling_manager.should_fill_slot(
+                    slot.name,  # type: ignore[union-attr]
+                    mapping,
+                )
+
+                if should_fill_slot:
+                    break
+
+            if not should_fill_slot:
+                structlogger.debug(
+                    "command_processor.check_commands_against_slot_mappings.slot_not_fillable",
+                    slot_name=slot.name,  # type: ignore[union-attr]
+                )
+                slots_to_be_removed.append(slot.name)  # type: ignore[union-attr]
+
+        if not slots_to_be_removed:
+            return commands
+
+        filtered_commands = [
+            command
+            for command in commands
+            if not (
+                isinstance(command, SetSlotCommand)
+                and command.name in slots_to_be_removed
+            )
+        ]
+
+        return filtered_commands
+
+    def _check_start_flow_command_overlap(
+        self,
+        prior_commands: List[Command],
+        commands: List[Command],
+        prior_start_flow_names: Set[str],
+        current_start_flow_names: Set[str],
+    ) -> List[Command]:
+        """Prioritize the prior commands over the LLM-issued commands."""
+        different_flow_names = current_start_flow_names.difference(
+            prior_start_flow_names
+        )
+
+        if not different_flow_names:
+            return prior_commands + commands
+
+        # discard the flow names that are different to prior start flow commands
+        filtered_commands = [
+            command
+            for command in commands
+            if not isinstance(command, StartFlowCommand)
+            or command.flow not in different_flow_names
+        ]
+        return prior_commands + filtered_commands
+
+    def _filter_slot_commands(
+        self,
+        prior_commands: List[Command],
+        commands: List[Command],
+        overlapping_slot_names: Set[str],
+    ) -> Tuple[List[Command], List[Command]]:
+        """Prioritize prior commands over LLM ones in the case of same slot."""
+        filtered_commands = (
+            rasa.dialogue_understanding.generator.utils.filter_slot_commands(
+                commands, overlapping_slot_names
+            )
+        )
+        return prior_commands, filtered_commands
+
+    @staticmethod
+    def get_default_llm_config() -> Dict[str, Any]:
+        """Get the default LLM config for the command generator."""
+        return DEFAULT_LLM_CONFIG

rasa/dialogue_understanding/generator/llm_command_generator.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional, Text
+from typing import Any, Dict, List, Optional, Text, Union
 
 import structlog
 from deprecated import deprecated  # type: ignore[import]
@@ -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()
@@ -24,7 +25,7 @@ structlogger = structlog.get_logger()
 @deprecated(
     reason=(
         "The LLMCommandGenerator is deprecated and will be removed in Rasa 4.0.0. "
-        "Please use use SingleStepLLMCommandGenerator instead."
+        "Please use SingleStepLLMCommandGenerator instead."
     )
 )
 class LLMCommandGenerator(SingleStepLLMCommandGenerator):
@@ -53,7 +54,9 @@ class LLMCommandGenerator(SingleStepLLMCommandGenerator):
             **kwargs,
         )
 
-    async def invoke_llm(self, prompt: Text) -> Optional[Text]:
+    async def invoke_llm(
+        self, prompt: Union[List[dict], List[str], str]
+    ) -> Optional[LLMResponse]:
         try:
             return await super().invoke_llm(prompt)
         except ProviderClientAPIException: