monocle-apptrace 0.3.1b1__py3-none-any.whl → 0.4.0b1__py3-none-any.whl

monocle_apptrace/instrumentation/common/tracing.md

@@ -0,0 +1,68 @@
+# Monocle tracing: concepts and principles
+
+## Span
+Span is an observation of a code/method executed. Each span has a unique ID. It records the start time and end time of the code's execution along with additional information relevant to that operation. Before the code execution starts, a span object is created in the memory of the host process executing this code. It'll capture the current time as start of time of span. At this stage the span is considered active. It'll stay active till the code execution ends. Once the code execution is complete, it'll record the current time as end time, capture any additional relevant information (eg argument, return value, environment setttings etc.). Now the span is marked as closed and it will be queued to be saved to some configured storage. 
+Note that the code that generated this span could in turn call other methods that are also instrumented. Those will generate spans of their own. These will be "child" spans which will refer to the span ID of the calling code as "parent" span. An initial span which has no parent is referred as "root" span.
+
+## Trace
+A trace is a collection of spans with a common ID called traceID. When the first active span gets created, a new unique traceID is generated and assigned to that span. All the child spans generated by execution of other instrumented code/methods will share the same traceID. Once this top span ends, this trace ends. This way all the code executed as part of the top level instrumented code will have a common traceID to group them together. For example, consider following sequence where `f1()` is the first instrumented method is executed, it calls other instrumented methods `f2(),f3(),f4() and f5()`
+```
+f1()--> f2() --> f3()
+    --> f4() --> f5()
+```
+In the above sequence, all method execution will generate a span each and they all will have a common traceID. Now if a new instrumented methods is executed after f1() finishes, it will be the first active span in the process's execution context and a will get a new traceID.
+
+### Trace ID propogation
+Each child span inherits the parent's trace ID. When spans running in same process, it captures it from process memory/context etc. But consider the above example again, where the `f4()-->f5()` code is not part of the process that executing f1(). It's a remote call, say over REST. From the overall application's point of view, the work done if `f4()` and `f5()` is part of `f1()` and you want same traceID associated with all spans. You want the instrumentation to seamlessly pass the tracedID over such remote calls and continue that instead of generating a new one. It's the responsibility of Monocle to provide such mechanism to make thsi trace ID propogation transparent to the application logic and architecture.
+
+## Propogation
+When the execution logic spans mulitple processes using remote calling mechanisms like REST, you want the trace ID also to propogate from process that originated it to the one that's continueing the remote execution. Monocle supports seamlessly propogating traceID over REST if both the sides for the trace execution are instrumented.
+
+## Types of spans in Monocle
+Monocle extends these generic span types by enriching additional attributes/data for genAI specific operations.
+### GenAI spans
+There are the core spans that capture details of genAI component operations like call to an LLM or vectore store. The purpose of these spans is to capture the details the applications interaction with core genAI comoponents. These spans are triggered by pre-instrumented methods that handle such operations. 
+- Inference span
+Represents interaction with LLMs, captures details like model, prompts, response and other metadata (eg tokens)
+- Retrieval span
+Represents interactions with vector stores like embedding creating, vector retrieval etc. Captures the model, search query, response, vector embedding etc.
+
+### anchor spans
+These are the spans that are created by a top level method that anchors a higher level of abstraction for underlying core genAI APIs. For example a langchain.invoke() which under the cover calls langchain.llm_invoke() or langchain.vector_retrieval(). Consider following psuedo code of a langchain rag pattern API,
+```
+response = rag_chain.invoke(prompt)
+            --> cleaned_prompt = llm1.chat(prompt)
+            --> context = vector_store.retrieve(cleaned_prompt)
+            --> response = llm2.chat(system_prompt+context+cleaned_prompt)
+            --> return response
+```
+If we only instrument the top level invoke call, then we'll trace the top level prompt and response interaction between application and langchain. But we'll miss the details like how a system prompt was added and send to mulitple LLMs and what context was extracted from a vector store etc. On the other hand, if we only instrument the low level calls to LLM and vector, then we'll miss the fact that those are part of same RAG. Hence we instrument all of them. This exaple would genearte an anchor spna for `invoke()` method, a retrieval span for `retrieve()` method and two inference spans for each `chat()` method. All of these will have common traceID.
+The anchor spans also provides an observation window of your application interaction with an high level SDK or service. It will illustrate facts such as how much time take by the genAI service invocation compared to other local logic.
+
+### Workflow spans
+Workflow spans are synthetic spans that are created to trace the full trace. It captures the summary of the full trace including the time window, the process running this code (set as `workflow_name` in the API to enab le Monocle instrumentation) and runtime environment details such as hosting service (Azure function, Lambda function etc). 
+The workflow spans is generated when a new trace starts or when a trace is propogated. They provide the base line observation window for the entire trace or a fragment of trace executed in a process.
+Consider following example,
+```
+setup_monocle_telemetry(workflow='bot')
+rag_chain.invoke()
+                --> context = retrieval() 
+                --> new_prompt = REST --> azure.func.chat(prompt) --> 
+                                                                setup_monocle_telemetry(workflow='moderator')
+                                                                return llm(moderator_system_prompt+prompt)
+                --> response = llm(new_prompt)
+```
+This will generate following spans:
+```
+Span{name='workflow.bot', type= workflow, traceID = xx1, spanID = yy0, parentID=None} ==> Workflow for new trace start
+Span{name='chain.invoke', type=anchor, traceID = xx1, spanID = yy1, parentID=yy0} ==> anchor span for chain invoke
+Span{name='chain.retrieval', type=retrieval, traceID = xx1, spanID = yy2, parentID = yy1} ==> Retrieval API span
+Span{name='workflow.moderator', type=workflow, traceID = xx1, spanID = zz1, parentID=yy1} ==> Workflow for propogated trace fragement
+Span{name='az.func.chat', type=anchor, traceID = xx1, spanID = zz2, parentID=zz1} ==> anchor span for az function invoke
+Span{name='chain.infer', type=inference, traceID = xx1, spanID = zz2, parentID=zz2} ==> inference
+Span{name='chain.infer',type=inference, traceID = xx1, spanID = yy3, parentID=yy1} ==> inference
+```
+
+## Scopes
+Scope is an way of grouping across traces. It's a tag with a value that can either be specified or auto generated (GUID) by Monocle. There can be any number of scopes active in an application code at a given point in time. All the active scopes are recorded in every span that's emmitted.
+

monocle_apptrace/instrumentation/common/utils.py

@@ -1,26 +1,32 @@
 import logging, json
 import os
+import traceback
 from typing import Callable, Generic, Optional, TypeVar, Mapping
-import threading, asyncio
 
 from opentelemetry.context import attach, detach, get_current, get_value, set_value, Context
-from opentelemetry.trace import NonRecordingSpan, Span, get_tracer
+from opentelemetry.trace import NonRecordingSpan, Span
 from opentelemetry.trace.propagation import _SPAN_KEY
 from opentelemetry.sdk.trace import id_generator, TracerProvider
-from opentelemetry.propagate import inject, extract
+from opentelemetry.propagate import extract
 from opentelemetry import baggage
-from monocle_apptrace.instrumentation.common.constants import MONOCLE_SCOPE_NAME_PREFIX, SCOPE_METHOD_FILE, SCOPE_CONFIG_PATH, llm_type_map
+from monocle_apptrace.instrumentation.common.constants import MONOCLE_SCOPE_NAME_PREFIX, SCOPE_METHOD_FILE, SCOPE_CONFIG_PATH, llm_type_map, MONOCLE_SDK_VERSION, ADD_NEW_WORKFLOW
+from importlib.metadata import version
 
 T = TypeVar('T')
 U = TypeVar('U')
 
 logger = logging.getLogger(__name__)
 
-monocle_tracer_provider: TracerProvider = None
 embedding_model_context = {}
 scope_id_generator = id_generator.RandomIdGenerator()
 http_scopes:dict[str:str] = {}
 
+try:
+    monocle_sdk_version = version("monocle_apptrace")
+except Exception as e:
+    monocle_sdk_version = "unknown"
+    logger.warning("Exception finding monocle-apptrace version.")
+
 class MonocleSpanException(Exception):
     def __init__(self, err_message:str):
         """
@@ -36,14 +42,6 @@ class MonocleSpanException(Exception):
         """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
-
-def get_tracer_provider() -> TracerProvider:
-    global monocle_tracer_provider
-    return monocle_tracer_provider
-
 def set_span_attribute(span, name, value):
     if value is not None:
         if value != "":
@@ -86,7 +84,12 @@ def with_tracer_wrapper(func):
             except Exception as e:
                 logger.error("Exception in attaching parent context: %s", e)
 
-            val = func(tracer, handler, to_wrap, wrapped, instance, args, kwargs)
+            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
 
         return wrapper
@@ -96,8 +99,8 @@ def with_tracer_wrapper(func):
 def resolve_from_alias(my_map, alias):
     """Find a alias that is not none from list of aliases"""
 
-    for i in alias and my_map[i] is not None:
-        if i in my_map.keys():
+    for i in alias:
+        if i in my_map.keys() and my_map[i] is not None:
             return my_map[i]
     return None
 
@@ -236,6 +239,7 @@ def set_scopes_from_baggage(baggage_context:Context):
 def extract_http_headers(headers) -> object:
     global http_scopes
     trace_context:Context = extract(headers, context=get_current())
+    trace_context = set_value(ADD_NEW_WORKFLOW, True, trace_context)
     imported_scope:dict[str, object] = {}
     for http_header, http_scope in http_scopes.items():
         if http_header in headers:
@@ -267,49 +271,57 @@ async def http_async_route_handler(func, *args, **kwargs):
         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
+        if headers is not None:
+            token = extract_http_headers(headers)
+        return await func(*args, **kwargs)
     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
+        if token is not None:
+            clear_http_scopes(token)
 
-    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)
+# 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)
+
+# async def async_wrapper(method, headers=None, *args, **kwargs):
+#     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)
+
+def get_monocle_version() -> str:
+    global monocle_sdk_version
+    return monocle_sdk_version
+
+def add_monocle_trace_state(headers:dict[str:str]) -> None:
+    if headers is None:
+        return
+    monocle_trace_state = f"{MONOCLE_SDK_VERSION}={get_monocle_version()}"
+    if 'tracestate' in headers:
+        headers['tracestate'] = f"{headers['tracestate']},{monocle_trace_state}"
+    else:
+        headers['tracestate'] = monocle_trace_state
 
 class Option(Generic[T]):
     def __init__(self, value: Optional[T]):
@@ -343,14 +355,25 @@ def try_option(func: Callable[..., T], *args, **kwargs) -> Option[T]:
 
 def get_llm_type(instance):
     try:
-        llm_type = llm_type_map.get(type(instance).__name__.lower())
+        t_name = type(instance).__name__.lower()
+        t_name = t_name.replace("async", "") if "async" in t_name else t_name
+        llm_type = llm_type_map.get(t_name)
         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:
-        if i in my_map.keys():
-            return my_map[i]
-    return None
+def patch_instance_method(obj, method_name, func):
+    """
+    Patch a special method (like __iter__) for a single instance.
+
+    Args:
+        obj: the instance to patch
+        method_name: the name of the method (e.g., '__iter__')
+        func: the new function, expecting (self, ...)
+    """
+    cls = obj.__class__
+    # Dynamically create a new class that inherits from obj's class
+    new_cls = type(f"Patched{cls.__name__}", (cls,), {
+        method_name: func
+    })
+    obj.__class__ = new_cls

monocle_apptrace/instrumentation/common/wrapper.py

@@ -6,76 +6,168 @@ 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,
+    set_scopes,
     with_tracer_wrapper,
     set_scope,
-    remove_scope,
-    async_wrapper
+    remove_scope
 )
-from monocle_apptrace.instrumentation.common.constants import WORKFLOW_TYPE_KEY
+from monocle_apptrace.instrumentation.common.constants import WORKFLOW_TYPE_KEY, ADD_NEW_WORKFLOW
 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)
+def get_auto_close_span(to_wrap, kwargs):
+    try:
+        if to_wrap.get("output_processor") and to_wrap.get("output_processor").get("is_auto_close"):
+            return to_wrap.get("output_processor").get("is_auto_close")(kwargs)
+        return True
+    except Exception as e:
+        logger.warning("Warning: Error occurred in get_auto_close_span: %s", str(e))
+        return True
+
+def pre_process_span(name, tracer, handler, add_workflow_span, to_wrap, wrapped, instance, args, kwargs, span, source_path):
+    SpanHandler.set_default_monocle_attributes(span, source_path)
+    if SpanHandler.is_root_span(span) or add_workflow_span:
+        # This is a direct API call of a non-framework type
+        SpanHandler.set_workflow_properties(span, to_wrap)
+    else:
+        SpanHandler.set_non_workflow_properties(span)
+        handler.pre_task_processing(to_wrap, wrapped, instance, args, kwargs, span)
 
-    if hasattr(instance, "name") and instance.name:
-        name = f"{to_wrap.get('span_name')}.{instance.name.lower()}"
-    elif to_wrap.get("span_name"):
+def post_process_span(handler, to_wrap, wrapped, instance, args, kwargs, return_value, span):
+    if not (SpanHandler.is_root_span(span) or get_value(ADD_NEW_WORKFLOW) == True):
+        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)
+
+def get_span_name(to_wrap, instance):
+    if to_wrap.get("span_name"):
         name = to_wrap.get("span_name")
     else:
         name = get_fully_qualified_class_name(instance)
+    return name
+
+def monocle_wrapper_span_processor(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, source_path, add_workflow_span, args, kwargs):
+    # Main span processing logic
+    name = get_span_name(to_wrap, instance)
+    return_value = None
+    span_status = None
+    if(get_auto_close_span(to_wrap, kwargs)):
+        with tracer.start_as_current_span(name) as span:
+            pre_process_span(name, tracer, handler, add_workflow_span, to_wrap, wrapped, instance, args, kwargs, span, source_path)
+            
+            if SpanHandler.is_root_span(span) or add_workflow_span:
+                # Recursive call for the actual span
+                return_value, span_status = monocle_wrapper_span_processor(tracer, handler, to_wrap, wrapped, instance, source_path, False, args, kwargs)
+                span.set_status(span_status)
+            else:
+                with SpanHandler.workflow_type(to_wrap):
+                    return_value = wrapped(*args, **kwargs)
+                post_process_span(handler, to_wrap, wrapped, instance, args, kwargs, return_value, span)
+                span_status = span.status
+    else:
+        span = tracer.start_span(name)
+        
+        pre_process_span(name, tracer, handler, add_workflow_span, to_wrap, wrapped, instance, args, kwargs, span, source_path)
+        
+        def post_process_span_internal(ret_val):
+            nonlocal handler, to_wrap, wrapped, instance, args, kwargs, span
+            post_process_span(handler, to_wrap, wrapped, instance, args, kwargs, ret_val, span)
+            span.end()
+                    
+        with SpanHandler.workflow_type(to_wrap):
+            return_value = wrapped(*args, **kwargs)
+        if to_wrap.get("output_processor") and to_wrap.get("output_processor").get("response_processor"):
+            # Process the stream
+            to_wrap.get("output_processor").get("response_processor")(to_wrap, return_value, post_process_span_internal)
+        else: 
+            span.end()
+        span_status = span.status
+    return return_value, span_status
 
+def monocle_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, source_path, args, kwargs):
     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)
-        if not to_wrap.get('skip_span'):
-            token = SpanHandler.attach_workflow_type(to_wrap=to_wrap)
-        if skip_scan:
-            if async_task:
-                return_value = async_wrapper(wrapped, None, None, None, *args, **kwargs)
-            else:
-                return_value = wrapped(*args, **kwargs)
+        if to_wrap.get('skip_span', False) or handler.skip_span(to_wrap, wrapped, instance, args, kwargs):
+            return_value = wrapped(*args, **kwargs)
         else:
-            return_value = span_processor(name, async_task, tracer, handler, to_wrap, wrapped, instance, args, kwargs)
+            add_workflow_span = get_value(ADD_NEW_WORKFLOW) == True
+            token = attach(set_value(ADD_NEW_WORKFLOW, False))
+            try:
+                return_value, span_status = monocle_wrapper_span_processor(tracer, handler, to_wrap, wrapped, instance, source_path, add_workflow_span, args, kwargs)
+            finally:
+                detach(token)
         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
+
+async def amonocle_wrapper_span_processor(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, source_path, add_workflow_span, args, kwargs):
+    # Main span processing logic
+    name = get_span_name(to_wrap, instance)
     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)
+    span_status = None
+    if(get_auto_close_span(to_wrap, kwargs)):
+        with tracer.start_as_current_span(name) as span:
+            pre_process_span(name, tracer, handler, add_workflow_span, to_wrap, wrapped, instance, args, kwargs, span, source_path)
+            
+            if SpanHandler.is_root_span(span) or add_workflow_span:
+                # Recursive call for the actual span
+                return_value, span_status = await amonocle_wrapper_span_processor(tracer, handler, to_wrap, wrapped, instance, source_path, False, args, kwargs)
+                span.set_status(span_status)
+            else:
+                with SpanHandler.workflow_type(to_wrap):
+                    return_value = await wrapped(*args, **kwargs)
+                span_status = span.status
+            post_process_span(handler, to_wrap, wrapped, instance, args, kwargs, return_value, span)
+    else:
+        span = tracer.start_span(name)
+        
+        pre_process_span(name, tracer, handler, add_workflow_span, to_wrap, wrapped, instance, args, kwargs, span, source_path)
+        
+        def post_process_span_internal(ret_val):
+            nonlocal handler, to_wrap, wrapped, instance, args, kwargs, span
+            post_process_span(handler, to_wrap, wrapped, instance, args, kwargs, ret_val, span)
+            span.end()
+
+        with SpanHandler.workflow_type(to_wrap):
+            return_value = await wrapped(*args, **kwargs)
+           
+        if to_wrap.get("output_processor") and to_wrap.get("output_processor").get("response_processor"):
+            # Process the stream
+            to_wrap.get("output_processor").get("response_processor")(to_wrap, return_value, post_process_span_internal)
+        else: 
+            span.end()
+        span_status = span.status
+    return return_value, span.status
+
+async def amonocle_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, source_path, args, kwargs):
+    return_value = None
+    token = None
+    try:
+        handler.pre_tracing(to_wrap, wrapped, instance, args, kwargs)
+        if to_wrap.get('skip_span', False) or handler.skip_span(to_wrap, wrapped, instance, args, kwargs):
+            return_value = await wrapped(*args, **kwargs)
         else:
-            handler.pre_task_processing(to_wrap, wrapped, instance, args, kwargs, span)
-            if async_task:
-                return_value = async_wrapper(wrapped, None, 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
+            add_workflow_span = get_value(ADD_NEW_WORKFLOW) == True
+            token = attach(set_value(ADD_NEW_WORKFLOW, False))
+            try:
+                return_value, span_status = await amonocle_wrapper_span_processor(tracer, handler, to_wrap, wrapped, instance, source_path, add_workflow_span, args, kwargs)    
+            finally:
+                detach(token)
+        return return_value
+    finally:
+        handler.post_tracing(to_wrap, wrapped, instance, args, kwargs, 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)
+def task_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, source_path, args, kwargs):
+    return monocle_wrapper(tracer, handler, to_wrap, wrapped, instance, source_path, 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)
+async def atask_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, source_path, args, kwargs):
+    return await amonocle_wrapper(tracer, handler, to_wrap, wrapped, instance, source_path, args, kwargs)
 
 @with_tracer_wrapper
-def scope_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, args, kwargs):
+def scope_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, source_path, args, kwargs):
     scope_name = to_wrap.get('scope_name', None)
     if scope_name:
         token = set_scope(scope_name)
@@ -85,8 +177,55 @@ def scope_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instan
     return return_value
 
 @with_tracer_wrapper
-async def ascope_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, args, kwargs):
+async def ascope_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, source_path, args, kwargs):
     scope_name = to_wrap.get('scope_name', None)
     scope_value = to_wrap.get('scope_value', None)
-    return_value = async_wrapper(wrapped, scope_name, scope_value, None, *args, **kwargs)
-    return return_value
+    token = None
+    try:
+        if scope_name:
+            token = set_scope(scope_name, scope_value)
+        return_value = await wrapped(*args, **kwargs)
+        return return_value
+    finally:
+        if token:
+            remove_scope(token)
+            
+@with_tracer_wrapper
+def scopes_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, source_path, args, kwargs):
+    scope_values = to_wrap.get('scope_values', None)
+    scope_values = evaluate_scope_values(args, kwargs, scope_values)
+    token = None
+    try:
+        if scope_values:
+            token = set_scopes(scope_values)
+        return_value = wrapped(*args, **kwargs)
+        return return_value
+    finally:
+        if token:
+            remove_scope(token)
+
+@with_tracer_wrapper
+async def ascopes_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, source_path, args, kwargs):
+    scope_values = to_wrap.get('scope_values', None)
+    scope_values = evaluate_scope_values(args, kwargs, scope_values)
+    token = None
+    try:
+        if scope_values:
+            token = set_scopes(scope_values)
+        return_value = await wrapped(*args, **kwargs)
+        return return_value
+    finally:
+        if token:
+            remove_scope(token)
+
+def evaluate_scope_values(args, kwargs, scope_values):
+    if callable(scope_values):
+        try:
+            scope_values = scope_values(args, kwargs)
+        except Exception as e:
+            logger.warning("Warning: Error occurred in evaluate_scope_values: %s", str(e))
+            scope_values = None
+    if isinstance(scope_values, dict):
+        return scope_values
+    return None
+    

monocle_apptrace/instrumentation/common/wrapper_method.py

@@ -12,12 +12,13 @@ from monocle_apptrace.instrumentation.metamodel.haystack.methods import (HAYSTAC
 from monocle_apptrace.instrumentation.metamodel.openai.methods import (OPENAI_METHODS,)
 from monocle_apptrace.instrumentation.metamodel.langgraph.methods import LANGGRAPH_METHODS
 from monocle_apptrace.instrumentation.metamodel.flask.methods import (FLASK_METHODS, )
-from monocle_apptrace.instrumentation.metamodel.flask._helper import FlaskSpanHandler
+from monocle_apptrace.instrumentation.metamodel.flask._helper import FlaskSpanHandler, FlaskResponseSpanHandler
 from monocle_apptrace.instrumentation.metamodel.requests.methods import (REQUESTS_METHODS, )
 from monocle_apptrace.instrumentation.metamodel.requests._helper import RequestSpanHandler
 from monocle_apptrace.instrumentation.metamodel.teamsai.methods import (TEAMAI_METHODS, )
 from monocle_apptrace.instrumentation.metamodel.anthropic.methods import (ANTHROPIC_METHODS, )
-
+from monocle_apptrace.instrumentation.metamodel.aiohttp.methods import (AIOHTTP_METHODS, )
+from monocle_apptrace.instrumentation.metamodel.aiohttp._helper import aiohttpSpanHandler
 class WrapperMethod:
     def __init__(
             self,
@@ -29,7 +30,8 @@ class WrapperMethod:
             wrapper_method = task_wrapper,
             span_handler = 'default',
             scope_name: str = None,
-            span_type: str = None
+            span_type: str = None,
+            scope_values = None,
             ):
         self.package = package
         self.object = object_name
@@ -37,10 +39,11 @@ class WrapperMethod:
         self.span_name = span_name
         self.output_processor=output_processor
         self.span_type = span_type
+        self.scope_values = scope_values
 
         self.span_handler:SpanHandler.__class__ = span_handler
         self.scope_name = scope_name
-        if scope_name:
+        if scope_name and not scope_values:
             self.wrapper_method = scope_wrapper
         else:
             self.wrapper_method = wrapper_method
@@ -56,19 +59,22 @@ class WrapperMethod:
             'wrapper_method': self.wrapper_method,
             'span_handler': self.span_handler,
             'scope_name': self.scope_name,
-            'span_type': self.span_type
+            'span_type': self.span_type,
+            'scope_values': self.scope_values,
         }
         return instance_dict
 
     def get_span_handler(self) -> SpanHandler:
         return self.span_handler()
 
-DEFAULT_METHODS_LIST = LANGCHAIN_METHODS + LLAMAINDEX_METHODS + HAYSTACK_METHODS + BOTOCORE_METHODS + FLASK_METHODS + REQUESTS_METHODS + LANGGRAPH_METHODS + OPENAI_METHODS + TEAMAI_METHODS + ANTHROPIC_METHODS
+DEFAULT_METHODS_LIST = LANGCHAIN_METHODS + LLAMAINDEX_METHODS + HAYSTACK_METHODS + BOTOCORE_METHODS + FLASK_METHODS + REQUESTS_METHODS + LANGGRAPH_METHODS + OPENAI_METHODS + TEAMAI_METHODS + ANTHROPIC_METHODS + AIOHTTP_METHODS
 
 MONOCLE_SPAN_HANDLERS: Dict[str, SpanHandler] = {
     "default": SpanHandler(),
+    "aiohttp_handler": aiohttpSpanHandler(),
     "botocore_handler": BotoCoreSpanHandler(),
     "flask_handler": FlaskSpanHandler(),
+    "flask_response_handler": FlaskResponseSpanHandler(),
     "request_handler": RequestSpanHandler(),
     "non_framework_handler": NonFrameworkSpanHandler()
 }