monocle-apptrace 0.3.0b5__tar.gz → 0.3.0b6__tar.gz

{monocle_apptrace-0.3.0b5 → monocle_apptrace-0.3.0b6}/CHANGELOG.md

@@ -1,4 +1,4 @@
-## Version 0.3.0b5 (2024-12-10)
+## Version 0.3.0b6 (2024-12-10)
 
 - Add dev dependency for Mistral AI integration ([#81](https://github.com/monocle2ai/monocle/pull/81))
 - Add VectorStore deployment URL capture support ([#80](https://github.com/monocle2ai/monocle/pull/80))  

{monocle_apptrace-0.3.0b5 → monocle_apptrace-0.3.0b6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: monocle_apptrace
-Version: 0.3.0b5
+Version: 0.3.0b6
 Summary: package with monocle genAI tracing
 Project-URL: Homepage, https://github.com/monocle2ai/monocle
 Project-URL: Issues, https://github.com/monocle2ai/monocle/issues

{monocle_apptrace-0.3.0b5 → monocle_apptrace-0.3.0b6}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "monocle_apptrace"
-version = "0.3.0b5"
+version = "0.3.0b6"
 authors = []
 description = "package with monocle genAI tracing"
 readme = "README.md"

{monocle_apptrace-0.3.0b5 → monocle_apptrace-0.3.0b6}/src/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-0.3.0b5 → monocle_apptrace-0.3.0b6}/src/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-0.3.0b5 → monocle_apptrace-0.3.0b6}/src/monocle_apptrace/instrumentation/common/constants.py

@@ -37,6 +37,19 @@ 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"
@@ -52,3 +65,7 @@ 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-0.3.0b5 → monocle_apptrace-0.3.0b6}/src/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,7 +15,7 @@ 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 (
@@ -22,11 +23,11 @@ from monocle_apptrace.instrumentation.common.wrapper_method import (
     WrapperMethod,
     MONOCLE_SPAN_HANDLERS
 )
-from monocle_apptrace.instrumentation.common.wrapper import scope_wrapper
+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
+    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__)
 
@@ -67,8 +68,8 @@ class MonocleInstrumentor(BaseInstrumentor):
                 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)
+                    handler.hydrate_span(to_wrap, wrapped=wrapped, instance=instance, args=args, kwargs=kwargs,
+                                         result=response, span=span)
                     return response
 
             return with_instrumentation
@@ -93,7 +94,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:
@@ -191,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= 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
+        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:
@@ -248,23 +242,38 @@ def monocle_trace_scope(scope_name: str, scope_value:str = None):
         yield
     finally:
         stop_scope(token)
-
+    
 def monocle_trace_scope_method(scope_name: str):
     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, 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)
+                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
+    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-0.3.0b5 → monocle_apptrace-0.3.0b6}/src/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, get_scopes
+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,18 +39,32 @@ 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'])
 
+    @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
 
@@ -61,10 +74,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:
@@ -124,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()

{monocle_apptrace-0.3.0b5 → monocle_apptrace-0.3.0b6}/src/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')
@@ -232,13 +232,55 @@ 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, headers, *args, **kwargs)
+
+def run_async_with_scope(method, scope_name, headers, *args, **kwargs):
+    token = None
+    if scope_name:
+        token = set_scope(scope_name)
+    elif headers:
+        token = extract_http_headers(headers)
+    try:
+        return asyncio.run(method(*args, **kwargs))
+    finally:
+        if token:
+            remove_scope(token)
+
+def async_wrapper(method, scope_name=None, headers=None, *args, **kwargs):
+    try:
+        run_loop = asyncio.get_running_loop()
+    except RuntimeError:
+        run_loop = None
+
+    if run_loop and run_loop.is_running():
+        results = []
+        thread = threading.Thread(target=lambda: results.append(run_async_with_scope(method, scope_name, headers, *args, **kwargs)))
+        thread.start()
+        thread.join()
+        return_value = results[0] if len(results) > 0 else None
+        return return_value
+    else:
+        return run_async_with_scope(method, scope_name, headers, *args, **kwargs)
 
 class Option(Generic[T]):
     def __init__(self, value: Optional[T]):
@@ -270,6 +312,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:

monocle_apptrace-0.3.0b6/src/monocle_apptrace/instrumentation/common/wrapper.py

@@ -0,0 +1,90 @@
+# pylint: disable=protected-access
+import logging
+from opentelemetry.trace import Tracer
+from opentelemetry.context import set_value, attach, detach, get_value
+
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.common.utils import (
+    get_fully_qualified_class_name,
+    with_tracer_wrapper,
+    set_scope,
+    remove_scope,
+    async_wrapper
+)
+from monocle_apptrace.instrumentation.common.constants import WORKFLOW_TYPE_KEY
+logger = logging.getLogger(__name__)
+
+def wrapper_processor(async_task: bool, tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, args, kwargs):
+    # Some Langchain objects are wrapped elsewhere, so we ignore them here
+    if instance.__class__.__name__ in ("AgentExecutor"):
+        return wrapped(*args, **kwargs)
+
+    if hasattr(instance, "name") and instance.name:
+        name = f"{to_wrap.get('span_name')}.{instance.name.lower()}"
+    elif to_wrap.get("span_name"):
+        name = to_wrap.get("span_name")
+    else:
+        name = get_fully_qualified_class_name(instance)
+
+    return_value = None
+    token = None
+    try:
+        handler.pre_tracing(to_wrap, wrapped, instance, args, kwargs)
+        skip_scan:bool = to_wrap.get('skip_span') or handler.skip_span(to_wrap, wrapped, instance, args, kwargs)
+        token = SpanHandler.attach_workflow_type(to_wrap=to_wrap)
+        if skip_scan:
+            if async_task:
+                return_value = async_wrapper(wrapped, None, None, *args, **kwargs)
+            else:
+                return_value = wrapped(*args, **kwargs)
+        else:
+            return_value = span_processor(name, async_task, tracer, handler, to_wrap, wrapped, instance, args, kwargs)
+        return return_value
+    finally:
+        handler.detach_workflow_type(token)
+        handler.post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
+
+def span_processor(name: str, async_task: bool, tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, args, kwargs):
+    # For singleton spans, eg OpenAI inference generate a workflow span to format the workflow specific attributes
+    return_value = None
+    with tracer.start_as_current_span(name) as span:
+        # Since Spanhandler can be overridden, ensure we set default monocle attributes.
+        SpanHandler.set_default_monocle_attributes(span)
+        if SpanHandler.is_root_span(span):
+            SpanHandler.set_workflow_properties(span, to_wrap)
+        if handler.is_non_workflow_root_span(span, to_wrap):
+            # This is a direct API call of a non-framework type, call the span_processor recursively for the actual span
+            return_value = span_processor(name, async_task, tracer, handler, to_wrap, wrapped, instance, args, kwargs)
+        else:
+            handler.pre_task_processing(to_wrap, wrapped, instance, args, kwargs, span)
+            if async_task:
+                return_value = async_wrapper(wrapped, None, None, *args, **kwargs)
+            else:                    
+                return_value = wrapped(*args, **kwargs)
+            handler.hydrate_span(to_wrap, wrapped, instance, args, kwargs, return_value, span)
+            handler.post_task_processing(to_wrap, wrapped, instance, args, kwargs, return_value, span)
+    return return_value
+
+@with_tracer_wrapper
+def task_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, args, kwargs):
+    return wrapper_processor(False, tracer, handler, to_wrap, wrapped, instance, args, kwargs)
+
+@with_tracer_wrapper
+async def atask_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, args, kwargs):
+    return wrapper_processor(True, tracer, handler, to_wrap, wrapped, instance, args, kwargs)
+
+@with_tracer_wrapper
+def scope_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, args, kwargs):
+    scope_name = to_wrap.get('scope_name', None)
+    if scope_name:
+        token = set_scope(scope_name)
+    return_value = wrapped(*args, **kwargs)
+    if scope_name:
+        remove_scope(token)
+    return return_value
+
+@with_tracer_wrapper
+async def ascope_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, args, kwargs):
+    scope_name = to_wrap.get('scope_name', None)
+    return_value = async_wrapper(wrapped, scope_name, None, *args, **kwargs)
+    return return_value

{monocle_apptrace-0.3.0b5 → monocle_apptrace-0.3.0b6}/src/monocle_apptrace/instrumentation/common/wrapper_method.py

@@ -1,7 +1,7 @@
 # pylint: disable=too-few-public-methods
 from typing import Any, Dict
 from monocle_apptrace.instrumentation.common.wrapper import task_wrapper, scope_wrapper
-from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler, NonFrameworkSpanHandler
 from monocle_apptrace.instrumentation.metamodel.botocore.methods import BOTOCORE_METHODS
 from monocle_apptrace.instrumentation.metamodel.botocore.handlers.botocore_span_handler import BotoCoreSpanHandler
 from monocle_apptrace.instrumentation.metamodel.langchain.methods import (
@@ -26,13 +26,15 @@ class WrapperMethod:
             output_processor : str = None,
             wrapper_method = task_wrapper,
             span_handler = 'default',
-            scope_name: str = None
+            scope_name: str = None,
+            span_type: str = None
             ):
         self.package = package
         self.object = object_name
         self.method = method
         self.span_name = span_name
         self.output_processor=output_processor
+        self.span_type = span_type
 
         self.span_handler:SpanHandler.__class__ = span_handler
         self.scope_name = scope_name
@@ -51,7 +53,8 @@ class WrapperMethod:
             'output_processor': self.output_processor,
             'wrapper_method': self.wrapper_method,
             'span_handler': self.span_handler,
-            'scope_name': self.scope_name
+            'scope_name': self.scope_name,
+            'span_type': self.span_type
         }
         return instance_dict
 
@@ -64,5 +67,6 @@ MONOCLE_SPAN_HANDLERS: Dict[str, SpanHandler] = {
     "default": SpanHandler(),
     "botocore_handler": BotoCoreSpanHandler(),
     "flask_handler": FlaskSpanHandler(),
-    "request_handler": RequestSpanHandler()
+    "request_handler": RequestSpanHandler(),
+    "non_framework_handler": NonFrameworkSpanHandler()
 }

{monocle_apptrace-0.3.0b5 → monocle_apptrace-0.3.0b6}/src/monocle_apptrace/instrumentation/metamodel/haystack/entities/inference.py

@@ -9,7 +9,7 @@ INFERENCE = {
             {
                 "_comment": "provider type ,name , deployment , inference_endpoint",
                 "attribute": "type",
-                "accessor": lambda arguments: 'inference.azure_oai'
+                "accessor": lambda arguments: 'inference.azure_openai'
             },
             {
                 "attribute": "provider_name",