monocle-apptrace 0.0.1__py3-none-any.whl → 0.1.1__py3-none-any.whl

monocle_apptrace/metamodel/spans/span_example.json

@@ -0,0 +1,140 @@
+{
+    "name": "llamaindex.retrieve",
+    "context": {
+        "trace_id": "0x93cd0bf865b3ffcc3cf9c075dc3e3797",
+        "span_id": "0x5d3f839e900bda24",
+        "trace_state": "[]"
+    },
+    "kind": "SpanKind.CLIENT",
+    "parent_id": "0x7a63d63e42ccac60",
+    "start_time": "2024-09-09T14:38:45.237182Z",
+    "end_time": "2024-09-09T14:38:45.620112Z",
+    "status": {
+        "status_code": "OK"
+    },
+    "attributes": {
+        "span.type": "Retrieval", 
+        "entity.count": 2,
+        "entity.1.name": "ChromaVectorStore",
+        "entity.1.type": "vectorstore.chroma",
+        "entity.1.embedding-model-name": "BAAI/bge-small-en-v1.5",
+        "entity.2.name": "BAAI/bge-small-en-v1.5",
+        "entity.2.type": "model.embedding",
+        "entity.2.model_name": "BAAI/bge-small-en-v1.5"
+    },
+    "events": [
+        {
+            "name": "data.input",
+            "timestamp": "timestamp",
+            "attributes": {
+                "context_input": "question: What is an americano?"
+            }
+        },
+        {
+            "name": "data.output",
+            "timestamp": "timestamp",
+            "attributes": {
+                "context_output": "Coffee is a hot drink made from the roasted and ground seeds (coffee beans) of a tropical shrub\nA latte consists of one or more shots of espresso, served in a glass (or sometimes a cup), into which hot steamed milk is added\nAmericano is a type of coffee drink prepared by diluting an espresso shot with hot water at a 1:3 to 1:4 ratio, resulting in a drink that retains the complex flavors of espresso, but in a lighter way"
+            }
+        }
+
+    ],
+    "links": [],
+    "resource": {
+        "attributes": {
+            "service.name": "coffee-bot"
+        },
+        "schema_url": ""
+    }
+},
+{
+    "name": "llamaindex.openai",
+    "context": {
+        "trace_id": "0x93cd0bf865b3ffcc3cf9c075dc3e3797",
+        "span_id": "0x8b6363e1937a4d7b",
+        "trace_state": "[]"
+    },
+    "kind": "SpanKind.CLIENT",
+    "parent_id": "0x7a63d63e42ccac60",
+    "start_time": "2024-09-09T14:38:45.622174Z",
+    "end_time": "2024-09-09T14:38:46.514120Z",
+    "status": {
+        "status_code": "OK"
+    },
+    "attributes": {
+        "span.type": "inference",
+        "entity.count": 2,
+        "entity.1.name": "AzureOpenAI",
+        "entity.1.type": "inference.azure_oai",
+        "entity.1.provider_name": "openai.azure.com",
+        "entity.1.deployment": "kshitiz-gpt",
+        "entity.1.inference_endpoint": "https://okahu-openai-dev.openai.azure.com/",
+
+        "entity.2.name": "gpt-35-turbo",
+        "entity.2.type": "model.llm",
+        "entity.2.model_name": "gpt-35-turbo"
+    },
+    "events": [
+        {
+            "name": "data.input",
+            "timestamp": "timestamp",
+            "attributes": {
+                "question": "What is an americano?",
+            }
+        },
+        {
+            "name": "data.output",
+            "timestamp": "timestamp",
+            "attributes": {
+                "response": "An americano is a type of coffee drink that is made by diluting an espresso shot with hot water at a 1:3 to 1:4 ratio, resulting in a drink that retains the complex flavors of espresso, but in a lighter way.",
+            }
+        },
+        {
+            "name": "metadata",
+            "timestamp": "timestamp",
+            "attributes": {
+                "temperature": 0.1,
+                "completion_tokens": 52,
+                "prompt_tokens": 233,
+                "total_tokens": 285
+            }
+        }
+    ],
+    "links": [],
+    "resource": {
+        "attributes": {
+            "service.name": "coffee-bot"
+        },
+        "schema_url": ""
+    }
+}
+{
+    "name": "llamaindex.query",
+    "context": {
+        "trace_id": "0x93cd0bf865b3ffcc3cf9c075dc3e3797",
+        "span_id": "0x7a63d63e42ccac60",
+        "trace_state": "[]"
+    },
+    "kind": "SpanKind.CLIENT",
+    "parent_id": null,
+    "start_time": "2024-09-09T14:38:45.236627Z",
+    "end_time": "2024-09-09T14:38:46.514442Z",
+    "status": {
+        "status_code": "OK"
+    },
+    "attributes": {
+        "span.type": "workflow",
+        "entity.count": 1,
+        "entity.1.name": "coffee-bot",
+        "entity.1.type": "workflow.llama_index"
+    },
+    "events": [
+    ],
+    "links": [],
+    "resource": {
+        "attributes": {
+            "service.name": "coffee-bot"
+        },
+        "schema_url": ""
+    }
+}

monocle_apptrace/metamodel/spans/span_format.json

@@ -0,0 +1,55 @@
+{
+    "name": "span-name",
+    "context": {
+        "trace_id": "trace-id",
+        "span_id": "span-id",
+        "trace_state": "[]"
+    },
+    "kind": "SpanKind.CLIENT",
+    "parent_id": "parent-id or None (for root span)",
+    "start_time": "UTC timestamp",
+    "end_time": "UTC timestamp",
+    "status": {
+        "status_code": "OK or Error"
+    },
+    "attributes": {
+        "description": "List of AI component entities used in this span, eg Model, Inference hosting service. Needs to be one of the supported entity types.",
+    
+        "span.type": "Monocle-span-type",
+        "entity.count": "count-of-entities",
+    
+        "entity.<index>.name": "Monocle-Entity-name",
+        "entity.<index>.type": "Monocle-Entity-Type",
+        "entity.<index>.<attribute>": "Value"
+    },
+    "events" : [
+        {
+            "name": "data.input",
+            "timestamp": "UTC timestamp",
+            "attributes": {
+                "input_attribute": "value"
+           }
+        },
+        {
+            "name": "data.output",
+            "timestamp": "UTC timestamp",
+            "attributes": {
+                "output_attribute": "value"
+            }
+        },
+        {
+            "name": "metadata",
+            "timestamp": "UTC timestamp",
+            "attributes": {
+                "metadata_attribute": "value"
+            }
+        }
+    ],
+    "links": [],
+    "resource": {
+        "attributes": {
+            "service.name": "top-workflow-name"
+        },
+        "schema_url": ""
+    }
+}

monocle_apptrace/utils.py

@@ -1,17 +1,16 @@
-
-
 import logging
-logger = logging.getLogger(__name__)
+import json
+from importlib import import_module
+import os
+from opentelemetry.trace import Span
+from monocle_apptrace.constants import azure_service_map, aws_service_map
 
-class Config:
-    exception_logger = None
+embedding_model_context = {}
 
 def set_span_attribute(span, name, value):
     if value is not None:
         if value != "":
             span.set_attribute(name, value)
-    return
-
 
 def dont_throw(func):
     """
@@ -22,15 +21,12 @@ def dont_throw(func):
     """
     # 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 e:
-            logger.warning("Failed to execute %s, error: %s", func.__name__, str(e))
-            if Config.exception_logger:
-                Config.exception_logger(e)
-
+        except Exception as ex:
+            logger.warning("Failed to execute %s, error: %s", func.__name__, str(ex))
     return wrapper
 
 def with_tracer_wrapper(func):
@@ -44,10 +40,54 @@ def with_tracer_wrapper(func):
 
     return _with_tracer
 
-def resolve_from_alias(map, alias):
+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 map.keys():
-            return map[i]
+        if i in my_map.keys():
+            return my_map[i]
     return None
+
+def load_wrapper_from_config(config_file_path: str, module_name: str = None):
+    wrapper_methods = []
+    with open(config_file_path, encoding='UTF-8') as config_file:
+        json_data = json.load(config_file)
+        wrapper_methods = json_data["wrapper_methods"]
+        for wrapper_method in wrapper_methods:
+            wrapper_method["wrapper"] = get_wrapper_method(
+                wrapper_method["wrapper_package"], wrapper_method["wrapper_method"])
+            if "span_name_getter_method" in wrapper_method :
+                wrapper_method["span_name_getter"] = get_wrapper_method(
+                    wrapper_method["span_name_getter_package"],
+                    wrapper_method["span_name_getter_method"])
+        return wrapper_methods
+
+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_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')

monocle_apptrace/wrap_common.py

@@ -1,9 +1,10 @@
+#pylint: disable=protected-access
 import logging
 import os
+from urllib.parse import urlparse
 
 from opentelemetry.trace import Span, Tracer
-
-from monocle_apptrace.utils import resolve_from_alias, with_tracer_wrapper
+from monocle_apptrace.utils import resolve_from_alias, update_span_with_infra_name, with_tracer_wrapper, get_embedding_model
 
 logger = logging.getLogger(__name__)
 WORKFLOW_TYPE_KEY = "workflow_type"
@@ -14,7 +15,12 @@ PROMPT_OUTPUT_KEY = "output"
 QUERY = "question"
 RESPONSE = "response"
 TAGS = "tags"
-CONTEXT_PROPERTIES_KEY = "workflow_context_properties"
+SESSION_PROPERTIES_KEY = "session"
+INFRA_SERVICE_KEY = "infra_service_name"
+TYPE = "type"
+PROVIDER = "provider_name"
+EMBEDDING_MODEL = "embedding_model"
+VECTOR_STORE = 'vector_store'
 
 
 WORKFLOW_TYPE_MAP = {
@@ -23,6 +29,24 @@ WORKFLOW_TYPE_MAP = {
     "haystack": "workflow.haystack"
 }
 
+framework_vector_store_mapping = {
+    'langchain_core.retrievers': lambda instance: {
+        'provider': instance.tags[0],
+        'embedding_model': instance.tags[1],
+        'type': VECTOR_STORE,
+    },
+    'llama_index.core.indices.base_retriever': lambda instance: {
+        'provider': type(instance._vector_store).__name__,
+        'embedding_model': instance._embed_model.model_name,
+        'type': VECTOR_STORE,
+    },
+    'haystack.components.retrievers': lambda instance: {
+        'provider': instance.__dict__.get("document_store").__class__.__name__,
+        'embedding_model': get_embedding_model(),
+        'type': VECTOR_STORE,
+    },
+}
+
 @with_tracer_wrapper
 def task_wrapper(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
     """Instruments and calls every function defined in TO_WRAP."""
@@ -37,29 +61,39 @@ def task_wrapper(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
         name = to_wrap.get("span_name")
     else:
         name = f"langchain.task.{instance.__class__.__name__}"
-    kind = to_wrap.get("kind")
 
     with tracer.start_as_current_span(name) as span:
-        if is_root_span(span):
-            update_span_with_prompt_input(to_wrap=to_wrap, wrapped_args=args, span=span)
-
-        #capture the tags attribute of the instance if present, else ignore 
-        try:
-            span.set_attribute(TAGS, getattr(instance, TAGS))
-        except AttributeError:
-            pass    
-        update_span_with_context_input(to_wrap=to_wrap, wrapped_args=args, span=span)
+        pre_task_processing(to_wrap, instance, args, span)
         return_value = wrapped(*args, **kwargs)
-        update_span_with_context_output(to_wrap=to_wrap, return_value=return_value, span=span)
-
-        if is_root_span(span):
-            workflow_name = span.resource.attributes.get("service.name")
-            span.set_attribute("workflow_name",workflow_name)
-            update_span_with_prompt_output(to_wrap=to_wrap, wrapped_args=return_value, span=span)
-            update_workflow_type(to_wrap, span)
+        post_task_processing(to_wrap, span, return_value)
 
     return return_value
 
+def post_task_processing(to_wrap, span, return_value):
+    update_span_with_context_output(to_wrap=to_wrap, return_value=return_value, span=span)
+
+    if is_root_span(span):
+        workflow_name = span.resource.attributes.get("service.name")
+        span.set_attribute("workflow_name",workflow_name)
+        update_span_with_prompt_output(to_wrap=to_wrap, wrapped_args=return_value, span=span)
+        update_workflow_type(to_wrap, span)
+
+def pre_task_processing(to_wrap, instance, args, span):
+    if is_root_span(span):
+        update_span_with_prompt_input(to_wrap=to_wrap, wrapped_args=args, span=span)
+
+        update_span_with_infra_name(span, INFRA_SERVICE_KEY)
+
+    #capture the tags attribute of the instance if present, else ignore
+    try:
+        update_tags(instance, span)
+        update_vectorstore_attributes(to_wrap, instance, span)
+    except AttributeError:
+        pass
+    update_span_with_context_input(to_wrap=to_wrap, wrapped_args=args, span=span)
+
+
+
 @with_tracer_wrapper
 async def atask_wrapper(tracer, to_wrap, wrapped, instance, args, kwargs):
     """Instruments and calls every function defined in TO_WRAP."""
@@ -74,9 +108,10 @@ async def atask_wrapper(tracer, to_wrap, wrapped, instance, args, kwargs):
         name = to_wrap.get("span_name")
     else:
         name = f"langchain.task.{instance.__class__.__name__}"
-    kind = to_wrap.get("kind")
     with tracer.start_as_current_span(name) as span:
+        pre_task_processing(to_wrap, instance, args, span)
         return_value = await wrapped(*args, **kwargs)
+        post_task_processing(to_wrap, span, return_value)
 
     return return_value
 
@@ -86,7 +121,7 @@ async def allm_wrapper(tracer, to_wrap, wrapped, instance, args, kwargs):
     if instance.__class__.__name__ in ("AgentExecutor"):
         return wrapped(*args, **kwargs)
 
-    if to_wrap.get("span_name_getter"):
+    if callable(to_wrap.get("span_name_getter")):
         name = to_wrap.get("span_name_getter")(instance)
 
     elif hasattr(instance, "name") and instance.name:
@@ -95,11 +130,11 @@ async def allm_wrapper(tracer, to_wrap, wrapped, instance, args, kwargs):
         name = to_wrap.get("span_name")
     else:
         name = f"langchain.task.{instance.__class__.__name__}"
-    kind = to_wrap.get("kind")
     with tracer.start_as_current_span(name) as span:
         update_llm_endpoint(curr_span= span, instance=instance)
 
         return_value = await wrapped(*args, **kwargs)
+        update_span_from_llm_response(response = return_value, span = span)
 
     return return_value
 
@@ -119,8 +154,9 @@ def llm_wrapper(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
         name = to_wrap.get("span_name")
     else:
         name = f"langchain.task.{instance.__class__.__name__}"
-    kind = to_wrap.get("kind")
     with tracer.start_as_current_span(name) as span:
+        if 'haystack.components.retrievers' in to_wrap['package'] and 'haystack.retriever' in span.name:
+            update_vectorstore_attributes(to_wrap, instance, span)
         update_llm_endpoint(curr_span= span, instance=instance)
 
         return_value = wrapped(*args, **kwargs)
@@ -136,22 +172,45 @@ def update_llm_endpoint(curr_span: Span, instance):
         if 'temperature' in instance.__dict__:
             temp_val = instance.__dict__.get("temperature")
             curr_span.set_attribute("temperature", temp_val)
-        # handling for model name
-        model_name =  resolve_from_alias(instance.__dict__ , ["model","model_name"])
-        curr_span.set_attribute("openai_model_name", model_name)
+            # handling for model name
+        model_name = resolve_from_alias(instance.__dict__ , ["model","model_name"])
+        curr_span.set_attribute("model_name", model_name)
+        set_provider_name(curr_span, instance)
         # handling AzureOpenAI deployment
-        deployment_name =  resolve_from_alias(instance.__dict__ , [ "engine", "azure_deployment", 
+        deployment_name = resolve_from_alias(instance.__dict__ , [ "engine", "azure_deployment",
                                                                    "deployment_name", "deployment_id", "deployment"])
         curr_span.set_attribute("az_openai_deployment", deployment_name)
         # handling the inference endpoint
         inference_ep = resolve_from_alias(instance.__dict__,["azure_endpoint","api_base"])
         curr_span.set_attribute("inference_endpoint",inference_ep)
 
+def set_provider_name(curr_span, instance):
+    provider_url = ""
+
+    try :
+        if isinstance(instance.client._client.base_url.host, str) :
+            provider_url = instance. client._client.base_url.host
+    except:
+        pass
+
+    try :
+        if isinstance(instance.api_base, str):
+            provider_url = instance.api_base
+    except:
+        pass
+
+    try :
+        if len(provider_url) > 0:
+            parsed_provider_url = urlparse(provider_url)
+            curr_span.set_attribute("provider_name", parsed_provider_url.hostname or provider_url)
+    except:
+        pass
+
 def is_root_span(curr_span: Span) -> bool:
-    return curr_span.parent == None
+    return curr_span.parent is None
 
 def get_input_from_args(chain_args):
-    if len(chain_args) > 0 and type(chain_args[0]) == str:
+    if len(chain_args) > 0 and isinstance(chain_args[0], str):
         return chain_args[0]
     return ""
 
@@ -166,15 +225,19 @@ def update_span_from_llm_response(response, span: Span):
             span.set_attribute("total_tokens", token_usage.get("total_tokens"))
     # extract token usage from llamaindex openai
     if(response is not None and hasattr(response, "raw")):
-        if response.raw is not None:
-            token_usage = response.raw.get("usage")
-            if token_usage is not None:
-                if(hasattr(token_usage, "completion_tokens")):
-                    span.set_attribute("completion_tokens", token_usage.completion_tokens)
-                if(hasattr(token_usage, "prompt_tokens")):
-                    span.set_attribute("prompt_tokens", token_usage.prompt_tokens)
-                if(hasattr(token_usage, "total_tokens")):
-                    span.set_attribute("total_tokens", token_usage.total_tokens)
+        try:
+            if response.raw is not None:
+                token_usage = response.raw.get("usage") if isinstance(response.raw, dict) else getattr(response.raw, "usage", None)
+                if token_usage is not None:
+                    if getattr(token_usage, "completion_tokens", None):
+                        span.set_attribute("completion_tokens", getattr(token_usage, "completion_tokens"))
+                    if getattr(token_usage, "prompt_tokens", None):
+                        span.set_attribute("prompt_tokens", getattr(token_usage, "prompt_tokens"))
+                    if getattr(token_usage, "total_tokens", None):
+                        span.set_attribute("total_tokens", getattr(token_usage, "total_tokens"))
+        except AttributeError:
+            token_usage = None
+
 
 def update_workflow_type(to_wrap, span: Span):
     package_name = to_wrap.get('package')
@@ -185,30 +248,64 @@ def update_workflow_type(to_wrap, span: Span):
 
 def update_span_with_context_input(to_wrap, wrapped_args ,span: Span):
     package_name: str = to_wrap.get('package')
-    if("langchain_core.retrievers" in package_name):
+    if "langchain_core.retrievers" in package_name:
         input_arg_text = wrapped_args[0]
         span.add_event(CONTEXT_INPUT_KEY, {QUERY:input_arg_text})
-    if("llama_index.core.indices.base_retriever" in package_name):
+    if "llama_index.core.indices.base_retriever" in package_name:
         input_arg_text = wrapped_args[0].query_str
         span.add_event(CONTEXT_INPUT_KEY, {QUERY:input_arg_text})
 
 def update_span_with_context_output(to_wrap, return_value ,span: Span):
     package_name: str = to_wrap.get('package')
-    if("llama_index.core.indices.base_retriever" in package_name):
+    if "llama_index.core.indices.base_retriever" in package_name:
         output_arg_text = return_value[0].text
         span.add_event(CONTEXT_OUTPUT_KEY, {RESPONSE:output_arg_text})
 
 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)
-    else:    
+    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)):
+    if isinstance(wrapped_args, str):
         span.add_event(PROMPT_OUTPUT_KEY, {RESPONSE:wrapped_args})
-    if("llama_index.core.base.base_query_engine" in package_name):
+    if "llama_index.core.base.base_query_engine" in package_name:
         span.add_event(PROMPT_OUTPUT_KEY, {RESPONSE:wrapped_args.response})
+
+def update_tags(instance, span):
+    try:
+        # copy tags as is from langchain
+        span.set_attribute(TAGS, getattr(instance, TAGS))
+    except:
+        pass
+    try:
+        # extract embed model and vector store names for llamaindex
+        model_name = instance.retriever._embed_model.model_name
+        vector_store_name = type(instance.retriever._vector_store).__name__
+        span.set_attribute(TAGS, [model_name, vector_store_name])
+    except:
+        pass
+
+
+def update_vectorstore_attributes(to_wrap, instance, span):
+    """
+       Updates the telemetry span attributes for vector store retrieval tasks.
+    """
+    try:
+        package = to_wrap.get('package')
+        if package in framework_vector_store_mapping:
+            attributes = framework_vector_store_mapping[package](instance)
+            span._attributes.update({
+                TYPE: attributes['type'],
+                PROVIDER: attributes['provider'],
+                EMBEDDING_MODEL: attributes['embedding_model']
+            })
+        else:
+            logger.warning(f"Package '{package}' not recognized for vector store telemetry.")
+
+    except Exception as e:
+        logger.error(f"Error updating span attributes: {e}")

monocle_apptrace/wrapper.py

@@ -5,20 +5,20 @@ from monocle_apptrace.langchain import LANGCHAIN_METHODS
 from monocle_apptrace.llamaindex import LLAMAINDEX_METHODS
 from monocle_apptrace.wrap_common import task_wrapper
 
+# pylint: disable=too-few-public-methods
 class WrapperMethod:
     def __init__(
             self,
             package: str,
-            object: str,
+            object_name: str,
             method: str,
             span_name: str = None,
             wrapper = task_wrapper
             ):
         self.package = package
-        self.object = object
+        self.object = object_name
         self.method = method
         self.span_name = span_name
         self.wrapper = wrapper
 
 INBUILT_METHODS_LIST = LANGCHAIN_METHODS + LLAMAINDEX_METHODS + HAYSTACK_METHODS
-