rasa-pro 3.11.0__py3-none-any.whl → 3.11.0a2__py3-none-any.whl

rasa/telemetry.py

@@ -32,13 +32,7 @@ from rasa.constants import (
     CONFIG_TELEMETRY_ENABLED,
     CONFIG_TELEMETRY_ID,
 )
-from rasa.shared.constants import (
-    PROMPT_CONFIG_KEY,
-    PROMPT_TEMPLATE_CONFIG_KEY,
-    LLM_API_HEALTH_CHECK_ENV_VAR,
-    LLM_API_HEALTH_CHECK_DEFAULT_VALUE,
-    MODEL_GROUP_CONFIG_KEY,
-)
+from rasa.shared.constants import PROMPT_CONFIG_KEY, PROMPT_TEMPLATE_CONFIG_KEY
 from rasa.engine.storage.local_model_storage import LocalModelStorage
 from rasa.shared.constants import DOCS_URL_TELEMETRY, UTTER_ASK_PREFIX
 from rasa.shared.core.flows import Flow
@@ -112,7 +106,6 @@ TELEMETRY_INTERACTIVE_LEARNING_STARTED_EVENT = "Interactive Learning Started"
 TELEMETRY_SERVER_STARTED_EVENT = "Server Started"
 TELEMETRY_PROJECT_CREATED_EVENT = "Project Created"
 TELEMETRY_SHELL_STARTED_EVENT = "Shell Started"
-TELEMETRY_INSPECT_STARTED_EVENT = "Inspect Started"
 TELEMETRY_VISUALIZATION_STARTED_EVENT = "Story Visualization Started"
 TELEMETRY_TEST_CORE_EVENT = "Model Core Tested"
 TELEMETRY_TEST_NLU_EVENT = "Model NLU Tested"
@@ -164,7 +157,6 @@ NUM_LINK_STEPS = "num_link_steps"
 NUM_CALL_STEPS = "num_call_steps"
 NUM_SHARED_SLOTS_BETWEEN_FLOWS = "num_shared_slots_between_flows"
 LLM_COMMAND_GENERATOR_MODEL_NAME = "llm_command_generator_model_name"
-LLM_COMMAND_GENERATOR_MODEL_GROUP_ID = "llm_command_generator_model_group_id"
 LLM_COMMAND_GENERATOR_CUSTOM_PROMPT_USED = "llm_command_generator_custom_prompt_used"
 MULTI_STEP_LLM_COMMAND_GENERATOR_HANDLE_FLOWS_PROMPT_USED = (
     "multi_step_llm_command_generator_custom_handle_flows_prompt_used"
@@ -174,7 +166,6 @@ MULTI_STEP_LLM_COMMAND_GENERATOR_FILL_SLOTS_PROMPT_USED = (
 )
 FLOW_RETRIEVAL_ENABLED = "flow_retrieval_enabled"
 FLOW_RETRIEVAL_EMBEDDING_MODEL_NAME = "flow_retrieval_embedding_model_name"
-FLOW_RETRIEVAL_EMBEDDING_MODEL_GROUP_ID = "flow_retrieval_embedding_model_group_id"
 TRACING_BACKEND = "tracing_backend"
 METRICS_BACKEND = "metrics_backend"
 VERSION = "version"
@@ -969,13 +960,6 @@ def track_model_training(
         "policies": config.get("policies"),
         "train_schema": config.get("train_schema"),
         "predict_schema": config.get("predict_schema"),
-        "model_groups": rasa.core.utils.AvailableEndpoints.get_instance().model_groups,
-        "api_health_check_enabled": (
-            os.getenv(
-                LLM_API_HEALTH_CHECK_ENV_VAR, LLM_API_HEALTH_CHECK_DEFAULT_VALUE
-            ).lower()
-            == "true"
-        ),
         "num_intent_examples": len(nlu_data.intent_examples),
         "num_entity_examples": len(nlu_data.entity_examples),
         "num_actions": len(domain.action_names_or_texts),
@@ -1130,78 +1114,46 @@ def _get_llm_command_generator_config(config: Dict[str, Any]) -> Optional[Dict]:
                 return component
         return None
 
-    def extract_llm_command_generator_llm_client_settings(component: Dict) -> Dict:
-        """Extracts settings related to LLM command generator."""
+    def extract_settings(component: Dict) -> Dict:
+        """Extracts the settings from the command generator component."""
         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)
-        llm_model_name = llm_config.get(MODEL_CONFIG_KEY) or llm_config.get(
-            MODEL_NAME_CONFIG_KEY
+        llm_model_name = (
+            llm_config.get(MODEL_CONFIG_KEY)
+            or llm_config.get(MODEL_NAME_CONFIG_KEY)
+            or DEFAULT_LLM_CONFIG[MODEL_CONFIG_KEY]
         )
-        if llm_model_group_id is None and llm_model_name is None:
-            llm_model_name = DEFAULT_LLM_CONFIG[MODEL_CONFIG_KEY]
-
-        custom_prompt_used = (
-            PROMPT_CONFIG_KEY in component or PROMPT_TEMPLATE_CONFIG_KEY in component
-        )
-        return {
-            LLM_COMMAND_GENERATOR_MODEL_NAME: llm_model_name,
-            LLM_COMMAND_GENERATOR_MODEL_GROUP_ID: llm_model_group_id,
-            LLM_COMMAND_GENERATOR_CUSTOM_PROMPT_USED: custom_prompt_used,
-        }
-
-    def extract_multistep_command_generator_prompt_settings(component: Dict) -> Dict:
-        """Extracts settings related to multistep command generator."""
-        prompt_templates = component.get("prompt_templates", {})
-        handle_flows_prompt_used = HANDLE_FLOWS_KEY in prompt_templates
-        fill_slots_prompt_used = FILL_SLOTS_KEY in prompt_templates
-        return {
-            MULTI_STEP_LLM_COMMAND_GENERATOR_HANDLE_FLOWS_PROMPT_USED: handle_flows_prompt_used,  # noqa: E501
-            MULTI_STEP_LLM_COMMAND_GENERATOR_FILL_SLOTS_PROMPT_USED: fill_slots_prompt_used,  # noqa: E501
-        }
-
-    def extract_flow_retrieval_settings(component: Dict) -> Dict:
-        """Extracts settings related to flow retrieval."""
         flow_retrieval_config = component.get(FLOW_RETRIEVAL_KEY, {})
         flow_retrieval_enabled = flow_retrieval_config.get("active", True)
-        embeddings_config = flow_retrieval_config.get(
+        flow_retrieval_embeddings_config = flow_retrieval_config.get(
             EMBEDDINGS_CONFIG_KEY, DEFAULT_EMBEDDINGS_CONFIG
         )
         flow_retrieval_embedding_model_name = (
             (
-                embeddings_config.get(MODEL_NAME_CONFIG_KEY)
-                or embeddings_config.get(MODEL_CONFIG_KEY)
+                flow_retrieval_embeddings_config.get(MODEL_NAME_CONFIG_KEY)
+                or flow_retrieval_embeddings_config.get(MODEL_CONFIG_KEY)
             )
             if flow_retrieval_enabled
             else None
         )
-        # Config at this stage is not yet resolved, so read from `model_group`
-        flow_retrieval_embedding_model_group_id = embeddings_config.get(
-            MODEL_GROUP_CONFIG_KEY
-        )
         return {
+            LLM_COMMAND_GENERATOR_MODEL_NAME: llm_model_name,
+            LLM_COMMAND_GENERATOR_CUSTOM_PROMPT_USED: PROMPT_CONFIG_KEY in component
+            or PROMPT_TEMPLATE_CONFIG_KEY in component,
+            MULTI_STEP_LLM_COMMAND_GENERATOR_HANDLE_FLOWS_PROMPT_USED: HANDLE_FLOWS_KEY
+            in component.get("prompt_templates", {}),
+            MULTI_STEP_LLM_COMMAND_GENERATOR_FILL_SLOTS_PROMPT_USED: FILL_SLOTS_KEY
+            in component.get("prompt_templates", {}),
             FLOW_RETRIEVAL_ENABLED: flow_retrieval_enabled,
             FLOW_RETRIEVAL_EMBEDDING_MODEL_NAME: flow_retrieval_embedding_model_name,
-            FLOW_RETRIEVAL_EMBEDDING_MODEL_GROUP_ID: flow_retrieval_embedding_model_group_id,  # noqa: E501
         }
 
-    def extract_settings(component: Dict) -> Dict:
-        """Extracts the settings from the command generator component."""
-        settings = {}
-        settings.update(extract_llm_command_generator_llm_client_settings(component))
-        settings.update(extract_multistep_command_generator_prompt_settings(component))
-        settings.update(extract_flow_retrieval_settings(component))
-        return settings
-
     command_generator_config = {
         LLM_COMMAND_GENERATOR_MODEL_NAME: None,
-        LLM_COMMAND_GENERATOR_MODEL_GROUP_ID: None,
         LLM_COMMAND_GENERATOR_CUSTOM_PROMPT_USED: None,
         MULTI_STEP_LLM_COMMAND_GENERATOR_HANDLE_FLOWS_PROMPT_USED: None,
         MULTI_STEP_LLM_COMMAND_GENERATOR_FILL_SLOTS_PROMPT_USED: None,
         FLOW_RETRIEVAL_ENABLED: None,
         FLOW_RETRIEVAL_EMBEDDING_MODEL_NAME: None,
-        FLOW_RETRIEVAL_EMBEDDING_MODEL_GROUP_ID: None,
     }
 
     pipeline = config.get("pipeline", [])
@@ -1381,17 +1333,6 @@ def track_shell_started(model_type: Text) -> None:
     _track(TELEMETRY_SHELL_STARTED_EVENT, {"type": model_type})
 
 
-@ensure_telemetry_enabled
-def track_inspect_started(model_type: Text) -> None:
-    """Track when a user starts a bot using rasa inspect.
-
-    Args:
-        channel: Channel name `socketio` (used for chat assistants)
-         or `browser_audio` (used for voice).
-    """
-    _track(TELEMETRY_INSPECT_STARTED_EVENT, {"type": model_type})
-
-
 @ensure_telemetry_enabled
 def track_visualization() -> None:
     """Track when a user runs the visualization."""
@@ -1612,7 +1553,6 @@ def track_response_rephrase(
     custom_prompt_template: Optional[str],
     llm_type: Optional[str],
     llm_model: Optional[str],
-    llm_model_group_id: Optional[str],
 ) -> None:
     """Track when a user rephrases a response."""
     _track(
@@ -1622,7 +1562,6 @@ def track_response_rephrase(
             "custom_prompt_template": custom_prompt_template,
             "llm_type": llm_type,
             "llm_model": llm_model,
-            "llm_model_group_id": llm_model_group_id,
         },
     )
 
@@ -1637,10 +1576,8 @@ def track_intentless_policy_train() -> None:
 def track_intentless_policy_train_completed(
     embeddings_type: Optional[str],
     embeddings_model: Optional[str],
-    embeddings_model_group_id: Optional[str],
     llm_type: Optional[str],
     llm_model: Optional[str],
-    llm_model_group_id: Optional[str],
 ) -> None:
     """Track when a user trains a policy."""
     _track(
@@ -1648,10 +1585,8 @@ def track_intentless_policy_train_completed(
         {
             "embeddings_type": embeddings_type,
             "embeddings_model": embeddings_model,
-            "embeddings_model_group_id": embeddings_model_group_id,
             "llm_type": llm_type,
             "llm_model": llm_model,
-            "llm_model_group_id": llm_model_group_id,
         },
     )
 
@@ -1660,10 +1595,8 @@ def track_intentless_policy_train_completed(
 def track_intentless_policy_predict(
     embeddings_type: Optional[str],
     embeddings_model: Optional[str],
-    embeddings_model_group_id: Optional[str],
     llm_type: Optional[str],
     llm_model: Optional[str],
-    llm_model_group_id: Optional[str],
     score: float,
 ) -> None:
     """Track when a user trains a policy."""
@@ -1672,10 +1605,8 @@ def track_intentless_policy_predict(
         {
             "embeddings_type": embeddings_type,
             "embeddings_model": embeddings_model,
-            "embeddings_model_group_id": embeddings_model_group_id,
             "llm_type": llm_type,
             "llm_model": llm_model,
-            "llm_model_group_id": llm_model_group_id,
             "score": score,
         },
     )
@@ -1765,10 +1696,8 @@ def track_enterprise_search_policy_train_completed(
     vector_store_type: Optional[str],
     embeddings_type: Optional[str],
     embeddings_model: Optional[str],
-    embeddings_model_group_id: Optional[str],
     llm_type: Optional[str],
     llm_model: Optional[str],
-    llm_model_group_id: Optional[str],
     citation_enabled: Optional[bool],
 ) -> None:
     """Track when a user completes training Enterprise Search policy."""
@@ -1778,10 +1707,8 @@ def track_enterprise_search_policy_train_completed(
             "vector_store_type": vector_store_type,
             "embeddings_type": embeddings_type,
             "embeddings_model": embeddings_model,
-            "embeddings_model_group_id": embeddings_model_group_id,
             "llm_type": llm_type,
             "llm_model": llm_model,
-            "llm_model_group_id": llm_model_group_id,
             "citation_enabled": citation_enabled,
         },
     )
@@ -1792,10 +1719,8 @@ def track_enterprise_search_policy_predict(
     vector_store_type: Optional[str],
     embeddings_type: Optional[str],
     embeddings_model: Optional[str],
-    embeddings_model_group_id: Optional[str],
     llm_type: Optional[str],
     llm_model: Optional[str],
-    llm_model_group_id: Optional[str],
     citation_enabled: Optional[bool],
 ) -> None:
     """Track when a user predicts the next action using Enterprise Search policy."""
@@ -1805,10 +1730,8 @@ def track_enterprise_search_policy_predict(
             "vector_store_type": vector_store_type,
             "embeddings_type": embeddings_type,
             "embeddings_model": embeddings_model,
-            "embeddings_model_group_id": embeddings_model_group_id,
             "llm_type": llm_type,
             "llm_model": llm_model,
-            "llm_model_group_id": llm_model_group_id,
             "citation_enabled": citation_enabled,
         },
     )

rasa/tracing/config.py

@@ -33,7 +33,6 @@ from rasa.dialogue_understanding.generator import (
     SingleStepLLMCommandGenerator,
     MultiStepLLMCommandGenerator,
 )
-from rasa.dialogue_understanding.generator.flow_retrieval import FlowRetrieval
 from rasa.dialogue_understanding.generator.nlu_command_adapter import NLUCommandAdapter
 from rasa.engine.graph import GraphNode
 from rasa.engine.training.graph_trainer import GraphTrainer
@@ -112,7 +111,6 @@ def configure_tracing(tracer_provider: Optional[TracerProvider]) -> None:
         single_step_llm_command_generator_class=SingleStepLLMCommandGenerator,
         multi_step_llm_command_generator_class=MultiStepLLMCommandGenerator,
         custom_action_executor_subclasses=custom_action_executor_subclasses,
-        flow_retrieval_class=FlowRetrieval,
     )
 
 
@@ -131,7 +129,7 @@ def get_tracer_provider(endpoints_file: Text) -> Optional[TracerProvider]:
 
     if not cfg:
         logger.info(
-            f"No endpoint for tracing type available in {endpoints_file}, "
+            f"No endpoint for tracing type available in {endpoints_file},"
             f"tracing will not be configured."
         )
         return None

rasa/tracing/instrumentation/attribute_extractors.py

@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Text, Tuple, Union
 
 import tiktoken
 from numpy import ndarray
-from rasa.dialogue_understanding.generator.constants import FLOW_RETRIEVAL_KEY
 from rasa_sdk.grpc_py import action_webhook_pb2
 
 from rasa.core.actions.action import DirectCustomActionExecutor
@@ -20,7 +19,6 @@ from rasa.core.processor import MessageProcessor
 from rasa.core.tracker_store import TrackerStore
 from rasa.dialogue_understanding.commands import Command
 from rasa.dialogue_understanding.stack.dialogue_stack import DialogueStack
-from rasa.dialogue_understanding.generator import LLMBasedCommandGenerator
 from rasa.engine.graph import ExecutionContext, GraphModelConfiguration, GraphNode
 from rasa.engine.training.graph_trainer import GraphTrainer
 from rasa.shared.constants import (
@@ -29,8 +27,6 @@ from rasa.shared.constants import (
     PROVIDER_CONFIG_KEY,
     TIMEOUT_CONFIG_KEY,
     DEPLOYMENT_CONFIG_KEY,
-    MODEL_GROUP_ID_CONFIG_KEY,
-    LLM_CONFIG_KEY,
 )
 from rasa.shared.core.constants import REQUESTED_SLOT
 from rasa.shared.core.domain import Domain
@@ -54,7 +50,10 @@ if TYPE_CHECKING:
     from rasa.core.policies.enterprise_search_policy import EnterpriseSearchPolicy
     from rasa.core.policies.intentless_policy import IntentlessPolicy
     from rasa.core.policies.policy import PolicyPrediction
-    from rasa.dialogue_understanding.generator import CommandGenerator
+    from rasa.dialogue_understanding.generator import (
+        CommandGenerator,
+        LLMBasedCommandGenerator,
+    )
 
 # This file contains all attribute extractors for tracing instrumentation.
 # These are functions that are applied to the arguments of the wrapped function to be
@@ -301,49 +300,22 @@ def extract_attrs_for_command(
     }
 
 
-def extract_llm_config(
-    self: Any,
-    default_llm_config: Dict[str, Any],
-    default_embeddings_config: Dict[str, Any],
-) -> Dict[str, Any]:
+def extract_llm_config(self: Any, default_llm_config: Dict[str, Any]) -> Dict[str, Any]:
     if isinstance(self, ContextualResponseRephraser):
-        # ContextualResponseRephraser is not a graph component, so it's
-        # not having a full config.
-        config = {"llm": self.llm_config}
+        config = self.nlg_endpoint.kwargs
     else:
         config = self.config
 
     llm_property = combine_custom_and_default_config(
-        config.get(LLM_CONFIG_KEY), default_llm_config
+        config.get("llm"), default_llm_config
     )
 
-    if isinstance(self, LLMBasedCommandGenerator):
-        flow_retrieval_config = config.get(FLOW_RETRIEVAL_KEY, {}) or {}
-        embeddings_property = combine_custom_and_default_config(
-            flow_retrieval_config.get(EMBEDDINGS_CONFIG_KEY),
-            default_embeddings_config,
-        )
-    else:
-        embeddings_property = combine_custom_and_default_config(
-            config.get(EMBEDDINGS_CONFIG_KEY), default_embeddings_config
-        )
-
     attributes = {
         "class_name": self.__class__.__name__,
-        # llm client attributes
         "llm_model": str(llm_property.get(MODEL_CONFIG_KEY)),
         "llm_type": str(llm_property.get(PROVIDER_CONFIG_KEY)),
-        "llm_model_group_id": str(llm_property.get(MODEL_GROUP_ID_CONFIG_KEY)),
+        "embeddings": json.dumps(config.get(EMBEDDINGS_CONFIG_KEY, {})),
         "llm_temperature": str(llm_property.get("temperature")),
-        "llm_request_timeout": str(llm_property.get(TIMEOUT_CONFIG_KEY)),
-        # embedding client attributes
-        "embeddings_model": str(embeddings_property.get(MODEL_CONFIG_KEY)),
-        "embeddings_type": str(embeddings_property.get(PROVIDER_CONFIG_KEY)),
-        "embeddings_model_group_id": str(
-            embeddings_property.get(MODEL_GROUP_ID_CONFIG_KEY)
-        ),
-        # TODO: Keeping this to avoid potential breaking changes
-        "embeddings": json.dumps(embeddings_property, sort_keys=True),
         "request_timeout": str(llm_property.get(TIMEOUT_CONFIG_KEY)),
     }
 
@@ -357,16 +329,11 @@ 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,
+    from rasa.dialogue_understanding.generator.constants import (
+        DEFAULT_LLM_CONFIG,
     )
 
-    attributes = extract_llm_config(
-        self,
-        default_llm_config=DEFAULT_LLM_CONFIG,
-        default_embeddings_config=DEFAULT_EMBEDDINGS_CONFIG,
-    )
+    attributes = extract_llm_config(self, default_llm_config=DEFAULT_LLM_CONFIG)
 
     return extend_attributes_with_prompt_tokens_length(self, attributes, prompt)
 
@@ -377,12 +344,7 @@ def extract_attrs_for_contextual_response_rephraser(
 ) -> Dict[str, Any]:
     from rasa.core.nlg.contextual_response_rephraser import DEFAULT_LLM_CONFIG
 
-    attributes = extract_llm_config(
-        self,
-        default_llm_config=DEFAULT_LLM_CONFIG,
-        # rephraser is not using embeddings
-        default_embeddings_config={},
-    )
+    attributes = extract_llm_config(self, default_llm_config=DEFAULT_LLM_CONFIG)
 
     return extend_attributes_with_prompt_tokens_length(self, attributes, prompt)
 
@@ -393,12 +355,7 @@ def extract_attrs_for_create_history(
 ) -> Dict[str, Any]:
     from rasa.core.nlg.contextual_response_rephraser import DEFAULT_LLM_CONFIG
 
-    return extract_llm_config(
-        self,
-        default_llm_config=DEFAULT_LLM_CONFIG,
-        # rephraser is not using embeddings
-        default_embeddings_config={},
-    )
+    return extract_llm_config(self, default_llm_config=DEFAULT_LLM_CONFIG)
 
 
 def extract_attrs_for_generate(
@@ -414,26 +371,6 @@ def extract_attrs_for_generate(
     }
 
 
-def extract_attrs_for_performing_health_check(
-    custom_config: Optional[Dict[str, Any]],
-    default_config: Dict[str, Any],
-    log_source_method: str,
-    log_source_component: str,
-) -> Dict[str, Any]:
-    from rasa.shared.utils.health_check.health_check import is_api_health_check_enabled
-
-    attrs = {
-        "api_health_check_enabled": is_api_health_check_enabled(),
-        "health_check_trigger_component": log_source_component,
-        "health_check_trigger_method": log_source_method,
-    }
-    if is_api_health_check_enabled():
-        attrs["config"] = json.dumps(
-            combine_custom_and_default_config(custom_config, default_config)
-        )
-    return attrs
-
-
 def extract_attrs_for_execute_commands(
     tracker: DialogueStateTracker,
     all_flows: FlowsList,
@@ -643,16 +580,9 @@ def extract_attrs_for_intentless_policy_find_closest_response(
 def extract_attrs_for_intentless_policy_generate_llm_answer(
     self: "IntentlessPolicy", llm: "BaseLLM", prompt: str
 ) -> Dict[str, Any]:
-    from rasa.core.policies.intentless_policy import (
-        DEFAULT_LLM_CONFIG,
-        DEFAULT_EMBEDDINGS_CONFIG,
-    )
+    from rasa.core.policies.intentless_policy import DEFAULT_LLM_CONFIG
 
-    attributes = extract_llm_config(
-        self,
-        default_llm_config=DEFAULT_LLM_CONFIG,
-        default_embeddings_config=DEFAULT_EMBEDDINGS_CONFIG,
-    )
+    attributes = extract_llm_config(self, default_llm_config=DEFAULT_LLM_CONFIG)
 
     return extend_attributes_with_prompt_tokens_length(self, attributes, prompt)
 
@@ -660,16 +590,9 @@ def extract_attrs_for_intentless_policy_generate_llm_answer(
 def extract_attrs_for_enterprise_search_generate_llm_answer(
     self: "EnterpriseSearchPolicy", llm: "BaseLLM", prompt: str
 ) -> Dict[str, Any]:
-    from rasa.core.policies.enterprise_search_policy import (
-        DEFAULT_LLM_CONFIG,
-        DEFAULT_EMBEDDINGS_CONFIG,
-    )
+    from rasa.core.policies.enterprise_search_policy import DEFAULT_LLM_CONFIG
 
-    attributes = extract_llm_config(
-        self,
-        default_llm_config=DEFAULT_LLM_CONFIG,
-        default_embeddings_config=DEFAULT_EMBEDDINGS_CONFIG,
-    )
+    attributes = extract_llm_config(self, default_llm_config=DEFAULT_LLM_CONFIG)
 
     return extend_attributes_with_prompt_tokens_length(self, attributes, prompt)
 

rasa/tracing/instrumentation/instrumentation.py

@@ -45,7 +45,6 @@ from rasa.dialogue_understanding.generator import (
     MultiStepLLMCommandGenerator,
     SingleStepLLMCommandGenerator,
 )
-from rasa.dialogue_understanding.generator.flow_retrieval import FlowRetrieval
 from rasa.dialogue_understanding.generator.nlu_command_adapter import NLUCommandAdapter
 from rasa.engine.graph import GraphNode
 from rasa.engine.training.graph_trainer import GraphTrainer
@@ -284,7 +283,6 @@ SingleStepLLMCommandGeneratorType = TypeVar(
 MultiStepLLMCommandGeneratorType = TypeVar(
     "MultiStepLLMCommandGeneratorType", bound=MultiStepLLMCommandGenerator
 )
-FlowRetrievalType = TypeVar("FlowRetrievalType", bound=FlowRetrieval)
 CommandType = TypeVar("CommandType", bound=Command)
 PolicyType = TypeVar("PolicyType", bound=Policy)
 InformationRetrievalType = TypeVar(
@@ -319,7 +317,6 @@ def instrument(
     custom_action_executor_subclasses: Optional[
         List[Type[CustomActionExecutor]]
     ] = None,
-    flow_retrieval_class: Optional[Type[FlowRetrievalType]] = None,
 ) -> None:
     """Substitute methods to be traced by their traced counterparts.
 
@@ -448,12 +445,6 @@ def instrument(
             "_check_commands_against_startable_flows",
             attribute_extractors.extract_attrs_for_check_commands_against_startable_flows,
         )
-        _instrument_perform_health_check_method_for_component(
-            tracer_provider.get_tracer(llm_command_generator_class.__module__),
-            llm_command_generator_class,
-            "perform_llm_health_check",
-            attribute_extractors.extract_attrs_for_performing_health_check,
-        )
         mark_class_as_instrumented(llm_command_generator_class)
 
     if (
@@ -477,14 +468,6 @@ def instrument(
             "_check_commands_against_startable_flows",
             attribute_extractors.extract_attrs_for_check_commands_against_startable_flows,
         )
-        _instrument_perform_health_check_method_for_component(
-            tracer_provider.get_tracer(
-                single_step_llm_command_generator_class.__module__
-            ),
-            single_step_llm_command_generator_class,
-            "perform_llm_health_check",
-            attribute_extractors.extract_attrs_for_performing_health_check,
-        )
         mark_class_as_instrumented(single_step_llm_command_generator_class)
 
     if multi_step_llm_command_generator_class is not None and not class_is_instrumented(
@@ -505,36 +488,8 @@ def instrument(
             ),
             multi_step_llm_command_generator_class,
         )
-        _instrument_perform_health_check_method_for_component(
-            tracer_provider.get_tracer(
-                multi_step_llm_command_generator_class.__module__
-            ),
-            multi_step_llm_command_generator_class,
-            "perform_llm_health_check",
-            attribute_extractors.extract_attrs_for_performing_health_check,
-        )
         mark_class_as_instrumented(multi_step_llm_command_generator_class)
 
-    if (
-        any(
-            llm_based_command_generator_class is not None
-            for llm_based_command_generator_class in (
-                llm_command_generator_class,
-                single_step_llm_command_generator_class,
-                multi_step_llm_command_generator_class,
-            )
-        )
-        and flow_retrieval_class is not None
-        and not class_is_instrumented(flow_retrieval_class)
-    ):
-        _instrument_perform_health_check_method_for_component(
-            tracer_provider.get_tracer(flow_retrieval_class.__module__),
-            flow_retrieval_class,
-            "perform_embeddings_health_check",
-            attribute_extractors.extract_attrs_for_performing_health_check,
-        )
-        mark_class_as_instrumented(flow_retrieval_class)
-
     if command_subclasses:
         for command_subclass in command_subclasses:
             if command_subclass is not None and not class_is_instrumented(
@@ -569,12 +524,6 @@ def instrument(
             "generate",
             attribute_extractors.extract_attrs_for_generate,
         )
-        _instrument_perform_health_check_method_for_component(
-            tracer_provider.get_tracer(contextual_response_rephraser_class.__module__),
-            contextual_response_rephraser_class,
-            "perform_llm_health_check",
-            attribute_extractors.extract_attrs_for_performing_health_check,
-        )
         mark_class_as_instrumented(contextual_response_rephraser_class)
 
     if not module_is_instrumented(COMMAND_PROCESSOR_MODULE_NAME):
@@ -806,18 +755,6 @@ def _instrument_enterprise_search_policy(
         "_generate_llm_answer",
         attribute_extractors.extract_attrs_for_enterprise_search_generate_llm_answer,
     )
-    _instrument_perform_health_check_method_for_component(
-        tracer_provider.get_tracer(policy_class.__module__),
-        policy_class,
-        "perform_embeddings_health_check",
-        attribute_extractors.extract_attrs_for_performing_health_check,
-    )
-    _instrument_perform_health_check_method_for_component(
-        tracer_provider.get_tracer(policy_class.__module__),
-        policy_class,
-        "perform_llm_health_check",
-        attribute_extractors.extract_attrs_for_performing_health_check,
-    )
 
 
 def _instrument_intentless_policy(
@@ -850,18 +787,6 @@ def _instrument_intentless_policy(
         "_generate_llm_answer",
         attribute_extractors.extract_attrs_for_intentless_policy_generate_llm_answer,
     )
-    _instrument_perform_health_check_method_for_component(
-        tracer_provider.get_tracer(policy_class.__module__),
-        policy_class,
-        "perform_embeddings_health_check",
-        attribute_extractors.extract_attrs_for_performing_health_check,
-    )
-    _instrument_perform_health_check_method_for_component(
-        tracer_provider.get_tracer(policy_class.__module__),
-        policy_class,
-        "perform_llm_health_check",
-        attribute_extractors.extract_attrs_for_performing_health_check,
-    )
 
 
 def _instrument_processor(
@@ -1214,52 +1139,6 @@ def _instrument_grpc_custom_action_executor(
     logger.debug(f"Instrumented '{grpc_custom_action_executor_class.__name__}.run.")
 
 
-def _instrument_perform_health_check_method_for_component(
-    tracer: Tracer,
-    instrumented_class: Type,
-    method_name: Text,
-    attr_extractor: Optional[Callable] = None,
-    return_value_attr_extractor: Optional[Callable] = None,
-) -> None:
-    def tracing_perform_health_check_for_component(
-        fn: Callable[..., S],
-    ) -> Callable[..., S]:
-        @functools.wraps(fn)
-        def wrapper(*args: Any, **kwargs: Any) -> S:
-            # Check the first argument to adjust for self/cls depending on how
-            # the static method from LLMHealthCheckMixin / EmbeddingsLLMHealthCheckMixin
-            # is called.
-            if args and isinstance(
-                args[0], (instrumented_class, type(instrumented_class))
-            ):
-                # The first argument is self/cls; align args to match the signature
-                args = args[1:]
-
-            span_name = f"{instrumented_class.__name__}.{fn.__name__}"
-            extracted_attrs = attr_extractor(*args, **kwargs) if attr_extractor else {}
-
-            with tracer.start_as_current_span(span_name) as span:
-                result = fn(*args, **kwargs)
-
-                # Extract attributes from the return value, if an extractor is provided
-                return_value_attributes = (
-                    return_value_attr_extractor(result, *args, **kwargs)
-                    if return_value_attr_extractor
-                    else {}
-                )
-
-                span.set_attributes({**extracted_attrs, **return_value_attributes})
-                return result
-
-        return wrapper
-
-    method_to_trace = getattr(instrumented_class, method_name)
-    traced_method = tracing_perform_health_check_for_component(method_to_trace)
-    setattr(instrumented_class, method_name, traced_method)
-
-    logger.debug(f"Instrumented '{instrumented_class.__name__}.{method_name}'.")
-
-
 def _mangled_instrumented_boolean_attribute_name(instrumented_class: Type) -> Text:
     # see https://peps.python.org/pep-0008/#method-names-and-instance-variables
     # and https://stackoverflow.com/a/50401073