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

monocle_apptrace/instrumentation/common/wrapper.py

@@ -6,78 +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, 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)
 
+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)
         if to_wrap.get('skip_span', False) or handler.skip_span(to_wrap, wrapped, instance, args, kwargs):
-            if async_task:
-                return_value = async_wrapper(wrapped, None, None, None, *args, **kwargs)
-            else:
-                return_value = wrapped(*args, **kwargs)
+            return_value = 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_processor(name, async_task, tracer, handler, add_workflow_span,
-                                        to_wrap, wrapped, instance, args, kwargs)
+                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.post_tracing(to_wrap, wrapped, instance, args, kwargs, return_value)
 
-def span_processor(name: str, async_task: bool, tracer: Tracer, handler: SpanHandler, add_workflow_span:bool,
-                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) or add_workflow_span:
-            # This is a direct API call of a non-framework type, call the span_processor recursively for the actual span
-            SpanHandler.set_workflow_properties(span, to_wrap)
-            return_value = span_processor(name, async_task, tracer, handler, False, 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:
-            with SpanHandler.workflow_type(to_wrap):
-                SpanHandler.set_non_workflow_properties(span)
-                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)
@@ -87,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

@@ -17,7 +17,8 @@ from monocle_apptrace.instrumentation.metamodel.requests.methods import (REQUEST
 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,17 +59,19 @@ 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(),

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/flask/_helper.py

@@ -2,6 +2,8 @@ import logging
 from threading import local
 from monocle_apptrace.instrumentation.common.utils import extract_http_headers, clear_http_scopes
 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
@@ -28,12 +30,15 @@ def get_body(args) -> dict:
 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:   
+    else:
         response = ""
     return response
 
 def extract_status(instance) -> str:
-    status = instance.status if hasattr(instance, 'status') else ""
+    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):

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

@@ -18,7 +18,6 @@ FLASK_HTTP_PROCESSOR = {
 }
 
 FLASK_RESPONSE_PROCESSOR = {
-    "type": "http.process",
     "events": [
         {
             "name": "data.input",

monocle_apptrace/instrumentation/metamodel/haystack/_helper.py

@@ -1,4 +1,5 @@
 import logging
+
 from monocle_apptrace.instrumentation.common.utils import (
     Option,
     get_keys_as_tuple,
@@ -11,13 +12,19 @@ logger = logging.getLogger(__name__)
 def extract_messages(kwargs):
     try:
         messages = []
+        system_message, user_message = None,None
         if isinstance(kwargs, dict):
             if 'system_prompt' in kwargs and kwargs['system_prompt']:
                 system_message = kwargs['system_prompt']
-                messages.append({"system" : system_message})
             if 'prompt' in kwargs and kwargs['prompt']:
                 user_message = extract_question_from_prompt(kwargs['prompt'])
+            if 'messages' in kwargs and len(kwargs['messages'])>1:
+                system_message = kwargs['messages'][0].text
+                user_message = kwargs['messages'][1].text
+            if system_message and user_message:
+                messages.append({"system": system_message})
                 messages.append({"user": user_message})
+
         return [str(message) for message in messages]
     except Exception as e:
         logger.warning("Warning: Error occurred in extract_messages: %s", str(e))
@@ -52,6 +59,8 @@ def extract_assistant_message(response):
             reply = response["replies"][0]
             if hasattr(reply, 'content'):
                 return [reply.content]
+            if hasattr(reply, 'text'):
+                return [reply.text]
             return [reply]
     except Exception as e:
         logger.warning("Warning: Error occurred in extract_assistant_message: %s", str(e))
@@ -108,15 +117,19 @@ def extract_embeding_model(instance):
 
 def update_span_from_llm_response(response, instance):
     meta_dict = {}
-    if response is not None and isinstance(response, dict) and "meta" in response:
-        token_usage = response["meta"][0]["usage"]
+    token_usage = None
+    if response is not None and isinstance(response, dict):
+        if "meta" in response:
+            token_usage = response["meta"][0]["usage"]
+        if "replies" in response:
+            token_usage = response["replies"][0].meta["usage"]
         if token_usage is not None:
             temperature = instance.__dict__.get("temperature", None)
             meta_dict.update({"temperature": temperature})
             meta_dict.update(
                 {"completion_tokens": token_usage.get("completion_tokens") or token_usage.get("output_tokens")})
             meta_dict.update({"prompt_tokens": token_usage.get("prompt_tokens") or token_usage.get("input_tokens")})
-            meta_dict.update({"total_tokens": token_usage.get("total_tokens")})
+            meta_dict.update({"total_tokens": token_usage.get("total_tokens") or token_usage.get("completion_tokens")+token_usage.get("prompt_tokens")})
     return meta_dict
 
 

monocle_apptrace/instrumentation/metamodel/haystack/methods.py

@@ -36,5 +36,12 @@ HAYSTACK_METHODS = [
         "object": "Pipeline",
         "method": "run",
         "wrapper_method": task_wrapper
-    }
+    },
+    {
+        "package": "haystack_integrations.components.generators.anthropic",
+        "object": "AnthropicChatGenerator",
+        "method": "run",
+        "wrapper_method": task_wrapper,
+        "output_processor": INFERENCE
+    },
 ]

monocle_apptrace/instrumentation/metamodel/llamaindex/_helper.py

@@ -96,12 +96,19 @@ def extract_query_from_content(content):
 
 
 def extract_provider_name(instance):
-    provider_url = try_option(getattr, instance, 'api_base').and_then(lambda url: urlparse(url).hostname)
-    return provider_url
+    if hasattr(instance,'api_base'):
+        provider_url: Option[str]= try_option(getattr, instance, 'api_base').and_then(lambda url: urlparse(url).hostname)
+    if hasattr(instance,'_client'):
+        provider_url:Option[str] = try_option(getattr, instance._client.base_url,'host')
+    return provider_url.unwrap_or(None)
 
 
 def extract_inference_endpoint(instance):
-    inference_endpoint = try_option(getattr, instance._client.sdk_configuration, 'server_url').map(str)
+    if hasattr(instance,'_client'):
+        if hasattr(instance._client,'sdk_configuration'):
+            inference_endpoint: Option[str] = try_option(getattr, instance._client.sdk_configuration, 'server_url').map(str)
+        if hasattr(instance._client,'base_url'):
+            inference_endpoint: Option[str] = try_option(getattr, instance._client, 'base_url').map(str)
     return inference_endpoint.unwrap_or(extract_provider_name(instance))
 
 
@@ -163,10 +170,7 @@ def update_span_from_llm_response(response, instance):
             if token_usage is not None:
                 temperature = instance.__dict__.get("temperature", None)
                 meta_dict.update({"temperature": temperature})
-                if getattr(token_usage, "completion_tokens", None):
-                    meta_dict.update({"completion_tokens": getattr(token_usage, "completion_tokens")})
-                if getattr(token_usage, "prompt_tokens", None):
-                    meta_dict.update({"prompt_tokens": getattr(token_usage, "prompt_tokens")})
-                if getattr(token_usage, "total_tokens", None):
-                    meta_dict.update({"total_tokens": getattr(token_usage, "total_tokens")})
+                meta_dict.update({"completion_tokens": getattr(token_usage, "completion_tokens",None) or getattr(token_usage,"output_tokens",None)})
+                meta_dict.update({"prompt_tokens": getattr(token_usage, "prompt_tokens",None) or getattr(token_usage,"input_tokens",None)})
+                meta_dict.update({"total_tokens": getattr(token_usage, "total_tokens",None) or getattr(token_usage,"output_tokens",None)+getattr(token_usage,"input_tokens",None)})
     return meta_dict

monocle_apptrace/instrumentation/metamodel/llamaindex/methods.py

@@ -84,5 +84,19 @@ LLAMAINDEX_METHODS = [
         "method": "chat",
         "wrapper_method": task_wrapper,
         "output_processor": AGENT
+    },
+    {
+        "package": "llama_index.llms.anthropic",
+        "object": "Anthropic",
+        "method": "chat",
+        "wrapper_method": task_wrapper,
+        "output_processor": INFERENCE
+    },
+    {
+        "package": "llama_index.llms.anthropic",
+        "object": "Anthropic",
+        "method": "achat",
+        "wrapper_method": atask_wrapper,
+        "output_processor": INFERENCE
     }
 ]