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

monocle_apptrace/instrumentation/common/wrapper.py

@@ -6,76 +6,185 @@ 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
+    get_parent_span
 )
-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
 
-    if hasattr(instance, "name") and instance.name:
-        name = f"{to_wrap.get('span_name')}.{instance.name.lower()}"
-    elif to_wrap.get("span_name"):
+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)
+        try:
+            handler.pre_task_processing(to_wrap, wrapped, instance, args, kwargs, span)
+        except Exception as e:
+            logger.info(f"Warning: Error occurred in pre_task_processing: {e}")
+
+def post_process_span(handler, to_wrap, wrapped, instance, args, kwargs, return_value, span, parent_span=None, ex = None):
+    if not (SpanHandler.is_root_span(span) or get_value(ADD_NEW_WORKFLOW) == True):
+        try:
+            handler.hydrate_span(to_wrap, wrapped, instance, args, kwargs, return_value, span, parent_span, ex)
+        except Exception as e:
+            logger.info(f"Warning: Error occurred in hydrate_span: {e}")
+        
+        try:
+            handler.post_task_processing(to_wrap, wrapped, instance, args, kwargs, return_value, span)
+        except Exception as e:
+            logger.info(f"Warning: Error occurred in post_task_processing: {e}")
+
+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
+    auto_close_span = get_auto_close_span(to_wrap, kwargs)
+    parent_span = get_parent_span()
+    with tracer.start_as_current_span(name, end_on_exit=auto_close_span) 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)
+            if not auto_close_span:
+                span.end()
+        else:
+            ex:Exception = None
+            try:
+                with SpanHandler.workflow_type(to_wrap, span):
+                    return_value = wrapped(*args, **kwargs)
+            except Exception as e:
+                ex = e
+                raise
+            finally:
+                def post_process_span_internal(ret_val):
+                    post_process_span(handler, to_wrap, wrapped, instance, args, kwargs, ret_val, span, parent_span ,ex)
+                    if not auto_close_span:
+                        span.end()
+                if ex is None and not auto_close_span and to_wrap.get("output_processor") and to_wrap.get("output_processor").get("response_processor"):
+                    to_wrap.get("output_processor").get("response_processor")(to_wrap, return_value, post_process_span_internal)
+                else:
+                    post_process_span_internal(return_value)
+            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)
+        try:
+            handler.pre_tracing(to_wrap, wrapped, instance, args, kwargs)
+        except Exception as e:
+            logger.info(f"Warning: Error occurred in pre_tracing: {e}")
+        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)
+        try:
+            handler.post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
+        except Exception as e:
+            logger.info(f"Warning: Error occurred in post_tracing: {e}")
 
-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
+    auto_close_span = get_auto_close_span(to_wrap, kwargs)
+    parent_span = get_parent_span()
+    with tracer.start_as_current_span(name, end_on_exit=auto_close_span) 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)
+            if not auto_close_span:
+                span.end()
         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
+            ex:Exception = None
+            try:
+                with SpanHandler.workflow_type(to_wrap, span):
+                    return_value = await wrapped(*args, **kwargs)
+            except Exception as e:
+                ex = e
+                raise
+            finally:
+                def post_process_span_internal(ret_val):
+                    post_process_span(handler, to_wrap, wrapped, instance, args, kwargs, ret_val, span, parent_span,ex)
+                    if not auto_close_span:
+                        span.end()
+                if ex is None and not auto_close_span and to_wrap.get("output_processor") and to_wrap.get("output_processor").get("response_processor"):
+                    to_wrap.get("output_processor").get("response_processor")(to_wrap, return_value, post_process_span_internal)
+                else:
+                    post_process_span_internal(return_value)
+        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:
+        try:
+            handler.pre_tracing(to_wrap, wrapped, instance, args, kwargs)
+        except Exception as e:
+            logger.info(f"Warning: Error occurred in pre_tracing: {e}")
+        if to_wrap.get('skip_span', False) or handler.skip_span(to_wrap, wrapped, instance, args, kwargs):
+            return_value = await wrapped(*args, **kwargs)
+        else:
+            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:
+        try:
+            handler.post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
+        except Exception as e:
+            logger.info(f"Warning: Error occurred in post_tracing: {e}")
 
 @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 +194,54 @@ 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, to_wrap, 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, to_wrap, 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, to_wrap, 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

@@ -10,14 +10,16 @@ from monocle_apptrace.instrumentation.metamodel.langchain.methods import (
 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.openai._helper import OpenAISpanHandler
 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 +31,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 +40,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 +60,23 @@ 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()
+    "non_framework_handler": NonFrameworkSpanHandler(),
+    "openai_handler": OpenAISpanHandler(),
 }

monocle_apptrace/instrumentation/metamodel/aiohttp/_helper.py

@@ -0,0 +1,66 @@
+import logging
+from threading import local
+from monocle_apptrace.instrumentation.common.utils import extract_http_headers, clear_http_scopes, try_option, Option, MonocleSpanException
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.common.constants import HTTP_SUCCESS_CODES
+from urllib.parse import unquote
+
+logger = logging.getLogger(__name__)
+MAX_DATA_LENGTH = 1000
+token_data = local()
+token_data.current_token = None
+
+def get_route(args) -> str:
+    route_path: Option[str] = try_option(getattr, args[0], 'path')
+    return route_path.unwrap_or("")
+
+def get_method(args) -> str:
+#    return args[0]['method'] if 'method' in args[0] else ""
+    http_method: Option[str] = try_option(getattr, args[0], 'method')
+    return http_method.unwrap_or("")
+
+def get_params(args) -> dict:
+    params: Option[str] = try_option(getattr, args[0], 'query_string')
+    return unquote(params.unwrap_or(""))
+
+def get_body(args) -> dict:
+    return ""
+
+def extract_response(result) -> str:
+    if hasattr(result, 'text'):
+        response = result.text[0:max(result.text.__len__(), MAX_DATA_LENGTH)]
+    else:
+        response = ""
+    return response
+
+def extract_status(result) -> str:
+    status = f"{result.status}" if hasattr(result, 'status') else ""
+    if status not in HTTP_SUCCESS_CODES:
+        error_message = extract_response(result)
+        raise MonocleSpanException(f"error: {status} - {error_message}")
+    return status
+
+def aiohttp_pre_tracing(args):
+    token_data.current_token = extract_http_headers(args[0].headers)
+
+def aiohttp_post_tracing():
+    clear_http_scopes(token_data.current_token)
+    token_data.current_token = None
+
+def aiohttp_skip_span(args) -> bool:
+    if get_method(args) == "HEAD":
+        return True
+    return False
+
+class aiohttpSpanHandler(SpanHandler):
+
+    def pre_tracing(self, to_wrap, wrapped, instance, args, kwargs):
+        aiohttp_pre_tracing(args)
+        return super().pre_tracing(to_wrap, wrapped, instance, args, kwargs)
+    
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+        aiohttp_post_tracing()
+        return super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
+
+    def skip_span(self, to_wrap, wrapped, instance, args, kwargs) -> bool:
+        return aiohttp_skip_span(args)

monocle_apptrace/instrumentation/metamodel/aiohttp/entities/http.py

@@ -0,0 +1,51 @@
+from monocle_apptrace.instrumentation.metamodel.aiohttp import _helper
+AIO_HTTP_PROCESSOR = {
+    "type": "http.process",
+    "attributes": [
+        [
+            {
+                "_comment": "request method, request URI",
+                "attribute": "method",
+                "accessor": lambda arguments: _helper.get_method(arguments['args'])
+            },
+            {
+                "_comment": "request method, request URI",
+                "attribute": "route",
+                "accessor": lambda arguments: _helper.get_route(arguments['args'])
+            },
+            {
+                "_comment": "request method, request URI",
+                "attribute": "body",
+                "accessor": lambda arguments: _helper.get_body(arguments['args'])
+            },
+        ]
+    ],
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "_comment": "route params",
+                    "attribute": "params",
+                    "accessor": lambda arguments: _helper.get_params(arguments['args'])
+                }
+            ]
+        },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "status from HTTP response",
+                    "attribute": "status",
+                    "accessor": lambda arguments: _helper.extract_status(arguments['result'])
+                },
+                {
+                    "_comment": "this is result from LLM",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_response(arguments['result'])
+                }
+            ]
+        }
+
+    ]
+}

monocle_apptrace/instrumentation/metamodel/aiohttp/methods.py

@@ -0,0 +1,13 @@
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper
+from monocle_apptrace.instrumentation.metamodel.aiohttp.entities.http import AIO_HTTP_PROCESSOR
+
+AIOHTTP_METHODS = [
+    {
+        "package": "aiohttp.web_app",
+        "object": "Application",
+        "method": "_handle",
+        "wrapper_method": atask_wrapper,
+        "span_handler": "aiohttp_handler",
+        "output_processor": AIO_HTTP_PROCESSOR
+    }
+]

monocle_apptrace/instrumentation/metamodel/anthropic/methods.py

@@ -5,17 +5,19 @@ from monocle_apptrace.instrumentation.metamodel.anthropic.entities.inference imp
 
 ANTHROPIC_METHODS = [
     {
-        "package": "anthropic.resources.messages.messages",
+        "package": "anthropic.resources",
         "object": "Messages",
         "method": "create",
         "wrapper_method": task_wrapper,
+        "span_handler": "non_framework_handler",
         "output_processor": INFERENCE
     },
     {
-        "package": "anthropic.resources.messages.messages",
+        "package": "anthropic.resources",
         "object": "AsyncMessages",
         "method": "create",
         "wrapper_method": atask_wrapper,
+        "span_handler": "non_framework_handler",
         "output_processor": INFERENCE
     },
 

monocle_apptrace/instrumentation/metamodel/flask/_helper.py

@@ -1,11 +1,46 @@
+import logging
 from threading import local
 from monocle_apptrace.instrumentation.common.utils import extract_http_headers, clear_http_scopes
-from opentelemetry.propagate import extract
-from opentelemetry.context import Context, attach, detach
 from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.common.constants import HTTP_SUCCESS_CODES
+from monocle_apptrace.instrumentation.common.utils import MonocleSpanException
+from urllib.parse import unquote
+from opentelemetry.context import get_current
+from opentelemetry.trace import Span, get_current_span
+from opentelemetry.trace.propagation import _SPAN_KEY
+
+logger = logging.getLogger(__name__)
+MAX_DATA_LENGTH = 1000
 token_data = local()
 token_data.current_token = None
 
+def get_route(args) -> str:
+    return args[0]['PATH_INFO'] if 'PATH_INFO' in args[0] else ""
+
+def get_method(args) -> str:
+    return args[0]['REQUEST_METHOD'] if 'REQUEST_METHOD' in args[0] else ""
+
+def get_params(args) -> dict:
+    params = args[0]['QUERY_STRING'] if 'QUERY_STRING' in args[0] else ""
+    return unquote(params)
+
+def get_body(args) -> dict:
+    return ""
+
+def extract_response(instance) -> str:
+    if hasattr(instance, 'data') and hasattr(instance, 'content_length'):
+        response = instance.data[0:max(instance.content_length, MAX_DATA_LENGTH)]
+    else:
+        response = ""
+    return response
+
+def extract_status(instance) -> str:
+    status = f"{instance.status_code}" if hasattr(instance, 'status_code') else ""
+    if status not in HTTP_SUCCESS_CODES:
+        error_message = extract_response(instance)
+        raise MonocleSpanException(f"error: {status} - {error_message}")
+    return status
+
 def flask_pre_tracing(args):
     headers = dict()
     for key, value in args[0].items():
@@ -26,4 +61,16 @@ class FlaskSpanHandler(SpanHandler):
     
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
         flask_post_tracing()
-        return super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
+        return super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
+
+class FlaskResponseSpanHandler(SpanHandler):
+    def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
+        try:
+            _parent_span_context = get_current()
+            if _parent_span_context is not None:
+                parent_span: Span = _parent_span_context.get(_SPAN_KEY, None)
+                if parent_span is not None:
+                    self.hydrate_events(to_wrap, wrapped, instance, args, kwargs, return_value, parent_span=parent_span)
+        except Exception as e:
+            logger.info(f"Failed to propogate flask response: {e}")
+        super().post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)

monocle_apptrace/instrumentation/metamodel/flask/entities/http.py

@@ -0,0 +1,48 @@
+from monocle_apptrace.instrumentation.metamodel.flask import _helper
+FLASK_HTTP_PROCESSOR = {
+    "type": "http.process",
+    "attributes": [
+        [
+            {
+                "_comment": "request method, request URI",
+                "attribute": "method",
+                "accessor": lambda arguments: _helper.get_method(arguments['args'])
+            },
+            {
+                "_comment": "request method, request URI",
+                "attribute": "route",
+                "accessor": lambda arguments: _helper.get_route(arguments['args'])
+            },
+        ]
+    ]
+}
+
+FLASK_RESPONSE_PROCESSOR = {
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "_comment": "route params",
+                    "attribute": "params",
+                    "accessor": lambda arguments: _helper.get_params(arguments['args'])
+                }
+         ]
+         },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "status from HTTP response",
+                    "attribute": "status",
+                    "accessor": lambda arguments: _helper.extract_status(arguments['instance'])
+                },
+                {
+                    "_comment": "this is result from LLM",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_response(arguments['instance'])
+                }
+            ]
+        }
+    ]
+}

monocle_apptrace/instrumentation/metamodel/flask/methods.py

@@ -1,13 +1,22 @@
 from monocle_apptrace.instrumentation.common.wrapper import task_wrapper
+from monocle_apptrace.instrumentation.metamodel.flask.entities.http import FLASK_HTTP_PROCESSOR, FLASK_RESPONSE_PROCESSOR
 
 FLASK_METHODS = [
     {
         "package": "flask.app",
         "object": "Flask",
         "method": "wsgi_app",
-        "span_name": "Flask.wsgi_app",
         "wrapper_method": task_wrapper,
         "span_handler": "flask_handler",
+        "output_processor": FLASK_HTTP_PROCESSOR,
+    },
+    {
+        "package": "werkzeug.wrappers.response",
+        "object": "Response",
+        "method": "__call__",
+        "wrapper_method": task_wrapper,
+        "span_handler": "flask_response_handler",
+        "output_processor": FLASK_RESPONSE_PROCESSOR,
         "skip_span": True
     }
 ]