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

monocle_apptrace/__main__.py

@@ -0,0 +1,19 @@
+import sys, os
+import runpy
+from monocle_apptrace.instrumentation.common.instrumentor import setup_monocle_telemetry
+
+def main():
+    if len(sys.argv) < 2 or not sys.argv[1].endswith(".py"):
+        print("Usage: python -m monocle_apptrace <your-main-module-file> <args>")
+        sys.exit(1)
+    file_name = os.path.basename(sys.argv[1])
+    workflow_name = file_name[:-3]
+    setup_monocle_telemetry(workflow_name=workflow_name)
+    sys.argv.pop(0)
+
+    try:
+        runpy.run_path(path_name=sys.argv[0], run_name="__main__")
+    except Exception as e:
+        print(e)
+if __name__ == "__main__":
+    main()

monocle_apptrace/exporters/monocle_exporters.py

@@ -1,9 +1,10 @@
 from typing import Dict, Any, List
 import os
-import warnings
+import logging
 from importlib import import_module
 from opentelemetry.sdk.trace.export import SpanExporter, ConsoleSpanExporter
 from monocle_apptrace.exporters.file_exporter import FileSpanExporter
+logger = logging.getLogger(__name__)
 
 monocle_exporters: Dict[str, Any] = {
     "s3": {"module": "monocle_apptrace.exporters.aws.s3_exporter", "class": "S3SpanExporter"},
@@ -25,21 +26,21 @@ def get_monocle_exporter() -> List[SpanExporter]:
         try:
             exporter_class_path = monocle_exporters[exporter_name]
         except KeyError:
-            warnings.warn(f"Unsupported Monocle span exporter '{exporter_name}', skipping.")
+            logger.debug(f"Unsupported Monocle span exporter '{exporter_name}', skipping.")
             continue
         try:
             exporter_module = import_module(exporter_class_path["module"])
             exporter_class = getattr(exporter_module, exporter_class_path["class"])
             exporters.append(exporter_class())
         except Exception as ex:
-            warnings.warn(
+            logger.debug(
                 f"Unable to initialize Monocle span exporter '{exporter_name}', error: {ex}. Using ConsoleSpanExporter as a fallback.")
             exporters.append(ConsoleSpanExporter())
             continue
 
     # If no exporters were created, default to FileSpanExporter
     if not exporters:
-        warnings.warn("No valid Monocle span exporters configured. Defaulting to FileSpanExporter.")
+        logger.debug("No valid Monocle span exporters configured. Defaulting to FileSpanExporter.")
         exporters.append(FileSpanExporter())
 
     return exporters

monocle_apptrace/instrumentation/common/constants.py

@@ -37,6 +37,7 @@ service_name_map = {
     GITHUB_CODESPACE_SERVICE_NAME: GITHUB_CODESPACE_IDENTIFIER_ENV_NAME
 }
 
+MONOCLE_INSTRUMENTOR = "monocle_apptrace"
 WORKFLOW_TYPE_KEY = "workflow_type"
 DATA_INPUT_KEY = "data.input"
 DATA_OUTPUT_KEY = "data.output"
@@ -47,3 +48,7 @@ RESPONSE = "response"
 SESSION_PROPERTIES_KEY = "session"
 INFRA_SERVICE_KEY = "infra_service_name"
 META_DATA = 'metadata'
+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'

monocle_apptrace/instrumentation/common/instrumentor.py

@@ -3,6 +3,7 @@ from typing import Collection, Dict, List, Union
 import random
 import uuid
 from opentelemetry import trace
+from contextlib import contextmanager
 from opentelemetry.context import attach, get_value, set_value, get_current, detach
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import unwrap
@@ -19,8 +20,14 @@ from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
 from monocle_apptrace.instrumentation.common.wrapper_method import (
     DEFAULT_METHODS_LIST,
     WrapperMethod,
+    MONOCLE_SPAN_HANDLERS
 )
-
+from monocle_apptrace.instrumentation.common.wrapper import scope_wrapper
+from monocle_apptrace.instrumentation.common.utils import (
+    set_scope, remove_scope, http_route_handler, load_scopes
+)
+from monocle_apptrace.instrumentation.common.constants import MONOCLE_INSTRUMENTOR
+from functools import wraps
 logger = logging.getLogger(__name__)
 
 SESSION_PROPERTIES_KEY = "session"
@@ -29,8 +36,6 @@ _instruments = ()
 
 monocle_tracer_provider: TracerProvider = None
 
-MONOCLE_INSTRUMENTOR = "monocle_apptrace"
-
 class MonocleInstrumentor(BaseInstrumentor):
     workflow_name: str = ""
     user_wrapper_methods: list[Union[dict,WrapperMethod]] = []
@@ -45,17 +50,36 @@ class MonocleInstrumentor(BaseInstrumentor):
             union_with_default_methods: bool = True
             ) -> None:
         self.user_wrapper_methods = user_wrapper_methods or []
-        self.handlers = handlers or {'default':SpanHandler()}
+        self.handlers = handlers
+        if self.handlers is not None:
+            for key, val in MONOCLE_SPAN_HANDLERS.items():
+                if key not in self.handlers:
+                    self.handlers[key] = val
+        else:
+            self.handlers = MONOCLE_SPAN_HANDLERS
         self.union_with_default_methods = union_with_default_methods
         super().__init__()
 
+    def get_instrumentor(self, tracer):
+        def instrumented_endpoint_invoke(to_wrap,wrapped, span_name, instance,fn):
+            @wraps(fn)
+            def with_instrumentation(*args, **kwargs):
+                handler = SpanHandler()
+                with tracer.start_as_current_span(span_name) as span:
+                    response = fn(*args, **kwargs)
+                    handler.hydrate_span(to_wrap, span=span, wrapped=wrapped, instance=instance, args=args, kwargs=kwargs,
+                                         result=response)
+                    return response
+
+            return with_instrumentation
+        return instrumented_endpoint_invoke
+
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
     def _instrument(self, **kwargs):
         tracer_provider: TracerProvider = kwargs.get("tracer_provider")
-        global monocle_tracer_provider
-        monocle_tracer_provider = tracer_provider
+        set_tracer_provider(tracer_provider)
         tracer = get_tracer(instrumenting_module_name=MONOCLE_INSTRUMENTOR, tracer_provider=tracer_provider)
 
         final_method_list = []
@@ -67,6 +91,10 @@ class MonocleInstrumentor(BaseInstrumentor):
                 final_method_list.append(method)
             elif isinstance(method, WrapperMethod):
                 final_method_list.append(method.to_dict())
+
+        for method in load_scopes():
+            method['wrapper_method'] = scope_wrapper
+            final_method_list.append(method)
         
         for method_config in final_method_list:
             target_package = method_config.get("package", None)
@@ -77,6 +105,10 @@ class MonocleInstrumentor(BaseInstrumentor):
             handler_key = method_config.get("span_handler",'default')
             try:
                 handler =  self.handlers.get(handler_key)
+                if not handler:
+                    logger.warning("incorrect or empty handler falling back to default handler")
+                    handler = self.handlers.get('default')
+                handler.set_instrumentor(self.get_instrumentor(tracer))
                 wrap_function_wrapper(
                     target_package,
                     f"{target_object}.{target_method}" if target_object else target_method,
@@ -108,6 +140,14 @@ class MonocleInstrumentor(BaseInstrumentor):
                              object:{wrap_object},
                              method:{wrap_method}""")
 
+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 setup_monocle_telemetry(
         workflow_name: str,
         span_processors: List[SpanProcessor] = None,
@@ -119,7 +159,7 @@ def setup_monocle_telemetry(
     })
     exporters = get_monocle_exporter()
     span_processors = span_processors or [BatchSpanProcessor(exporter) for exporter in exporters]
-    trace_provider = TracerProvider(resource=resource)
+    set_tracer_provider(TracerProvider(resource=resource))
     attach(set_value("workflow_name", workflow_name))
     tracer_provider_default = trace.get_tracer_provider()
     provider_type = type(tracer_provider_default).__name__
@@ -129,14 +169,14 @@ def setup_monocle_telemetry(
         if not is_proxy_provider:
             tracer_provider_default.add_span_processor(processor)
         else:
-            trace_provider.add_span_processor(processor)
+            get_tracer_provider().add_span_processor(processor)
     if is_proxy_provider:
-        trace.set_tracer_provider(trace_provider)
+        trace.set_tracer_provider(get_tracer_provider())
     instrumentor = MonocleInstrumentor(user_wrapper_methods=wrapper_methods or [], 
                                        handlers=span_handlers, union_with_default_methods = union_with_default_methods)
     # instrumentor.app_name = workflow_name
     if not instrumentor.is_instrumented_by_opentelemetry:
-        instrumentor.instrument(trace_provider=trace_provider)
+        instrumentor.instrument(trace_provider=get_tracer_provider())
 
     return instrumentor
 
@@ -156,7 +196,7 @@ def propagate_trace_id(traceId = "", use_trace_context = False):
     try:
         if traceId.startswith("0x"):
             traceId = traceId.lstrip("0x")
-        tracer = get_tracer(instrumenting_module_name= MONOCLE_INSTRUMENTOR, tracer_provider= monocle_tracer_provider)
+        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):
@@ -194,6 +234,37 @@ def is_valid_trace_id_uuid(traceId: str) -> bool:
         pass
     return False
 
+def start_scope(scope_name: str, scope_value:str = None) -> object:
+    return set_scope(scope_name, scope_value)
+
+def stop_scope(token:object) -> None:
+    remove_scope(token)
+    return
+
+@contextmanager
+def monocle_trace_scope(scope_name: str, scope_value:str = None):
+    token = start_scope(scope_name, scope_value)
+    try:
+        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)
+                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
 
 class FixedIdGenerator(id_generator.IdGenerator):
     def __init__(

monocle_apptrace/instrumentation/common/span_handler.py

@@ -10,7 +10,7 @@ from monocle_apptrace.instrumentation.common.constants import (
     service_name_map,
     service_type_map,
 )
-from monocle_apptrace.instrumentation.common.utils import set_attribute
+from monocle_apptrace.instrumentation.common.utils import set_attribute, get_scopes
 
 logger = logging.getLogger(__name__)
 
@@ -24,10 +24,25 @@ WORKFLOW_TYPE_MAP = {
 
 class SpanHandler:
 
+    def __init__(self,instrumentor=None):
+        self.instrumentor=instrumentor
+
+    def set_instrumentor(self, instrumentor):
+        self.instrumentor = instrumentor
+
     def validate(self, to_wrap, wrapped, instance, args, kwargs):
         pass
 
-    def pre_task_processing(self, to_wrap, wrapped, instance, args, span):
+    def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
+        pass
+
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+        pass
+
+    def skip_span(self, to_wrap, wrapped, instance, args, kwargs) -> bool:
+        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")
@@ -37,14 +52,9 @@ class SpanHandler:
         if "pipeline" in to_wrap['package']:
             set_attribute(QUERY, args[0]['prompt_builder']['question'])
 
-
-
     def post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, span):
         pass
 
-    def set_context_properties(self, to_wrap, wrapped, instance, args, kwargs):
-        pass
-
     def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span):
         self.hydrate_attributes(to_wrap, wrapped, instance, args, kwargs, result, span)
         self.hydrate_events(to_wrap, wrapped, instance, args, kwargs, result, span)
@@ -72,16 +82,21 @@ class SpanHandler:
                             try:
                                 arguments = {"instance":instance, "args":args, "kwargs":kwargs, "result":result}
                                 result = accessor(arguments)
-                                if result and isinstance(result, str):
+                                if result and isinstance(result, (str, list)):
                                     span.set_attribute(attribute_name, result)
                             except Exception as e:
                                 logger.debug(f"Error processing accessor: {e}")
                         else:
-                            logger.warning(f"{' and '.join([key for key in ['attribute', 'accessor'] if not processor.get(key)])} not found or incorrect in entity JSON")
+                            logger.debug(f"{' and '.join([key for key in ['attribute', 'accessor'] if not processor.get(key)])} not found or incorrect in entity JSON")
                     span_index += 1
             else:
-                logger.warning("attributes not found or incorrect written in entity json")
+                logger.debug("attributes not found or incorrect written in entity json")
 
+        # set scopes as attributes by calling get_scopes()
+        # scopes is a Mapping[str:object], iterate directly with .items()
+        for scope_key, scope_value in get_scopes().items():
+            span.set_attribute(f"scope.{scope_key}", scope_value)
+        
         if span_index > 0:
             span.set_attribute("entity.count", span_index)
 

monocle_apptrace/instrumentation/common/utils.py

@@ -1,16 +1,33 @@
-import logging
-from typing import Callable, Generic, Optional, TypeVar
+import logging, json
+import os
+from typing import Callable, Generic, Optional, TypeVar, Mapping
+from threading import local
 
-from opentelemetry.context import attach, detach, get_current, get_value, set_value
-from opentelemetry.trace import NonRecordingSpan, Span
+from opentelemetry.context import attach, detach, get_current, get_value, set_value, Context
+from opentelemetry.trace import NonRecordingSpan, Span, get_tracer
 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
 
 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] = {}
+
+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:
@@ -139,6 +156,89 @@ def get_nested_value(data, keys):
 def get_keys_as_tuple(dictionary, *keys):
     return tuple(next((value for key, value in dictionary.items() if key.endswith(k) and value is not None), None) for k in keys)
 
+def load_scopes() -> dict:
+    methods_data = []
+    scope_methods = []
+    try:
+        scope_config_file_path = os.environ.get(SCOPE_CONFIG_PATH, 
+                                        os.path.join(os.getcwd(), SCOPE_METHOD_FILE))
+        with open(scope_config_file_path) as f:
+            methods_data = json.load(f)
+            for method in methods_data:
+                if method.get('http_header'):
+                    http_scopes[method.get('http_header')] = method.get('scope_name')
+                else:
+                    scope_methods.append(method)
+    except Exception as e:
+        logger.debug(f"Error loading scope methods from file: {e}")
+    return scope_methods
+
+def __generate_scope_id() -> str:
+    global scope_id_generator
+    return f"{hex(scope_id_generator.generate_trace_id())}"
+
+def set_scope(scope_name: str, scope_value:str = None) -> object:
+    return set_scopes({scope_name: scope_value})
+
+def set_scopes(scopes:dict[str, object], baggage_context:Context = None) -> object:
+    if baggage_context is None:
+        baggage_context:Context = get_current()
+    for scope_name, scope_value in scopes.items():
+        if scope_value is None:
+            scope_value = __generate_scope_id()
+        baggage_context = baggage.set_baggage(f"{MONOCLE_SCOPE_NAME_PREFIX}{scope_name}", scope_value, baggage_context)
+        
+    token:object = attach(baggage_context)
+    return token
+
+def remove_scope(token:object) -> None:
+    remove_scopes(token)
+
+def remove_scopes(token:object) -> None:
+    if token is not None:
+        detach(token)
+
+def get_scopes() -> dict[str, object]:
+    monocle_scopes:dict[str, object] = {}
+    for key, val in baggage.get_all().items():
+        if key.startswith(MONOCLE_SCOPE_NAME_PREFIX):
+            monocle_scopes[key[len(MONOCLE_SCOPE_NAME_PREFIX):]] = val
+    return monocle_scopes
+
+def get_baggage_for_scopes():
+    baggage_context:Context = None
+    for scope_key, scope_value in get_scopes():
+        monocle_scope_name = f"{MONOCLE_SCOPE_NAME_PREFIX}{scope_key}"
+        baggage_context = baggage.set_baggage(monocle_scope_name, scope_value, context=baggage_context)
+    return baggage_context
+
+def set_scopes_from_baggage(baggage_context:Context):
+    for scope_key, scope_value in baggage.get_all(baggage_context):
+        if scope_key.startswith(MONOCLE_SCOPE_NAME_PREFIX):
+            scope_name = scope_key[len(MONOCLE_SCOPE_NAME_PREFIX):]
+            set_scope(scope_name, scope_value)
+
+def extract_http_headers(headers) -> object:
+    global http_scopes
+    trace_context:Context = extract(headers, context=get_current())
+    imported_scope:dict[str, object] = {}
+    for http_header, http_scope in http_scopes.items():
+        if http_header in headers:
+            imported_scope[http_scope] = f"{http_header}: {headers[http_header]}"
+    token = set_scopes(imported_scope, trace_context)
+    return token
+
+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)
+    try:
+        result = func(req)
+        return result
+    finally:
+        clear_http_scopes(token)
 
 class Option(Generic[T]):
     def __init__(self, value: Optional[T]):
@@ -168,4 +268,11 @@ def try_option(func: Callable[..., T], *args, **kwargs) -> Option[T]:
     try:
         return Option(func(*args, **kwargs))
     except Exception:
-        return Option(None)
+        return Option(None)
+
+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

monocle_apptrace/instrumentation/common/wrapper.py

@@ -7,8 +7,9 @@ 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
 )
-from monocle_apptrace.instrumentation.metamodel.botocore import _helper
 logger = logging.getLogger(__name__)
 
 
@@ -26,22 +27,20 @@ def task_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instanc
     else:
         name = get_fully_qualified_class_name(instance)
 
-    handler.validate(to_wrap, wrapped, instance, args, kwargs)
-    handler.set_context_properties(to_wrap, wrapped, instance, args, kwargs)
-
-    if to_wrap.get('skip_span'):
-        return_value = wrapped(*args, **kwargs)
-        _helper.botocore_processor(tracer, to_wrap, wrapped, instance, args, kwargs, return_value)
+    return_value = None
+    try:
+        handler.pre_tracing(to_wrap, wrapped, instance, args, kwargs)
+        if to_wrap.get('skip_span') or handler.skip_span(to_wrap, wrapped, instance, args, kwargs):
+            return_value = wrapped(*args, **kwargs)
+        else:
+            with tracer.start_as_current_span(name) as span:
+                handler.pre_task_processing(to_wrap, wrapped, instance, args, kwargs, span)
+                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.start_as_current_span(name) as span:
-        handler.pre_task_processing(to_wrap, wrapped, instance, args, span)
-        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
-
+    finally:
+        handler.post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
 
 @with_tracer_wrapper
 async def atask_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, args, kwargs):
@@ -58,12 +57,38 @@ async def atask_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped,
     else:
         name = get_fully_qualified_class_name(instance)
 
-    handler.validate(to_wrap, wrapped, instance, args, kwargs)
-    handler.set_context_properties(to_wrap, wrapped, instance, args, kwargs)
-    with tracer.start_as_current_span(name) as span:
-        handler.pre_task_processing(to_wrap, wrapped, instance, args, span)
-        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_value = None
+    try:
+        handler.pre_tracing(to_wrap, wrapped, instance, args, kwargs)
+        if to_wrap.get('skip_span') or handler.skip_span(to_wrap, wrapped, instance, args, kwargs):
+            return_value = wrapped(*args, **kwargs)
+        else:
+            with tracer.start_as_current_span(name) as span:
+                handler.pre_task_processing(to_wrap, wrapped, instance, args, kwargs, span)
+                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
+    finally:
+        handler.post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
 
+@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)
+    if scope_name:
+        token = set_scope(scope_name)
+    return_value = wrapped(*args, **kwargs)
+    if scope_name:
+        remove_scope(token)
+    return return_value

monocle_apptrace/instrumentation/common/wrapper_method.py

@@ -1,12 +1,20 @@
 # pylint: disable=too-few-public-methods
-from monocle_apptrace.instrumentation.common.wrapper import task_wrapper
+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.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 (
     LANGCHAIN_METHODS,
 )
 from monocle_apptrace.instrumentation.metamodel.llamaindex.methods import (LLAMAINDEX_METHODS, )
 from monocle_apptrace.instrumentation.metamodel.haystack.methods import (HAYSTACK_METHODS, )
-
+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.requests.methods import (REQUESTS_METHODS, )
+from monocle_apptrace.instrumentation.metamodel.requests._helper import RequestSpanHandler
 
 class WrapperMethod:
     def __init__(
@@ -17,16 +25,21 @@ class WrapperMethod:
             span_name: str = None,
             output_processor : str = None,
             wrapper_method = task_wrapper,
-            span_handler = 'default'
+            span_handler = 'default',
+            scope_name: str = None
             ):
         self.package = package
         self.object = object_name
         self.method = method
         self.span_name = span_name
         self.output_processor=output_processor
-        self.span_handler = span_handler
 
-        self.wrapper_method = wrapper_method
+        self.span_handler:SpanHandler.__class__ = span_handler
+        self.scope_name = scope_name
+        if scope_name:
+            self.wrapper_method = scope_wrapper
+        else:
+            self.wrapper_method = wrapper_method
 
     def to_dict(self) -> dict:
         # Create a dictionary representation of the instance
@@ -37,9 +50,19 @@ class WrapperMethod:
             'span_name': self.span_name,
             'output_processor': self.output_processor,
             'wrapper_method': self.wrapper_method,
-            'span_handler': self.span_handler
+            'span_handler': self.span_handler,
+            'scope_name': self.scope_name
         }
         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
 
-DEFAULT_METHODS_LIST = LANGCHAIN_METHODS + LLAMAINDEX_METHODS + HAYSTACK_METHODS + BOTOCORE_METHODS
+MONOCLE_SPAN_HANDLERS: Dict[str, SpanHandler] = {
+    "default": SpanHandler(),
+    "botocore_handler": BotoCoreSpanHandler(),
+    "flask_handler": FlaskSpanHandler(),
+    "request_handler": RequestSpanHandler()
+}

monocle_apptrace/instrumentation/metamodel/botocore/_helper.py

@@ -81,37 +81,6 @@ def resolve_from_alias(my_map, alias):
             return my_map[i]
     return None
 
-
-def botocore_processor(tracer, to_wrap, wrapped, instance, args, kwargs,return_value):
-    service_name = kwargs.get("service_name")
-    service_method_mapping = {
-        "sagemaker-runtime": "invoke_endpoint",
-        "bedrock-runtime": "converse",
-    }
-    if service_name in service_method_mapping:
-        method_name = service_method_mapping[service_name]
-        original_method = getattr(return_value, method_name, None)
-
-        if original_method:
-            instrumented_method = _instrumented_endpoint_invoke(
-                to_wrap, wrapped,return_value, original_method, tracer, service_name
-            )
-            setattr(return_value, method_name, instrumented_method)
-
-
-def _instrumented_endpoint_invoke(to_wrap,wrapped, instance, fn, tracer,service_name):
-    @wraps(fn)
-    def with_instrumentation(*args, **kwargs):
-        span_name="botocore-"+service_name+"-invoke-endpoint"
-        handler = SpanHandler()
-        with tracer.start_as_current_span(span_name) as span:
-            response = fn(*args, **kwargs)
-            handler.hydrate_span(to_wrap, span=span,wrapped=wrapped, instance=instance,args=args, kwargs=kwargs, result=response)
-            return response
-
-    return with_instrumentation
-
-
 def update_span_from_llm_response(response, instance):
     meta_dict = {}
     if response is not None and isinstance(response, dict) and "usage" in response: