rasa-pro 3.12.0rc2__py3-none-any.whl → 3.12.0rc3__py3-none-any.whl

rasa/llm_fine_tuning/llm_data_preparation_module.py

@@ -4,8 +4,10 @@ from typing import Any, Dict, List, Optional
 import structlog
 from tqdm import tqdm
 
+from rasa.dialogue_understanding.commands.prompt_command import PromptCommand
 from rasa.llm_fine_tuning.conversations import Conversation, ConversationStep
 from rasa.llm_fine_tuning.storage import StorageContext
+from rasa.llm_fine_tuning.utils import commands_as_string
 
 LLM_DATA_PREPARATION_MODULE_STORAGE_LOCATION = "3_llm_finetune_data/llm_ft_data.jsonl"
 
@@ -15,7 +17,7 @@ structlogger = structlog.get_logger()
 @dataclass
 class LLMDataExample:
     prompt: str
-    output: str
+    output: List[PromptCommand]
     original_test_name: str
     original_user_utterance: str
     rephrased_user_utterance: str
@@ -23,7 +25,7 @@ class LLMDataExample:
     def as_dict(self) -> Dict[str, Any]:
         return {
             "prompt": self.prompt,
-            "output": self.output,
+            "output": commands_as_string(self.output),
             "original_test_name": self.original_test_name,
             "original_user_utterance": self.original_user_utterance,
             "rephrased_user_utterance": self.rephrased_user_utterance,
@@ -38,7 +40,7 @@ def _create_data_point(
 ) -> LLMDataExample:
     return LLMDataExample(
         prompt,
-        step.commands_as_string(),
+        step.llm_commands,
         conversation.get_full_name(),
         step.original_test_step.text,
         rephrased_user_message,

rasa/llm_fine_tuning/paraphrasing/rephrase_validator.py

@@ -1,18 +1,18 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Type
 
 import structlog
 
 from rasa.dialogue_understanding.commands import Command, SetSlotCommand
-from rasa.dialogue_understanding.generator import SingleStepLLMCommandGenerator
+from rasa.dialogue_understanding.generator.llm_based_command_generator import (
+    LLMBasedCommandGenerator,
+)
 from rasa.llm_fine_tuning.conversations import Conversation, ConversationStep
 from rasa.llm_fine_tuning.paraphrasing.rephrased_user_message import (
     RephrasedUserMessage,
 )
 from rasa.shared.core.flows import FlowsList
 from rasa.shared.exceptions import ProviderClientAPIException
-from rasa.shared.utils.llm import (
-    llm_factory,
-)
+from rasa.shared.utils.llm import llm_factory
 
 structlogger = structlog.get_logger()
 
@@ -26,6 +26,7 @@ class RephraseValidator:
         self,
         rephrasings: List[RephrasedUserMessage],
         conversation: Conversation,
+        llm_command_generator: Type[LLMBasedCommandGenerator],
     ) -> List[RephrasedUserMessage]:
         """Split rephrased user messages into passing and failing.
 
@@ -38,6 +39,7 @@ class RephraseValidator:
         Args:
             rephrasings: The rephrased user messages.
             conversation: The conversation.
+            llm_command_generator: A LLM based command generator class.
 
         Returns:
             A list of rephrased user messages including the passing and failing
@@ -49,7 +51,9 @@ class RephraseValidator:
             current_rephrasings = rephrasings[i]
 
             for rephrase in current_rephrasings.rephrasings:
-                if await self._validate_rephrase_is_passing(rephrase, step):
+                if await self._validate_rephrase_is_passing(
+                    rephrase, step, llm_command_generator
+                ):
                     current_rephrasings.passed_rephrasings.append(rephrase)
                 else:
                     current_rephrasings.failed_rephrasings.append(rephrase)
@@ -60,25 +64,26 @@ class RephraseValidator:
         self,
         rephrase: str,
         step: ConversationStep,
+        llm_command_generator: Type[LLMBasedCommandGenerator],
     ) -> bool:
         prompt = self._update_prompt(
             rephrase, step.original_test_step.text, step.llm_prompt
         )
 
-        action_list = await self._invoke_llm(prompt)
+        action_list = await self._invoke_llm(
+            prompt, llm_command_generator.get_default_llm_config()
+        )
 
         commands_from_original_utterance = step.llm_commands
-        commands_from_rephrased_utterance = (
-            SingleStepLLMCommandGenerator.parse_commands(action_list, None, self.flows)
+        commands_from_rephrased_utterance = llm_command_generator.parse_commands(  # type: ignore
+            action_list, None, self.flows
         )
         return self._check_commands_match(
             commands_from_original_utterance, commands_from_rephrased_utterance
         )
 
-    async def _invoke_llm(self, prompt: str) -> str:
-        from rasa.dialogue_understanding.generator.constants import DEFAULT_LLM_CONFIG
-
-        llm = llm_factory(self.llm_config, DEFAULT_LLM_CONFIG)
+    async def _invoke_llm(self, prompt: str, default_llm_config: Dict[str, Any]) -> str:
+        llm = llm_factory(self.llm_config, default_llm_config)
 
         try:
             llm_response = await llm.acompletion(prompt)

rasa/llm_fine_tuning/paraphrasing_module.py

@@ -1,8 +1,11 @@
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Tuple, Type
 
 import structlog
 from tqdm import tqdm
 
+from rasa.dialogue_understanding.generator.llm_based_command_generator import (
+    LLMBasedCommandGenerator,
+)
 from rasa.llm_fine_tuning.conversations import Conversation
 from rasa.llm_fine_tuning.paraphrasing.conversation_rephraser import (
     ConversationRephraser,
@@ -25,6 +28,7 @@ async def create_paraphrased_conversations(
     rephrase_config: Dict[str, Any],
     num_rephrases: int,
     flows: FlowsList,
+    llm_command_generator: Type[LLMBasedCommandGenerator],
     llm_command_generator_config: Dict[str, Any],
     storage_context: StorageContext,
 ) -> Tuple[List[Conversation], Dict[str, Any]]:
@@ -71,7 +75,7 @@ async def create_paraphrased_conversations(
             rephrasings = _filter_rephrasings(rephrasings, conversations[i])
             # check if the rephrasings are still producing the same commands
             rephrasings = await validator.validate_rephrasings(
-                rephrasings, current_conversation
+                rephrasings, current_conversation, llm_command_generator
             )
         except ProviderClientAPIException as e:
             structlogger.error(

rasa/llm_fine_tuning/train_test_split_module.py

@@ -1,27 +1,18 @@
 import random
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, Dict, List, Protocol, Set, Tuple
+from typing import Any, Dict, List, Protocol, Set, Tuple, Type
 
 import structlog
 
+from rasa.dialogue_understanding.commands.prompt_command import PromptCommand
 from rasa.e2e_test.e2e_test_case import TestSuite
 from rasa.llm_fine_tuning.llm_data_preparation_module import LLMDataExample
 from rasa.llm_fine_tuning.storage import StorageContext
+from rasa.llm_fine_tuning.utils import commands_as_string
 
 TRAIN_TEST_MODULE_STORAGE_LOCATION = "4_train_test_split"
 
-SUPPORTED_COMMANDS = [
-    "SetSlot",
-    "StartFlow",
-    "CancelFlow",
-    "ChitChat",
-    "SkipQuestion",
-    "SearchAndReply",
-    "HumanHandoff",
-    "Clarify",
-]
-
 INSTRUCTION_DATA_FORMAT = "instruction"
 CONVERSATIONAL_DATA_FORMAT = "conversational"
 
@@ -77,17 +68,19 @@ class ConversationalDataFormat(DataExampleFormat):
         }
 
 
-def _get_command_types_covered_by_llm_data_point(commands: LLMDataExample) -> Set[str]:
+def _get_command_types_covered_by_llm_data_point(
+    data_point: LLMDataExample,
+) -> Set[Type[PromptCommand]]:
     """Get the command types covered by the LLM data point.
 
     This function returns the set of command types from the output present in a
-    LLMDataExample object. Eg: The function returns {'SetSlot', 'StartFlow'} when the
-    LLMDataExample.output is 'SetSlot(slot, abc), SetSlot(slot, cde), StartFlow(xyz)'.
+    LLMDataExample object. Eg: The function returns {'SetSlotCommand',
+    'StartFlowCommand'} when the LLMDataExample.output is 'SetSlotCommand(slot, abc),
+    SetSlotCommand(slot, cde), StartFlowCommand(xyz)'.
     """
     commands_covered = set()
-    for command in SUPPORTED_COMMANDS:
-        if command in commands.output:
-            commands_covered.add(command)
+    for command in data_point.output:
+        commands_covered.add(command.__class__)
     return commands_covered
 
 
@@ -146,14 +139,18 @@ def _get_minimum_test_case_groups_to_cover_all_commands(
         {
             "test_case_name": "t1",
             "data_examples": [],
-            "commands": {"SetSlot", "CancelFlow"}
+            "commands": {"SetSlotCommand", "CancelFlowCommand"}
         },
-        {"test_case_name": "t2", "data_examples": [], "commands": {"CancelFlow"}},
-        {"test_case_name": "t3", "data_examples": [], "commands": {"StartFlow"}},
+        {
+            "test_case_name": "t2",
+            "data_examples": [],
+            "commands": {"CancelFlowCommand"}
+        },
+        {"test_case_name": "t3", "data_examples": [], "commands": {"StartFlowCommand"}},
         {
             "test_case_name": "t4",
             "data_examples": [],
-            "commands": {"SetSlot", "StartFlow"}
+            "commands": {"SetSlotCommand", "StartFlowCommand"}
         },
     ]
 
@@ -166,7 +163,7 @@ def _get_minimum_test_case_groups_to_cover_all_commands(
         command for test_group in grouped_data for command in test_group[KEY_COMMANDS]
     )
     selected_test_cases = []
-    covered_commands: Set[str] = set()
+    covered_commands: Set[Type[PromptCommand]] = set()
 
     while covered_commands != all_commands:
         # Find the test case group that covers the most number of uncovered commands
@@ -187,7 +184,7 @@ def _get_minimum_test_case_groups_to_cover_all_commands(
 
     structlogger.info(
         "llm_fine_tuning.train_test_split_module.command_coverage_in_train_dataset",
-        covered_commands=covered_commands,
+        covered_commands=[command.__name__ for command in covered_commands],
     )
     return selected_test_cases
 
@@ -205,7 +202,10 @@ def _get_finetuning_data_in_instruction_data_format(
         data: List[Dict[str, Any]],
     ) -> List[DataExampleFormat]:
         return [
-            InstructionDataFormat(llm_data_example.prompt, llm_data_example.output)
+            InstructionDataFormat(
+                llm_data_example.prompt,
+                commands_as_string(llm_data_example.output),
+            )
             for test_group in data
             for llm_data_example in test_group[KEY_DATA_EXAMPLES]
         ]
@@ -232,7 +232,7 @@ def _get_finetuning_data_in_conversational_data_format(
                 [
                     ConversationalMessageDataFormat("user", llm_data_example.prompt),
                     ConversationalMessageDataFormat(
-                        "assistant", llm_data_example.output
+                        "assistant", commands_as_string(llm_data_example.output)
                     ),
                 ]
             )
@@ -271,7 +271,7 @@ def _check_and_log_missing_validation_dataset_command_coverage(
         structlogger.warning(
             "llm_fine_tuning.train_test_split_module.missing_commands_in_validation_dat"
             "aset",
-            missing_commands=missing_commands,
+            missing_commands=[command.__name__ for command in missing_commands],
         )
 
 

rasa/llm_fine_tuning/utils.py

@@ -0,0 +1,7 @@
+from typing import List
+
+from rasa.dialogue_understanding.commands.prompt_command import PromptCommand
+
+
+def commands_as_string(commands: List[PromptCommand], delimiter: str = "\n") -> str:
+    return delimiter.join([command.to_dsl() for command in commands])

rasa/shared/constants.py

@@ -194,6 +194,9 @@ PROVIDER_CONFIG_KEY = "provider"
 REQUEST_TIMEOUT_CONFIG_KEY = "request_timeout"  # deprecated
 TIMEOUT_CONFIG_KEY = "timeout"
 
+TEMPERATURE_CONFIG_KEY = "temperature"
+MAX_TOKENS_CONFIG_KEY = "max_tokens"
+
 DEPLOYMENT_NAME_CONFIG_KEY = "deployment_name"
 DEPLOYMENT_CONFIG_KEY = "deployment"
 EMBEDDINGS_CONFIG_KEY = "embeddings"
@@ -264,6 +267,7 @@ LITELLM_SSL_CERTIFICATE_ENV_VAR = "SSL_CERTIFICATE"
 
 OPENAI_PROVIDER = "openai"
 AZURE_OPENAI_PROVIDER = "azure"
+ANTHROPIC_PROVIDER = "anthropic"
 SELF_HOSTED_PROVIDER = "self-hosted"
 HUGGINGFACE_LOCAL_EMBEDDING_PROVIDER = "huggingface_local"
 RASA_PROVIDER = "rasa"

rasa/shared/core/domain.py

@@ -52,6 +52,7 @@ from rasa.shared.core.constants import (
     SlotMappingType,
 )
 from rasa.shared.core.events import SlotSet, UserUttered
+from rasa.shared.core.flows.constants import KEY_TRANSLATION
 from rasa.shared.core.slots import (
     AnySlot,
     CategoricalSlot,
@@ -117,6 +118,7 @@ RESPONSE_KEYS_TO_INTERPOLATE = [
     KEY_RESPONSES_BUTTONS,
     KEY_RESPONSES_ATTACHMENT,
     KEY_RESPONSES_QUICK_REPLIES,
+    KEY_TRANSLATION,
 ]
 
 ALL_DOMAIN_KEYS = [

rasa/shared/providers/_configs/azure_entra_id_config.py

@@ -8,7 +8,7 @@ from functools import lru_cache
 from typing import Any, Callable, Dict, List, Optional, Set, Type
 
 import structlog
-from azure.core.credentials import TokenProvider
+from azure.core.credentials import TokenCredential
 from azure.identity import (
     CertificateCredential,
     ClientSecretCredential,
@@ -77,7 +77,7 @@ class AzureEntraIDTokenProviderConfig(abc.ABC):
     """Interface for Azure Entra ID OAuth credential configuration."""
 
     @abc.abstractmethod
-    def create_azure_token_provider(self) -> TokenProvider:
+    def create_azure_token_provider(self) -> TokenCredential:
         """Create an Azure Entra ID token provider."""
         ...
 
@@ -159,7 +159,7 @@ class AzureEntraIDClientCredentialsConfig(AzureEntraIDTokenProviderConfig, BaseM
             ),
         )
 
-    def create_azure_token_provider(self) -> TokenProvider:
+    def create_azure_token_provider(self) -> TokenCredential:
         """Create a ClientSecretCredential for Azure Entra ID."""
         return create_azure_entra_id_client_credentials(
             client_id=self.client_id,
@@ -286,7 +286,7 @@ class AzureEntraIDClientCertificateConfig(AzureEntraIDTokenProviderConfig, BaseM
             ),
         )
 
-    def create_azure_token_provider(self) -> TokenProvider:
+    def create_azure_token_provider(self) -> TokenCredential:
         """Creates a CertificateCredential for Azure Entra ID."""
         return create_azure_entra_id_certificate_credentials(
             client_id=self.client_id,
@@ -369,7 +369,7 @@ class AzureEntraIDDefaultCredentialsConfig(AzureEntraIDTokenProviderConfig, Base
         """
         return cls(authority_host=config.pop(AZURE_AUTHORITY_FIELD, None))
 
-    def create_azure_token_provider(self) -> TokenProvider:
+    def create_azure_token_provider(self) -> TokenCredential:
         """Creates a DefaultAzureCredential."""
         return create_azure_entra_id_default_credentials(
             authority_host=self.authority_host
@@ -530,12 +530,12 @@ class AzureEntraIDOAuthConfig(OAuth, BaseModel):
         azure_oauth_class = AzureEntraIDOAuthConfig._get_azure_oauth_by_type(oauth_type)
         return azure_oauth_class.from_dict(oauth_config)
 
-    def _create_azure_credential(
+    def create_azure_credential(
         self,
-    ) -> TokenProvider:
+    ) -> TokenCredential:
         """Create an Azure Entra ID client which can be used to get a bearer token."""
         return self.azure_entra_id_token_provider_config.create_azure_token_provider()
 
     def get_bearer_token(self) -> str:
         """Returns a bearer token."""
-        return self._create_azure_credential().get_token(*self.scopes).token  # type: ignore
+        return self.create_azure_credential().get_token(*self.scopes).token

rasa/shared/providers/llm/litellm_router_llm_client.py

@@ -198,6 +198,7 @@ class LiteLLMRouterLLMClient(_BaseLiteLLMRouterClient, _BaseLiteLLMClient):
         """Returns the completion arguments for invoking a call through
         LiteLLM's completion functions.
         """
+
         return {
             **self._litellm_extra_parameters,
             LITE_LLM_MODEL_FIELD: self.model_group_id,

rasa/shared/providers/router/_base_litellm_router_client.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+from copy import deepcopy
 from typing import Any, Dict, List
 
 import structlog
@@ -18,6 +19,7 @@ from rasa.shared.constants import (
     USE_CHAT_COMPLETIONS_ENDPOINT_CONFIG_KEY,
 )
 from rasa.shared.exceptions import ProviderClientValidationError
+from rasa.shared.providers._configs.azure_entra_id_config import AzureEntraIDOAuthConfig
 from rasa.shared.providers._configs.litellm_router_client_config import (
     LiteLLMRouterClientConfig,
 )
@@ -61,12 +63,8 @@ class _BaseLiteLLMRouterClient:
         self._extra_parameters = kwargs or {}
         self.additional_client_setup()
         try:
-            resolved_model_configurations = (
-                self._resolve_env_vars_in_model_configurations()
-            )
-            self._router_client = Router(
-                model_list=resolved_model_configurations, **router_settings
-            )
+            # We instantiate a router client here to validate the configuration.
+            self._router_client = self._create_router_client()
         except Exception as e:
             event_info = "Cannot instantiate a router client."
             structlogger.error(
@@ -145,6 +143,14 @@ class _BaseLiteLLMRouterClient:
     @property
     def router_client(self) -> Router:
         """Returns the instantiated LiteLLM Router client."""
+        # In ca se oauth is used, due to a bug in LiteLLM,
+        # azure_ad_token_provider is not working as expected.
+        # To work around this, we create a new client every
+        # time we need to make a call which will
+        # ensure that the token is always fresh.
+        # GitHub issue for LiteLLm: https://github.com/BerriAI/litellm/issues/4417
+        if self._has_oauth():
+            return self._create_router_client()
         return self._router_client
 
     @property
@@ -175,11 +181,36 @@ class _BaseLiteLLMRouterClient:
             **self._litellm_extra_parameters,
         }
 
+    def _create_router_client(self) -> Router:
+        resolved_model_configurations = self._resolve_env_vars_in_model_configurations()
+        return Router(model_list=resolved_model_configurations, **self.router_settings)
+
+    def _has_oauth(self) -> bool:
+        for model_configuration in self.model_configurations:
+            if model_configuration.get("litellm_params", {}).get("oauth", None):
+                return True
+        return False
+
     def _resolve_env_vars_in_model_configurations(self) -> List:
         model_configuration_with_resolved_keys = []
         for model_configuration in self.model_configurations:
             resolved_model_configuration = resolve_environment_variables(
-                model_configuration
+                deepcopy(model_configuration)
             )
+
+            if not isinstance(resolved_model_configuration, dict):
+                continue
+
+            lite_llm_params = resolved_model_configuration.get("litellm_params", {})
+            if lite_llm_params.get("oauth", None):
+                oauth_config_dict = lite_llm_params.pop("oauth")
+                oauth_config = AzureEntraIDOAuthConfig.from_dict(oauth_config_dict)
+                credential = oauth_config.create_azure_credential()
+                # token_provider = get_bearer_token_provider(
+                #     credential, *oauth_config.scopes
+                # )
+                resolved_model_configuration["litellm_params"]["azure_ad_token"] = (
+                    credential.get_token(*oauth_config.scopes).token
+                )
             model_configuration_with_resolved_keys.append(resolved_model_configuration)
         return model_configuration_with_resolved_keys

rasa/shared/utils/llm.py

@@ -667,38 +667,94 @@ def get_prompt_template(
     """
     try:
         if jinja_file_path is not None:
-            return rasa.shared.utils.io.read_file(jinja_file_path)
+            prompt_template = rasa.shared.utils.io.read_file(jinja_file_path)
+            structlogger.info(
+                "utils.llm.get_prompt_template.custom_prompt_template_read_successfull",
+                event_info=(
+                    f"Custom prompt template read successfully from "
+                    f"`{jinja_file_path}`."
+                ),
+                prompt_file_path=jinja_file_path,
+            )
+            return prompt_template
     except (FileIOException, FileNotFoundException):
         structlogger.warning(
-            "Failed to read custom prompt template. Using default template instead.",
-            jinja_file_path=jinja_file_path,
+            "utils.llm.get_prompt_template.failed_to_read_custom_prompt_template",
+            event_info=(
+                "Failed to read custom prompt template. Using default template instead."
+            ),
         )
     return default_prompt_template
 
 
 def get_default_prompt_template_based_on_model(
-    config: Dict[str, Any],
+    llm_config: Dict[str, Any],
     model_prompt_mapping: Dict[str, Any],
+    default_prompt_path: str,
     fallback_prompt_path: str,
 ) -> Text:
     """Returns the default prompt template based on the model name.
 
     Args:
-        config: The model config.
+        llm_config: The model config.
         model_prompt_mapping: The mapping of model name to prompt template.
-        fallback_prompt_path: The fallback prompt path.
+        default_prompt_path: The default prompt path of the component.
+        fallback_prompt_path: The fallback prompt path for all other models
+            that do not have a mapping in the model_prompt_mapping.
 
     Returns:
         The default prompt template.
     """
-    _config = deepcopy(config)
-    if MODELS_CONFIG_KEY in _config:
-        _config = _config[MODELS_CONFIG_KEY][0]
-    provider = _config.get(PROVIDER_CONFIG_KEY)
-    model = _config.get(MODEL_CONFIG_KEY, "")
+    _llm_config = deepcopy(llm_config)
+    if MODELS_CONFIG_KEY in _llm_config:
+        _llm_config = _llm_config[MODELS_CONFIG_KEY][0]
+    provider = _llm_config.get(PROVIDER_CONFIG_KEY)
+    model = _llm_config.get(MODEL_CONFIG_KEY)
+    if not model:
+        # If the model is not defined, we default to the default prompt template.
+        structlogger.info(
+            "utils.llm.get_default_prompt_template_based_on_model.using_default_prompt_template",
+            event_info=(
+                f"Model not defined in the config. Default prompt template read from"
+                f" - `{default_prompt_path}`."
+            ),
+            default_prompt_path=default_prompt_path,
+        )
+        return importlib.resources.read_text(
+            DEFAULT_PROMPT_PACKAGE_NAME, default_prompt_path
+        )
+
     model_name = model if provider and provider in model else f"{provider}/{model}"
-    prompt_file_path = model_prompt_mapping.get(model_name, fallback_prompt_path)
-    return importlib.resources.read_text(DEFAULT_PROMPT_PACKAGE_NAME, prompt_file_path)
+    if prompt_file_path := model_prompt_mapping.get(model_name):
+        # If the model is found in the mapping, we use the model-specific prompt
+        # template.
+        structlogger.info(
+            "utils.llm.get_default_prompt_template_based_on_model.using_model_specific_prompt_template",
+            event_info=(
+                f"Using model-specific default prompt template. Default prompt "
+                f"template read from - `{prompt_file_path}`."
+            ),
+            default_prompt_path=prompt_file_path,
+            model_name=model_name,
+        )
+        return importlib.resources.read_text(
+            DEFAULT_PROMPT_PACKAGE_NAME, prompt_file_path
+        )
+
+    # If the model is not found in the mapping, we default to the fallback prompt
+    # template.
+    structlogger.info(
+        "utils.llm.get_default_prompt_template_based_on_model.using_fallback_prompt_template",
+        event_info=(
+            f"Model not found in the model prompt mapping. Fallback prompt template "
+            f"read from - `{fallback_prompt_path}`."
+        ),
+        fallback_prompt_path=fallback_prompt_path,
+        model_name=model_name,
+    )
+    return importlib.resources.read_text(
+        DEFAULT_PROMPT_PACKAGE_NAME, fallback_prompt_path
+    )
 
 
 def allowed_values_for_slot(slot: Slot) -> Union[str, None]:

rasa/telemetry.py

@@ -15,7 +15,7 @@ from collections import defaultdict
 from datetime import datetime
 from functools import wraps
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Text, Tuple
+from typing import Any, Callable, Dict, List, Optional, Text, Tuple, Type, cast
 
 import importlib_resources
 import requests
@@ -1126,12 +1126,12 @@ def _get_llm_command_generator_config(config: Dict[str, Any]) -> Optional[Dict]:
     """
     from rasa.dialogue_understanding.generator import (
         CompactLLMCommandGenerator,
+        LLMBasedCommandGenerator,
         LLMCommandGenerator,
         MultiStepLLMCommandGenerator,
         SingleStepLLMCommandGenerator,
     )
     from rasa.dialogue_understanding.generator.constants import (
-        DEFAULT_LLM_CONFIG,
         FLOW_RETRIEVAL_KEY,
         LLM_CONFIG_KEY,
     )
@@ -1162,6 +1162,12 @@ def _get_llm_command_generator_config(config: Dict[str, Any]) -> Optional[Dict]:
 
     def extract_llm_command_generator_llm_client_settings(component: Dict) -> Dict:
         """Extracts settings related to LLM command generator."""
+        component_class_lookup = {
+            LLMCommandGenerator.__name__: LLMCommandGenerator,
+            SingleStepLLMCommandGenerator.__name__: SingleStepLLMCommandGenerator,
+            MultiStepLLMCommandGenerator.__name__: MultiStepLLMCommandGenerator,
+            CompactLLMCommandGenerator.__name__: CompactLLMCommandGenerator,
+        }
         llm_config = component.get(LLM_CONFIG_KEY, {})
         # Config at this stage is not yet resolved, so read from `model_group`
         llm_model_group_id = llm_config.get(MODEL_GROUP_CONFIG_KEY)
@@ -1169,7 +1175,11 @@ def _get_llm_command_generator_config(config: Dict[str, Any]) -> Optional[Dict]:
             MODEL_NAME_CONFIG_KEY
         )
         if llm_model_group_id is None and llm_model_name is None:
-            llm_model_name = DEFAULT_LLM_CONFIG[MODEL_CONFIG_KEY]
+            component_clz = cast(
+                Type[LLMBasedCommandGenerator],
+                component_class_lookup[component["name"]],
+            )
+            llm_model_name = component_clz.get_default_llm_config()[MODEL_CONFIG_KEY]
 
         custom_prompt_used = (
             PROMPT_CONFIG_KEY in component or PROMPT_TEMPLATE_CONFIG_KEY in component

rasa/tracing/instrumentation/attribute_extractors.py

@@ -58,9 +58,7 @@ from rasa.shared.core.trackers import DialogueStateTracker
 from rasa.shared.core.training_data.structures import StoryGraph
 from rasa.shared.importers.importer import TrainingDataImporter
 from rasa.shared.nlu.constants import INTENT_NAME_KEY, SET_SLOT_COMMAND
-from rasa.shared.utils.llm import (
-    combine_custom_and_default_config,
-)
+from rasa.shared.utils.llm import combine_custom_and_default_config
 from rasa.tracing.constants import (
     PROMPT_TOKEN_LENGTH_ATTRIBUTE_NAME,
     REQUEST_BODY_SIZE_IN_BYTES_ATTRIBUTE_NAME,
@@ -375,14 +373,13 @@ def extract_attrs_for_llm_based_command_generator(
     self: "LLMBasedCommandGenerator",
     prompt: str,
 ) -> Dict[str, Any]:
-    from rasa.dialogue_understanding.generator.constants import DEFAULT_LLM_CONFIG
     from rasa.dialogue_understanding.generator.flow_retrieval import (
         DEFAULT_EMBEDDINGS_CONFIG,
     )
 
     attributes = extract_llm_config(
         self,
-        default_llm_config=DEFAULT_LLM_CONFIG,
+        default_llm_config=self.get_default_llm_config(),
         default_embeddings_config=DEFAULT_EMBEDDINGS_CONFIG,
     )