monocle-apptrace 0.2.0__py3-none-any.whl → 0.3.0b2__py3-none-any.whl

monocle_apptrace/metamodel/maps/llamaindex_methods.json

@@ -69,6 +69,24 @@
         "wrapper_package": "wrap_common",
         "wrapper_method": "allm_wrapper",
         "output_processor": ["metamodel/maps/attributes/inference/llamaindex_entities.json"]
+    },
+    {
+        "package": "llama_index.llms.mistralai.base",
+        "object": "MistralAI",
+        "method": "chat",
+        "span_name": "llamaindex.mistralai",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "llm_wrapper",
+        "output_processor": ["metamodel/maps/attributes/inference/llamaindex_entities.json"]
+    },
+    {
+        "package": "llama_index.llms.mistralai.base",
+        "object": "MistralAI",
+        "method": "achat",
+        "span_name": "llamaindex.mistralai",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "allm_wrapper",
+        "output_processor": ["metamodel/maps/attributes/inference/llamaindex_entities.json"]
     }
 ]
 }

monocle_apptrace/utils.py

@@ -2,10 +2,13 @@ import logging
 import json
 from importlib import import_module
 import os
-from opentelemetry.trace import Span
+from opentelemetry.trace import NonRecordingSpan,Span
+from opentelemetry.trace.propagation import _SPAN_KEY
+from opentelemetry.context import (attach, detach,get_current)
 from opentelemetry.context import attach, set_value, get_value
-from monocle_apptrace.constants import azure_service_map, aws_service_map
+from monocle_apptrace.constants import service_name_map, service_type_map
 from json.decoder import JSONDecodeError
+
 logger = logging.getLogger(__name__)
 
 embedding_model_context = {}
@@ -39,8 +42,25 @@ def with_tracer_wrapper(func):
 
     def _with_tracer(tracer, to_wrap):
         def wrapper(wrapped, instance, args, kwargs):
-            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
-
+            token = None
+            try:
+                _parent_span_context = get_current()
+                if _parent_span_context is not None and _parent_span_context.get(_SPAN_KEY, None):
+                    parent_span: Span = _parent_span_context.get(_SPAN_KEY, None)
+                    is_invalid_span = isinstance(parent_span, NonRecordingSpan)
+                    if is_invalid_span:
+                        token = attach(context={})
+            except Exception as e:
+                logger.error("Exception in attaching parent context: %s", e)
+
+            val = func(tracer, to_wrap, wrapped, instance, args, kwargs)
+            # Detach the token if it was set
+            if token:
+                try:
+                    detach(token=token)
+                except Exception as e:
+                    logger.error("Exception in detaching parent context: %s", e)
+            return val
         return wrapper
 
     return _with_tracer
@@ -118,13 +138,16 @@ def get_wrapper_method(package_name: str, method_name: str):
     wrapper_module = import_module("monocle_apptrace." + package_name)
     return getattr(wrapper_module, method_name)
 
-def update_span_with_infra_name(span: Span, span_key: str):
-    for key, val in azure_service_map.items():
-        if key in os.environ:
-            span.set_attribute(span_key, val)
-    for key, val in aws_service_map.items():
-        if key in os.environ:
-            span.set_attribute(span_key, val)
+def set_app_hosting_identifier_attribute(span, span_index):
+    return_value = 0
+    # Search env to indentify the infra service type, if found check env for service name if possible
+    for type_env, type_name in service_type_map.items():
+        if type_env in os.environ:
+            return_value = 1
+            span.set_attribute(f"entity.{span_index}.type", f"app_hosting.{type_name}")
+            entity_name_env = service_name_map.get(type_name, "unknown")
+            span.set_attribute(f"entity.{span_index}.name", os.environ.get(entity_name_env, "generic"))
+    return return_value
 
 def set_embedding_model(model_name: str):
     """
@@ -164,9 +187,66 @@ def get_attribute(key: str) -> str:
     """
     return get_value(key)
 
+def flatten_dict(d, parent_key='', sep='_'):
+    items = []
+    for k, v in d.items():
+        new_key = f"{parent_key}{sep}{k}" if parent_key else k
+        if isinstance(v, dict):
+            items.extend(flatten_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)
+
+def get_fully_qualified_class_name(instance):
+    if instance is None:
+        return None
+    module_name = instance.__class__.__module__
+    qualname = instance.__class__.__qualname__
+    return f"{module_name}.{qualname}"
+
+# returns json path like key probe in a dictionary
+def get_nested_value(data, keys):
+    for key in keys:
+        if isinstance(data, dict) and key in data:
+            data = data[key]
+        elif hasattr(data, key):
+            data = getattr(data, key)
+        else:
+            return None
+    return data
+
 def get_workflow_name(span: Span) -> str:
     try:
         return get_value("workflow_name") or span.resource.attributes.get("service.name")
     except Exception as e:
         logger.exception(f"Error getting workflow name: {e}")
-        return None
+        return None
+
+def get_vectorstore_deployment(my_map):
+    if isinstance(my_map, dict):
+        if '_client_settings' in my_map:
+            client = my_map['_client_settings'].__dict__
+            host, port = get_keys_as_tuple(client, 'host', 'port')
+            if host:
+                return f"{host}:{port}" if port else host
+        keys_to_check = ['client', '_client']
+        host = get_host_from_map(my_map, keys_to_check)
+        if host:
+            return host
+    else:
+        if hasattr(my_map, 'client') and '_endpoint' in my_map.client.__dict__:
+            return my_map.client.__dict__['_endpoint']
+        host, port = get_keys_as_tuple(my_map.__dict__, 'host', 'port')
+        if host:
+            return f"{host}:{port}" if port else host
+    return None
+
+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 get_host_from_map(my_map, keys_to_check):
+    for key in keys_to_check:
+        seed_connections = get_nested_value(my_map, [key, 'transport', 'seed_connections'])
+        if seed_connections and 'host' in seed_connections[0].__dict__:
+            return seed_connections[0].__dict__['host']
+    return None

monocle_apptrace/wrap_common.py

@@ -2,18 +2,23 @@
 import logging
 import os
 import inspect
+from importlib.metadata import version
 from urllib.parse import urlparse
-from opentelemetry.trace import Span, Tracer
-from monocle_apptrace.utils import resolve_from_alias, update_span_with_infra_name, with_tracer_wrapper, get_embedding_model, get_attribute, get_workflow_name
-from monocle_apptrace.utils import set_attribute
-from opentelemetry.context import get_value, attach, set_value
+from opentelemetry.trace import Tracer
+from opentelemetry.sdk.trace import Span
+from monocle_apptrace.utils import resolve_from_alias, with_tracer_wrapper, get_embedding_model, get_attribute, get_workflow_name, set_embedding_model, set_app_hosting_identifier_attribute
+from monocle_apptrace.utils import set_attribute, get_vectorstore_deployment
+from monocle_apptrace.utils import get_fully_qualified_class_name, get_nested_value
+from monocle_apptrace.message_processing import extract_messages, extract_assistant_message
+from functools import wraps
+
 logger = logging.getLogger(__name__)
 WORKFLOW_TYPE_KEY = "workflow_type"
 DATA_INPUT_KEY = "data.input"
 DATA_OUTPUT_KEY = "data.output"
 PROMPT_INPUT_KEY = "data.input"
 PROMPT_OUTPUT_KEY = "data.output"
-QUERY = "question"
+QUERY = "input"
 RESPONSE = "response"
 SESSION_PROPERTIES_KEY = "session"
 INFRA_SERVICE_KEY = "infra_service_name"
@@ -30,52 +35,18 @@ WORKFLOW_TYPE_MAP = {
     "haystack": "workflow.haystack"
 }
 
-def get_embedding_model_for_vectorstore(instance):
-    # Handle Langchain or other frameworks where vectorstore exists
-    if hasattr(instance, 'vectorstore'):
-        vectorstore_dict = instance.vectorstore.__dict__
-
-        # Use inspect to check if the embedding function is from Sagemaker
-        if 'embedding_func' in vectorstore_dict:
-            embedding_func = vectorstore_dict['embedding_func']
-            class_name = embedding_func.__class__.__name__
-            file_location = inspect.getfile(embedding_func.__class__)
-
-            # Check if the class is SagemakerEndpointEmbeddings
-            if class_name == 'SagemakerEndpointEmbeddings' and 'langchain_community' in file_location:
-                # Set embedding_model as endpoint_name if it's Sagemaker
-                if hasattr(embedding_func, 'endpoint_name'):
-                    return embedding_func.endpoint_name
-
-        # Default to the regular embedding model if not Sagemaker
-        return instance.vectorstore.embeddings.model
-
-    # Handle llama_index where _embed_model is present
-    if hasattr(instance, '_embed_model') and hasattr(instance._embed_model, 'model_name'):
-        return instance._embed_model.model_name
-
-    # Fallback if no specific model is found
-    return "Unknown Embedding Model"
-
-
-framework_vector_store_mapping = {
-    'langchain_core.retrievers': lambda instance: {
-        'provider': type(instance.vectorstore).__name__,
-        'embedding_model': get_embedding_model_for_vectorstore(instance),
-        'type': VECTOR_STORE,
-    },
-    'llama_index.core.indices.base_retriever': lambda instance: {
-        'provider': type(instance._vector_store).__name__,
-        'embedding_model': get_embedding_model_for_vectorstore(instance),
-        'type': VECTOR_STORE,
-    },
-    'haystack.components.retrievers.in_memory': lambda instance: {
-        'provider': instance.__dict__.get("document_store").__class__.__name__,
-        'embedding_model': get_embedding_model(),
-        'type': VECTOR_STORE,
-    },
-}
 
+def get_embedding_model_haystack(instance):
+    try:
+        if hasattr(instance, 'get_component'):
+            text_embedder = instance.get_component('text_embedder')
+            if text_embedder and hasattr(text_embedder, 'model'):
+                # Set the embedding model attribute
+                return text_embedder.model
+    except:
+        pass
+
+    return None
 
 @with_tracer_wrapper
 def task_wrapper(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
@@ -90,30 +61,64 @@ def task_wrapper(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
     elif to_wrap.get("span_name"):
         name = to_wrap.get("span_name")
     else:
-        name = f"langchain.task.{instance.__class__.__name__}"
+        name = get_fully_qualified_class_name(instance)
+
+    if 'haystack.core.pipeline.pipeline' in to_wrap['package']:
+        embedding_model = get_embedding_model_haystack(instance)
+        set_embedding_model(embedding_model)
+        inputs = set()
+        workflow_input = get_workflow_input(args, inputs)
+        set_attribute(DATA_INPUT_KEY, workflow_input)
+
+    if to_wrap.get('skip_span'):
+        return_value = wrapped(*args, **kwargs)
+        botocore_processor(tracer, to_wrap, wrapped, instance, args, kwargs, return_value)
+        return return_value
 
     with tracer.start_as_current_span(name) as span:
-        process_span(to_wrap, span, instance, args)
         pre_task_processing(to_wrap, instance, args, span)
         return_value = wrapped(*args, **kwargs)
+        process_span(to_wrap, span, instance, args, kwargs, return_value)
         post_task_processing(to_wrap, span, return_value)
 
     return return_value
 
+def botocore_processor(tracer, to_wrap, wrapped, instance, args, kwargs,return_value):
+    if kwargs.get("service_name") == "sagemaker-runtime":
+        return_value.invoke_endpoint = _instrumented_endpoint_invoke(to_wrap,return_value,return_value.invoke_endpoint,tracer)
+
+def _instrumented_endpoint_invoke(to_wrap, instance, fn, tracer):
+    @wraps(fn)
+    def with_instrumentation(*args, **kwargs):
+
+        with tracer.start_as_current_span("botocore-sagemaker-invoke-endpoint") as span:
+            response = fn(*args, **kwargs)
+            process_span(to_wrap, span, instance=instance,args=args, kwargs=kwargs, return_value=response)
+            return response
+
+    return with_instrumentation
+
+def get_workflow_input(args, inputs):
+    if args is not None and len(args) > 0:
+        for value in args[0].values():
+            for text in value.values():
+                inputs.add(text)
+
+    workflow_input: str = ""
+
+    if inputs is not None and len(inputs) > 0:
+        for input_str in inputs:
+            workflow_input = workflow_input + input_str
+    return workflow_input
 
-def process_span(to_wrap, span, instance, args):
+def process_span(to_wrap, span, instance, args, kwargs, return_value):
     # Check if the output_processor is a valid JSON (in Python, that means it's a dictionary)
-    span_index = 1
+    instance_args = {}
+    set_provider_name(instance, instance_args)
+    span_index = 0
     if is_root_span(span):
-        workflow_name = get_workflow_name(span)
-        if workflow_name:
-            span.set_attribute(f"entity.{span_index}.name", workflow_name)
-        # workflow type
-        package_name = to_wrap.get('package')
-        for (package, workflow_type) in WORKFLOW_TYPE_MAP.items():
-            if (package_name is not None and package in package_name):
-                span.set_attribute(f"entity.{span_index}.type", workflow_type)
-        span_index += 1
+        span_index += set_workflow_attributes(to_wrap, span, span_index+1)
+        span_index += set_app_hosting_identifier_attribute(span, span_index+1)
     if 'output_processor' in to_wrap:
         output_processor=to_wrap['output_processor']
         if isinstance(output_processor, dict) and len(output_processor) > 0:
@@ -121,20 +126,17 @@ def process_span(to_wrap, span, instance, args):
                 span.set_attribute("span.type", output_processor['type'])
             else:
                 logger.warning("type of span not found or incorrect written in entity json")
-            count = 0
             if 'attributes' in output_processor:
-                count = len(output_processor["attributes"])
-                span.set_attribute("entity.count", count)
-                span_index = 1
                 for processors in output_processor["attributes"]:
                     for processor in processors:
                         attribute = processor.get('attribute')
                         accessor = processor.get('accessor')
 
                         if attribute and accessor:
-                            attribute_name = f"entity.{span_index}.{attribute}"
+                            attribute_name = f"entity.{span_index+1}.{attribute}"
                             try:
-                                result = eval(accessor)(instance, args)
+                                arguments = {"instance":instance, "args":args, "kwargs":kwargs, "output":return_value}
+                                result = eval(accessor)(arguments)
                                 if result and isinstance(result, str):
                                     span.set_attribute(attribute_name, result)
                             except Exception as e:
@@ -144,11 +146,54 @@ def process_span(to_wrap, span, instance, args):
                     span_index += 1
             else:
                 logger.warning("attributes not found or incorrect written in entity json")
-                span.set_attribute("span.count", count)
+            if 'events' in output_processor:
+                events = output_processor['events']
+                arguments = {"instance": instance, "args": args, "kwargs": kwargs, "output": return_value}
+                accessor_mapping = {
+                    "arguments": arguments,
+                    "response": return_value
+                }
+                for event in events:
+                    event_name = event.get("name")
+                    event_attributes = {}
+                    attributes = event.get("attributes", [])
+                    for attribute in attributes:
+                        attribute_key = attribute.get("attribute")
+                        accessor = attribute.get("accessor")
+                        if accessor:
+                            try:
+                                accessor_function = eval(accessor)
+                                for keyword, value in accessor_mapping.items():
+                                    if keyword in accessor:
+                                        evaluated_val = accessor_function(value)
+                                        if isinstance(evaluated_val, list):
+                                            evaluated_val = [str(d) for d in evaluated_val]
+                                        event_attributes[attribute_key] = evaluated_val
+                            except Exception as e:
+                                logger.error(f"Error evaluating accessor for attribute '{attribute_key}': {e}")
+                    span.add_event(name=event_name, attributes=event_attributes)
 
         else:
             logger.warning("empty or entities json is not in correct format")
-
+    if span_index > 0:
+        span.set_attribute("entity.count", span_index)
+
+def set_workflow_attributes(to_wrap, span: Span, span_index):
+    return_value = 1
+    workflow_name = get_workflow_name(span=span)
+    if workflow_name:
+        span.set_attribute("span.type", "workflow")
+        span.set_attribute(f"entity.{span_index}.name", workflow_name)
+        # workflow type
+    package_name = to_wrap.get('package')
+    workflow_type_set = False
+    for (package, workflow_type) in WORKFLOW_TYPE_MAP.items():
+        if (package_name is not None and package in package_name):
+            span.set_attribute(f"entity.{span_index}.type", workflow_type)
+            workflow_type_set = True
+    if not workflow_type_set:
+        span.set_attribute(f"entity.{span_index}.type", "workflow.generic")
+    return return_value
 
 def post_task_processing(to_wrap, span, return_value):
     try:
@@ -163,9 +208,12 @@ def post_task_processing(to_wrap, span, return_value):
 def pre_task_processing(to_wrap, instance, args, span):
     try:
         if is_root_span(span):
+            try:
+                sdk_version = version("monocle_apptrace")
+                span.set_attribute("monocle_apptrace.version", sdk_version)
+            except:
+                logger.warning(f"Exception finding monocle-apptrace version.")
             update_span_with_prompt_input(to_wrap=to_wrap, wrapped_args=args, span=span)
-            update_span_with_infra_name(span, INFRA_SERVICE_KEY)
-
         update_span_with_context_input(to_wrap=to_wrap, wrapped_args=args, span=span)
     except:
         logger.exception("exception in pre_task_processing")
@@ -184,11 +232,18 @@ async def atask_wrapper(tracer, to_wrap, wrapped, instance, args, kwargs):
     elif to_wrap.get("span_name"):
         name = to_wrap.get("span_name")
     else:
-        name = f"langchain.task.{instance.__class__.__name__}"
+        name = get_fully_qualified_class_name(instance)
+    if 'haystack.core.pipeline.pipeline' in to_wrap['package']:
+        embedding_model = get_embedding_model_haystack(instance)
+        set_embedding_model(embedding_model)
+        inputs = set()
+        workflow_input = get_workflow_input(args, inputs)
+        set_attribute(DATA_INPUT_KEY, workflow_input)
+
     with tracer.start_as_current_span(name) as span:
-        process_span(to_wrap, span, instance, args)
         pre_task_processing(to_wrap, instance, args, span)
         return_value = await wrapped(*args, **kwargs)
+        process_span(to_wrap, span, instance, args, kwargs, return_value)
         post_task_processing(to_wrap, span, return_value)
 
     return return_value
@@ -208,19 +263,12 @@ async def allm_wrapper(tracer, to_wrap, wrapped, instance, args, kwargs):
     elif to_wrap.get("span_name"):
         name = to_wrap.get("span_name")
     else:
-        name = f"langchain.task.{instance.__class__.__name__}"
+        name = get_fully_qualified_class_name(instance)
     with tracer.start_as_current_span(name) as span:
-        if 'haystack.components.retrievers' in to_wrap['package'] and 'haystack.retriever' in span.name:
-            input_arg_text = get_attribute(DATA_INPUT_KEY)
-            span.add_event(DATA_INPUT_KEY, {QUERY: input_arg_text})
         provider_name, inference_endpoint = get_provider_name(instance)
-        instance_args = {"provider_name": provider_name, "inference_endpoint": inference_endpoint}
-
-        process_span(to_wrap, span, instance, instance_args)
-
         return_value = await wrapped(*args, **kwargs)
-        if 'haystack.components.retrievers' in to_wrap['package'] and 'haystack.retriever' in span.name:
-            update_span_with_context_output(to_wrap=to_wrap, return_value=return_value, span=span)
+        kwargs.update({"provider_name": provider_name, "inference_endpoint": inference_endpoint or getattr(instance, 'endpoint', None)})
+        process_span(to_wrap, span, instance, args, kwargs, return_value)
         update_span_from_llm_response(response=return_value, span=span, instance=instance)
 
     return return_value
@@ -240,20 +288,13 @@ def llm_wrapper(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
     elif to_wrap.get("span_name"):
         name = to_wrap.get("span_name")
     else:
-        name = f"langchain.task.{instance.__class__.__name__}"
+        name = get_fully_qualified_class_name(instance)
 
     with tracer.start_as_current_span(name) as span:
-        if 'haystack.components.retrievers' in to_wrap['package'] and 'haystack.retriever' in span.name:
-            input_arg_text = get_attribute(DATA_INPUT_KEY)
-            span.add_event(DATA_INPUT_KEY, {QUERY: input_arg_text})
         provider_name, inference_endpoint = get_provider_name(instance)
-        instance_args = {"provider_name": provider_name, "inference_endpoint": inference_endpoint}
-
-        process_span(to_wrap, span, instance, instance_args)
-
         return_value = wrapped(*args, **kwargs)
-        if 'haystack.components.retrievers' in to_wrap['package'] and 'haystack.retriever' in span.name:
-            update_span_with_context_output(to_wrap=to_wrap, return_value=return_value, span=span)
+        kwargs.update({"provider_name": provider_name, "inference_endpoint": inference_endpoint or getattr(instance, 'endpoint', None)})
+        process_span(to_wrap, span, instance, args, kwargs, return_value)
         update_span_from_llm_response(response=return_value, span=span, instance=instance)
 
     return return_value
@@ -292,13 +333,46 @@ def update_llm_endpoint(curr_span: Span, instance):
 def get_provider_name(instance):
     provider_url = ""
     inference_endpoint = ""
+    parsed_provider_url = ""
+    try:
+        base_url = getattr(instance.client._client, "base_url", None)
+        if base_url:
+            if isinstance(getattr(base_url, "host", None), str):
+                provider_url = base_url.host
+            inference_endpoint = base_url if isinstance(base_url, str) else str(base_url)
+    except:
+        pass
+
+    try:
+        if isinstance(instance.client.meta.endpoint_url, str):
+            inference_endpoint = instance.client.meta.endpoint_url
+    except:
+        pass
+
+    api_base = getattr(instance, "api_base", None)
+    if isinstance(api_base, str):
+        provider_url = api_base
+
+    # Handle inference endpoint for Mistral AI (llamaindex)
+    sdk_config = getattr(instance, "_client", None)
+    if sdk_config and hasattr(sdk_config, "sdk_configuration"):
+        inference_endpoint = getattr(sdk_config.sdk_configuration, "server_url", inference_endpoint)
+
+    if provider_url:
+        try:
+            parsed_provider_url = urlparse(provider_url)
+        except:
+            pass
+
+    return parsed_provider_url.hostname if parsed_provider_url else provider_url, inference_endpoint
+
+
+def set_provider_name(instance, instance_args: dict):
+    provider_url = ""
+    parsed_provider_url = ""
     try:
         if isinstance(instance.client._client.base_url.host, str):
             provider_url = instance.client._client.base_url.host
-        if isinstance(instance.client._client.base_url, str):
-            inference_endpoint = instance.client._client.base_url
-        else:
-            inference_endpoint = str(instance.client._client.base_url)
     except:
         pass
 
@@ -307,13 +381,13 @@ def get_provider_name(instance):
             provider_url = instance.api_base
     except:
         pass
-
     try:
         if len(provider_url) > 0:
-            parsed_provider_url = urlparse(provider_url)
+            parsed_provider_url = urlparse(provider_url).hostname
     except:
         pass
-    return parsed_provider_url.hostname or provider_url,inference_endpoint
+    if parsed_provider_url or provider_url:
+        instance_args[PROVIDER] = parsed_provider_url or provider_url
 
 
 def is_root_span(curr_span: Span) -> bool:
@@ -327,16 +401,25 @@ def get_input_from_args(chain_args):
 
 
 def update_span_from_llm_response(response, span: Span, instance):
-    # extract token uasge from langchain openai
-    if (response is not None and hasattr(response, "response_metadata")):
-        response_metadata = response.response_metadata
-        token_usage = response_metadata.get("token_usage")
+    if (response is not None and isinstance(response, dict) and "meta" in response) or (
+            response is not None and hasattr(response, "response_metadata")):
+        token_usage = None
+        if (response is not None and isinstance(response, dict) and "meta" in response):  # haystack
+            token_usage = response["meta"][0]["usage"]
+
+        if (response is not None and hasattr(response, "response_metadata")):
+            if hasattr(response, "usage_metadata") and response.usage_metadata is not None:
+                token_usage = response.usage_metadata
+            else:
+                response_metadata = response.response_metadata
+                token_usage = response_metadata.get("token_usage")
+
         meta_dict = {}
         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")})
-            meta_dict.update({"prompt_tokens": token_usage.get("prompt_tokens")})
+            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")})
             span.add_event(META_DATA, meta_dict)
     # extract token usage from llamaindex openai
@@ -401,17 +484,28 @@ def update_span_with_context_output(to_wrap, return_value, span: Span):
 def update_span_with_prompt_input(to_wrap, wrapped_args, span: Span):
     input_arg_text = wrapped_args[0]
 
-    if isinstance(input_arg_text, dict):
-        span.add_event(PROMPT_INPUT_KEY, input_arg_text)
+    prompt_inputs = get_nested_value(input_arg_text, ['prompt_builder', 'question'])
+    if prompt_inputs is not None:  # haystack
+        span.add_event(PROMPT_INPUT_KEY, {QUERY: prompt_inputs})
+    elif isinstance(input_arg_text, dict):
+        span.add_event(PROMPT_INPUT_KEY, {QUERY: input_arg_text['input']})
     else:
         span.add_event(PROMPT_INPUT_KEY, {QUERY: input_arg_text})
 
 
 def update_span_with_prompt_output(to_wrap, wrapped_args, span: Span):
     package_name: str = to_wrap.get('package')
-    if isinstance(wrapped_args, str):
-        span.add_event(PROMPT_OUTPUT_KEY, {RESPONSE: wrapped_args})
-    if isinstance(wrapped_args, dict):
-        span.add_event(PROMPT_OUTPUT_KEY, wrapped_args)
+
     if "llama_index.core.base.base_query_engine" in package_name:
         span.add_event(PROMPT_OUTPUT_KEY, {RESPONSE: wrapped_args.response})
+    elif "haystack.core.pipeline.pipeline" in package_name:
+        resp = get_nested_value(wrapped_args, ['llm', 'replies'])
+        if resp is not None:
+            if isinstance(resp, list) and hasattr(resp[0], 'content'):
+                span.add_event(PROMPT_OUTPUT_KEY, {RESPONSE: resp[0].content})
+            else:
+                span.add_event(PROMPT_OUTPUT_KEY, {RESPONSE: resp[0]})
+    elif isinstance(wrapped_args, str):
+        span.add_event(PROMPT_OUTPUT_KEY, {RESPONSE: wrapped_args})
+    elif isinstance(wrapped_args, dict):
+        span.add_event(PROMPT_OUTPUT_KEY, wrapped_args)

monocle_apptrace/wrapper.py

@@ -2,6 +2,7 @@
 from monocle_apptrace.haystack import HAYSTACK_METHODS
 from monocle_apptrace.langchain import LANGCHAIN_METHODS
 from monocle_apptrace.llamaindex import LLAMAINDEX_METHODS
+from monocle_apptrace.botocore import BOTOCORE_METHODS
 from monocle_apptrace.wrap_common import task_wrapper
 
 # pylint: disable=too-few-public-methods
@@ -23,4 +24,4 @@ class WrapperMethod:
 
         self.wrapper = wrapper
 
-INBUILT_METHODS_LIST = LANGCHAIN_METHODS + LLAMAINDEX_METHODS + HAYSTACK_METHODS
+INBUILT_METHODS_LIST = LANGCHAIN_METHODS + LLAMAINDEX_METHODS + HAYSTACK_METHODS + BOTOCORE_METHODS

{monocle_apptrace-0.2.0.dist-info → monocle_apptrace-0.3.0b2.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: monocle_apptrace
-Version: 0.2.0
+Version: 0.3.0b2
 Summary: package with monocle genAI tracing
 Project-URL: Homepage, https://github.com/monocle2ai/monocle
 Project-URL: Issues, https://github.com/monocle2ai/monocle/issues
@@ -25,11 +25,14 @@ Requires-Dist: faiss-cpu==1.8.0; extra == 'dev'
 Requires-Dist: instructorembedding==1.0.1; extra == 'dev'
 Requires-Dist: langchain-chroma==0.1.1; extra == 'dev'
 Requires-Dist: langchain-community==0.2.5; extra == 'dev'
+Requires-Dist: langchain-mistralai==0.1.13; extra == 'dev'
 Requires-Dist: langchain-openai==0.1.8; extra == 'dev'
 Requires-Dist: langchain==0.2.5; extra == 'dev'
 Requires-Dist: llama-index-embeddings-huggingface==0.2.0; extra == 'dev'
+Requires-Dist: llama-index-llms-mistralai==0.1.20; extra == 'dev'
 Requires-Dist: llama-index-vector-stores-chroma==0.1.9; extra == 'dev'
 Requires-Dist: llama-index==0.10.30; extra == 'dev'
+Requires-Dist: mistral-haystack==0.0.2; extra == 'dev'
 Requires-Dist: numpy==1.26.4; extra == 'dev'
 Requires-Dist: parameterized==0.9.0; extra == 'dev'
 Requires-Dist: pytest==8.0.0; extra == 'dev'