monocle-apptrace 0.3.0b4__py3-none-any.whl → 0.3.0b6__py3-none-any.whl

monocle_apptrace/__main__.py

@@ -0,0 +1,19 @@
+import sys, os
+import runpy
+from monocle_apptrace.instrumentation.common.instrumentor import setup_monocle_telemetry
+
+def main():
+    if len(sys.argv) < 2 or not sys.argv[1].endswith(".py"):
+        print("Usage: python -m monocle_apptrace <your-main-module-file> <args>")
+        sys.exit(1)
+    file_name = os.path.basename(sys.argv[1])
+    workflow_name = file_name[:-3]
+    setup_monocle_telemetry(workflow_name=workflow_name)
+    sys.argv.pop(0)
+
+    try:
+        runpy.run_path(path_name=sys.argv[0], run_name="__main__")
+    except Exception as e:
+        print(e)
+if __name__ == "__main__":
+    main()

monocle_apptrace/exporters/azure/blob_exporter.py

@@ -11,6 +11,7 @@ from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from typing import Sequence
 from monocle_apptrace.exporters.base_exporter import SpanExporterBase
 import json
+from monocle_apptrace.instrumentation.common.constants import MONOCLE_SDK_VERSION
 logger = logging.getLogger(__name__)
 
 class AzureBlobSpanExporter(SpanExporterBase):
@@ -72,6 +73,12 @@ class AzureBlobSpanExporter(SpanExporterBase):
         """The actual async export logic is run here."""
         # Add spans to the export queue
         for span in spans:
+            # Azure blob library has a check to generate it's own span if OpenTelemetry is loaded and Azure trace package is installed (just pip install azure-trace-opentelemetry)
+            # With Monocle,OpenTelemetry is always loaded. If the Azure trace package is installed, then it triggers the blob trace generation on every blob operation.
+            # Thus, the Monocle span write ends up generating a blob span which again comes back to the exporter .. and would result in an infinite loop.
+            # To avoid this, we check if the span has the Monocle SDK version attribute and skip it if it doesn't. That way the blob span genearted by Azure library are not exported.
+            if not span.attributes.get(MONOCLE_SDK_VERSION):
+                continue # TODO: All exporters to use same base class and check it there
             self.export_queue.append(span)
             if len(self.export_queue) >= self.max_batch_size:
                 await self.__export_spans()

monocle_apptrace/exporters/monocle_exporters.py

@@ -1,9 +1,10 @@
 from typing import Dict, Any, List
 import os
-import warnings
+import logging
 from importlib import import_module
 from opentelemetry.sdk.trace.export import SpanExporter, ConsoleSpanExporter
 from monocle_apptrace.exporters.file_exporter import FileSpanExporter
+logger = logging.getLogger(__name__)
 
 monocle_exporters: Dict[str, Any] = {
     "s3": {"module": "monocle_apptrace.exporters.aws.s3_exporter", "class": "S3SpanExporter"},
@@ -25,21 +26,21 @@ def get_monocle_exporter() -> List[SpanExporter]:
         try:
             exporter_class_path = monocle_exporters[exporter_name]
         except KeyError:
-            warnings.warn(f"Unsupported Monocle span exporter '{exporter_name}', skipping.")
+            logger.debug(f"Unsupported Monocle span exporter '{exporter_name}', skipping.")
             continue
         try:
             exporter_module = import_module(exporter_class_path["module"])
             exporter_class = getattr(exporter_module, exporter_class_path["class"])
             exporters.append(exporter_class())
         except Exception as ex:
-            warnings.warn(
+            logger.debug(
                 f"Unable to initialize Monocle span exporter '{exporter_name}', error: {ex}. Using ConsoleSpanExporter as a fallback.")
             exporters.append(ConsoleSpanExporter())
             continue
 
     # If no exporters were created, default to FileSpanExporter
     if not exporters:
-        warnings.warn("No valid Monocle span exporters configured. Defaulting to FileSpanExporter.")
+        logger.debug("No valid Monocle span exporters configured. Defaulting to FileSpanExporter.")
         exporters.append(FileSpanExporter())
 
     return exporters

monocle_apptrace/exporters/okahu/okahu_exporter.py

@@ -83,7 +83,7 @@ class OkahuSpanExporter(SpanExporter):
                         result.text,
                     )
                     return SpanExportResult.FAILURE
-                logger.warning("spans successfully exported to okahu")
+                logger.debug("spans successfully exported to okahu")
                 return SpanExportResult.SUCCESS
             except ReadTimeout as e:
                 logger.warning("Trace export timed out: %s", str(e))

monocle_apptrace/instrumentation/common/constants.py

@@ -37,6 +37,20 @@ service_name_map = {
     GITHUB_CODESPACE_SERVICE_NAME: GITHUB_CODESPACE_IDENTIFIER_ENV_NAME
 }
 
+
+llm_type_map = {
+    "sagemakerendpoint": "aws_sagemaker",
+    "azureopenai": "azure_openai",
+    "openai": "openai",
+    "chatopenai": "openai",
+    "azurechatopenai": "azure_openai",
+    "bedrock": "aws_bedrock",
+    "sagemakerllm": "aws_sagemaker",
+    "chatbedrock": "aws_bedrock",
+    "openaigenerator": "openai",
+}
+
+MONOCLE_INSTRUMENTOR = "monocle_apptrace"
 WORKFLOW_TYPE_KEY = "workflow_type"
 DATA_INPUT_KEY = "data.input"
 DATA_OUTPUT_KEY = "data.output"
@@ -47,3 +61,11 @@ RESPONSE = "response"
 SESSION_PROPERTIES_KEY = "session"
 INFRA_SERVICE_KEY = "infra_service_name"
 META_DATA = 'metadata'
+MONOCLE_SCOPE_NAME_PREFIX = "monocle.scope."
+SCOPE_METHOD_LIST = 'MONOCLE_SCOPE_METHODS'
+SCOPE_METHOD_FILE = 'monocle_scopes.json'
+SCOPE_CONFIG_PATH = 'MONOCLE_SCOPE_CONFIG_PATH'
+TRACE_PROPOGATION_URLS = "MONOCLE_TRACE_PROPAGATATION_URLS"
+WORKFLOW_TYPE_KEY = "monocle.workflow_type"
+WORKFLOW_TYPE_GENERIC = "workflow.generic"
+MONOCLE_SDK_VERSION = "monocle_apptrace.version"

monocle_apptrace/instrumentation/common/instrumentor.py

@@ -1,8 +1,10 @@
 import logging
+import inspect
 from typing import Collection, Dict, List, Union
 import random
 import uuid
 from opentelemetry import trace
+from contextlib import contextmanager
 from opentelemetry.context import attach, get_value, set_value, get_current, detach
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import unwrap
@@ -13,14 +15,20 @@ from opentelemetry.sdk.trace import Span, TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanProcessor
 from opentelemetry.trace import get_tracer
 from wrapt import wrap_function_wrapper
-from opentelemetry.trace.propagation import set_span_in_context
+from opentelemetry.trace.propagation import set_span_in_context, _SPAN_KEY
 from monocle_apptrace.exporters.monocle_exporters import get_monocle_exporter
 from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
 from monocle_apptrace.instrumentation.common.wrapper_method import (
     DEFAULT_METHODS_LIST,
     WrapperMethod,
+    MONOCLE_SPAN_HANDLERS
 )
-
+from monocle_apptrace.instrumentation.common.wrapper import scope_wrapper, ascope_wrapper
+from monocle_apptrace.instrumentation.common.utils import (
+    set_scope, remove_scope, http_route_handler, load_scopes, async_wrapper, http_async_route_handler
+)
+from monocle_apptrace.instrumentation.common.constants import MONOCLE_INSTRUMENTOR, WORKFLOW_TYPE_KEY
+from functools import wraps
 logger = logging.getLogger(__name__)
 
 SESSION_PROPERTIES_KEY = "session"
@@ -29,8 +37,6 @@ _instruments = ()
 
 monocle_tracer_provider: TracerProvider = None
 
-MONOCLE_INSTRUMENTOR = "monocle_apptrace"
-
 class MonocleInstrumentor(BaseInstrumentor):
     workflow_name: str = ""
     user_wrapper_methods: list[Union[dict,WrapperMethod]] = []
@@ -45,17 +51,36 @@ class MonocleInstrumentor(BaseInstrumentor):
             union_with_default_methods: bool = True
             ) -> None:
         self.user_wrapper_methods = user_wrapper_methods or []
-        self.handlers = handlers or {'default':SpanHandler()}
+        self.handlers = handlers
+        if self.handlers is not None:
+            for key, val in MONOCLE_SPAN_HANDLERS.items():
+                if key not in self.handlers:
+                    self.handlers[key] = val
+        else:
+            self.handlers = MONOCLE_SPAN_HANDLERS
         self.union_with_default_methods = union_with_default_methods
         super().__init__()
 
+    def get_instrumentor(self, tracer):
+        def instrumented_endpoint_invoke(to_wrap,wrapped, span_name, instance,fn):
+            @wraps(fn)
+            def with_instrumentation(*args, **kwargs):
+                handler = SpanHandler()
+                with tracer.start_as_current_span(span_name) as span:
+                    response = fn(*args, **kwargs)
+                    handler.hydrate_span(to_wrap, wrapped=wrapped, instance=instance, args=args, kwargs=kwargs,
+                                         result=response, span=span)
+                    return response
+
+            return with_instrumentation
+        return instrumented_endpoint_invoke
+
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
     def _instrument(self, **kwargs):
         tracer_provider: TracerProvider = kwargs.get("tracer_provider")
-        global monocle_tracer_provider
-        monocle_tracer_provider = tracer_provider
+        set_tracer_provider(tracer_provider)
         tracer = get_tracer(instrumenting_module_name=MONOCLE_INSTRUMENTOR, tracer_provider=tracer_provider)
 
         final_method_list = []
@@ -67,6 +92,13 @@ class MonocleInstrumentor(BaseInstrumentor):
                 final_method_list.append(method)
             elif isinstance(method, WrapperMethod):
                 final_method_list.append(method.to_dict())
+
+        for method in load_scopes():
+            if method.get('async', False):
+                method['wrapper_method'] = ascope_wrapper
+            else:
+                method['wrapper_method'] = scope_wrapper
+            final_method_list.append(method)
         
         for method_config in final_method_list:
             target_package = method_config.get("package", None)
@@ -77,6 +109,10 @@ class MonocleInstrumentor(BaseInstrumentor):
             handler_key = method_config.get("span_handler",'default')
             try:
                 handler =  self.handlers.get(handler_key)
+                if not handler:
+                    logger.warning("incorrect or empty handler falling back to default handler")
+                    handler = self.handlers.get('default')
+                handler.set_instrumentor(self.get_instrumentor(tracer))
                 wrap_function_wrapper(
                     target_package,
                     f"{target_object}.{target_method}" if target_object else target_method,
@@ -108,6 +144,14 @@ class MonocleInstrumentor(BaseInstrumentor):
                              object:{wrap_object},
                              method:{wrap_method}""")
 
+def set_tracer_provider(tracer_provider: TracerProvider):
+    global monocle_tracer_provider
+    monocle_tracer_provider = tracer_provider
+
+def get_tracer_provider() -> TracerProvider:
+    global monocle_tracer_provider
+    return monocle_tracer_provider
+
 def setup_monocle_telemetry(
         workflow_name: str,
         span_processors: List[SpanProcessor] = None,
@@ -119,7 +163,7 @@ def setup_monocle_telemetry(
     })
     exporters = get_monocle_exporter()
     span_processors = span_processors or [BatchSpanProcessor(exporter) for exporter in exporters]
-    trace_provider = TracerProvider(resource=resource)
+    set_tracer_provider(TracerProvider(resource=resource))
     attach(set_value("workflow_name", workflow_name))
     tracer_provider_default = trace.get_tracer_provider()
     provider_type = type(tracer_provider_default).__name__
@@ -129,14 +173,14 @@ def setup_monocle_telemetry(
         if not is_proxy_provider:
             tracer_provider_default.add_span_processor(processor)
         else:
-            trace_provider.add_span_processor(processor)
+            get_tracer_provider().add_span_processor(processor)
     if is_proxy_provider:
-        trace.set_tracer_provider(trace_provider)
+        trace.set_tracer_provider(get_tracer_provider())
     instrumentor = MonocleInstrumentor(user_wrapper_methods=wrapper_methods or [], 
                                        handlers=span_handlers, union_with_default_methods = union_with_default_methods)
     # instrumentor.app_name = workflow_name
     if not instrumentor.is_instrumented_by_opentelemetry:
-        instrumentor.instrument(trace_provider=trace_provider)
+        instrumentor.instrument(trace_provider=get_tracer_provider())
 
     return instrumentor
 
@@ -151,40 +195,30 @@ def on_processor_start(span: Span, parent_context):
 def set_context_properties(properties: dict) -> None:
     attach(set_value(SESSION_PROPERTIES_KEY, properties))
 
-
-def propagate_trace_id(traceId = "", use_trace_context = False):
+def start_trace():
     try:
-        if traceId.startswith("0x"):
-            traceId = traceId.lstrip("0x")
-        tracer = get_tracer(instrumenting_module_name= MONOCLE_INSTRUMENTOR, tracer_provider= monocle_tracer_provider)
-        initial_id_generator = tracer.id_generator
-        _parent_span_context = get_current() if use_trace_context else None
-        if traceId and is_valid_trace_id_uuid(traceId):
-            tracer.id_generator = FixedIdGenerator(uuid.UUID(traceId).int)
-
-        span = tracer.start_span(name = "parent_placeholder_span", context= _parent_span_context)
-        updated_span_context = set_span_in_context(span=span, context= _parent_span_context)
-        updated_span_context = set_value("root_span_id", span.get_span_context().span_id, updated_span_context)
-        token = attach(updated_span_context)
-
-        span.end()
-        tracer.id_generator = initial_id_generator
+        tracer = get_tracer(instrumenting_module_name= MONOCLE_INSTRUMENTOR, tracer_provider= get_tracer_provider())
+        span = tracer.start_span(name = "workflow")
+        updated_span_context = set_span_in_context(span=span)
+        SpanHandler.set_default_monocle_attributes(span)
+        SpanHandler.set_workflow_properties(span)
+        token = SpanHandler.attach_workflow_type(context=updated_span_context)
         return token
     except:
-        logger.warning("Failed to propagate trace id")
-    return
+        logger.warning("Failed to start trace")
+        return None
 
-
-def propagate_trace_id_from_traceparent():
-    propagate_trace_id(use_trace_context = True)
-
-
-def stop_propagate_trace_id(token) -> None:
+def stop_trace(token) -> None:
     try:
-        detach(token)
+        _parent_span_context = get_current()
+        if _parent_span_context is not None:
+            parent_span: Span = _parent_span_context.get(_SPAN_KEY, None)
+            if parent_span is not None:
+                parent_span.end()
+        if token is not None:
+            SpanHandler.detach_workflow_type(token)
     except:
-        logger.warning("Failed to stop propagating trace id")
-
+        logger.warning("Failed to stop trace")
 
 def is_valid_trace_id_uuid(traceId: str) -> bool:
     try:
@@ -194,6 +228,52 @@ def is_valid_trace_id_uuid(traceId: str) -> bool:
         pass
     return False
 
+def start_scope(scope_name: str, scope_value:str = None) -> object:
+    return set_scope(scope_name, scope_value)
+
+def stop_scope(token:object) -> None:
+    remove_scope(token)
+    return
+
+@contextmanager
+def monocle_trace_scope(scope_name: str, scope_value:str = None):
+    token = start_scope(scope_name, scope_value)
+    try:
+        yield
+    finally:
+        stop_scope(token)
+    
+def monocle_trace_scope_method(scope_name: str):
+    def decorator(func):
+        if inspect.iscoroutinefunction(func):
+            @wraps(func)
+            async def wrapper(*args, **kwargs):
+                result = async_wrapper(func, scope_name, None, *args, **kwargs)
+                return result
+            return wrapper
+        else:
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                token = start_scope(scope_name)
+                try:
+                    result = func(*args, **kwargs)
+                    return result
+                finally:
+                    stop_scope(token)
+            return wrapper
+    return decorator
+
+def monocle_trace_http_route(func):
+    if inspect.iscoroutinefunction(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            return http_async_route_handler(func, *args, **kwargs)
+        return wrapper
+    else:
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            return http_route_handler(func, *args, **kwargs)
+        return wrapper
 
 class FixedIdGenerator(id_generator.IdGenerator):
     def __init__(

monocle_apptrace/instrumentation/common/span_handler.py

@@ -1,16 +1,17 @@
 import logging
 import os
 from importlib.metadata import version
-from opentelemetry.context import get_current
-from opentelemetry.context import get_value
+from opentelemetry.context import get_value, set_value, attach, detach
 from opentelemetry.sdk.trace import Span
 
 from monocle_apptrace.instrumentation.common.constants import (
     QUERY,
     service_name_map,
     service_type_map,
+    MONOCLE_SDK_VERSION
 )
-from monocle_apptrace.instrumentation.common.utils import set_attribute
+from monocle_apptrace.instrumentation.common.utils import set_attribute, get_scopes
+from monocle_apptrace.instrumentation.common.constants import WORKFLOW_TYPE_KEY, WORKFLOW_TYPE_GENERIC
 
 logger = logging.getLogger(__name__)
 
@@ -20,41 +21,61 @@ WORKFLOW_TYPE_MAP = {
     "haystack": "workflow.haystack"
 }
 
+class SpanHandler:
 
+    def __init__(self,instrumentor=None):
+        self.instrumentor=instrumentor
 
-class SpanHandler:
+    def set_instrumentor(self, instrumentor):
+        self.instrumentor = instrumentor
 
     def validate(self, to_wrap, wrapped, instance, args, kwargs):
         pass
 
-    def pre_task_processing(self, to_wrap, wrapped, instance, args, span):
-        if self.__is_root_span(span):
-            try:
-                sdk_version = version("monocle_apptrace")
-                span.set_attribute("monocle_apptrace.version", sdk_version)
-            except Exception as e:
-                logger.warning("Exception finding monocle-apptrace version.")
+    def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
+        pass
+
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+        pass
+
+    def skip_span(self, to_wrap, wrapped, instance, args, kwargs) -> bool:
+        # If this is a workflow span type and a workflow span is already generated, then skip generating this span
+        if to_wrap.get('span_type') == "workflow" and self.is_workflow_span_active():
+            return True
+        return False
+
+    def pre_task_processing(self, to_wrap, wrapped, instance, args,kwargs, span):
         if "pipeline" in to_wrap['package']:
             set_attribute(QUERY, args[0]['prompt_builder']['question'])
 
+    @staticmethod
+    def set_default_monocle_attributes(span: Span):
+        """ Set default monocle attributes for all spans """
+        try:
+            sdk_version = version("monocle_apptrace")
+            span.set_attribute(MONOCLE_SDK_VERSION, sdk_version)
+        except Exception as e:
+            logger.warning("Exception finding monocle-apptrace version.")
+        for scope_key, scope_value in get_scopes().items():
+            span.set_attribute(f"scope.{scope_key}", scope_value)
 
+    @staticmethod
+    def set_workflow_properties(span: Span, to_wrap = None):
+        """ Set attributes of workflow if this is a root span"""
+        SpanHandler.set_workflow_attributes(to_wrap, span)
+        SpanHandler.set_app_hosting_identifier_attribute(span)
 
     def post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, span):
         pass
 
-    def set_context_properties(self, to_wrap, wrapped, instance, args, kwargs):
-        pass
-
     def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span):
         self.hydrate_attributes(to_wrap, wrapped, instance, args, kwargs, result, span)
         self.hydrate_events(to_wrap, wrapped, instance, args, kwargs, result, span)
 
     def hydrate_attributes(self, to_wrap, wrapped, instance, args, kwargs, result, span):
         span_index = 0
-        if self.__is_root_span(span):
-            span_index += self.set_workflow_attributes(to_wrap, span, span_index+1)
-            span_index += self.set_app_hosting_identifier_attribute(span, span_index+1)
-
+        if SpanHandler.is_root_span(span):
+            span_index = 2 # root span will have workflow and hosting entities pre-populated
         if 'output_processor' in to_wrap and to_wrap["output_processor"] is not None:    
             output_processor=to_wrap['output_processor']
             if 'type' in output_processor:
@@ -72,16 +93,21 @@ class SpanHandler:
                             try:
                                 arguments = {"instance":instance, "args":args, "kwargs":kwargs, "result":result}
                                 result = accessor(arguments)
-                                if result and isinstance(result, str):
+                                if result and isinstance(result, (str, list)):
                                     span.set_attribute(attribute_name, result)
                             except Exception as e:
                                 logger.debug(f"Error processing accessor: {e}")
                         else:
-                            logger.warning(f"{' and '.join([key for key in ['attribute', 'accessor'] if not processor.get(key)])} not found or incorrect in entity JSON")
+                            logger.debug(f"{' and '.join([key for key in ['attribute', 'accessor'] if not processor.get(key)])} not found or incorrect in entity JSON")
                     span_index += 1
             else:
-                logger.warning("attributes not found or incorrect written in entity json")
+                logger.debug("attributes not found or incorrect written in entity json")
 
+        # set scopes as attributes by calling get_scopes()
+        # scopes is a Mapping[str:object], iterate directly with .items()
+        for scope_key, scope_value in get_scopes().items():
+            span.set_attribute(f"scope.{scope_key}", scope_value)
+        
         if span_index > 0:
             span.set_attribute("entity.count", span_index)
 
@@ -109,46 +135,79 @@ class SpanHandler:
                                 logger.debug(f"Error evaluating accessor for attribute '{attribute_key}': {e}")
                     span.add_event(name=event_name, attributes=event_attributes)
 
-
-
-    def set_workflow_attributes(self, to_wrap, span: Span, span_index):
-        return_value = 1
-        workflow_name = self.get_workflow_name(span=span)
+    @staticmethod
+    def set_workflow_attributes(to_wrap, span: Span):
+        span_index = 1
+        workflow_name = SpanHandler.get_workflow_name(span=span)
         if workflow_name:
             span.set_attribute("span.type", "workflow")
             span.set_attribute(f"entity.{span_index}.name", workflow_name)
-            # workflow type
-        package_name = to_wrap.get('package')
-        workflow_type_set = False
-        for (package, workflow_type) in WORKFLOW_TYPE_MAP.items():
-            if (package_name is not None and package in package_name):
-                span.set_attribute(f"entity.{span_index}.type", workflow_type)
-                workflow_type_set = True
-        if not workflow_type_set:
-            span.set_attribute(f"entity.{span_index}.type", "workflow.generic")
-        return return_value
-
-    def set_app_hosting_identifier_attribute(self, span, span_index):
-        return_value = 0
+        workflow_type = SpanHandler.get_workflow_type(to_wrap)
+        span.set_attribute(f"entity.{span_index}.type", workflow_type)
+
+    @staticmethod
+    def get_workflow_type(to_wrap):
+        # workflow type
+        workflow_type = WORKFLOW_TYPE_GENERIC
+        if to_wrap is not None:
+            package_name = to_wrap.get('package')
+            for (package, framework_workflow_type) in WORKFLOW_TYPE_MAP.items():
+                if (package_name is not None and package in package_name):
+                    workflow_type = framework_workflow_type
+                    break
+        return workflow_type
+
+    def set_app_hosting_identifier_attribute(span):
+        span_index = 2
         # Search env to indentify the infra service type, if found check env for service name if possible
+        span.set_attribute(f"entity.{span_index}.type", f"app_hosting.generic")
+        span.set_attribute(f"entity.{span_index}.name", "generic")
         for type_env, type_name in service_type_map.items():
             if type_env in os.environ:
-                return_value = 1
                 span.set_attribute(f"entity.{span_index}.type", f"app_hosting.{type_name}")
                 entity_name_env = service_name_map.get(type_name, "unknown")
                 span.set_attribute(f"entity.{span_index}.name", os.environ.get(entity_name_env, "generic"))
-        return return_value
 
-    def get_workflow_name(self, span: Span) -> str:
+    @staticmethod
+    def get_workflow_name(span: Span) -> str:
         try:
             return get_value("workflow_name") or span.resource.attributes.get("service.name")
         except Exception as e:
             logger.exception(f"Error getting workflow name: {e}")
             return None
 
-    def __is_root_span(self, curr_span: Span) -> bool:
+    @staticmethod
+    def is_root_span(curr_span: Span) -> bool:
         try:
             if curr_span is not None and hasattr(curr_span, "parent"):
-                return curr_span.parent is None or get_current().get("root_span_id") == curr_span.parent.span_id
+                return curr_span.parent is None
         except Exception as e:
             logger.warning(f"Error finding root span: {e}")
+
+    def is_non_workflow_root_span(self, curr_span: Span, to_wrap) -> bool:
+        return SpanHandler.is_root_span(curr_span) and to_wrap.get("span_type") != "workflow"
+    
+    def is_workflow_span_active(self):
+        return get_value(WORKFLOW_TYPE_KEY) is not None
+
+    @staticmethod
+    def attach_workflow_type(to_wrap=None, context=None): 
+        token = None
+        if to_wrap:
+            if to_wrap.get('span_type') == "workflow":
+                token = attach(set_value(WORKFLOW_TYPE_KEY,
+                                        SpanHandler.get_workflow_type(to_wrap), context))
+        else:
+            token = attach(set_value(WORKFLOW_TYPE_KEY, WORKFLOW_TYPE_GENERIC, context))
+        return token
+
+    @staticmethod
+    def detach_workflow_type(token):
+        if token:
+            return detach(token)
+
+class NonFrameworkSpanHandler(SpanHandler):
+
+    # If the language framework is being executed, then skip generating direct openAI spans
+    def skip_span(self, to_wrap, wrapped, instance, args, kwargs) -> bool:
+        return get_value(WORKFLOW_TYPE_KEY) in WORKFLOW_TYPE_MAP.values()