langtrace-python-sdk 1.0.9__py3-none-any.whl

instrumentation/chroma/instrumentation.py

@@ -0,0 +1,33 @@
+import importlib.metadata
+from typing import Collection
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from instrumentation.chroma.lib.apis import APIS
+from instrumentation.chroma.patch import collection_patch
+
+
+class ChromaInstrumentation(BaseInstrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["chromadb >= 0.4.23"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version('chromadb')
+
+        for operation, details in APIS.items():
+            wrap_function_wrapper(
+                'chromadb.api.models.Collection',
+                f'Collection.{operation.lower()}',
+                collection_patch(operation, version, tracer)
+            )
+
+    def _instrument_module(self, module_name):
+        print(module_name)
+
+    def _uninstrument(self, **kwargs):
+        print(kwargs)

instrumentation/chroma/lib/apis.py

@@ -0,0 +1,40 @@
+from langtrace.trace_attributes import ChromaDBMethods
+
+APIS = {
+    "ADD": {
+        "METHOD": ChromaDBMethods.ADD.value,
+        "OPERATION": "add",
+    },
+    "GET": {
+        "METHOD": ChromaDBMethods.GET.value,
+        "OPERATION": "get",
+    },
+    "QUERY": {
+        "METHOD": ChromaDBMethods.QUERY.value,
+        "OPERATION": "query",
+    },
+    "DELETE": {
+        "METHOD": ChromaDBMethods.DELETE.value,
+        "OPERATION": "delete",
+    },
+    "PEEK": {
+        "METHOD": ChromaDBMethods.PEEK.value,
+        "OPERATION": "peek",
+    },
+    "UPDATE": {
+        "METHOD": ChromaDBMethods.UPDATE.value,
+        "OPERATION": "update",
+    },
+    "UPSERT": {
+        "METHOD": ChromaDBMethods.UPSERT.value,
+        "OPERATION": "upsert",
+    },
+    "MODIFY": {
+        "METHOD": ChromaDBMethods.MODIFY.value,
+        "OPERATION": "modify",
+    },
+    "COUNT": {
+        "METHOD": ChromaDBMethods.COUNT.value,
+        "OPERATION": "count",
+    },
+}

instrumentation/chroma/patch.py

@@ -0,0 +1,46 @@
+from langtrace.trace_attributes import DatabaseSpanAttributes
+from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace.status import Status, StatusCode
+
+from instrumentation.constants import SERVICE_PROVIDERS
+from instrumentation.chroma.lib.apis import APIS
+
+
+def collection_patch(method, version, tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+        api = APIS[method]
+        service_provider = SERVICE_PROVIDERS['CHROMA']
+        span_attributes = {
+            'langtrace.service.name': service_provider,
+            'langtrace.service.type': 'vectordb',
+            'langtrace.service.version': version,
+            'langtrace.version': '1.0.0',
+            "db.system": "chromadb",
+            "db.operation": api['OPERATION'],
+        }
+
+        if hasattr(instance, 'name') and instance.name is not None:
+            span_attributes["db.collection.name"] = instance.name
+
+        attributes = DatabaseSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(api["METHOD"], kind=SpanKind.CLIENT) as span:
+            for field, value in attributes.model_dump(by_alias=True).items():
+                if value is not None:
+                    span.set_attribute(field, value)
+            try:
+                # Attempt to call the original method
+                result = wrapped(*args, **kwargs)
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as e:
+                # Record the exception in the span
+                span.record_exception(e)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method

instrumentation/constants.py

@@ -0,0 +1,18 @@
+
+TRACE_NAMESPACES = {
+    "OPENAI": "Langtrace OpenAI SDK",
+    "LANGCHAIN": "Langtrace Langchain SDK",
+    "PINECONE": "Langtrace Pinecone SDK",
+    "LLAMAINDEX": "Langtrace LlamaIndex SDK",
+    "CHROMA": "Langtrace Chroma SDK",
+}
+
+SERVICE_PROVIDERS = {
+    "OPENAI": "OpenAI",
+    "AZURE": "Azure",
+    "LANGCHAIN": "Langchain",
+    "LANGCHAIN_COMMUNITY": "Langchain Community",
+    "PINECONE": "Pinecone",
+    "LLAMAINDEX": "LlamaIndex",
+    "CHROMA": "Chroma",
+}

instrumentation/langchain/instrumentation.py

@@ -0,0 +1,74 @@
+"""
+Instrumentation for langchain.
+"""
+import importlib.metadata
+import inspect
+from typing import Collection
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from instrumentation.langchain.patch import generic_patch
+
+
+def patch_module_classes(module_name, tracer, version, task, trace_output=True, trace_input=True):
+    """
+    Generic function to patch all public methods of all classes in a given module.
+
+    Parameters:
+    - module: The module object containing the classes to patch.
+    - module_name: The name of the module, used in the prefix for `wrap_function_wrapper`.
+    - tracer: The tracer object used in `generic_patch`.
+    - version: The version parameter used in `generic_patch`.
+    - task: The name used to identify the type of task in `generic_patch`.
+    - exclude_private: Whether to exclude private methods (those starting with '_').
+    - trace_output: Whether to trace the output of the patched methods.
+    - trace_input: Whether to trace the input of the patched methods.
+    """
+    # import the module
+    module = importlib.import_module(module_name)
+    # loop through all public classes in the module
+    for name, obj in inspect.getmembers(module, lambda member: inspect.isclass(member) and
+                                        member.__module__ == module.__name__):
+        # loop through all public methods of the class
+        for method_name, _ in inspect.getmembers(obj, predicate=inspect.isfunction):
+            # Skip private methods
+            if method_name.startswith('_'):
+                continue
+            try:
+                method_path = f'{name}.{method_name}'
+                wrap_function_wrapper(
+                    module_name,
+                    method_path,
+                    generic_patch(
+                        method_path, task, tracer, version, trace_output, trace_input)
+                )
+            # pylint: disable=broad-except
+            except Exception:
+                pass
+
+
+class LangchainInstrumentation(BaseInstrumentor):
+    """
+    Instrumentor for langchain.
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["langchain >= 0.1.9"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version('langchain')
+
+        modules_to_patch = [
+            ('langchain.text_splitter', 'split_text', True, True),
+        ]
+
+        for module_name, task, trace_output, trace_input in modules_to_patch:
+            patch_module_classes(module_name, tracer, version,
+                                 task, trace_output, trace_input)
+
+    def _uninstrument(self, **kwargs):
+        pass

instrumentation/langchain/patch.py

@@ -0,0 +1,84 @@
+"""
+This module contains the patching logic for the langchain package.
+"""
+import json
+
+from langtrace.trace_attributes import FrameworkSpanAttributes
+from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace.status import Status
+
+from instrumentation.constants import SERVICE_PROVIDERS
+
+
+def generic_patch(method_name, task, tracer, version, trace_output=True, trace_input=True):
+    """
+    patch method for generic methods.
+    """
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS['LANGCHAIN']
+        span_attributes = {
+            'langtrace.service.name': service_provider,
+            'langtrace.service.type': 'framework',
+            'langtrace.service.version': version,
+            'langtrace.version': '1.0.0',
+            'langchain.task.name': task,
+        }
+
+        if len(args) > 0 and trace_input:
+            span_attributes['langchain.inputs'] = to_json_string(args)
+
+        attributes = FrameworkSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(method_name, kind=SpanKind.CLIENT) as span:
+            for field, value in attributes.model_dump(by_alias=True).items():
+                if value is not None:
+                    span.set_attribute(field, value)
+            try:
+                # Attempt to call the original method
+                result = wrapped(*args, **kwargs)
+                if trace_output:
+                    span.set_attribute(
+                        'langchain.outputs', to_json_string(result))
+
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as e:
+                # Record the exception in the span
+                span.record_exception(e)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method
+
+
+def clean_empty(d):
+    """Recursively remove empty lists, empty dicts, or None elements from a dictionary."""
+    if not isinstance(d, (dict, list)):
+        return d
+    if isinstance(d, list):
+        return [v for v in (clean_empty(v) for v in d) if v != [] and v is not None]
+    return {
+        k: v
+        for k, v in ((k, clean_empty(v)) for k, v in d.items())
+        if v is not None and v != {}
+    }
+
+
+def custom_serializer(obj):
+    """Fallback function to convert unserializable objects."""
+    if hasattr(obj, "__dict__"):
+        # Attempt to serialize custom objects by their __dict__ attribute.
+        return clean_empty(obj.__dict__)
+    else:
+        # For other types, just convert to string
+        return str(obj)
+
+
+def to_json_string(any_object):
+    """Converts any object to a JSON-parseable string, omitting empty or None values."""
+    cleaned_object = clean_empty(any_object)
+    return json.dumps(cleaned_object, default=custom_serializer, indent=2)

instrumentation/langchain_community/instrumentation.py

@@ -0,0 +1,99 @@
+"""
+Instrumentation for langchain-community.
+"""
+import importlib.metadata
+import inspect
+from typing import Collection
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from instrumentation.langchain_community.patch import generic_patch
+
+
+def patch_module_classes(module_name, tracer, version, task, trace_output=True, trace_input=True):
+    """
+    Generic function to patch all public methods of all classes in a given module.
+
+    Parameters:
+    - module: The module object containing the classes to patch.
+    - module_name: The name of the module, used in the prefix for `wrap_function_wrapper`.
+    - tracer: The tracer object used in `generic_patch`.
+    - version: The version parameter used in `generic_patch`.
+    - task: The name used to identify the type of patch in `generic_patch`.
+    - exclude_private: Whether to exclude private methods (those starting with '_').
+    - trace_output: Whether to trace the output of the patched methods.
+    - trace_input: Whether to trace the input of the patched methods.
+    """
+    # import the module
+    module = importlib.import_module(module_name)
+    # loop through all public classes in the module
+    for name, obj in inspect.getmembers(module, lambda member: inspect.isclass(member) and
+                                        member.__module__ == module.__name__):
+        # loop through all public methods of the class
+        for method_name, _ in inspect.getmembers(obj, predicate=inspect.isfunction):
+            # Skip private methods
+            if method_name.startswith('_'):
+                continue
+            try:
+                method_path = f'{name}.{method_name}'
+                wrap_function_wrapper(
+                    module_name,
+                    method_path,
+                    generic_patch(
+                        method_path, task, tracer, version, trace_output, trace_input)
+                )
+            # pylint: disable=broad-except
+            except Exception:
+                pass
+
+
+class LangchainCommunityInstrumentation(BaseInstrumentor):
+    """
+    Instrumentor for langchain-community.
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["langchain-community >= 0.0.24"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version('langchain-community')
+
+        # List of modules to patch, with their corresponding patch names
+        modules_to_patch = [
+            ('langchain_community.document_loaders.pdf', 'load_pdf', True, True),
+            ('langchain_community.vectorstores.faiss', 'vector_store', False, False),
+            ('langchain_community.vectorstores.pgvector',
+             'vector_store', False, False),
+            ('langchain_community.vectorstores.pinecone',
+             'vector_store', False, False),
+            ('langchain_community.vectorstores.mongodb_atlas',
+             'vector_store', False, False),
+            ('langchain_community.vectorstores.azuresearch',
+             'vector_store', False, False),
+            ('langchain_community.vectorstores.azure_cosmos_db',
+             'vector_store', False, False),
+            ('langchain_community.vectorstores.cassandra',
+             'vector_store', False, False),
+            ('langchain_community.vectorstores.chroma', 'vector_store', False, False),
+            ('langchain_community.vectorstores.clickhouse',
+             'vector_store', False, False),
+            ('langchain_community.vectorstores.elasticsearch',
+             'vector_store', False, False),
+            ('langchain_community.vectorstores.supabase',
+             'vector_store', False, False),
+            ('langchain_community.vectorstores.weaviate',
+             'vector_store', False, False),
+            ('langchain_community.vectorstores.vectara',
+             'vector_store', False, False),
+        ]
+
+        for module_name, task, trace_output, trace_input in modules_to_patch:
+            patch_module_classes(module_name, tracer, version,
+                                 task, trace_output, trace_input)
+
+    def _uninstrument(self, **kwargs):
+        pass

instrumentation/langchain_community/patch.py

@@ -0,0 +1,78 @@
+import json
+
+from langtrace.trace_attributes import FrameworkSpanAttributes
+from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace.status import Status, StatusCode
+
+from instrumentation.constants import SERVICE_PROVIDERS
+
+
+def generic_patch(method_name, task, tracer, version, trace_output=True, trace_input=True):
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS['LANGCHAIN_COMMUNITY']
+        span_attributes = {
+            'langtrace.service.name': service_provider,
+            'langtrace.service.type': 'framework',
+            'langtrace.service.version': version,
+            'langtrace.version': '1.0.0',
+            'langchain.task.name': task,
+        }
+
+        if trace_input and len(args) > 0:
+            span_attributes['langchain.inputs'] = to_json_string(args)
+
+        attributes = FrameworkSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(method_name, kind=SpanKind.CLIENT) as span:
+            for field, value in attributes.model_dump(by_alias=True).items():
+                if value is not None:
+                    span.set_attribute(field, value)
+            try:
+                # Attempt to call the original method
+                result = wrapped(*args, **kwargs)
+                if trace_output:
+                    span.set_attribute(
+                        'langchain.outputs', to_json_string(result))
+
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as e:
+                # Record the exception in the span
+                span.record_exception(e)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method
+
+
+def clean_empty(d):
+    """Recursively remove empty lists, empty dicts, or None elements from a dictionary."""
+    if not isinstance(d, (dict, list)):
+        return d
+    if isinstance(d, list):
+        return [v for v in (clean_empty(v) for v in d) if v != [] and v is not None]
+    return {
+        k: v
+        for k, v in ((k, clean_empty(v)) for k, v in d.items())
+        if v is not None and v != {}
+    }
+
+
+def custom_serializer(obj):
+    """Fallback function to convert unserializable objects."""
+    if hasattr(obj, "__dict__"):
+        # Attempt to serialize custom objects by their __dict__ attribute.
+        return clean_empty(obj.__dict__)
+    else:
+        # For other types, just convert to string
+        return str(obj)
+
+
+def to_json_string(any_object):
+    """Converts any object to a JSON-parseable string, omitting empty or None values."""
+    cleaned_object = clean_empty(any_object)
+    return json.dumps(cleaned_object, default=custom_serializer, indent=2)

instrumentation/langchain_core/instrumentation.py

@@ -0,0 +1,101 @@
+"""
+Instrumentation for langchain-core.
+"""
+import importlib.metadata
+import inspect
+from typing import Collection
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from instrumentation.langchain_core.patch import generic_patch, runnable_patch
+
+
+# pylint: disable=dangerous-default-value
+def patch_module_classes(module_name, tracer, version, task, patch_method,
+                         exclude_methods=[], exclude_classes=[],
+                         trace_output=True, trace_input=True):
+    """
+    Generic function to patch all public methods of all classes in a given module.
+
+    Parameters:
+    - module: The module object containing the classes to patch.
+    - module_name: The name of the module, used in the prefix for `wrap_function_wrapper`.
+    - tracer: The tracer object used in `generic_patch`.
+    - version: The version parameter used in `generic_patch`.
+    - task: The name used to identify the type of task in `generic_patch`.
+    - patch_method: The patch method to use.
+    - exclude_methods: A list of methods to exclude from patching.
+    - exclude_classes: A list of classes to exclude from patching.
+    - trace_output: Whether to trace the output of the patched methods.
+    - trace_input: Whether to trace the input of the patched methods.
+    """
+    # import the module
+    module = importlib.import_module(module_name)
+    # loop through all public classes in the module
+    for name, obj in inspect.getmembers(module, lambda member: inspect.isclass(member) and
+                                        member.__module__ == module.__name__):
+        # Skip private classes
+        if name.startswith('_') or name in exclude_classes:
+            continue
+        # loop through all public methods of the class
+        for method_name, _ in inspect.getmembers(obj, predicate=inspect.isfunction):
+            # Skip private methods
+            if method_name.startswith('_') or method_name in exclude_methods:
+                continue
+            try:
+                method_path = f'{name}.{method_name}'
+                wrap_function_wrapper(
+                    module_name,
+                    method_path,
+                    patch_method(
+                        method_path, task, tracer, version, trace_output, trace_input)
+                )
+            # pylint: disable=broad-except
+            except Exception:
+                pass
+
+
+class LangchainCoreInstrumentation(BaseInstrumentor):
+    """
+    Instrumentor for langchain.
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["langchain-core >= 0.1.27"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version('langchain-core')
+
+        exclude_methods = ['get_name', 'get_output_schema',
+                           'get_input_schema', 'get_graph', 'to_json']
+        exclude_classes = ['BaseChatPromptTemplate']
+        modules_to_patch = [
+            ('langchain_core.retrievers', 'retriever',
+             generic_patch, True, True),
+            ('langchain_core.prompts.chat', 'chatprompt',
+             generic_patch, True, False),
+            ('langchain_core.runnables.base',
+             'runnableparallel', runnable_patch, True, True),
+            ('langchain_core.runnables.passthrough',
+             'runnablepassthrough', runnable_patch, True, True),
+            ('langchain_core.output_parsers.string',
+             'stroutputparser', runnable_patch, True, True),
+            ('langchain_core.output_parsers.json',
+             'jsonoutputparser', runnable_patch, True, True),
+            ('langchain_core.output_parsers.list',
+             'listoutputparser', runnable_patch, True, True),
+            ('langchain_core.output_parsers.xml',
+             'xmloutputparser', runnable_patch, True, True)
+        ]
+
+        for module_name, task, patch_method, trace_output, trace_input in modules_to_patch:
+            patch_module_classes(module_name, tracer, version,
+                                 task, patch_method, exclude_methods,
+                                 exclude_classes, trace_output, trace_input)
+
+    def _uninstrument(self, **kwargs):
+        pass