monocle-apptrace 0.4.2__py3-none-any.whl → 0.5.0b1__py3-none-any.whl

monocle_apptrace/instrumentation/common/method_wrappers.py

@@ -0,0 +1,380 @@
+import logging
+import inspect
+from typing import Dict, List, Optional, Any
+from functools import wraps
+import inspect
+from opentelemetry.context import attach, get_current, detach
+from opentelemetry.sdk.trace import Span
+from opentelemetry.sdk.trace import Span
+from opentelemetry.trace import get_tracer
+from opentelemetry.trace.propagation import set_span_in_context, _SPAN_KEY
+from contextlib import contextmanager, asynccontextmanager
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper, task_wrapper
+from monocle_apptrace.instrumentation.common.utils import (
+    set_scope, remove_scope, http_route_handler, http_async_route_handler
+)
+from monocle_apptrace.instrumentation.common.constants import MONOCLE_INSTRUMENTOR
+from monocle_apptrace.instrumentation.common.instrumentor import get_tracer_provider
+
+logger = logging.getLogger(__name__)
+
+
+
+
+def start_trace(
+    span_name: Optional[str] = None,
+    attributes: Optional[Dict[str, Any]] = None,
+    events: Optional[List[Dict[str, Any]]] = None
+):
+    """
+    Starts a new trace. All the spans created after this call will be part of the same trace. 
+    
+    Args:
+        span_name: Optional custom span name. If None, uses the default span name.
+        attributes: Optional dictionary of custom attributes to set on the span.
+        events: Optional list of events to add to the span. Each event should be a dict with 
+                'name' and optionally 'attributes' keys.
+    
+    Returns:
+        Token: A token representing the attached context for the 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:
+        tracer = get_tracer(instrumenting_module_name= MONOCLE_INSTRUMENTOR, tracer_provider= get_tracer_provider())
+        span_name = span_name or "custom_span"
+        span = tracer.start_span(name=span_name)
+        updated_span_context = set_span_in_context(span=span)
+        
+        # Set default monocle attributes
+        SpanHandler.set_default_monocle_attributes(span)
+        if SpanHandler.is_root_span(span):
+            SpanHandler.set_workflow_properties(span)
+        
+        # Set custom attributes and events using common method
+        _setup_span_attributes_and_events(span, attributes, events)
+        
+        token = attach(updated_span_context)
+        return token
+    except Exception as e:
+        logger.warning(f"Failed to start trace: {e}")
+        return None
+
+def stop_trace(
+    token,
+    final_attributes: Optional[Dict[str, Any]] = None,
+    final_events: Optional[List[Dict[str, Any]]] = None
+) -> None:
+    """
+    Stop the active trace. 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. Can be None in which case only the span is ended.
+        final_attributes: Optional dictionary of final attributes to set on the span before ending.
+        final_events: Optional list of final events to add to the span before ending.
+    
+    Returns:
+        None
+    """
+    try:
+        _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:
+                # Set final attributes and events using common method
+                _setup_span_attributes_and_events(parent_span, final_attributes, final_events)
+                
+                parent_span.end()
+        if token is not None:
+            detach(token)
+    except Exception as e:
+        logger.warning(f"Failed to stop trace: {e}")
+
+def start_scope(
+    scope_name: str, 
+    scope_value: Optional[str] = None
+) -> object:
+    """
+    Start a new scope with the given name 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.
+    """
+    try:
+        # Set the scope using existing utility
+        token = set_scope(scope_name, scope_value)
+        return token
+    except Exception as e:
+        logger.warning(f"Failed to start scope: {e}")
+        return None
+
+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
+    """
+    try:
+        # Remove the scope
+        remove_scope(token)
+    except Exception as e:
+        logger.warning(f"Failed to stop scope: {e}")
+    return
+
+
+@contextmanager
+def monocle_trace(
+    span_name: Optional[str] = None,
+    attributes: Optional[Dict[str, Any]] = None,
+    events: Optional[List[Dict[str, Any]]] = None
+):
+    """
+    Context manager to start and stop a trace. All the spans, across traces created within the encapsulated code will have same trace ID
+    
+    Args:
+        span_name: Optional custom span name.
+        attributes: Optional dictionary of custom attributes to set on the span.
+        events: Optional list of events to add to the span at start.
+    """
+    try:
+        tracer = get_tracer(instrumenting_module_name=MONOCLE_INSTRUMENTOR, tracer_provider=get_tracer_provider())
+        span_name = span_name or "custom_span"
+                
+        with tracer.start_as_current_span(span_name) as span:
+            # Set default monocle attributes
+            SpanHandler.set_default_monocle_attributes(span)
+            if SpanHandler.is_root_span(span):
+                SpanHandler.set_workflow_properties(span)
+            
+            # Set custom attributes and events using common method
+            _setup_span_attributes_and_events(span, attributes, events)
+            
+            try:
+                yield
+            finally:
+                pass
+
+            
+    except Exception as e:
+        logger.warning(f"Failed in monocle_trace: {e}")
+        yield  # Still yield to not break the context manager
+
+@asynccontextmanager
+async def amonocle_trace(
+    span_name: Optional[str] = None,
+    attributes: Optional[Dict[str, Any]] = None,
+    events: Optional[List[Dict[str, Any]]] = None
+):
+    """
+    Async context manager to start and stop a trace. All the spans, across traces created within the encapsulated code will have same trace ID
+    
+    Args:
+        span_name: Optional custom span name.
+        attributes: Optional dictionary of custom attributes to set on the span.
+        events: Optional list of events to add to the span at start.
+    """
+    try:
+        tracer = get_tracer(instrumenting_module_name=MONOCLE_INSTRUMENTOR, tracer_provider=get_tracer_provider())
+        span_name = span_name or "custom_span"
+
+        with tracer.start_as_current_span(span_name) as span:
+            # Set default monocle attributes
+            SpanHandler.set_default_monocle_attributes(span)
+            if SpanHandler.is_root_span(span):
+                SpanHandler.set_workflow_properties(span)
+            
+            # Set custom attributes and events using common method
+            _setup_span_attributes_and_events(span, attributes, events)
+            
+            try:
+                yield
+            finally:
+                pass
+                
+            
+    except Exception as e:
+        logger.warning(f"Failed in amonocle_trace: {e}")
+        yield  # Still yield to not break the context manager
+
+@contextmanager
+def monocle_trace_scope(
+    scope_name: str, 
+    scope_value: Optional[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)
+
+@asynccontextmanager
+async def amonocle_trace_scope(
+    scope_name: str, 
+    scope_value: Optional[str] = None
+):
+    """
+    Async 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, 
+    scope_value: Optional[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.
+    
+    Args:
+        scope_name: The name of the scope.
+        scope_value: Optional value of the scope. If None, a random UUID will be generated.
+    """
+    def decorator(func):
+        if inspect.iscoroutinefunction(func):
+            @wraps(func)
+            async def wrapper(*args, **kwargs):
+                async with amonocle_trace_scope(
+                    scope_name, scope_value
+                ):
+                    result = await func(*args, **kwargs)
+                    return result
+            return wrapper
+        else:
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                with monocle_trace_scope(
+                    scope_name, scope_value
+                ):
+                    result = func(*args, **kwargs)
+                    return result
+            return wrapper
+    return decorator
+
+def monocle_trace_method(
+    span_name: Optional[str] = None
+):
+    """
+    Decorator to start and stop a trace for a method. All the spans created in the method will be part of the same trace.
+    
+    Args:
+        span_name: Optional custom span name. If None, uses the decorated function's name.
+    """
+    
+    def decorator(func):
+        tracer = get_tracer(instrumenting_module_name=MONOCLE_INSTRUMENTOR, tracer_provider=get_tracer_provider())
+        handler = SpanHandler()
+        source_path= func.__code__.co_filename + ":" + str(func.__code__.co_firstlineno)
+        # Use function name as span name if not provided
+        effective_span_name = span_name or func.__name__ or "custom_span"
+
+        if inspect.iscoroutinefunction(func):
+            @wraps(func)
+            async def wrapper(*args, **kwargs):
+                return await atask_wrapper(
+                    tracer=tracer,
+                    handler=handler,
+                    to_wrap={
+                        "span_name": effective_span_name,
+                        "output_processor":{
+                            "type": "custom",
+                        }
+                    }
+                )(  wrapped=func,                        
+                    instance=None,
+                    source_path=source_path,
+                    args=args,
+                    kwargs=kwargs)
+            return wrapper
+        else:
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                return task_wrapper(
+                    tracer=tracer,
+                    handler=handler,
+                    to_wrap={
+                        "span_name": effective_span_name,
+                        "output_processor":{
+                            "type": "custom",
+                        }
+                    }
+                )(  wrapped=func,                        
+                    instance=None,
+                    source_path=source_path,
+                    args=args,
+                    kwargs=kwargs)
+            return wrapper
+    return decorator
+
+def monocle_trace_http_route(func):
+    """
+    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 await 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
+
+
+def _setup_span_attributes_and_events(
+    span,
+    attributes: Optional[Dict[str, Any]] = None,
+    events: Optional[List[Dict[str, Any]]] = None
+) -> None:
+    """
+    Common method to set attributes and events on a span.
+    
+    Args:
+        span: The span to configure
+        attributes: Optional dictionary of custom attributes to set on the span
+        events: Optional list of events to add to the span
+    """
+    # Set custom attributes if provided
+    if attributes:
+        for key, value in attributes.items():
+            if value is not None:
+                span.set_attribute(key, value)
+    
+    # Add custom events if provided
+    if events:
+        for event in events:
+            event_name = event.get('name')
+            event_attributes = event.get('attributes', {})
+            if event_name:
+                span.add_event(event_name, event_attributes)
+

monocle_apptrace/instrumentation/common/span_handler.py

@@ -11,17 +11,28 @@ from monocle_apptrace.instrumentation.common.constants import (
     MONOCLE_SDK_VERSION, MONOCLE_SDK_LANGUAGE, MONOCLE_DETECTED_SPAN_ERROR
 )
 from monocle_apptrace.instrumentation.common.utils import set_attribute, get_scopes, MonocleSpanException, get_monocle_version
-from monocle_apptrace.instrumentation.common.constants import WORKFLOW_TYPE_KEY, WORKFLOW_TYPE_GENERIC
+from monocle_apptrace.instrumentation.common.constants import WORKFLOW_TYPE_KEY, WORKFLOW_TYPE_GENERIC, CHILD_ERROR_CODE
 
 logger = logging.getLogger(__name__)
 
 WORKFLOW_TYPE_MAP = {
+    "llama_index.core.agent.workflow": WORKFLOW_TYPE_GENERIC,
     "llama_index": "workflow.llamaindex",
     "langchain": "workflow.langchain",
     "haystack": "workflow.haystack",
     "teams.ai": "workflow.teams_ai",
+    "langgraph": "workflow.langgraph",
+    "openai": "workflow.openai",
+    "anthropic": "workflow.anthropic",
+    "gemini": "workflow.gemini",
 }
 
+FRAMEWORK_WORKFLOW_LIST = [
+    "workflow.llamaindex",
+    "workflow.langchain",
+    "workflow.haystack",
+    "workflow.teams_ai",
+]
 class SpanHandler:
 
     def __init__(self,instrumentor=None):
@@ -36,7 +47,7 @@ class SpanHandler:
     def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
         pass
 
-    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value, token=None):
         pass
 
     def skip_span(self, to_wrap, wrapped, instance, args, kwargs) -> bool:
@@ -66,6 +77,9 @@ class SpanHandler:
         span.set_attribute("span_source", source_path)
         for scope_key, scope_value in get_scopes().items():
             span.set_attribute(f"scope.{scope_key}", scope_value)
+        workflow_name = SpanHandler.get_workflow_name(span=span)
+        if workflow_name:
+            span.set_attribute("workflow.name", workflow_name)
 
     @staticmethod
     def set_workflow_properties(span: Span, to_wrap = None):
@@ -75,17 +89,14 @@ class SpanHandler:
 
     @staticmethod
     def set_non_workflow_properties(span: Span, to_wrap = None):
-        workflow_name = SpanHandler.get_workflow_name(span=span)
-        if workflow_name:
-            span.set_attribute("workflow.name", workflow_name)
         span.set_attribute("span.type", "generic")
 
-    def post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, span:Span):
+    def post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, ex, span:Span, parent_span:Span):
         pass
 
     def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None) -> bool:
         try:
-            detected_error_in_attribute = self.hydrate_attributes(to_wrap, wrapped, instance, args, kwargs, result, span)
+            detected_error_in_attribute = self.hydrate_attributes(to_wrap, wrapped, instance, args, kwargs, result, span, parent_span)
             detected_error_in_event = self.hydrate_events(to_wrap, wrapped, instance, args, kwargs, result, span, parent_span, ex)
             if detected_error_in_attribute or detected_error_in_event:
                 span.set_attribute(MONOCLE_DETECTED_SPAN_ERROR, True)
@@ -93,7 +104,7 @@ class SpanHandler:
             if span.status.status_code == StatusCode.UNSET and ex is None:
                 span.set_status(StatusCode.OK)
 
-    def hydrate_attributes(self, to_wrap, wrapped, instance, args, kwargs, result, span:Span) -> bool:
+    def hydrate_attributes(self, to_wrap, wrapped, instance, args, kwargs, result, span:Span, parent_span:Span) -> bool:
         detected_error:bool = False
         span_index = 0
         if SpanHandler.is_root_span(span):
@@ -104,6 +115,7 @@ class SpanHandler:
             skip_processors:list[str] = self.skip_processor(to_wrap, wrapped, instance, span, args, kwargs) or []
 
             if 'attributes' in output_processor and 'attributes' not in skip_processors:
+                arguments = {"instance":instance, "args":args, "kwargs":kwargs, "result":result, "parent_span":parent_span, "span":span}
                 for processors in output_processor["attributes"]:
                     for processor in processors:
                         attribute = processor.get('attribute')
@@ -112,10 +124,9 @@ class SpanHandler:
                         if attribute and accessor:
                             attribute_name = f"entity.{span_index+1}.{attribute}"
                             try:
-                                arguments = {"instance":instance, "args":args, "kwargs":kwargs, "result":result}
-                                result = accessor(arguments)
-                                if result and isinstance(result, (str, list)):
-                                    span.set_attribute(attribute_name, result)
+                                processor_result = accessor(arguments)
+                                if processor_result and isinstance(processor_result, (str, list)):
+                                    span.set_attribute(attribute_name, processor_result)
                             except MonocleSpanException as e:
                                 span.set_status(StatusCode.ERROR, e.message)
                                 detected_error = True
@@ -140,7 +151,7 @@ class SpanHandler:
             output_processor=to_wrap['output_processor']
             skip_processors:list[str] = self.skip_processor(to_wrap, wrapped, instance, span, args, kwargs) or []
 
-            arguments = {"instance": instance, "args": args, "kwargs": kwargs, "result": ret_result, "exception":ex}
+            arguments = {"instance": instance, "args": args, "kwargs": kwargs, "result": ret_result, "exception":ex, "parent_span":parent_span, "span": span}
             # Process events if they are defined in the output_processor.
             # In case of inference.modelapi skip the event processing unless the span has an exception
             if 'events' in output_processor and ('events' not in skip_processors or ex is not None):
@@ -187,6 +198,16 @@ class SpanHandler:
         workflow_type = SpanHandler.get_workflow_type(to_wrap)
         span.set_attribute(f"entity.{span_index}.type", workflow_type)
 
+    def get_workflow_name_in_progress(self) -> str:
+        return get_value(WORKFLOW_TYPE_KEY)
+
+    @staticmethod
+    def is_framework_workflow(workflow_type) -> bool:
+        return workflow_type in FRAMEWORK_WORKFLOW_LIST
+
+    def is_framework_span_in_progress(self) -> bool:
+        return SpanHandler.is_framework_workflow(self.get_workflow_name_in_progress())
+
     @staticmethod
     def get_workflow_type(to_wrap):
         # workflow type
@@ -231,7 +252,7 @@ class SpanHandler:
         token = None
         if to_wrap:
             workflow_type = SpanHandler.get_workflow_type(to_wrap)
-            if workflow_type != WORKFLOW_TYPE_GENERIC:
+            if SpanHandler.is_framework_workflow(workflow_type):
                 token = attach(set_value(WORKFLOW_TYPE_KEY,
                                         SpanHandler.get_workflow_type(to_wrap), context))
         return token
@@ -252,21 +273,15 @@ class SpanHandler:
 
 
 class NonFrameworkSpanHandler(SpanHandler):
-
-    def get_workflow_name_in_progress(self) -> str:
-        return get_value(WORKFLOW_TYPE_KEY)
-
-    def is_framework_span_in_progess(self) -> bool:
-        return self.get_workflow_name_in_progress() in WORKFLOW_TYPE_MAP.values()
-
     # If the language framework is being executed, then skip generating direct openAI attributes and events
     def skip_processor(self, to_wrap, wrapped, instance, span, args, kwargs) -> list[str]:
-        if self.is_framework_span_in_progess():
+        if super().is_framework_span_in_progress():
             return ["attributes", "events"]
     
     def set_span_type(self, to_wrap, wrapped, instance, output_processor, span:Span, args, kwargs) -> str:
         span_type = super().set_span_type(to_wrap, wrapped, instance, output_processor, span, args, kwargs)
-        if self.is_framework_span_in_progess() and span_type is not None:
+        if self.is_framework_span_in_progress() and span_type is not None:
             span_type = span_type+".modelapi"
             span.set_attribute("span.type", span_type)
-        return span_type
+        return span_type
+

monocle_apptrace/instrumentation/common/utils.py

@@ -70,7 +70,7 @@ def with_tracer_wrapper(func):
     """Helper for providing tracer for wrapper functions."""
 
     def _with_tracer(tracer, handler, to_wrap):
-        def wrapper(wrapped, instance, args, kwargs):
+        def wrapper(wrapped, instance, args, kwargs, source_path=None):
             try:
                 # get and log the parent span context if injected by the application
                 # This is useful for debugging and tracing of Azure functions
@@ -83,12 +83,12 @@ def with_tracer_wrapper(func):
                             f"Parent span is found with trace id {hex(parent_span.get_span_context().trace_id)}")
             except Exception as e:
                 logger.error("Exception in attaching parent context: %s", e)
-
-            if traceback.extract_stack().__len__() > 2:
-                filename, line_number, _, _ = traceback.extract_stack()[-2]
-                source_path = f"{filename}:{line_number}"
-            else:
-                source_path = ""
+            if not source_path:
+                if traceback.extract_stack().__len__() > 2:
+                    filename, line_number, _, _ = traceback.extract_stack()[-2]
+                    source_path = f"{filename}:{line_number}"
+                else:
+                    source_path = ""
             val = func(tracer, handler, to_wrap, wrapped, instance, source_path, args, kwargs)
             return val
 
@@ -152,13 +152,6 @@ def flatten_dict(d, parent_key='', sep='_'):
             items.append((new_key, v))
     return dict(items)
 
-def get_fully_qualified_class_name(instance):
-    if instance is None:
-        return None
-    module_name = instance.__class__.__module__
-    qualname = instance.__class__.__qualname__
-    return f"{module_name}.{qualname}"
-
 # returns json path like key probe in a dictionary
 def get_nested_value(data, keys):
     for key in keys:
@@ -330,6 +323,12 @@ def add_monocle_trace_state(headers:dict[str:str]) -> None:
     else:
         headers['tracestate'] = monocle_trace_state
 
+def get_json_dumps(obj) -> str:
+    try:
+        return json.dumps(obj)
+    except TypeError as e:
+        return str(obj)
+
 class Option(Generic[T]):
     def __init__(self, value: Optional[T]):
         self.value = value
@@ -392,6 +391,13 @@ def get_exception_message(arguments):
     else:
         return ''
 
+def get_error_message(arguments):
+    status_code = get_status_code(arguments)
+    if status_code == 'success':
+        return ''
+    else:
+        return status_code
+
 def get_status_code(arguments):
     if arguments["exception"] is not None:
         return get_exception_status_code(arguments)