rasa-pro 3.8.17__py3-none-any.whl → 3.9.14__py3-none-any.whl

rasa/tracing/instrumentation/instrumentation.py

@@ -13,7 +13,6 @@ from typing import (
     Dict,
     List,
     Optional,
-    TYPE_CHECKING,
     Text,
     Type,
     TypeVar,
@@ -21,15 +20,19 @@ from typing import (
 
 from multidict import MultiDict
 from opentelemetry.context import Context
-
-import rasa.shared.utils.io
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.trace import SpanKind, Tracer
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
-from rasa.core.actions.action import Action, RemoteAction
+
+from rasa.core.actions.action import Action, RemoteAction, CustomActionExecutor
+from rasa.core.actions.custom_action_executor import RetryCustomActionExecutor
+from rasa.core.actions.grpc_custom_action_executor import GRPCCustomActionExecutor
 from rasa.core.agent import Agent
 from rasa.core.channels import OutputChannel
-from rasa.core.information_retrieval.information_retrieval import InformationRetrieval
+from rasa.core.information_retrieval.information_retrieval import (
+    InformationRetrieval,
+    SearchResultList,
+)
 from rasa.core.lock_store import LockStore
 from rasa.core.nlg import NaturalLanguageGenerator
 from rasa.core.policies.flows.flow_step_result import FlowActionPrediction
@@ -37,16 +40,21 @@ from rasa.core.policies.policy import Policy, PolicyPrediction
 from rasa.core.processor import MessageProcessor
 from rasa.core.tracker_store import TrackerStore
 from rasa.dialogue_understanding.commands import Command
-from rasa.dialogue_understanding.generator.llm_command_generator import (
+from rasa.dialogue_understanding.generator import (
     LLMCommandGenerator,
+    MultiStepLLMCommandGenerator,
+    SingleStepLLMCommandGenerator,
 )
 from rasa.dialogue_understanding.generator.nlu_command_adapter import NLUCommandAdapter
 from rasa.engine.graph import GraphNode
 from rasa.engine.training.graph_trainer import GraphTrainer
-from rasa.shared.constants import DOCS_BASE_URL
+from rasa.shared.core.domain import Domain
 from rasa.shared.core.flows import FlowsList
 from rasa.shared.core.trackers import DialogueStateTracker
+from rasa.shared.nlu.constants import SET_SLOT_COMMAND
 from rasa.shared.nlu.training_data.message import Message
+from rasa.tracing.constants import REQUEST_BODY_SIZE_IN_BYTES_ATTRIBUTE_NAME
+from rasa.tracing.instrumentation import attribute_extractors
 from rasa.tracing.instrumentation.intentless_policy_instrumentation import (
     _instrument_extract_ai_responses,
     _instrument_generate_answer,
@@ -55,15 +63,12 @@ from rasa.tracing.instrumentation.intentless_policy_instrumentation import (
 )
 from rasa.tracing.instrumentation.metrics import (
     record_llm_command_generator_metrics,
+    record_single_step_llm_command_generator_metrics,
+    record_multi_step_llm_command_generator_metrics,
     record_callable_duration_metrics,
     record_request_size_in_bytes,
 )
-from rasa.utils.endpoints import EndpointConfig
-
-from rasa.tracing.instrumentation import attribute_extractors
-
-if TYPE_CHECKING:
-    from langchain.schema import Document
+from rasa.utils.endpoints import concat_url, EndpointConfig
 
 # The `TypeVar` representing the return type for a function to be wrapped.
 S = TypeVar("S")
@@ -272,6 +277,12 @@ GraphNodeType = TypeVar("GraphNodeType", bound=GraphNode)
 LockStoreType = TypeVar("LockStoreType", bound=LockStore)
 GraphTrainerType = TypeVar("GraphTrainerType", bound=GraphTrainer)
 LLMCommandGeneratorType = TypeVar("LLMCommandGeneratorType", bound=LLMCommandGenerator)
+SingleStepLLMCommandGeneratorType = TypeVar(
+    "SingleStepLLMCommandGeneratorType", bound=SingleStepLLMCommandGenerator
+)
+MultiStepLLMCommandGeneratorType = TypeVar(
+    "MultiStepLLMCommandGeneratorType", bound=MultiStepLLMCommandGenerator
+)
 CommandType = TypeVar("CommandType", bound=Command)
 PolicyType = TypeVar("PolicyType", bound=Policy)
 InformationRetrievalType = TypeVar(
@@ -296,6 +307,16 @@ def instrument(
     vector_store_subclasses: Optional[List[Type[InformationRetrievalType]]] = None,
     nlu_command_adapter_class: Optional[Type[NLUCommandAdapterType]] = None,
     endpoint_config_class: Optional[Type[EndpointConfigType]] = None,
+    grpc_custom_action_executor_class: Optional[Type[GRPCCustomActionExecutor]] = None,
+    single_step_llm_command_generator_class: Optional[
+        Type[SingleStepLLMCommandGeneratorType]
+    ] = None,
+    multi_step_llm_command_generator_class: Optional[
+        Type[MultiStepLLMCommandGeneratorType]
+    ] = None,
+    custom_action_executor_subclasses: Optional[
+        List[Type[CustomActionExecutor]]
+    ] = None,
 ) -> None:
     """Substitute methods to be traced by their traced counterparts.
 
@@ -333,6 +354,18 @@ def instrument(
         `None` is given, no `NLUCommandAdapter` will be instrumented.
     :param endpoint_config_class: The `EndpointConfig` to be instrumented. If
         `None` is given, no `EndpointConfig` will be instrumented.
+    :param grpc_custom_action_executor_class: The `GRPCCustomActionExecution` to be
+        instrumented. If `None` is given, no `GRPCCustomActionExecution`
+        will be instrumented.
+    :param single_step_llm_command_generator_class: The `SingleStepLLMCommandGenerator`
+        to be instrumented. If `None` is given, no `SingleStepLLMCommandGenerator` will
+        be instrumented.
+    :param multi_step_llm_command_generator_class: The `MultiStepLLMCommandGenerator`
+        to be instrumented. If `None` is given, no `MultiStepLLMCommandGenerator` will
+        be instrumented.
+    :param custom_action_executor_subclasses: The subclasses of `CustomActionExecutor`
+    to be instrumented. If `None` is given, no subclass of `CustomActionExecutor`
+    will be instrumented.
     """
     if agent_class is not None and not class_is_instrumented(agent_class):
         _instrument_method(
@@ -379,7 +412,7 @@ def instrument(
             tracer_provider.get_tracer(lock_store_class.__module__),
             attribute_extractors.extract_attrs_for_lock_store,
         )
-        lock_store_class.lock = contextlib.asynccontextmanager(traced_lock_method)  # type: ignore[assignment]  # noqa: E501
+        lock_store_class.lock = contextlib.asynccontextmanager(traced_lock_method)  # type: ignore[assignment]
 
         logger.debug(f"Instrumented '{lock_store_class.__name__}.lock'.")
 
@@ -402,8 +435,8 @@ def instrument(
         _instrument_method(
             tracer_provider.get_tracer(llm_command_generator_class.__module__),
             llm_command_generator_class,
-            "_generate_action_list_using_llm",
-            attribute_extractors.extract_attrs_for_llm_command_generator,
+            "invoke_llm",
+            attribute_extractors.extract_attrs_for_llm_based_command_generator,
             metrics_recorder=record_llm_command_generator_metrics,
         )
         _instrument_method(
@@ -414,6 +447,49 @@ def instrument(
         )
         mark_class_as_instrumented(llm_command_generator_class)
 
+    if (
+        single_step_llm_command_generator_class is not None
+        and not class_is_instrumented(single_step_llm_command_generator_class)
+    ):
+        _instrument_method(
+            tracer_provider.get_tracer(
+                single_step_llm_command_generator_class.__module__
+            ),
+            single_step_llm_command_generator_class,
+            "invoke_llm",
+            attribute_extractors.extract_attrs_for_llm_based_command_generator,
+            metrics_recorder=record_single_step_llm_command_generator_metrics,
+        )
+        _instrument_method(
+            tracer_provider.get_tracer(
+                single_step_llm_command_generator_class.__module__
+            ),
+            single_step_llm_command_generator_class,
+            "_check_commands_against_startable_flows",
+            attribute_extractors.extract_attrs_for_check_commands_against_startable_flows,
+        )
+        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(
+        multi_step_llm_command_generator_class
+    ):
+        _instrument_method(
+            tracer_provider.get_tracer(
+                multi_step_llm_command_generator_class.__module__
+            ),
+            multi_step_llm_command_generator_class,
+            "invoke_llm",
+            attribute_extractors.extract_attrs_for_llm_based_command_generator,
+            metrics_recorder=record_multi_step_llm_command_generator_metrics,
+        )
+        _instrument_multi_step_llm_command_generator_parse_commands(
+            tracer_provider.get_tracer(
+                multi_step_llm_command_generator_class.__module__
+            ),
+            multi_step_llm_command_generator_class,
+        )
+        mark_class_as_instrumented(multi_step_llm_command_generator_class)
+
     if command_subclasses:
         for command_subclass in command_subclasses:
             if command_subclass is not None and not class_is_instrumented(
@@ -503,14 +579,48 @@ def instrument(
     if endpoint_config_class is not None and not class_is_instrumented(
         endpoint_config_class
     ):
-        _instrument_method(
+        _instrument_endpoint_config(
             tracer_provider.get_tracer(endpoint_config_class.__module__),
             endpoint_config_class,
-            "request",
-            attribute_extractors.extract_attrs_for_endpoint_config,
-            metrics_recorder=record_request_size_in_bytes,
         )
 
+    if grpc_custom_action_executor_class is not None and not class_is_instrumented(
+        grpc_custom_action_executor_class
+    ):
+        _instrument_grpc_custom_action_executor(
+            tracer_provider.get_tracer(grpc_custom_action_executor_class.__module__),
+            grpc_custom_action_executor_class,
+        )
+
+    if custom_action_executor_subclasses:
+        for custom_action_executor_subclass in custom_action_executor_subclasses:
+            if (
+                custom_action_executor_subclass is not None
+                and not class_is_instrumented(custom_action_executor_subclass)
+            ):
+                _instrument_method(
+                    tracer_provider.get_tracer(
+                        custom_action_executor_subclass.__module__
+                    ),
+                    custom_action_executor_subclass,
+                    "run",
+                    attribute_extractors.extract_attrs_for_custom_action_executor_run,
+                )
+
+                if issubclass(
+                    custom_action_executor_subclass, GRPCCustomActionExecutor
+                ):
+                    _instrument_method(
+                        tracer=tracer_provider.get_tracer(
+                            custom_action_executor_subclass.__module__
+                        ),
+                        instrumented_class=custom_action_executor_subclass,
+                        method_name="_request",
+                        attr_extractor=attribute_extractors.extract_attrs_for_grpc_custom_action_executor_request,
+                    )
+
+                mark_class_as_instrumented(custom_action_executor_subclass)
+
 
 def _instrument_nlu_command_adapter_predict_commands(
     tracer: Tracer, nlu_command_adapter_class: Type[NLUCommandAdapterType]
@@ -522,11 +632,12 @@ def _instrument_nlu_command_adapter_predict_commands(
             message: Message,
             flows: FlowsList,
             tracker: Optional[DialogueStateTracker] = None,
+            **kwargs: Any,
         ) -> List[Command]:
             with tracer.start_as_current_span(
                 f"{self.__class__.__name__}.{fn.__name__}"
             ) as span:
-                commands = await fn(self, message, flows, tracker)
+                commands = await fn(self, message, flows, tracker, **kwargs)
 
                 span.set_attributes(
                     {
@@ -540,8 +651,10 @@ def _instrument_nlu_command_adapter_predict_commands(
 
         return wrapper
 
-    nlu_command_adapter_class.predict_commands = tracing_nlu_command_adapter_predict_commands_wrapper(  # type: ignore[assignment]  # noqa: E501
-        nlu_command_adapter_class.predict_commands
+    nlu_command_adapter_class.predict_commands = (  # type: ignore[assignment]
+        tracing_nlu_command_adapter_predict_commands_wrapper(
+            nlu_command_adapter_class.predict_commands
+        )
     )
 
     logger.debug(
@@ -549,23 +662,71 @@ def _instrument_nlu_command_adapter_predict_commands(
     )
 
 
+def _instrument_multi_step_llm_command_generator_parse_commands(
+    tracer: Tracer,
+    multi_step_llm_command_generator_class: Type[MultiStepLLMCommandGeneratorType],
+) -> None:
+    def tracing_multi_step_llm_command_generator_parse_commands_wrapper(
+        fn: Callable,
+    ) -> Callable:
+        @functools.wraps(fn)
+        def wrapper(
+            cls: MultiStepLLMCommandGenerator,
+            actions: Optional[str],
+            tracker: DialogueStateTracker,
+            flows: FlowsList,
+            is_start_or_end_prompt: bool = False,
+        ) -> List[Command]:
+            with tracer.start_as_current_span(
+                f"{cls.__class__.__name__}.{fn.__name__}"
+            ) as span:
+                commands = fn(actions, tracker, flows, is_start_or_end_prompt)
+
+                commands_list = []
+                for command in commands:
+                    command_as_dict = command.as_dict()
+                    command_type = command.command()
+
+                    if command_type == SET_SLOT_COMMAND:
+                        slot_value = command_as_dict.pop("value", None)
+                        command_as_dict["is_slot_value_missing_or_none"] = (
+                            slot_value is None
+                        )
+
+                    commands_list.append(command_as_dict)
+
+                span.set_attributes({"commands": json.dumps(commands_list)})
+                return commands
+
+        return wrapper
+
+    multi_step_llm_command_generator_class.parse_commands = (  # type: ignore[assignment]
+        tracing_multi_step_llm_command_generator_parse_commands_wrapper(
+            multi_step_llm_command_generator_class.parse_commands
+        )
+    )
+
+
 def _instrument_information_retrieval_search(
     tracer: Tracer, vector_store_class: Type[InformationRetrievalType]
 ) -> None:
     def tracing_information_retrieval_search_wrapper(fn: Callable) -> Callable:
         @functools.wraps(fn)
         async def wrapper(
-            self: InformationRetrieval, query: Text, threshold: float = 0.0
-        ) -> List["Document"]:
+            self: InformationRetrieval,
+            query: Text,
+            tracker_state: Dict[str, Any],
+            threshold: float = 0.0,
+        ) -> SearchResultList:
             with tracer.start_as_current_span(
                 f"{self.__class__.__name__}.{fn.__name__}"
             ) as span:
-                documents = await fn(self, query, threshold)
+                documents = await fn(self, query, tracker_state, threshold)
                 span.set_attributes(
                     {
                         "query": query,
                         "document_metadata": json.dumps(
-                            [document.metadata for document in documents]
+                            [document.metadata for document in documents.results]
                         ),
                     }
                 )
@@ -573,7 +734,7 @@ def _instrument_information_retrieval_search(
 
         return wrapper
 
-    vector_store_class.search = tracing_information_retrieval_search_wrapper(  # type: ignore[assignment]  # noqa: E501
+    vector_store_class.search = tracing_information_retrieval_search_wrapper(  # type: ignore[assignment]
         vector_store_class.search
     )
 
@@ -682,7 +843,7 @@ def _instrument_get_tracker(
 
         return wrapper
 
-    processor_class.get_tracker = tracing_get_tracker_wrapper(  # type: ignore[assignment]  # noqa: E501
+    processor_class.get_tracker = tracing_get_tracker_wrapper(  # type: ignore[assignment]
         processor_class.get_tracker
     )
 
@@ -862,37 +1023,122 @@ def _instrument_run_action(
             attrs = {
                 "action_name": action.name(),
                 "sender_id": tracker.sender_id,
-                "message_id": tracker.latest_message.message_id or "default",  # type: ignore[union-attr]  # noqa: E501
+                "message_id": tracker.latest_message.message_id or "default",  # type: ignore[union-attr]
             }
+            if isinstance(action, RemoteAction):
+                if isinstance(action.executor, RetryCustomActionExecutor):
+                    attrs["executor_class_name"] = type(
+                        action.executor._custom_action_executor
+                    ).__name__
+                else:
+                    attrs["executor_class_name"] = type(action.executor).__name__
             with tracer.start_as_current_span(
                 f"{self.__class__.__name__}.{fn.__name__}",
                 kind=SpanKind.CLIENT,
                 attributes=attrs,
             ):
-                if isinstance(action, RemoteAction):
-                    if not isinstance(action.action_endpoint, EndpointConfig):
-                        rasa.shared.utils.io.raise_warning(
-                            f"No endpoint is configured to propagate the trace of this "
-                            f"custom action {action.name()}. Please take a look at "
-                            f"the docs and set an endpoint configuration in the "
-                            f"endpoints.yml file",
-                            docs=f"{DOCS_BASE_URL}/custom-actions",
-                        )
-                    else:
-                        propagator = TraceContextTextMapPropagator()
-                        propagator.inject(action.action_endpoint.headers)
-
                 return await fn(self, action, tracker, output_channel, nlg, prediction)
 
         return wrapper
 
-    processor_class._run_action = tracing_run_action_wrapper(  # type: ignore[assignment]  # noqa: E501
+    processor_class._run_action = tracing_run_action_wrapper(  # type: ignore[assignment]
         processor_class._run_action
     )
 
     logger.debug(f"Instrumented '{processor_class.__name__}._run_action'.")
 
 
+def _instrument_endpoint_config(
+    tracer: Tracer, endpoint_config_class: Type[EndpointConfigType]
+) -> None:
+    """Instrument the `request` method of the `EndpointConfig` class.
+
+    Args:
+        tracer: The `Tracer` that shall be used for tracing.
+        endpoint_config_class: The `EndpointConfig` to be instrumented.
+    """
+
+    def tracing_endpoint_config_wrapper(fn: Callable) -> Callable:
+        @functools.wraps(fn)
+        async def wrapper(
+            self: Type[EndpointConfigType],
+            method: Text = "post",
+            subpath: Optional[Text] = None,
+            content_type: Optional[Text] = "application/json",
+            compress: bool = False,
+            **kwargs: Any,
+        ) -> bool:
+            request_body = kwargs.get("json")
+            attrs: Dict[str, Any] = {"url": concat_url(self.url, subpath)}
+
+            if not request_body:
+                attrs.update({REQUEST_BODY_SIZE_IN_BYTES_ATTRIBUTE_NAME: 0})
+            else:
+                attrs.update(
+                    {
+                        REQUEST_BODY_SIZE_IN_BYTES_ATTRIBUTE_NAME: len(
+                            json.dumps(request_body).encode("utf-8")
+                        )
+                    }
+                )
+            with tracer.start_as_current_span(
+                f"{self.__class__.__name__}.{fn.__name__}",
+                attributes=attrs,
+            ):
+                TraceContextTextMapPropagator().inject(self.headers)
+
+                start_time = time.perf_counter_ns()
+                result = await fn(
+                    self, method, subpath, content_type, compress, **kwargs
+                )
+                end_time = time.perf_counter_ns()
+
+                record_callable_duration_metrics(self, start_time, end_time)
+                record_request_size_in_bytes(attrs)
+
+                return result
+
+        return wrapper
+
+    endpoint_config_class.request = tracing_endpoint_config_wrapper(  # type: ignore[assignment]
+        endpoint_config_class.request
+    )
+
+    logger.debug(f"Instrumented '{endpoint_config_class.__name__}.request'.")
+
+
+def _instrument_grpc_custom_action_executor(
+    tracer: Tracer, grpc_custom_action_executor_class: Type[GRPCCustomActionExecutor]
+) -> None:
+    """Instrument the `run` method of the `GRPCCustomActionExecutor` class.
+
+    Args:
+        tracer: The `Tracer` that shall be used for tracing.
+        grpc_custom_action_executor_class: The `GRPCCustomActionExecutor` to
+            be instrumented.
+    """
+
+    def tracing_grpc_custom_action_executor_wrapper(fn: Callable) -> Callable:
+        @functools.wraps(fn)
+        async def wrapper(
+            self: Type[GRPCCustomActionExecutor],
+            tracker: Type[DialogueStateTracker],
+            domain: Type[Domain],
+            include_domain: bool = False,
+        ) -> bool:
+            TraceContextTextMapPropagator().inject(self.action_endpoint.headers)
+            result = await fn(self, tracker, domain, include_domain)
+            return result
+
+        return wrapper
+
+    grpc_custom_action_executor_class.run = tracing_grpc_custom_action_executor_wrapper(  # type: ignore[assignment]
+        grpc_custom_action_executor_class.run
+    )
+
+    logger.debug(f"Instrumented '{grpc_custom_action_executor_class.__name__}.run.")
+
+
 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

rasa/tracing/instrumentation/intentless_policy_instrumentation.py

@@ -37,7 +37,7 @@ def _instrument_select_response_examples(
 
         return wrapper
 
-    policy_class.select_response_examples = tracing_select_response_examples_wrapper(  # type: ignore[assignment]  # noqa: E501
+    policy_class.select_response_examples = tracing_select_response_examples_wrapper(  # type: ignore[assignment]
         policy_class.select_response_examples
     )
 
@@ -72,8 +72,10 @@ def _instrument_select_few_shot_conversations(
 
         return wrapper
 
-    policy_class.select_few_shot_conversations = tracing_select_few_shot_conversations_wrapper(  # type: ignore[assignment]  # noqa: E501
-        policy_class.select_few_shot_conversations
+    policy_class.select_few_shot_conversations = (  # type: ignore[assignment]
+        tracing_select_few_shot_conversations_wrapper(
+            policy_class.select_few_shot_conversations
+        )
     )
 
     logger.debug(
@@ -102,7 +104,7 @@ def _instrument_extract_ai_responses(
 
         return wrapper
 
-    policy_class.extract_ai_responses = tracing_extract_ai_responses_wrapper(  # type: ignore[assignment]  # noqa: E501
+    policy_class.extract_ai_responses = tracing_extract_ai_responses_wrapper(  # type: ignore[assignment]
         policy_class.extract_ai_responses
     )
 
@@ -135,7 +137,7 @@ def _instrument_generate_answer(
 
         return wrapper
 
-    policy_class.generate_answer = tracing_generate_answer_wrapper(  # type: ignore[assignment]  # noqa: E501
+    policy_class.generate_answer = tracing_generate_answer_wrapper(  # type: ignore[assignment]
         policy_class.generate_answer
     )