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

monocle_apptrace/README.md

@@ -62,13 +62,13 @@ setup_monocle_telemetry(
         wrapper_methods=[
             WrapperMethod(
                 package="langchain.schema.runnable",
-                object="RunnableParallel",
+                object_name="RunnableParallel",
                 method="invoke",
                 span_name="langchain.workflow",
                 wrapper=task_wrapper),
             WrapperMethod(
                 package="langchain.schema.runnable",
-                object="RunnableParallel",
+                object_name="RunnableParallel",
                 method="ainvoke",
                 span_name="langchain.workflow",
                 wrapper=atask_wrapper)

monocle_apptrace/__init__.py

@@ -1,2 +0,0 @@
-
-

monocle_apptrace/constants.py

@@ -0,0 +1,22 @@
+# Azure environment constants
+AZURE_ML_ENDPOINT_ENV_NAME = "AZUREML_ENTRY_SCRIPT"
+AZURE_FUNCTION_WORKER_ENV_NAME = "FUNCTIONS_WORKER_RUNTIME"
+AZURE_APP_SERVICE_ENV_NAME = "WEBSITE_SITE_NAME"
+AWS_LAMBDA_ENV_NAME = "AWS_LAMBDA_RUNTIME_API"
+
+# Azure naming reference can be found here
+# https://learn.microsoft.com/en-us/azure/cloud-adoption-framework/ready/azure-best-practices/resource-abbreviations
+AZURE_FUNCTION_NAME = "azure.func"
+AZURE_APP_SERVICE_NAME = "azure.asp"
+AZURE_ML_SERVICE_NAME = "azure.mlw"
+AWS_LAMBDA_SERVICE_NAME = "aws.lambda"
+
+azure_service_map = {
+    AZURE_ML_ENDPOINT_ENV_NAME: AZURE_ML_SERVICE_NAME,
+    AZURE_APP_SERVICE_ENV_NAME: AZURE_APP_SERVICE_NAME,
+    AZURE_FUNCTION_WORKER_ENV_NAME: AZURE_FUNCTION_NAME
+}
+
+aws_service_map = {
+    AWS_LAMBDA_ENV_NAME: AWS_LAMBDA_SERVICE_NAME
+}

monocle_apptrace/exporters/file_exporter.py

@@ -0,0 +1,63 @@
+#pylint: disable=consider-using-with
+
+from os import linesep, path
+from io import TextIOWrapper
+from datetime import datetime
+from typing import Optional, Callable, Sequence
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from opentelemetry.sdk.resources import SERVICE_NAME
+
+DEFAULT_FILE_PREFIX:str = "monocle_trace_"
+DEFAULT_TIME_FORMAT:str = "%Y-%m-%d_%H.%M.%S"
+
+class FileSpanExporter(SpanExporter):
+    current_trace_id: int = None
+    current_file_path: str = None
+
+    def __init__(
+        self,
+        service_name: Optional[str] = None,
+        out_path:str = ".",
+        file_prefix = DEFAULT_FILE_PREFIX,
+        time_format = DEFAULT_TIME_FORMAT,
+        formatter: Callable[
+            [ReadableSpan], str
+        ] = lambda span: span.to_json()
+        + linesep,
+    ):
+        self.out_handle:TextIOWrapper = None
+        self.formatter = formatter
+        self.service_name = service_name
+        self.output_path = out_path
+        self.file_prefix = file_prefix
+        self.time_format = time_format
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        for span in spans:
+            if span.context.trace_id != self.current_trace_id:
+                self.rotate_file(span.resource.attributes[SERVICE_NAME],
+                                span.context.trace_id)
+            self.out_handle.write(self.formatter(span))
+        self.out_handle.flush()
+        return SpanExportResult.SUCCESS
+
+    def rotate_file(self, trace_name:str, trace_id:int) -> None:
+        self.reset_handle()
+        self.current_file_path = path.join(self.output_path,
+                        self.file_prefix + trace_name + "_" + hex(trace_id) + "_"
+                        + datetime.now().strftime(self.time_format) + ".json")
+        self.out_handle = open(self.current_file_path, "w", encoding='UTF-8')
+        self.current_trace_id = trace_id
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        self.out_handle.flush()
+        return True
+
+    def reset_handle(self) -> None:
+        if self.out_handle is not None:
+            self.out_handle.close()
+            self.out_handle = None
+
+    def shutdown(self) -> None:
+        self.reset_handle()

monocle_apptrace/haystack/__init__.py

@@ -1,28 +1,9 @@
 
-
+import os
 import logging
-from monocle_apptrace.haystack.wrap_openai import wrap_openai
-from monocle_apptrace.haystack.wrap_pipeline import wrap as wrap_pipeline
+from monocle_apptrace.utils import load_wrapper_from_config
 
 logger = logging.getLogger(__name__)
-
-HAYSTACK_METHODS = [
-    {
-        "package": "haystack.components.generators.openai",
-        "object": "OpenAIGenerator",
-        "method": "run",
-        "wrapper": wrap_openai,
-    },
-    {
-        "package": "haystack.components.generators.chat.openai",
-        "object": "OpenAIChatGenerator",
-        "method": "run",
-        "wrapper": wrap_openai,
-    },
-    {
-        "package": "haystack.core.pipeline.pipeline",
-        "object": "Pipeline",
-        "method": "run",
-        "wrapper": wrap_pipeline,
-    },
-]
+parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+HAYSTACK_METHODS = load_wrapper_from_config(
+    os.path.join(parent_dir, 'wrapper_config', 'haystack_methods.json'))

monocle_apptrace/haystack/wrap_node.py

@@ -21,7 +21,7 @@ def wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
         workflow_name = span.resource.attributes.get("service.name")
         span.set_attribute("workflow_name",workflow_name)
         span.set_attribute("workflow_type", WORKFLOW_TYPE_MAP["haystack"])
-        
+
         response = wrapped(*args, **kwargs)
 
     return response

monocle_apptrace/haystack/wrap_openai.py

@@ -1,5 +1,3 @@
-
-
 import logging
 from opentelemetry import context as context_api
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
@@ -17,9 +15,7 @@ def _set_input_attributes(span, kwargs, instance, args):
 
     if 'model' in instance.__dict__:
         model_name = instance.__dict__.get("model")
-        set_span_attribute(span, "openai_model_name", model_name)
-
-    return
+        set_span_attribute(span, "model_name", model_name)
 
 @dont_throw
 def _set_response_attributes(span, response):
@@ -39,14 +35,10 @@ def wrap_openai(tracer, to_wrap, wrapped, instance, args, kwargs):
     with tracer.start_as_current_span("haystack.openai") as span:
         if span.is_recording():
             _set_input_attributes(span, kwargs, instance, args)
-
-
-        
         response = wrapped(*args, **kwargs)
 
         if response:
             if span.is_recording():
                 _set_response_attributes(span, response)
-                
 
         return response

monocle_apptrace/haystack/wrap_pipeline.py

@@ -1,5 +1,3 @@
-
-
 import logging
 from opentelemetry import context as context_api
 from opentelemetry.context import attach, set_value
@@ -18,13 +16,12 @@ def wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
     name = "haystack_pipeline"
     attach(set_value("workflow_name", name))
     inputs = set()
-    input = get_workflow_input(args, inputs)
+    workflow_input = get_workflow_input(args, inputs)
 
     with tracer.start_as_current_span(f"{name}.workflow") as span:
-        span.set_attribute(PROMPT_INPUT_KEY, input)
+        span.set_attribute(PROMPT_INPUT_KEY, workflow_input)
         workflow_name = span.resource.attributes.get("service.name")
         set_workflow_attributes(span, workflow_name)
-        
         response = wrapped(*args, **kwargs)
         set_workflow_output(span, response)
     return response
@@ -37,12 +34,12 @@ def get_workflow_input(args, inputs):
     for value in args[0].values():
         for text in value.values():
             inputs.add(text)
-    
-    input: str = ""
+
+    workflow_input: str = ""
 
     for input_str in inputs:
-        input = input + input_str
-    return input
+        workflow_input = workflow_input + input_str
+    return workflow_input
 
 def set_workflow_attributes(span, workflow_name):
     span.set_attribute("workflow_name",workflow_name)

monocle_apptrace/instrumentor.py

@@ -10,25 +10,26 @@ from opentelemetry.sdk.trace import TracerProvider, Span
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanProcessor
 from opentelemetry.sdk.resources import SERVICE_NAME, Resource
 from opentelemetry import trace
+from opentelemetry.context import get_value, attach, set_value
 from monocle_apptrace.wrap_common import CONTEXT_PROPERTIES_KEY
 from monocle_apptrace.wrapper import INBUILT_METHODS_LIST, WrapperMethod
-from opentelemetry.context import get_value, attach, set_value
+from monocle_apptrace.exporters.file_exporter import FileSpanExporter
 
 
 logger = logging.getLogger(__name__)
 
-_instruments = ("langchain >= 0.0.346",)
+_instruments = ()
 
 class MonocleInstrumentor(BaseInstrumentor):
-    
+
     workflow_name: str = ""
     user_wrapper_methods: list[WrapperMethod] = []
     instrumented_method_list: list[object] = []
-    
+
     def __init__(
             self,
-            user_wrapper_methods: list[WrapperMethod] = []) -> None:
-        self.user_wrapper_methods = user_wrapper_methods
+            user_wrapper_methods: list[WrapperMethod] = None) -> None:
+        self.user_wrapper_methods = user_wrapper_methods or []
         super().__init__()
 
     def instrumentation_dependencies(self) -> Collection[str]:
@@ -63,11 +64,11 @@ class MonocleInstrumentor(BaseInstrumentor):
                 self.instrumented_method_list.append(wrapped_method)
             except Exception as ex:
                 if wrapped_method in user_method_list:
-                    logger.error(f"""_instrument wrap Exception: {str(ex)} 
+                    logger.error(f"""_instrument wrap Exception: {str(ex)}
                                 for package: {wrap_package},
                                 object:{wrap_object},
                                 method:{wrap_method}""")
-            
+
 
     def _uninstrument(self, **kwargs):
         for wrapped_method in self.instrumented_method_list:
@@ -80,33 +81,33 @@ class MonocleInstrumentor(BaseInstrumentor):
                     wrap_method,
                 )
             except Exception as ex:
-                logger.error(f"""_instrument unwrap Exception: {str(ex)} 
+                logger.error(f"""_instrument unwrap Exception: {str(ex)}
                              for package: {wrap_package},
                              object:{wrap_object},
                              method:{wrap_method}""")
-           
 
 def setup_monocle_telemetry(
         workflow_name: str,
-        span_processors: List[SpanProcessor] = [],
-        wrapper_methods: List[WrapperMethod] = []):
+        span_processors: List[SpanProcessor] = None,
+        wrapper_methods: List[WrapperMethod] = None):
     resource = Resource(attributes={
         SERVICE_NAME: workflow_name
     })
-    traceProvider = TracerProvider(resource=resource)
-    tracerProviderDefault = trace.get_tracer_provider()
-    providerType = type(tracerProviderDefault).__name__
-    isProxyProvider = "Proxy" in providerType
+    span_processors = span_processors or [BatchSpanProcessor(FileSpanExporter())]
+    trace_provider = TracerProvider(resource=resource)
+    tracer_provider_default = trace.get_tracer_provider()
+    provider_type = type(tracer_provider_default).__name__
+    is_proxy_provider = "Proxy" in provider_type
     for processor in span_processors:
         processor.on_start = on_processor_start
-        if not isProxyProvider:
-            tracerProviderDefault.add_span_processor(processor)
+        if not is_proxy_provider:
+            tracer_provider_default.add_span_processor(processor)
         else :
-            traceProvider.add_span_processor(processor)
-    if isProxyProvider :
-        trace.set_tracer_provider(traceProvider)
-    instrumentor = MonocleInstrumentor(user_wrapper_methods=wrapper_methods)
-    instrumentor.app_name = workflow_name
+            trace_provider.add_span_processor(processor)
+    if is_proxy_provider :
+        trace.set_tracer_provider(trace_provider)
+    instrumentor = MonocleInstrumentor(user_wrapper_methods=wrapper_methods or [])
+    # instrumentor.app_name = workflow_name
     if not instrumentor.is_instrumented_by_opentelemetry:
         instrumentor.instrument()
 
@@ -117,11 +118,7 @@ def on_processor_start(span: Span, parent_context):
         for key, value in context_properties.items():
             span.set_attribute(
                 f"{CONTEXT_PROPERTIES_KEY}.{key}", value
-            )    
+            )
 
 def set_context_properties(properties: dict) -> None:
     attach(set_value(CONTEXT_PROPERTIES_KEY, properties))
-
-
-
-

monocle_apptrace/langchain/__init__.py

@@ -1,95 +1,6 @@
+import os
+from monocle_apptrace.utils import load_wrapper_from_config
 
-
-from monocle_apptrace.wrap_common import allm_wrapper, atask_wrapper, llm_wrapper, task_wrapper
-
-LANGCHAIN_METHODS = [
-    {
-        "package": "langchain.prompts.base",
-        "object": "BasePromptTemplate",
-        "method": "invoke",
-        "wrapper": task_wrapper,
-    },
-    {
-        "package": "langchain.prompts.base",
-        "object": "BasePromptTemplate",
-        "method": "ainvoke",
-        "wrapper": atask_wrapper,
-    },
-    {
-        "package": "langchain.chat_models.base",
-        "object": "BaseChatModel",
-        "method": "invoke",
-        "wrapper": llm_wrapper,
-    },
-    {
-        "package": "langchain.chat_models.base",
-        "object": "BaseChatModel",
-        "method": "ainvoke",
-        "wrapper": allm_wrapper,
-    },
-    {
-        "package": "langchain_core.language_models.llms",
-        "object": "LLM",
-        "method": "_generate",
-        "wrapper": llm_wrapper,
-    },
-    {
-        "package": "langchain_core.language_models.llms",
-        "object": "LLM",
-        "method": "_agenerate",
-        "wrapper": llm_wrapper,
-    },
-    {
-        "package": "langchain_core.retrievers",
-        "object": "BaseRetriever",
-        "method": "invoke",
-        "wrapper": task_wrapper,
-    },
-    {
-        "package": "langchain_core.retrievers",
-        "object": "BaseRetriever",
-        "method": "ainvoke",
-        "wrapper": atask_wrapper,
-    },
-    {
-        "package": "langchain.schema",
-        "object": "BaseOutputParser",
-        "method": "invoke",
-        "wrapper": task_wrapper,
-    },
-    {
-        "package": "langchain.schema",
-        "object": "BaseOutputParser",
-        "method": "ainvoke",
-        "wrapper": atask_wrapper,
-    },
-    {
-        "package": "langchain.schema.runnable",
-        "object": "RunnableSequence",
-        "method": "invoke",
-        "span_name": "langchain.workflow",
-        "wrapper": task_wrapper,
-    },
-    {
-        "package": "langchain.schema.runnable",
-        "object": "RunnableSequence",
-        "method": "ainvoke",
-        "span_name": "langchain.workflow",
-        "wrapper": atask_wrapper,
-    },
-    {
-        "package": "langchain.schema.runnable",
-        "object": "RunnableParallel",
-        "method": "invoke",
-        "span_name": "langchain.workflow",
-        "wrapper": task_wrapper,
-    },
-    {
-        "package": "langchain.schema.runnable",
-        "object": "RunnableParallel",
-        "method": "ainvoke",
-        "span_name": "langchain.workflow",
-        "wrapper": atask_wrapper,
-    },
-    
-]
+parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+LANGCHAIN_METHODS = load_wrapper_from_config(
+    os.path.join(parent_dir, 'wrapper_config', 'lang_chain_methods.json'))

monocle_apptrace/llamaindex/__init__.py

@@ -1,71 +1,15 @@
 
-
-from monocle_apptrace.wrap_common import allm_wrapper, atask_wrapper, llm_wrapper, task_wrapper
+#pylint: disable=protected-access
+import os
+from monocle_apptrace.utils import load_wrapper_from_config
 
 def get_llm_span_name_for_openai(instance):
-    if (hasattr(instance, "_is_azure_client") 
+    if (hasattr(instance, "_is_azure_client")
         and callable(getattr(instance, "_is_azure_client"))
         and instance._is_azure_client()):
         return "llamaindex.azure_openai"
     return "llamaindex.openai"
 
-LLAMAINDEX_METHODS = [
-    {
-        "package": "llama_index.core.indices.base_retriever",
-        "object": "BaseRetriever",
-        "method": "retrieve",
-        "span_name": "llamaindex.retrieve",
-        "wrapper": task_wrapper
-    },
-    {
-        "package": "llama_index.core.indices.base_retriever",
-        "object": "BaseRetriever",
-        "method": "aretrieve",
-        "span_name": "llamaindex.retrieve",
-        "wrapper": atask_wrapper
-    },
-    {
-        "package": "llama_index.core.base.base_query_engine",
-        "object": "BaseQueryEngine",
-        "method": "query",
-        "span_name": "llamaindex.query",
-        "wrapper": task_wrapper,
-    },
-    {
-        "package": "llama_index.core.base.base_query_engine",
-        "object": "BaseQueryEngine",
-        "method": "aquery",
-        "span_name": "llamaindex.query",
-        "wrapper": atask_wrapper,
-    },
-    {
-        "package": "llama_index.core.llms.custom",
-        "object": "CustomLLM",
-        "method": "chat",
-        "span_name": "llamaindex.llmchat",
-        "wrapper": task_wrapper,
-    },
-    {
-        "package": "llama_index.core.llms.custom",
-        "object": "CustomLLM",
-        "method": "achat",
-        "span_name": "llamaindex.llmchat",
-        "wrapper": atask_wrapper,
-    },
-    {
-        "package": "llama_index.llms.openai.base",
-        "object": "OpenAI",
-        "method": "chat",
-        "span_name": "llamaindex.openai",
-        "span_name_getter" : get_llm_span_name_for_openai,
-        "wrapper": llm_wrapper,
-    },
-     {
-        "package": "llama_index.llms.openai.base",
-        "object": "OpenAI",
-        "method": "achat",
-        "span_name": "llamaindex.openai",
-        "wrapper": allm_wrapper,
-    }
-]
-
+parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+LLAMAINDEX_METHODS = load_wrapper_from_config(
+    os.path.join(parent_dir, 'wrapper_config', 'llama_index_methods.json'))

monocle_apptrace/utils.py

@@ -1,17 +1,14 @@
-
-
 import logging
-logger = logging.getLogger(__name__)
-
-class Config:
-    exception_logger = None
+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
 
 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 +19,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 +38,36 @@ 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)

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
 
 logger = logging.getLogger(__name__)
 WORKFLOW_TYPE_KEY = "workflow_type"
@@ -15,6 +16,8 @@ QUERY = "question"
 RESPONSE = "response"
 TAGS = "tags"
 CONTEXT_PROPERTIES_KEY = "workflow_context_properties"
+INFRA_SERVICE_KEY = "infra_service_name"
+
 
 
 WORKFLOW_TYPE_MAP = {
@@ -37,29 +40,38 @@ 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)
+    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 +86,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 +99,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 +108,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,7 +132,6 @@ 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:
         update_llm_endpoint(curr_span= span, instance=instance)
 
@@ -138,24 +150,48 @@ def update_llm_endpoint(curr_span: Span, instance):
             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)
+        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 ""
 
 def update_span_from_llm_response(response, span: Span):
+    
     # extract token uasge from langchain openai
     if (response is not None and hasattr(response, "response_metadata")):
         response_metadata = response.response_metadata
@@ -166,15 +202,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 +225,44 @@ 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

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
-

monocle_apptrace/wrapper_config/haystack_methods.json

@@ -0,0 +1,25 @@
+{
+"wrapper_methods" :  [
+    {
+        "package": "haystack.components.generators.openai",
+        "object": "OpenAIGenerator",
+        "method": "run",
+        "wrapper_package": "haystack.wrap_openai",
+        "wrapper_method": "wrap_openai"
+    },
+    {
+        "package": "haystack.components.generators.chat.openai",
+        "object": "OpenAIChatGenerator",
+        "method": "run",
+        "wrapper_package": "haystack.wrap_openai",
+        "wrapper_method": "wrap_openai"
+    },
+    {
+        "package": "haystack.core.pipeline.pipeline",
+        "object": "Pipeline",
+        "method": "run",
+        "wrapper_package": "haystack.wrap_pipeline",
+        "wrapper_method": "wrap"
+    }
+] 
+}

monocle_apptrace/wrapper_config/lang_chain_methods.json

@@ -0,0 +1,106 @@
+{ 
+"wrapper_methods" : [
+    {
+        "package": "langchain.prompts.base",
+        "object": "BasePromptTemplate",
+        "method": "invoke",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
+    },
+    {
+        "package": "langchain.prompts.base",
+        "object": "BasePromptTemplate",
+        "method": "ainvoke",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
+    },
+    {
+        "package": "langchain.chat_models.base",
+        "object": "BaseChatModel",
+        "method": "invoke",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "llm_wrapper"
+    },
+    {
+        "package": "langchain.chat_models.base",
+        "object": "BaseChatModel",
+        "method": "ainvoke",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "allm_wrapper"
+    },
+    {
+        "package": "langchain_core.language_models.llms",
+        "object": "LLM",
+        "method": "_generate",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "llm_wrapper"
+    },
+    {
+        "package": "langchain_core.language_models.llms",
+        "object": "LLM",
+        "method": "_agenerate",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "llm_wrapper"
+    },
+    {
+        "package": "langchain_core.retrievers",
+        "object": "BaseRetriever",
+        "method": "invoke",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
+    },
+    {
+        "package": "langchain_core.retrievers",
+        "object": "BaseRetriever",
+        "method": "ainvoke",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
+    },
+    {
+        "package": "langchain.schema",
+        "object": "BaseOutputParser",
+        "method": "invoke",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
+    },
+    {
+        "package": "langchain.schema",
+        "object": "BaseOutputParser",
+        "method": "ainvoke",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
+    },
+    {
+        "package": "langchain.schema.runnable",
+        "object": "RunnableSequence",
+        "method": "invoke",
+        "span_name": "langchain.workflow",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
+    },
+    {
+        "package": "langchain.schema.runnable",
+        "object": "RunnableSequence",
+        "method": "ainvoke",
+        "span_name": "langchain.workflow",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
+    },
+    {
+        "package": "langchain.schema.runnable",
+        "object": "RunnableParallel",
+        "method": "invoke",
+        "span_name": "langchain.workflow",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
+    },
+    {
+        "package": "langchain.schema.runnable",
+        "object": "RunnableParallel",
+        "method": "ainvoke",
+        "span_name": "langchain.workflow",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
+    }
+]
+}

monocle_apptrace/wrapper_config/llama_index_methods.json

@@ -0,0 +1,70 @@
+{
+"wrapper_methods" : [
+    {
+        "package": "llama_index.core.indices.base_retriever",
+        "object": "BaseRetriever",
+        "method": "retrieve",
+        "span_name": "llamaindex.retrieve",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
+    },
+    {
+        "package": "llama_index.core.indices.base_retriever",
+        "object": "BaseRetriever",
+        "method": "aretrieve",
+        "span_name": "llamaindex.retrieve",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
+    },
+    {
+        "package": "llama_index.core.base.base_query_engine",
+        "object": "BaseQueryEngine",
+        "method": "query",
+        "span_name": "llamaindex.query",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
+    },
+    {
+        "package": "llama_index.core.base.base_query_engine",
+        "object": "BaseQueryEngine",
+        "method": "aquery",
+        "span_name": "llamaindex.query",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
+    },
+    {
+        "package": "llama_index.core.llms.custom",
+        "object": "CustomLLM",
+        "method": "chat",
+        "span_name": "llamaindex.llmchat",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "task_wrapper"
+    },
+    {
+        "package": "llama_index.core.llms.custom",
+        "object": "CustomLLM",
+        "method": "achat",
+        "span_name": "llamaindex.llmchat",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "atask_wrapper"
+    },
+    {
+        "package": "llama_index.llms.openai.base",
+        "object": "OpenAI",
+        "method": "chat",
+        "span_name": "llamaindex.openai",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "llm_wrapper",
+        "span_name_getter_package" : "llamaindex",
+        "span_name_getter_mothod" : "get_llm_span_name_for_openai"
+    },
+    {
+        "package": "llama_index.llms.openai.base",
+        "object": "OpenAI",
+        "method": "achat",
+        "span_name": "llamaindex.openai",
+        "wrapper_package": "wrap_common",
+        "wrapper_method": "allm_wrapper"
+    }
+]
+}

{monocle_apptrace-0.0.1.dist-info → monocle_apptrace-0.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: monocle_apptrace
-Version: 0.0.1
+Version: 0.1.0
 Summary: package with monocle genAI tracing
 Project-URL: Homepage, https://github.com/monocle2ai/monocle
 Project-URL: Issues, https://github.com/monocle2ai/monocle/issues
@@ -27,6 +27,7 @@ Requires-Dist: llama-index-embeddings-huggingface==0.2.0; 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: numpy==1.26.4; extra == 'dev'
+Requires-Dist: parameterized==0.9.0; extra == 'dev'
 Requires-Dist: pytest==8.0.0; extra == 'dev'
 Requires-Dist: sentence-transformers==2.6.1; extra == 'dev'
 Requires-Dist: types-requests==2.31.0.20240106; extra == 'dev'
@@ -63,7 +64,7 @@ The traces are compatible with OpenTelemetry format. They are further enriched t
 ```
 - Enable Monocle tracing in your app by adding following
 ```
-    setup_okahu_telemetry(workflow_name="your-app-name")
+    setup_monocle_telemetry(workflow_name="your-app-name")
 ```
 Please refer to [Monocle user guide](Monocle_User_Guide.md) for more details
 

monocle_apptrace-0.1.0.dist-info/RECORD

@@ -0,0 +1,22 @@
+monocle_apptrace/README.md,sha256=OVJgf1_HAm5L2-FSDZ-3AYfBjSxiT4fhE_CgKVaDVuQ,3104
+monocle_apptrace/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+monocle_apptrace/constants.py,sha256=wjObbmMTFL201x-bf3EOXevYygwkFH_1ng5dDrpE3z0,810
+monocle_apptrace/instrumentor.py,sha256=IDCuqzvTLzYKBzHGX3q1gbiboq07LMkdOA8DhOUTEEU,5051
+monocle_apptrace/utils.py,sha256=VMiWscPCNjp29IQE3ahprXjkfiMw160DbCo8WrjQCXk,2658
+monocle_apptrace/wrap_common.py,sha256=zbYlhL7V655exfdE0h_aT6-THmmHL1ltzsm4cBZ4jq8,10875
+monocle_apptrace/wrapper.py,sha256=cNUdfciAXNYAhvtOA2O4ONRvuT2bbHb4ax_7pALijEI,734
+monocle_apptrace/exporters/file_exporter.py,sha256=gN9pJ_X5pcstVVsyivgHsjWhr443eRa6Y6Hx1rGLQAM,2280
+monocle_apptrace/haystack/__init__.py,sha256=zcluKUIu1M9g_B3s_ZfzRS_vr7yMz5gj-6rqQ7bJ5B0,318
+monocle_apptrace/haystack/wrap_node.py,sha256=IK07Wn3Lk1Os9URsyrmB1HXOH2FNdzK9fNLlR8TZdYc,908
+monocle_apptrace/haystack/wrap_openai.py,sha256=Yp916DhOl0WI6virRi3L43snfsQm7PhI28wlDsg19v8,1536
+monocle_apptrace/haystack/wrap_pipeline.py,sha256=xRYMRzxvFPdcJ64E0bbMdMuWO_p3V1T7eIvb3-Um5DE,1661
+monocle_apptrace/langchain/__init__.py,sha256=HhvRJ_rl9cX4M8ckiOkJC7QHbklrttaY9RvDC51m1l4,268
+monocle_apptrace/llamaindex/__init__.py,sha256=3zmSNoVDjB-hh_M4eUr-hUP0bup7HKmWFVv_3xPAwsA,570
+monocle_apptrace/wrapper_config/haystack_methods.json,sha256=JmngkaKICAzOyrWNTWEOLYFrp99l5wcERYKE_SQRNxE,698
+monocle_apptrace/wrapper_config/lang_chain_methods.json,sha256=HaOhhxb3PkI7tXPxXhWR4cnWrnEHU--k5pOY9RS0Uew,3119
+monocle_apptrace/wrapper_config/llama_index_methods.json,sha256=qpODnBHkaDjPBYZNd7clwmp_2subTu-fmI08Ky5OWdg,2192
+monocle_apptrace-0.1.0.dist-info/METADATA,sha256=sQfm_x1NFthDhfMKX_ZDE0E5bDLSxUOtV1v6yXt0Wpc,5699
+monocle_apptrace-0.1.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+monocle_apptrace-0.1.0.dist-info/licenses/LICENSE,sha256=ay9trLiP5I7ZsFXo6AqtkLYdRqe5S9r-DrPOvsNlZrg,9136
+monocle_apptrace-0.1.0.dist-info/licenses/NOTICE,sha256=9jn4xtwM_uUetJMx5WqGnhrR7MIhpoRlpokjSTlyt8c,112
+monocle_apptrace-0.1.0.dist-info/RECORD,,

monocle_apptrace-0.0.1.dist-info/RECORD

@@ -1,17 +0,0 @@
-monocle_apptrace/README.md,sha256=C0JfJtNC7LOlr_iHSxOwuZn89vuoZ2RbyeGB7ACu4iY,3094
-monocle_apptrace/__init__.py,sha256=daEdpEyAJIa8b2VkCqSKcw8PaExcB6Qro80XNes_sHA,2
-monocle_apptrace/instrumentor.py,sha256=T-g0OpR2VYm7lswyiFoBlBeRPKHX6a-mhH8q7rE-o2g,4927
-monocle_apptrace/utils.py,sha256=KZwwXK8NTbVjQo4BkFwPdccE6hzvBzawPp6an44Gmz8,1366
-monocle_apptrace/wrap_common.py,sha256=Qxwq2drHOovd0CMd7YiQbe_4FqGL07RLfu4Vek1Z-SU,9107
-monocle_apptrace/wrapper.py,sha256=95Yg_dwGv5T6qtNv-ozhXRGaDVI0P27exU2901Pw2ps,684
-monocle_apptrace/haystack/__init__.py,sha256=GspHWTmqMDxQhpTgl91xighNL8MHaWs6BF0YzzZayuE,714
-monocle_apptrace/haystack/wrap_node.py,sha256=96qUlDZtNOVIPrrDZzjJ7ZDCi70mYbd7bIEX41e-RH0,916
-monocle_apptrace/haystack/wrap_openai.py,sha256=9ecWsOUXNjTKrQjNZz1GCgkCAN9GChhasMRDmZMuANE,1585
-monocle_apptrace/haystack/wrap_pipeline.py,sha256=eVxJ4-yS4gFxInwaWy5oxNcNl44X2M7_dhBpKVFmI6Q,1622
-monocle_apptrace/langchain/__init__.py,sha256=PrZl6vxX5WsYfsPFVhumVmUGkwddAgPBFs1hpifbmZw,2519
-monocle_apptrace/llamaindex/__init__.py,sha256=iMaboVDeM5yGR01cEqCD8MotCj2FuY4QjSZmFBv4f2o,2094
-monocle_apptrace-0.0.1.dist-info/METADATA,sha256=uYd9O71FeKzrkLq6DmJYp9F62QxxZ7UfSAaGa5dPVm8,5645
-monocle_apptrace-0.0.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-monocle_apptrace-0.0.1.dist-info/licenses/LICENSE,sha256=ay9trLiP5I7ZsFXo6AqtkLYdRqe5S9r-DrPOvsNlZrg,9136
-monocle_apptrace-0.0.1.dist-info/licenses/NOTICE,sha256=9jn4xtwM_uUetJMx5WqGnhrR7MIhpoRlpokjSTlyt8c,112
-monocle_apptrace-0.0.1.dist-info/RECORD,,