rasa-pro 3.12.6.dev2__py3-none-any.whl → 3.12.7.dev2__py3-none-any.whl

rasa/hooks.py

@@ -1,20 +1,8 @@
 import argparse
 import logging
-import os
 from typing import TYPE_CHECKING, List, Optional, Text, Union
 
-import litellm
 import pluggy
-import structlog
-
-from rasa.shared.providers.constants import (
-    LANGFUSE_CALLBACK_NAME,
-    LANGFUSE_HOST_ENV_VAR,
-    LANGFUSE_PROJECT_ID_ENV_VAR,
-    LANGFUSE_PUBLIC_KEY_ENV_VAR,
-    LANGFUSE_SECRET_KEY_ENV_VAR,
-    RASA_LANGFUSE_INTEGRATION_ENABLED_ENV_VAR,
-)
 
 # IMPORTANT: do not import anything from rasa here - use scoped imports
 #  this avoids circular imports, as the hooks are used in different places
@@ -30,7 +18,6 @@ if TYPE_CHECKING:
 
 hookimpl = pluggy.HookimplMarker("rasa")
 logger = logging.getLogger(__name__)
-structlogger = structlog.get_logger()
 
 
 @hookimpl  # type: ignore[misc]
@@ -70,8 +57,6 @@ def configure_commandline(cmdline_arguments: argparse.Namespace) -> Optional[Tex
         config.configure_tracing(tracer_provider)
         config.configure_metrics(endpoints_file)
 
-    _init_langfuse_integration()
-
     return endpoints_file
 
 
@@ -130,43 +115,3 @@ def after_server_stop() -> None:
 
     if anon_pipeline is not None:
         anon_pipeline.stop()
-
-
-def _is_langfuse_integration_enabled() -> bool:
-    return (
-        os.environ.get(RASA_LANGFUSE_INTEGRATION_ENABLED_ENV_VAR, "false").lower()
-        == "true"
-    )
-
-
-def _init_langfuse_integration() -> None:
-    if not _is_langfuse_integration_enabled():
-        structlogger.info(
-            "hooks._init_langfuse_integration.disabled",
-            event_info="Langfuse integration is disabled.",
-        )
-        return
-
-    if (
-        not os.environ.get(LANGFUSE_HOST_ENV_VAR)
-        or not os.environ.get(LANGFUSE_PROJECT_ID_ENV_VAR)
-        or not os.environ.get(LANGFUSE_PUBLIC_KEY_ENV_VAR)
-        or not os.environ.get(LANGFUSE_SECRET_KEY_ENV_VAR)
-    ):
-        structlogger.warning(
-            "hooks._init_langfuse_integration.missing_langfuse_keys",
-            event_info=(
-                "Langfuse integration is enabled, but some environment variables"
-                "are missing. Please set LANGFUSE_HOST, LANGFUSE_PROJECT_ID, "
-                "LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY environment "
-                "variables to use Langfuse integration."
-            ),
-        )
-        return
-
-    litellm.success_callback = [LANGFUSE_CALLBACK_NAME]
-    litellm.failure_callback = [LANGFUSE_CALLBACK_NAME]
-    structlogger.info(
-        "hooks.langfuse_callbacks_initialized",
-        event_info="Langfuse integration initialized.",
-    )

rasa/shared/constants.py

@@ -338,8 +338,3 @@ ROLE_SYSTEM = "system"
 # Used for key values in ValidateSlotPatternFlowStackFrame
 REFILL_UTTER = "refill_utter"
 REJECTIONS = "rejections"
-
-LANGFUSE_METADATA_USER_ID = "trace_user_id"
-LANGFUSE_METADATA_SESSION_ID = "session_id"
-LANGFUSE_CUSTOM_METADATA_DICT = "trace_metadata"
-LANGFUSE_TAGS = "tags"

rasa/shared/core/constants.py

@@ -52,8 +52,6 @@ ACTION_TRIGGER_CHITCHAT = "action_trigger_chitchat"
 ACTION_RESET_ROUTING = "action_reset_routing"
 ACTION_HANGUP = "action_hangup"
 ACTION_REPEAT_BOT_MESSAGES = "action_repeat_bot_messages"
-ACTION_BLOCK_DIGRESSION = "action_block_digression"
-ACTION_CONTINUE_DIGRESSION = "action_continue_digression"
 
 ACTION_METADATA_EXECUTION_SUCCESS = "execution_success"
 ACTION_METADATA_EXECUTION_ERROR_MESSAGE = "execution_error_message"
@@ -84,8 +82,6 @@ DEFAULT_ACTION_NAMES = [
     ACTION_RESET_ROUTING,
     ACTION_HANGUP,
     ACTION_REPEAT_BOT_MESSAGES,
-    ACTION_BLOCK_DIGRESSION,
-    ACTION_CONTINUE_DIGRESSION,
 ]
 
 ACTION_SHOULD_SEND_DOMAIN = "send_domain"
@@ -205,8 +201,4 @@ CLASSIFIER_NAME_FALLBACK = "FallbackClassifier"
 
 POLICIES_THAT_EXTRACT_ENTITIES = {"TEDPolicy"}
 
-# digression constants
-KEY_ASK_CONFIRM_DIGRESSIONS = "ask_confirm_digressions"
-KEY_BLOCK_DIGRESSIONS = "block_digressions"
-
 ERROR_CODE_KEY = "error_code"

rasa/shared/core/domain.py

@@ -1678,6 +1678,14 @@ class Domain:
         """Write domain to a file."""
         as_yaml = self.as_yaml()
         rasa.shared.utils.io.write_text_file(as_yaml, filename)
+        # run the check again on the written domain to catch any errors
+        # that may have been missed in the user defined domain files
+        structlogger.info(
+            "domain.persist.domain_written_to_file",
+            event_info="The entire domain content has been written to file.",
+            filename=filename,
+        )
+        Domain.is_domain_file(filename)
 
     def as_yaml(self) -> Text:
         """Dump the `Domain` object as a YAML string.
@@ -1972,17 +1980,18 @@ class Domain:
 
         try:
             content = read_yaml_file(filename, expand_env_vars=cls.expand_env_vars)
-        except (RasaException, YamlSyntaxException):
-            structlogger.warning(
+        except (RasaException, YamlSyntaxException) as error:
+            structlogger.error(
                 "domain.cannot_load_domain_file",
                 file=filename,
+                error=error,
                 event_info=(
                     f"The file {filename} could not be loaded as domain file. "
                     f"You can use https://yamlchecker.com/ to validate "
                     f"the YAML syntax of your file."
                 ),
             )
-            return False
+            raise RasaException(f"Domain could not be loaded: {error}")
 
         return any(key in content for key in ALL_DOMAIN_KEYS)
 

rasa/shared/core/flows/flow.py

@@ -13,10 +13,6 @@ from pypred import Predicate
 import rasa.shared.utils.io
 from rasa.engine.language import Language
 from rasa.shared.constants import RASA_DEFAULT_FLOW_PATTERN_PREFIX
-from rasa.shared.core.constants import (
-    KEY_ASK_CONFIRM_DIGRESSIONS,
-    KEY_BLOCK_DIGRESSIONS,
-)
 from rasa.shared.core.flows.constants import (
     KEY_ALWAYS_INCLUDE_IN_PROMPT,
     KEY_DESCRIPTION,
@@ -52,7 +48,6 @@ from rasa.shared.core.flows.steps.constants import (
     START_STEP,
 )
 from rasa.shared.core.flows.steps.continuation import ContinueFlowStep
-from rasa.shared.core.flows.utils import extract_digression_prop
 from rasa.shared.core.slots import Slot
 
 structlogger = structlog.get_logger()
@@ -94,10 +89,6 @@ class Flow:
     """The path to the file where the flow is stored."""
     persisted_slots: List[str] = field(default_factory=list)
     """The list of slots that should be persisted after the flow ends."""
-    ask_confirm_digressions: List[str] = field(default_factory=list)
-    """The flow ids for which the assistant should ask for confirmation."""
-    block_digressions: List[str] = field(default_factory=list)
-    """The flow ids that the assistant should block from digressing to."""
     run_pattern_completed: bool = True
     """Whether the pattern_completed flow should be run after the flow ends."""
 
@@ -138,10 +129,6 @@ class Flow:
             # data. When the model is trained, take the provided file_path.
             file_path=data.get(KEY_FILE_PATH) if KEY_FILE_PATH in data else file_path,
             persisted_slots=data.get(KEY_PERSISTED_SLOTS, []),
-            ask_confirm_digressions=extract_digression_prop(
-                KEY_ASK_CONFIRM_DIGRESSIONS, data
-            ),
-            block_digressions=extract_digression_prop(KEY_BLOCK_DIGRESSIONS, data),
             run_pattern_completed=data.get(KEY_RUN_PATTERN_COMPLETED, True),
             translation=extract_translations(
                 translation_data=data.get(KEY_TRANSLATION, {})
@@ -220,10 +207,6 @@ class Flow:
             data[KEY_FILE_PATH] = self.file_path
         if self.persisted_slots:
             data[KEY_PERSISTED_SLOTS] = self.persisted_slots
-        if self.ask_confirm_digressions:
-            data[KEY_ASK_CONFIRM_DIGRESSIONS] = self.ask_confirm_digressions
-        if self.block_digressions:
-            data[KEY_BLOCK_DIGRESSIONS] = self.block_digressions
         if self.run_pattern_completed is not None:
             data["run_pattern_completed"] = self.run_pattern_completed
         if self.translation:

rasa/shared/core/flows/flows_yaml_schema.json

@@ -217,15 +217,12 @@
                             "reset_after_flow_ends": {
                                 "type": "boolean"
                             },
-                            "ask_confirm_digressions": {
-                                "$ref": "#/$defs/ask_confirm_digressions"
-                            },
-                            "block_digressions": {
-                                "$ref": "#/$defs/block_digressions"
-                            },
                             "utter": {
                                 "type": "string"
                             },
+                            "force_slot_filling": {
+                                "type": "boolean"
+                            },
                             "rejections": {
                                 "type": "array",
                                 "schema_name": "list of rejections",
@@ -253,32 +250,6 @@
                 }
             }
         },
-        "ask_confirm_digressions": {
-            "oneOf": [
-                {
-                    "type": "boolean"
-                },
-                {
-                    "type": "array",
-                    "items": {
-                        "type": "string"
-                    }
-                }
-            ]
-        },
-        "block_digressions": {
-            "oneOf": [
-                {
-                    "type": "boolean"
-                },
-                {
-                    "type": "array",
-                    "items": {
-                        "type": "string"
-                    }
-                }
-            ]
-        },
         "flow": {
             "required": [
                 "steps",
@@ -340,12 +311,6 @@
                 "persisted_slots": {
                     "$ref": "#/$defs/persisted_slots"
                 },
-                "ask_confirm_digressions": {
-                    "$ref": "#/$defs/ask_confirm_digressions"
-                },
-                "block_digressions": {
-                    "$ref": "#/$defs/block_digressions"
-                },
                 "run_pattern_completed": {
                     "type": "boolean"
                 }

rasa/shared/core/flows/steps/collect.py

@@ -1,15 +1,10 @@
 from __future__ import annotations
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Any, Dict, List, Set, Text
 
 from rasa.shared.constants import ACTION_ASK_PREFIX, UTTER_ASK_PREFIX
-from rasa.shared.core.constants import (
-    KEY_ASK_CONFIRM_DIGRESSIONS,
-    KEY_BLOCK_DIGRESSIONS,
-)
 from rasa.shared.core.flows.flow_step import FlowStep
-from rasa.shared.core.flows.utils import extract_digression_prop
 from rasa.shared.core.slots import SlotRejection
 
 
@@ -29,10 +24,8 @@ class CollectInformationFlowStep(FlowStep):
     """Whether to always ask the question even if the slot is already filled."""
     reset_after_flow_ends: bool = True
     """Whether to reset the slot value at the end of the flow."""
-    ask_confirm_digressions: List[str] = field(default_factory=list)
-    """The flow id digressions for which the assistant should ask for confirmation."""
-    block_digressions: List[str] = field(default_factory=list)
-    """The flow id digressions that should be blocked during the flow step."""
+    force_slot_filling: bool = False
+    """Whether to keep only the SetSlot command for the collected slot."""
 
     @classmethod
     def from_json(
@@ -60,10 +53,7 @@ class CollectInformationFlowStep(FlowStep):
                 SlotRejection.from_dict(rejection)
                 for rejection in data.get("rejections", [])
             ],
-            ask_confirm_digressions=extract_digression_prop(
-                KEY_ASK_CONFIRM_DIGRESSIONS, data
-            ),
-            block_digressions=extract_digression_prop(KEY_BLOCK_DIGRESSIONS, data),
+            force_slot_filling=data.get("force_slot_filling", False),
             **base.__dict__,
         )
 
@@ -79,10 +69,7 @@ class CollectInformationFlowStep(FlowStep):
         data["ask_before_filling"] = self.ask_before_filling
         data["reset_after_flow_ends"] = self.reset_after_flow_ends
         data["rejections"] = [rejection.as_dict() for rejection in self.rejections]
-        data["ask_confirm_digressions"] = self.ask_confirm_digressions
-        data["block_digressions"] = (
-            self.block_digressions if self.block_digressions else False
-        )
+        data["force_slot_filling"] = self.force_slot_filling
 
         return data
 

rasa/shared/core/flows/utils.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Dict, List, Set, Text
+from typing import TYPE_CHECKING, Any, Dict, Set, Text
 
 from rasa.shared.utils.io import raise_deprecation_warning
 
@@ -8,7 +8,6 @@ if TYPE_CHECKING:
 
 RESET_PROPERTY_NAME = "reset_after_flow_ends"
 PERSIST_PROPERTY_NAME = "persisted_slots"
-ALL_LABEL = "ALL"
 
 
 def warn_deprecated_collect_step_config() -> None:
@@ -45,20 +44,6 @@ def get_invalid_slot_persistence_config_error_message(
     )
 
 
-def extract_digression_prop(prop: str, data: Dict[str, Any]) -> List[str]:
-    """Extracts the digression property from the data.
-
-    There can be two types of properties: ask_confirm_digressions and
-    block_digressions.
-    """
-    digression_property = data.get(prop, [])
-
-    if isinstance(digression_property, bool):
-        digression_property = [ALL_LABEL] if digression_property else []
-
-    return digression_property
-
-
 def extract_translations(
     translation_data: Dict[Text, Any],
 ) -> Dict[Text, "FlowLanguageTranslation"]:

rasa/shared/core/slot_mappings.py

@@ -648,12 +648,14 @@ class SlotFillingManager:
         output_channel: "OutputChannel",
         nlg: "NaturalLanguageGenerator",
         recreate_tracker: bool = False,
+        slot_events: Optional[List[Event]] = None,
     ) -> List[Event]:
         from rasa.core.actions.action import RemoteAction
         from rasa.shared.core.trackers import DialogueStateTracker
         from rasa.utils.endpoints import ClientResponseError
 
-        slot_events: List[Event] = []
+        validated_slot_events: List[Event] = []
+        slot_events = slot_events if slot_events is not None else []
         remote_action = RemoteAction(custom_action, self._action_endpoint)
         disallowed_types = set()
 
@@ -673,9 +675,9 @@ class SlotFillingManager:
             )
             for event in custom_events:
                 if isinstance(event, SlotSet):
-                    slot_events.append(event)
+                    validated_slot_events.append(event)
                 elif isinstance(event, BotUttered):
-                    slot_events.append(event)
+                    validated_slot_events.append(event)
                 else:
                     disallowed_types.add(event.type_name)
         except (RasaException, ClientResponseError) as e:
@@ -699,7 +701,7 @@ class SlotFillingManager:
                 f"updated with this event.",
             )
 
-        return slot_events
+        return validated_slot_events
 
     async def execute_validation_action(
         self,
@@ -722,7 +724,11 @@ class SlotFillingManager:
             return cast(List[Event], slot_events)
 
         validate_events = await self._run_custom_action(
-            ACTION_VALIDATE_SLOT_MAPPINGS, output_channel, nlg, recreate_tracker=True
+            ACTION_VALIDATE_SLOT_MAPPINGS,
+            output_channel,
+            nlg,
+            recreate_tracker=True,
+            slot_events=cast(List[Event], slot_events),
         )
         validated_slot_names = [
             event.key for event in validate_events if isinstance(event, SlotSet)

rasa/shared/nlu/constants.py

@@ -55,4 +55,3 @@ SPLIT_ENTITIES_BY_COMMA_DEFAULT_VALUE = True
 SINGLE_ENTITY_ALLOWED_INTERLEAVING_CHARSET = {".", ",", " ", ";"}
 
 SET_SLOT_COMMAND = "set slot"
-HANDLE_DIGRESSIONS_COMMAND = "handle digressions"

rasa/shared/providers/constants.py

@@ -4,12 +4,3 @@ LITE_LLM_API_KEY_FIELD = "api_key"
 LITE_LLM_API_VERSION_FIELD = "api_version"
 LITE_LLM_MODEL_FIELD = "model"
 LITE_LLM_AZURE_AD_TOKEN = "azure_ad_token"
-
-# Enable or disable Langfuse integration
-RASA_LANGFUSE_INTEGRATION_ENABLED_ENV_VAR = "RASA_LANGFUSE_INTEGRATION_ENABLED"
-# Langfuse configuration
-LANGFUSE_CALLBACK_NAME = "langfuse"
-LANGFUSE_HOST_ENV_VAR = "LANGFUSE_HOST"
-LANGFUSE_PROJECT_ID_ENV_VAR = "LANGFUSE_PROJECT_ID"
-LANGFUSE_PUBLIC_KEY_ENV_VAR = "LANGFUSE_PUBLIC_KEY"
-LANGFUSE_SECRET_KEY_ENV_VAR = "LANGFUSE_SECRET_KEY"

rasa/shared/providers/llm/_base_litellm_client.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 from abc import abstractmethod
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Union, cast
 
 import structlog
 from litellm import acompletion, completion, validate_environment
@@ -120,11 +120,7 @@ class _BaseLiteLLMClient:
             raise ProviderClientValidationError(event_info)
 
     @suppress_logs(log_level=logging.WARNING)
-    def completion(
-        self,
-        messages: Union[List[dict], List[str], str],
-        metadata: Optional[Dict[str, Any]] = None,
-    ) -> LLMResponse:
+    def completion(self, messages: Union[List[dict], List[str], str]) -> LLMResponse:
         """Synchronously generate completions for given list of messages.
 
         Args:
@@ -136,7 +132,6 @@ class _BaseLiteLLMClient:
                 - 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.
-            metadata: Optional metadata to be passed to the LLM call.
 
         Returns:
             List of message completions.
@@ -154,9 +149,7 @@ class _BaseLiteLLMClient:
 
     @suppress_logs(log_level=logging.WARNING)
     async def acompletion(
-        self,
-        messages: Union[List[dict], List[str], str],
-        metadata: Optional[Dict[str, Any]] = None,
+        self, messages: Union[List[dict], List[str], str]
     ) -> LLMResponse:
         """Asynchronously generate completions for given list of messages.
 
@@ -169,7 +162,6 @@ class _BaseLiteLLMClient:
                 - 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.
-            metadata: Optional metadata to be passed to the LLM call.
 
         Returns:
             List of message completions.
@@ -180,9 +172,7 @@ class _BaseLiteLLMClient:
         try:
             formatted_messages = self._get_formatted_messages(messages)
             arguments = resolve_environment_variables(self._completion_fn_args)
-            response = await acompletion(
-                messages=formatted_messages, metadata=metadata, **arguments
-            )
+            response = await acompletion(messages=formatted_messages, **arguments)
             return self._format_response(response)
         except Exception as e:
             message = ""

rasa/shared/providers/llm/litellm_router_llm_client.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Union
 
 import structlog
 
@@ -122,12 +122,9 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
             raise ProviderClientAPIException(e)
 
     @suppress_logs(log_level=logging.WARNING)
-    def completion(
-        self,
-        messages: Union[List[dict], List[str], str],
-        metadata: Optional[Dict[str, Any]] = None,
-    ) -> LLMResponse:
-        """Synchronously generate completions for given list of messages.
+    def completion(self, messages: Union[List[dict], List[str], str]) -> LLMResponse:
+        """
+        Synchronously generate completions for given list of messages.
 
         Method overrides the base class method to call the appropriate
         completion method based on the configuration. If the chat completions
@@ -143,11 +140,8 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
                 - 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.
-            metadata: Optional metadata to be passed to the LLM call.
-
         Returns:
             List of message completions.
-
         Raises:
             ProviderClientAPIException: If the API request fails.
         """
@@ -164,11 +158,10 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
 
     @suppress_logs(log_level=logging.WARNING)
     async def acompletion(
-        self,
-        messages: Union[List[dict], List[str], str],
-        metadata: Optional[Dict[str, Any]] = None,
+        self, messages: Union[List[dict], List[str], str]
     ) -> LLMResponse:
-        """Asynchronously generate completions for given list of messages.
+        """
+        Asynchronously generate completions for given list of messages.
 
         Method overrides the base class method to call the appropriate
         completion method based on the configuration. If the chat completions
@@ -184,11 +177,8 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
                 - 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.
-            metadata: Optional metadata to be passed to the LLM call.
-
         Returns:
             List of message completions.
-
         Raises:
             ProviderClientAPIException: If the API request fails.
         """

rasa/shared/providers/llm/llm_client.py

@@ -1,19 +1,21 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional, Protocol, Union, runtime_checkable
+from typing import Dict, List, Protocol, Union, runtime_checkable
 
 from rasa.shared.providers.llm.llm_response import LLMResponse
 
 
 @runtime_checkable
 class LLMClient(Protocol):
-    """Protocol for an LLM client that specifies the interface for interacting
+    """
+    Protocol for an LLM client that specifies the interface for interacting
     with the API.
     """
 
     @classmethod
     def from_config(cls, config: dict) -> LLMClient:
-        """Initializes the llm client with the given configuration.
+        """
+        Initializes the llm client with the given configuration.
 
         This class method should be implemented to parse the given
         configuration and create an instance of an llm client.
@@ -22,24 +24,17 @@ class LLMClient(Protocol):
 
     @property
     def config(self) -> Dict:
-        """Returns the configuration for that the llm client is initialized with.
+        """
+        Returns the configuration for that the llm client is initialized with.
 
         This property should be implemented to return a dictionary containing
         the configuration settings for the llm client.
         """
         ...
 
-    def completion(
-        self,
-        messages: Union[List[dict], List[str], str],
-        metadata: Optional[Dict[str, Any]] = None,
-    ) -> LLMResponse:
-        """Synchronously generate completions for given list of messages.
-        def completion(
-        self,
-        messages: Union[List[dict], List[str], str],
-        metadata: Optional[Dict[str, Any]] = None,
-        ) -> LLMResponse:
+    def completion(self, messages: Union[List[dict], List[str], str]) -> LLMResponse:
+        """
+        Synchronously generate completions for given list of messages.
 
         This method should be implemented to take a list of messages (as
         strings) and return a list of completions (as strings).
@@ -53,19 +48,16 @@ class LLMClient(Protocol):
                 - 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.
-            metadata: Optional metadata to be passed to the LLM call.
-
         Returns:
             LLMResponse
         """
         ...
 
     async def acompletion(
-        self,
-        messages: Union[List[dict], List[str], str],
-        metadata: Optional[Dict[str, Any]] = None,
+        self, messages: Union[List[dict], List[str], str]
     ) -> LLMResponse:
-        """Asynchronously generate completions for given list of messages.
+        """
+        Asynchronously generate completions for given list of messages.
 
         This method should be implemented to take a list of messages (as
         strings) and return a list of completions (as strings).
@@ -79,15 +71,14 @@ class LLMClient(Protocol):
                 - 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.
-            metadata: Optional metadata to be passed to the LLM call.
-
         Returns:
             LLMResponse
         """
         ...
 
     def validate_client_setup(self, *args, **kwargs) -> None:  # type: ignore
-        """Perform client setup validation.
+        """
+        Perform client setup validation.
 
         This method should be implemented to validate whether the client can be
         used with the parameters provided through configuration or environment