refactor

rasa_sdk/__main__.py

@@ -3,6 +3,7 @@
 from rasa_sdk import utils
 from rasa_sdk.endpoint import create_argument_parser, run
 from rasa_sdk.constants import APPLICATION_ROOT_LOGGER_NAME
+from rasa_sdk.tracing.utils import get_tracer_provider
 
 
 def main_from_args(args):
@@ -17,7 +18,7 @@ def main_from_args(args):
         args.logging_config_file,
     )
     utils.update_sanic_log_level()
-    tracer_provider = utils.get_tracer_provider(args)
+    tracer_provider = get_tracer_provider(args)
 
     run(
         args.actions,

rasa_sdk/endpoint.py

@@ -4,9 +4,7 @@
 import types
 import zlib
 import json
-from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from typing import List, Text, Union, Optional
 from ssl import SSLContext
 
@@ -21,6 +19,7 @@
 from rasa_sdk.executor import ActionExecutor
 from rasa_sdk.interfaces import ActionExecutionRejection, ActionNotFoundException
 from rasa_sdk.plugin import plugin_manager
+from rasa_sdk.tracing.utils import get_tracer_and_context, set_span_attributes
 
 logger = logging.getLogger(__name__)
 
@@ -97,13 +96,7 @@ async def health(_) -> HTTPResponse:
     @app.post("/webhook")
     async def webhook(request: Request) -> HTTPResponse:
         """Webhook to retrieve action calls."""
-        span_name = "rasa_sdk.create_app.webhook"
-        if tracer_provider is None:
-            tracer = trace.get_tracer(span_name)
-            context = None
-        else:
-            tracer = tracer_provider.get_tracer(span_name)
-            context = TraceContextTextMapPropagator().extract(request.headers)
+        tracer, context, span_name = get_tracer_and_context(tracer_provider, request)
 
         with tracer.start_as_current_span(span_name, context=context) as span:
             if request.headers.get("Content-Encoding") == "deflate":
@@ -132,16 +125,7 @@ async def webhook(request: Request) -> HTTPResponse:
                 body = {"error": e.message, "action_name": e.action_name}
                 return response.json(body, status=404)
 
-            if span.is_recording():
-                span.set_attribute("http.method", "POST")
-                span.set_attribute("http.route", "/webhook")
-                span.set_attribute("next_action", action_call.get("next_action"))
-                span.set_attribute("version", action_call.get("version"))
-                span.set_attribute("sender_id", action_call.get("tracker")["sender_id"])
-                span.set_attribute(
-                    "message_id",
-                    action_call.get("tracker")["latest_message"]["message_id"],
-                )
+            set_span_attributes(span, action_call)
 
             return response.json(result, status=200)
 

rasa_sdk/tracing/utils.py

@@ -0,0 +1,56 @@
+import argparse
+from rasa_sdk.tracing import config
+from opentelemetry import trace
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+from opentelemetry.sdk.trace import TracerProvider
+from sanic.request import Request
+
+from typing import Optional, Tuple, Any, Text
+
+
+def get_tracer_provider(
+    cmdline_arguments: argparse.Namespace,
+) -> Optional[TracerProvider]:
+    """Gets the tracer provider from the command line arguments."""
+    tracer_provider = None
+    if "endpoints" in cmdline_arguments:
+        endpoints_file = cmdline_arguments.endpoints
+
+    if endpoints_file is not None:
+        tracer_provider = config.get_tracer_provider(endpoints_file)
+    return tracer_provider
+
+
+def get_tracer_and_context(
+    tracer_provider: Optional[TracerProvider], request: Request
+) -> Tuple[Any, Any, Text]:
+    """Gets tracer and context"""
+    span_name = "rasa_sdk.create_app.webhook"
+    if tracer_provider is None:
+        tracer = trace.get_tracer(span_name)
+        context = None
+    else:
+        tracer = tracer_provider.get_tracer(span_name)
+        context = TraceContextTextMapPropagator().extract(request.headers)
+    return (tracer, context, span_name)
+
+
+def set_span_attributes(span: Any, action_call: dict) -> None:
+    """Sets span attributes"""
+    set_span_attributes = {
+        "http.method": "POST",
+        "http.route": "/webhook",
+        "next_action": action_call.get("next_action"),
+        "version": action_call.get("version"),
+        "sender_id": action_call.get("tracker")["sender_id"],  # type: ignore
+        "message_id": action_call.get("tracker")["latest_message"][
+            "message_id"
+        ],  # type: ignore
+    }
+
+    if span.is_recording():
+        for key, value in set_span_attributes.items():
+            span.set_attribute(key, value)
+
+    return None

rasa_sdk/utils.py

@@ -1,4 +1,3 @@
-import argparse
 import asyncio
 import inspect
 import logging
@@ -13,8 +12,6 @@
 from typing import AbstractSet, Any, Dict, List, Text, Optional, Coroutine, Union
 
 import rasa_sdk
-from rasa_sdk.tracing import config
-from opentelemetry.sdk.trace import TracerProvider
 
 from rasa_sdk.constants import (
     DEFAULT_ENCODING,
@@ -367,15 +364,3 @@ def read_yaml_file(filename: Union[Text, Path]) -> Dict[Text, Any]:
         return read_yaml(read_file(filename, DEFAULT_ENCODING))
     except (YAMLError, DuplicateKeyError) as e:
         raise YamlSyntaxException(filename, e)
-
-
-def get_tracer_provider(
-    cmdline_arguments: argparse.Namespace,
-) -> Optional[TracerProvider]:
-    tracer_provider = None
-    if "endpoints" in cmdline_arguments:
-        endpoints_file = cmdline_arguments.endpoints
-
-    if endpoints_file is not None:
-        tracer_provider = config.get_tracer_provider(endpoints_file)
-    return tracer_provider