monocle-apptrace 0.3.0b5__py3-none-any.whl → 0.3.0b7__py3-none-any.whl

monocle_apptrace/instrumentation/common/instrumentor.py

@@ -1,4 +1,5 @@
 import logging
+import inspect
 from typing import Collection, Dict, List, Union
 import random
 import uuid
@@ -14,19 +15,19 @@ 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.span_handler import SpanHandler, NonFrameworkSpanHandler
 from monocle_apptrace.instrumentation.common.wrapper_method import (
     DEFAULT_METHODS_LIST,
     WrapperMethod,
     MONOCLE_SPAN_HANDLERS
 )
-from monocle_apptrace.instrumentation.common.wrapper import scope_wrapper
+from monocle_apptrace.instrumentation.common.wrapper import scope_wrapper, ascope_wrapper, wrapper_processor
 from monocle_apptrace.instrumentation.common.utils import (
-    set_scope, remove_scope, http_route_handler, load_scopes
+    set_scope, remove_scope, http_route_handler, load_scopes, async_wrapper, http_async_route_handler
 )
-from monocle_apptrace.instrumentation.common.constants import MONOCLE_INSTRUMENTOR
+from monocle_apptrace.instrumentation.common.constants import MONOCLE_INSTRUMENTOR, WORKFLOW_TYPE_KEY
 from functools import wraps
 logger = logging.getLogger(__name__)
 
@@ -64,13 +65,11 @@ class MonocleInstrumentor(BaseInstrumentor):
         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, span=span, wrapped=wrapped, instance=instance, args=args, kwargs=kwargs,
-                                         result=response)
-                    return response
-
+                async_task = inspect.iscoroutinefunction(fn)
+                boto_method_to_wrap = to_wrap.copy()
+                boto_method_to_wrap['skip_span'] = False
+                return wrapper_processor(async_task, tracer, NonFrameworkSpanHandler(),
+                            boto_method_to_wrap, fn, instance, args, kwargs)
             return with_instrumentation
         return instrumented_endpoint_invoke
 
@@ -93,7 +92,10 @@ class MonocleInstrumentor(BaseInstrumentor):
                 final_method_list.append(method.to_dict())
 
         for method in load_scopes():
-            method['wrapper_method'] = scope_wrapper
+            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:
@@ -154,6 +156,24 @@ def setup_monocle_telemetry(
         span_handlers: Dict[str,SpanHandler] = None,
         wrapper_methods: List[Union[dict,WrapperMethod]] = None,
         union_with_default_methods: bool = True) -> None:
+    """
+    Set up Monocle telemetry for the application.
+
+    Parameters
+    ----------
+    workflow_name : str
+        The name of the workflow to be used as the service name in telemetry.
+    span_processors : List[SpanProcessor], optional
+        Custom span processors to use instead of the default ones. If None, 
+        BatchSpanProcessors with Monocle exporters will be used.
+    span_handlers : Dict[str, SpanHandler], optional
+        Dictionary of span handlers to be used by the instrumentor, mapping handler names to handler objects.
+    wrapper_methods : List[Union[dict, WrapperMethod]], optional
+        Custom wrapper methods for instrumentation. If None, default methods will be used.
+    union_with_default_methods : bool, default=True
+        If True, combine the provided wrapper_methods with the default methods.
+        If False, only use the provided wrapper_methods.
+    """
     resource = Resource(attributes={
         SERVICE_NAME: workflow_name
     })
@@ -191,40 +211,48 @@ 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():
+    """
+    Starts a new trace. All the spans created after this call will be part of the same trace. 
+    Returns:
+        Token: A token representing the attached context for the workflow span.
+                      This token is to be used later to stop the current trace.
+                      Returns None if tracing fails.
+    
+    Raises:
+        Exception: The function catches all exceptions internally and logs a warning.
+    """
     try:
-        if traceId.startswith("0x"):
-            traceId = traceId.lstrip("0x")
         tracer = get_tracer(instrumenting_module_name= MONOCLE_INSTRUMENTOR, tracer_provider= get_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
+        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
-
-
-def propagate_trace_id_from_traceparent():
-    propagate_trace_id(use_trace_context = True)
-
-
-def stop_propagate_trace_id(token) -> None:
+        logger.warning("Failed to start trace")
+        return None
+
+def stop_trace(token) -> None:
+    """
+    Stop the active trace and detach workflow type if token is provided. All the spans created after this will not be part of the trace.
+    Args:
+        token: The token that was returned when the trace was started. Used to detach 
+               workflow type. Can be None in which case only the span is ended.
+    Returns:
+        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:
@@ -235,36 +263,90 @@ def is_valid_trace_id_uuid(traceId: str) -> bool:
     return False
 
 def start_scope(scope_name: str, scope_value:str = None) -> object:
+    """
+    Start a new scope with the given name and and optional value. If no value is provided, a random UUID will be generated.
+    All the spans, across traces created after this call will have the scope attached until the scope is stopped.
+    Args:
+        scope_name: The name of the scope.
+        scope_value: Optional value of the scope. If None, a random UUID will be generated.
+    Returns:
+        Token: A token representing the attached context for the scope. This token is to be used later to stop the current scope.
+    """
     return set_scope(scope_name, scope_value)
 
 def stop_scope(token:object) -> None:
+    """
+    Stop the active scope. All the spans created after this will not have the scope attached.
+    Args:
+        token: The token that was returned when the scope was started.
+    Returns:
+        None
+    """
     remove_scope(token)
     return
 
+@contextmanager
+def monocle_trace():
+    """
+    Context manager to start and stop a scope. All the spans, across traces created within the encapsulated code will have same trace ID
+    """
+    token = start_trace()
+    try:
+        yield
+    finally:
+        stop_trace(token)
+
 @contextmanager
 def monocle_trace_scope(scope_name: str, scope_value:str = None):
+    """
+    Context manager to start and stop a scope. All the spans, across traces created within the encapsulated code will have the scope attached.
+    Args:
+        scope_name: The name of the scope.
+        scope_value: Optional value of the scope. If None, a random UUID will be generated."""
     token = start_scope(scope_name, scope_value)
     try:
         yield
     finally:
         stop_scope(token)
-
-def monocle_trace_scope_method(scope_name: str):
+    
+def monocle_trace_scope_method(scope_name: str, scope_value:str=None):
+    """
+    Decorator to start and stop a scope for a method. All the spans, across traces created in the method will have the scope attached.
+    """
     def decorator(func):
-        def wrapper(*args, **kwargs):
-            token = start_scope(scope_name)
-            try:
-                result = func(*args, **kwargs)
+        if inspect.iscoroutinefunction(func):
+            @wraps(func)
+            async def wrapper(*args, **kwargs):
+                result = async_wrapper(func, scope_name, scope_value, None, *args, **kwargs)
                 return result
-            finally:
-                stop_scope(token)
-        return wrapper
+            return wrapper
+        else:
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                token = start_scope(scope_name, scope_value)
+                try:
+                    result = func(*args, **kwargs)
+                    return result
+                finally:
+                    stop_scope(token)
+            return wrapper
     return decorator
 
 def monocle_trace_http_route(func):
-    def wrapper(req):
-        return http_route_handler(req.headers, func, req)
-    return wrapper
+    """
+    Decorator to start and stop a continue traces and scope for a http route. It will also initiate new scopes from the http headers if configured in ``monocle_scopes.json``
+    All the spans, across traces created in the route will have the scope attached.
+    """
+    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__(
@@ -277,3 +359,4 @@ class FixedIdGenerator(id_generator.IdGenerator):
 
     def generate_trace_id(self) -> int:
         return self.trace_id
+

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 opentelemetry.trace.status import Status, StatusCode
 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, get_scopes
+from monocle_apptrace.instrumentation.common.utils import set_attribute, get_scopes, MonocleSpanException
+from monocle_apptrace.instrumentation.common.constants import WORKFLOW_TYPE_KEY, WORKFLOW_TYPE_GENERIC
 
 logger = logging.getLogger(__name__)
 
@@ -20,8 +21,6 @@ WORKFLOW_TYPE_MAP = {
     "haystack": "workflow.haystack"
 }
 
-
-
 class SpanHandler:
 
     def __init__(self,instrumentor=None):
@@ -40,20 +39,37 @@ class SpanHandler:
         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 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.")
         if "pipeline" in to_wrap['package']:
             set_attribute(QUERY, args[0]['prompt_builder']['question'])
 
-    def post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, span):
-        pass
+    @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)
+        span.set_status(StatusCode.OK)
+
+
+    def post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, span:Span):
+        if span.status.status_code == StatusCode.UNSET:
+            span.set_status(StatusCode.OK)
 
     def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span):
         self.hydrate_attributes(to_wrap, wrapped, instance, args, kwargs, result, span)
@@ -61,10 +77,8 @@ class SpanHandler:
 
     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:
@@ -84,6 +98,8 @@ class SpanHandler:
                                 result = accessor(arguments)
                                 if result and isinstance(result, (str, list)):
                                     span.set_attribute(attribute_name, result)
+                            except MonocleSpanException as e:
+                                span.set_status(StatusCode.ERROR, e.message)
                             except Exception as e:
                                 logger.debug(f"Error processing accessor: {e}")
                         else:
@@ -120,50 +136,85 @@ class SpanHandler:
                                     event_attributes[attribute_key] = accessor(arguments)
                                 else:
                                     event_attributes.update(accessor(arguments))
+                            except MonocleSpanException as e:
+                                span.set_status(StatusCode.ERROR, e.message)
                             except Exception as e:
                                 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()

monocle_apptrace/instrumentation/common/utils.py

@@ -1,7 +1,7 @@
 import logging, json
 import os
 from typing import Callable, Generic, Optional, TypeVar, Mapping
-from threading import local
+import threading, asyncio
 
 from opentelemetry.context import attach, detach, get_current, get_value, set_value, Context
 from opentelemetry.trace import NonRecordingSpan, Span, get_tracer
@@ -9,7 +9,7 @@ from opentelemetry.trace.propagation import _SPAN_KEY
 from opentelemetry.sdk.trace import id_generator, TracerProvider
 from opentelemetry.propagate import inject, extract
 from opentelemetry import baggage
-from monocle_apptrace.instrumentation.common.constants import MONOCLE_SCOPE_NAME_PREFIX, SCOPE_METHOD_FILE, SCOPE_CONFIG_PATH
+from monocle_apptrace.instrumentation.common.constants import MONOCLE_SCOPE_NAME_PREFIX, SCOPE_METHOD_FILE, SCOPE_CONFIG_PATH, llm_type_map
 
 T = TypeVar('T')
 U = TypeVar('U')
@@ -21,6 +21,21 @@ embedding_model_context = {}
 scope_id_generator = id_generator.RandomIdGenerator()
 http_scopes:dict[str:str] = {}
 
+class MonocleSpanException(Exception):
+    def __init__(self, err_message:str):
+        """
+        Monocle exeption to indicate error in span processing.        
+        Parameters:
+        - err_message (str): Error message.
+        - status (str): Status code
+        """
+        super().__init__(err_message)
+        self.message = err_message
+
+    def __str__(self):
+        """String representation of the exception."""
+        return f"[Monocle Span Error: {self.message} {self.status}"
+
 def set_tracer_provider(tracer_provider: TracerProvider):
     global monocle_tracer_provider
     monocle_tracer_provider = tracer_provider
@@ -232,13 +247,69 @@ def clear_http_scopes(token:object) -> None:
     global http_scopes
     remove_scopes(token)
 
-def http_route_handler(headers, func, req):
-    token = extract_http_headers(headers)
+def http_route_handler(func, *args, **kwargs):
+    if 'req' in kwargs and hasattr(kwargs['req'], 'headers'):
+        headers = kwargs['req'].headers
+    else:
+        headers = None
+    token = None
+    if headers is not None:
+        token = extract_http_headers(headers)
     try:
-        result = func(req)
+        result = func(*args, **kwargs)
         return result
     finally:
-        clear_http_scopes(token)
+        if token is not None:
+            clear_http_scopes(token)
+
+async def http_async_route_handler(func, *args, **kwargs):
+    if 'req' in kwargs and hasattr(kwargs['req'], 'headers'):
+        headers = kwargs['req'].headers
+    else:
+        headers = None
+    return async_wrapper(func, None, None, headers, *args, **kwargs)
+
+def run_async_with_scope(method, current_context, exceptions, *args, **kwargs):
+    token = None
+    try:
+        if current_context:
+            token = attach(current_context)
+        return asyncio.run(method(*args, **kwargs))
+    except Exception as e:
+        exceptions['exception'] = e
+        raise e
+    finally:
+        if token:
+            detach(token)
+
+def async_wrapper(method, scope_name=None, scope_value=None, headers=None, *args, **kwargs):
+    try:
+        run_loop = asyncio.get_running_loop()
+    except RuntimeError:
+        run_loop = None
+
+    token = None
+    exceptions = {}
+    if scope_name:
+        token = set_scope(scope_name, scope_value)
+    elif headers:
+        token = extract_http_headers(headers)
+    current_context = get_current()
+    try:
+        if run_loop and run_loop.is_running():
+            results = []
+            thread = threading.Thread(target=lambda: results.append(run_async_with_scope(method, current_context, exceptions, *args, **kwargs)))
+            thread.start()
+            thread.join()
+            if 'exception' in exceptions:
+                raise exceptions['exception']
+            return_value = results[0] if len(results) > 0 else None
+            return return_value
+        else:
+            return run_async_with_scope(method, None, exceptions, *args, **kwargs)
+    finally:
+        if token:
+            remove_scope(token)
 
 class Option(Generic[T]):
     def __init__(self, value: Optional[T]):
@@ -270,6 +341,13 @@ def try_option(func: Callable[..., T], *args, **kwargs) -> Option[T]:
     except Exception:
         return Option(None)
 
+def get_llm_type(instance):
+    try:
+        llm_type = llm_type_map.get(type(instance).__name__.lower())
+        return llm_type
+    except:
+        pass
+
 def resolve_from_alias(my_map, alias):
     """Find a alias that is not none from list of aliases"""
     for i in alias: