monocle-apptrace 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

monocle_apptrace/instrumentation/common/method_wrappers.py

@@ -0,0 +1,376 @@
+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.trace.span import INVALID_SPAN
+from opentelemetry.trace import get_tracer
+from contextlib import contextmanager, asynccontextmanager
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper, get_current_monocle_span, set_monocle_span_in_context, 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_monocle_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 = get_current_monocle_span()
+        if parent_span is not None and parent_span != INVALID_SPAN:
+            # 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,30 @@ 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",
+    "litellm": "workflow.litellm",
 }
 
+FRAMEWORK_WORKFLOW_LIST = [
+    "workflow.llamaindex",
+    "workflow.langchain",
+    "workflow.haystack",
+    "workflow.teams_ai",
+    "workflow.litellm",
+]
 class SpanHandler:
 
     def __init__(self,instrumentor=None):
@@ -36,7 +49,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:
@@ -55,8 +68,11 @@ class SpanHandler:
         return span_type
 
     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'])
+        try:
+            if "pipeline" in to_wrap['package']:
+                set_attribute(QUERY, args[0]['prompt_builder']['question'])
+        except Exception as e:
+            logger.warning("Warning: Error occurred in pre_task_processing: %s", str(e))
 
     @staticmethod
     def set_default_monocle_attributes(span: Span, source_path = "" ):
@@ -66,6 +82,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 +94,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 +109,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 +120,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 +129,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
@@ -134,13 +150,13 @@ class SpanHandler:
             span.set_attribute("entity.count", span_index)
         return detected_error
 
-    def hydrate_events(self, to_wrap, wrapped, instance, args, kwargs, ret_result, span, parent_span=None, ex:Exception=None) -> bool:
+    def hydrate_events(self, to_wrap, wrapped, instance, args, kwargs, ret_result, span: Span, parent_span=None, ex:Exception=None) -> bool:
         detected_error:bool = False
         if 'output_processor' in to_wrap and to_wrap["output_processor"] is not None:
             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):
@@ -170,10 +186,16 @@ class SpanHandler:
                             except Exception as e:
                                 logger.debug(f"Error evaluating accessor for attribute '{attribute_key}': {e}")
                     matching_timestamp = getattr(ret_result, "timestamps", {}).get(event_name, None)
-                    if isinstance(matching_timestamp, int):
-                        span.add_event(name=event_name, attributes=event_attributes, timestamp=matching_timestamp)
-                    else:
-                        span.add_event(name=event_name, attributes=event_attributes)
+                    alreadyExist = False
+                    for existing_event in span.events:
+                        if event_name == existing_event.name:
+                            existing_event.attributes._dict.update(event_attributes)
+                            alreadyExist = True
+                    if not alreadyExist:
+                        if isinstance(matching_timestamp, int):
+                            span.add_event(name=event_name, attributes=event_attributes, timestamp=matching_timestamp)
+                        else:
+                            span.add_event(name=event_name, attributes=event_attributes)
         return detected_error
 
     @staticmethod
@@ -187,6 +209,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
@@ -194,7 +226,7 @@ class SpanHandler:
         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):
+                if (package_name is not None and package_name.startswith(package)):
                     workflow_type = framework_workflow_type
                     break
         return workflow_type
@@ -231,7 +263,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 +284,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
+