monocle-apptrace 0.3.0b1__py3-none-any.whl → 0.3.0b3__py3-none-any.whl

monocle_apptrace/instrumentation/common/span_handler.py

@@ -0,0 +1,154 @@
+import logging
+import os
+from importlib.metadata import version
+from opentelemetry.context import get_current
+from opentelemetry.context import get_value
+from opentelemetry.sdk.trace import Span
+
+from monocle_apptrace.instrumentation.common.constants import (
+    QUERY,
+    service_name_map,
+    service_type_map,
+)
+from monocle_apptrace.instrumentation.common.utils import set_attribute
+
+logger = logging.getLogger(__name__)
+
+WORKFLOW_TYPE_MAP = {
+    "llama_index": "workflow.llamaindex",
+    "langchain": "workflow.langchain",
+    "haystack": "workflow.haystack"
+}
+
+
+
+class SpanHandler:
+
+    def validate(self, to_wrap, wrapped, instance, args, kwargs):
+        pass
+
+    def pre_task_processing(self, to_wrap, wrapped, instance, args, span):
+        if self.__is_root_span(span):
+            try:
+                sdk_version = version("monocle_apptrace")
+                span.set_attribute("monocle_apptrace.version", sdk_version)
+            except Exception as e:
+                logger.warning("Exception finding monocle-apptrace version.")
+        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)
+
+    def hydrate_attributes(self, to_wrap, wrapped, instance, args, kwargs, result, span):
+        span_index = 0
+        if self.__is_root_span(span):
+            span_index += self.set_workflow_attributes(to_wrap, span, span_index+1)
+            span_index += self.set_app_hosting_identifier_attribute(span, span_index+1)
+
+        if 'output_processor' in to_wrap and to_wrap["output_processor"] is not None:    
+            output_processor=to_wrap['output_processor']
+            if 'type' in output_processor:
+                        span.set_attribute("span.type", output_processor['type'])
+            else:
+                logger.warning("type of span not found or incorrect written in entity json")
+            if 'attributes' in output_processor:
+                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+1}.{attribute}"
+                            try:
+                                arguments = {"instance":instance, "args":args, "kwargs":kwargs, "result":result}
+                                result = accessor(arguments)
+                                if result and isinstance(result, str):
+                                    span.set_attribute(attribute_name, result)
+                            except Exception as e:
+                                logger.error(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")
+                    span_index += 1
+            else:
+                logger.warning("attributes not found or incorrect written in entity json")
+
+        if span_index > 0:
+            span.set_attribute("entity.count", span_index)
+
+
+    def hydrate_events(self, to_wrap, wrapped, instance, args, kwargs, result, span):
+        if 'output_processor' in to_wrap and to_wrap["output_processor"] is not None:
+            output_processor=to_wrap['output_processor']
+            arguments = {"instance": instance, "args": args, "kwargs": kwargs, "result": result}
+            if 'events' in output_processor:
+                events = output_processor['events']
+                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:
+                                if attribute_key is not None:
+                                    event_attributes[attribute_key] = accessor(arguments)
+                                else:
+                                    event_attributes.update(accessor(arguments))
+                            except Exception as e:
+                                logger.error(f"Error evaluating accessor for attribute '{attribute_key}': {e}")
+                    span.add_event(name=event_name, attributes=event_attributes)
+
+
+
+    def set_workflow_attributes(self, to_wrap, span: Span, span_index):
+        return_value = 1
+        workflow_name = self.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 set_app_hosting_identifier_attribute(self, 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 get_workflow_name(self, 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
+
+    def __is_root_span(self, curr_span: Span) -> bool:
+        try:
+            if curr_span is not None and hasattr(curr_span, "parent"):
+                return curr_span.parent is None or get_current().get("root_span_id") == curr_span.parent.span_id
+        except Exception as e:
+            logger.warning(f"Error finding root span: {e}")

monocle_apptrace/instrumentation/common/utils.py

@@ -0,0 +1,171 @@
+import logging
+from typing import Callable, Generic, Optional, TypeVar
+
+from opentelemetry.context import attach, detach, get_current, get_value, set_value
+from opentelemetry.trace import NonRecordingSpan, Span
+from opentelemetry.trace.propagation import _SPAN_KEY
+
+T = TypeVar('T')
+U = TypeVar('U')
+
+logger = logging.getLogger(__name__)
+
+embedding_model_context = {}
+
+def set_span_attribute(span, name, value):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+
+def dont_throw(func):
+    """
+    A decorator that wraps the passed in function and logs exceptions instead of throwing them.
+
+    @param func: The function to wrap
+    @return: The wrapper function
+    """
+    # Obtain a logger specific to the function's module
+    logger = logging.getLogger(func.__module__)
+
+    # pylint: disable=inconsistent-return-statements
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as ex:
+            logger.warning("Failed to execute %s, error: %s", func.__name__, str(ex))
+
+    return wrapper
+
+def with_tracer_wrapper(func):
+    """Helper for providing tracer for wrapper functions."""
+
+    def _with_tracer(tracer, handler, to_wrap):
+        def wrapper(wrapped, instance, args, kwargs):
+            try:
+                # get and log the parent span context if injected by the application
+                # This is useful for debugging and tracing of Azure functions
+                _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_span = isinstance(parent_span, NonRecordingSpan)
+                    if is_span:
+                        logger.debug(
+                            f"Parent span is found with trace id {hex(parent_span.get_span_context().trace_id)}")
+            except Exception as e:
+                logger.error("Exception in attaching parent context: %s", e)
+
+            val = func(tracer, handler, to_wrap, wrapped, instance, args, kwargs)
+            return val
+
+        return wrapper
+
+    return _with_tracer
+
+def resolve_from_alias(my_map, alias):
+    """Find a alias that is not none from list of aliases"""
+
+    for i in alias and my_map[i] is not None:
+        if i in my_map.keys():
+            return my_map[i]
+    return None
+
+def set_embedding_model(model_name: str):
+    """
+    Sets the embedding model in the global context.
+
+    @param model_name: The name of the embedding model to set
+    """
+    embedding_model_context['embedding_model'] = model_name
+
+def get_embedding_model() -> str:
+    """
+    Retrieves the embedding model from the global context.
+
+    @return: The name of the embedding model, or 'unknown' if not set
+    """
+    return embedding_model_context.get('embedding_model', 'unknown')
+
+def set_attribute(key: str, value: str):
+    """
+    Set a value in the global context for a given key.
+
+    Args:
+        key: The key for the context value to set.
+        value: The value to set for the given key.
+    """
+    attach(set_value(key, value))
+
+def get_attribute(key: str) -> str:
+    """
+    Retrieve a value from the global context for a given key.
+
+    Args:
+        key: The key for the context value to retrieve.
+
+    Returns:
+        The value associated with the given key.
+    """
+    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_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)
+
+
+class Option(Generic[T]):
+    def __init__(self, value: Optional[T]):
+        self.value = value
+
+    def is_some(self) -> bool:
+        return self.value is not None
+
+    def is_none(self) -> bool:
+        return self.value is None
+
+    def unwrap_or(self, default: T) -> T:
+        return self.value if self.is_some() else default
+
+    def map(self, func: Callable[[T], U]) -> 'Option[U]':
+        if self.is_some():
+            return Option(func(self.value))
+        return Option(None)
+
+    def and_then(self, func: Callable[[T], 'Option[U]']) -> 'Option[U]':
+        if self.is_some():
+            return func(self.value)
+        return Option(None)
+
+# Example usage
+def try_option(func: Callable[..., T], *args, **kwargs) -> Option[T]:
+    try:
+        return Option(func(*args, **kwargs))
+    except Exception:
+        return Option(None)

monocle_apptrace/instrumentation/common/wrapper.py

@@ -0,0 +1,69 @@
+# pylint: disable=protected-access
+import logging
+
+from opentelemetry.trace import Tracer
+
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.common.utils import (
+    get_fully_qualified_class_name,
+    with_tracer_wrapper,
+)
+from monocle_apptrace.instrumentation.metamodel.botocore import _helper
+logger = logging.getLogger(__name__)
+
+
+@with_tracer_wrapper
+def task_wrapper(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)
+
+    if hasattr(instance, "name") and instance.name:
+        name = f"{to_wrap.get('span_name')}.{instance.name.lower()}"
+    elif to_wrap.get("span_name"):
+        name = to_wrap.get("span_name")
+    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 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
+
+
+@with_tracer_wrapper
+async def atask_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, args, kwargs):
+    """Instruments and calls every function defined in TO_WRAP."""
+
+    # Some Langchain objects are wrapped elsewhere, so we ignore them here
+    if instance.__class__.__name__ in ("AgentExecutor"):
+        return wrapped(*args, **kwargs)
+
+    if hasattr(instance, "name") and instance.name:
+        name = f"{to_wrap.get('span_name')}.{instance.name.lower()}"
+    elif to_wrap.get("span_name"):
+        name = to_wrap.get("span_name")
+    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 return_value

monocle_apptrace/instrumentation/common/wrapper_method.py

@@ -0,0 +1,45 @@
+# pylint: disable=too-few-public-methods
+from monocle_apptrace.instrumentation.common.wrapper import task_wrapper
+from monocle_apptrace.instrumentation.metamodel.botocore.methods import BOTOCORE_METHODS
+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, )
+
+
+class WrapperMethod:
+    def __init__(
+            self,
+            package: str,
+            object_name: str,
+            method: str,
+            span_name: str = None,
+            output_processor : str = None,
+            wrapper_method = task_wrapper,
+            span_handler = 'default'
+            ):
+        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
+
+    def to_dict(self) -> dict:
+        # Create a dictionary representation of the instance
+        instance_dict = {
+            'package': self.package,
+            'object': self.object,
+            'method': self.method,
+            'span_name': self.span_name,
+            'output_processor': self.output_processor,
+            'wrapper_method': self.wrapper_method,
+            'span_handler': self.span_handler
+        }
+        return instance_dict
+
+
+DEFAULT_METHODS_LIST = LANGCHAIN_METHODS + LLAMAINDEX_METHODS + HAYSTACK_METHODS + BOTOCORE_METHODS

monocle_apptrace/instrumentation/metamodel/botocore/_helper.py

@@ -0,0 +1,126 @@
+"""
+This module provides utility functions for extracting system, user,
+and assistant messages from various input formats.
+"""
+
+import logging
+import json
+from io import BytesIO
+from functools import wraps
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+
+logger = logging.getLogger(__name__)
+
+
+def extract_messages(args):
+    """Extract system and user messages"""
+    try:
+        messages = []
+        if args and isinstance(args, dict) and len(args) > 0:
+            if 'Body' in args and isinstance(args['Body'], str):
+                data = json.loads(args['Body'])
+                question = data.get("question")
+                messages.append(question)
+            if 'messages' in args and isinstance(args['messages'], list):
+                role = args['messages'][0]['role']
+                user_message = extract_query_from_content(args['messages'][0]['content'][0]['text'])
+                messages.append({role: user_message})
+        return [str(d) for d in messages]
+    except Exception as e:
+        logger.warning("Warning: Error occurred in extract_messages: %s", str(e))
+        return []
+
+
+def extract_assistant_message(response):
+    try:
+        if "Body" in response and hasattr(response['Body'], "_raw_stream"):
+            raw_stream = getattr(response['Body'], "_raw_stream")
+            if hasattr(raw_stream, "data"):
+                response_bytes = getattr(raw_stream, "data")
+                response_str = response_bytes.decode('utf-8')
+                response_dict = json.loads(response_str)
+                response['Body'] = BytesIO(response_bytes)
+                return [response_dict["answer"]]
+        if "output" in response:
+            output = response.get("output", {})
+            message = output.get("message", {})
+            content = message.get("content", [])
+            if isinstance(content, list) and len(content) > 0 and "text" in content[0]:
+                reply = content[0]["text"]
+                return [reply]
+    except Exception as e:
+        logger.warning("Warning: Error occurred in extract_assistant_message: %s", str(e))
+        return []
+
+
+def extract_query_from_content(content):
+    try:
+        query_prefix = "Query:"
+        answer_prefix = "Answer:"
+        query_start = content.find(query_prefix)
+        if query_start == -1:
+            return None
+
+        query_start += len(query_prefix)
+        answer_start = content.find(answer_prefix, query_start)
+        if answer_start == -1:
+            query = content[query_start:].strip()
+        else:
+            query = content[query_start:answer_start].strip()
+        return query
+    except Exception as e:
+        logger.warning("Warning: Error occurred in extract_query_from_content: %s", str(e))
+        return ""
+
+
+def resolve_from_alias(my_map, alias):
+    """Find a alias that is not none from list of aliases"""
+
+    for i in alias:
+        if i in my_map.keys():
+            return my_map[i]
+    return None
+
+
+def 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:
+        token_usage = response["usage"]
+
+        if token_usage is not None:
+            temperature = instance.__dict__.get("temperature", None)
+            meta_dict.update({"temperature": temperature})
+            meta_dict.update({"completion_tokens": resolve_from_alias(token_usage,["completion_tokens","output_tokens","outputTokens"])})
+            meta_dict.update({"prompt_tokens": resolve_from_alias(token_usage,["prompt_tokens","input_tokens","inputTokens"])})
+            meta_dict.update({"total_tokens": resolve_from_alias(token_usage,["total_tokens","totalTokens"])})
+    return meta_dict

monocle_apptrace/instrumentation/metamodel/botocore/entities/inference.py

@@ -0,0 +1,65 @@
+from monocle_apptrace.instrumentation.metamodel.botocore import (
+    _helper,
+)
+
+INFERENCE = {
+    "type": "inference",
+    "attributes": [
+        [
+            {
+                "_comment": "provider type  , inference_endpoint",
+                "attribute": "type",
+                "accessor": lambda arguments: 'inference.aws_sagemaker'
+            },
+            {
+                "attribute": "inference_endpoint",
+                "accessor": lambda arguments: arguments['instance'].meta.endpoint_url
+            }
+        ],
+        [
+            {
+                "_comment": "LLM Model",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.resolve_from_alias(arguments['kwargs'],
+                                                                         ['EndpointName', 'modelId'])
+            },
+            {
+                "attribute": "type",
+                "accessor": lambda arguments: 'model.llm.' + _helper.resolve_from_alias(arguments['kwargs'],
+                                                                                        ['EndpointName', 'modelId'])
+            }
+        ]
+    ],
+    "events": [
+        {"name": "data.input",
+         "attributes": [
+
+             {
+                 "_comment": "this is instruction and user query to LLM",
+                 "attribute": "input",
+                 "accessor": lambda arguments: _helper.extract_messages(arguments['kwargs'])
+             }
+         ]
+         },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "this is response from LLM",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_assistant_message(arguments['result'])
+                }
+            ]
+        },
+        {
+            "name": "metadata",
+            "attributes": [
+                {
+                    "_comment": "this is metadata usage from LLM",
+                    "accessor": lambda arguments: _helper.update_span_from_llm_response(arguments['result'],
+                                                                                        arguments['instance'])
+                }
+            ]
+        }
+    ]
+}

monocle_apptrace/instrumentation/metamodel/botocore/methods.py

@@ -0,0 +1,16 @@
+from monocle_apptrace.instrumentation.common.wrapper import task_wrapper
+from monocle_apptrace.instrumentation.metamodel.botocore.entities.inference import (
+    INFERENCE,
+)
+
+BOTOCORE_METHODS = [{
+
+      "package": "botocore.client",
+      "object": "ClientCreator",
+      "method": "create_client",
+      "wrapper_method": task_wrapper,
+      "skip_span": True,
+      "output_processor": INFERENCE
+
+}
+]