feat: llamaindex auto instrumentation (#157)

.cz.toml

@@ -28,6 +28,8 @@ version_files = [
     "packages/opentelemetry-instrumentation-vertexai/opentelemetry/instrumentation/vertexai/version.py",
     "packages/opentelemetry-instrumentation-weaviate/pyproject.toml:^version",
     "packages/opentelemetry-instrumentation-weaviate/opentelemetry/instrumentation/weaviate/version.py",
+    "packages/opentelemetry-instrumentation-llamaindex/pyproject.toml:^version",
+    "packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/version.py",
     "packages/traceloop-sdk/pyproject.toml:^version",
     "packages/traceloop-sdk/traceloop/sdk/version.py"
 ]

packages/opentelemetry-instrumentation-llamaindex/.flake8

@@ -0,0 +1,11 @@
+[flake8]
+exclude =
+	.git,
+	__pycache__,
+	build,
+	dist,
+	.tox,
+	venv,
+	.venv,
+	.pytest_cache
+max-line-length = 120

packages/opentelemetry-instrumentation-llamaindex/.python-version

@@ -0,0 +1 @@
+3.9.5

packages/opentelemetry-instrumentation-llamaindex/README.md

@@ -0,0 +1,3 @@
+# opentelemetry-instrumentation-llamaindex
+
+Project description here.

packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/__init__.py

@@ -0,0 +1,105 @@
+"""OpenTelemetry LlamaIndex instrumentation"""
+import logging
+from typing import Collection
+from wrapt import wrap_function_wrapper
+
+from opentelemetry.trace import get_tracer
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+
+from opentelemetry.instrumentation.llamaindex.task_wrapper import task_wrapper
+from opentelemetry.instrumentation.llamaindex.workflow_wrapper import workflow_wrapper
+from opentelemetry.instrumentation.llamaindex.version import __version__
+
+logger = logging.getLogger(__name__)
+
+_instruments = ("llama-index >= 0.7.0",)
+
+WRAPPED_METHODS = [
+    {
+        "package": "llama_index.indices.query.base",
+        "object": "BaseQueryEngine",
+        "method": "query",
+        "wrapper": workflow_wrapper,
+    },
+    {
+        "package": "llama_index.indices.query.base",
+        "object": "BaseQueryEngine",
+        "method": "aquery",
+        "wrapper": workflow_wrapper,
+    },
+    {
+        "package": "llama_index.indices.base_retriever",
+        "object": "BaseRetriever",
+        "method": "retrieve",
+        "span_name": "retrieve",
+        "wrapper": task_wrapper
+    },
+    {
+        "package": "llama_index.indices.base_retriever",
+        "object": "BaseRetriever",
+        "method": "aretrieve",
+        "span_name": "retrieve",
+        "wrapper": task_wrapper
+    },
+    {
+        "package": "llama_index.embeddings.base",
+        "object": "BaseEmbedding",
+        "method": "get_query_embedding",
+        "span_name": "get_query_embedding",
+        "wrapper": task_wrapper
+    },
+    {
+        "package": "llama_index.embeddings.base",
+        "object": "BaseEmbedding",
+        "method": "aget_query_embedding",
+        "span_name": "get_query_embedding",
+        "wrapper": task_wrapper
+    },
+    {
+        "package": "llama_index.response_synthesizers",
+        "object": "BaseSynthesizer",
+        "method": "synthesize",
+        "span_name": "synthesize",
+        "wrapper": task_wrapper
+    },
+    {
+        "package": "llama_index.response_synthesizers",
+        "object": "BaseSynthesizer",
+        "method": "asynthesize",
+        "span_name": "synthesize",
+        "wrapper": task_wrapper
+    }
+]
+
+
+class LlamaIndexInstrumentor(BaseInstrumentor):
+    """An instrumentor for LlamaIndex SDK."""
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            wrap_method = wrapped_method.get("method")
+            wrapper = wrapped_method.get("wrapper")
+            wrap_function_wrapper(
+                wrap_package,
+                f"{wrap_object}.{wrap_method}" if wrap_object else wrap_method,
+                wrapper(tracer, wrapped_method),
+            )
+
+    def _uninstrument(self, **kwargs):
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            wrap_method = wrapped_method.get("method")
+            unwrap(
+                f"{wrap_package}.{wrap_object}" if wrap_object else wrap_package,
+                wrap_method,
+            )

packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/task_wrapper.py

@@ -0,0 +1,27 @@
+from opentelemetry import context as context_api
+
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+
+from opentelemetry.semconv.ai import SpanAttributes, TraceloopSpanKindValues
+
+from opentelemetry.instrumentation.llamaindex.utils import _with_tracer_wrapper
+
+
+@_with_tracer_wrapper
+def task_wrapper(tracer, to_wrap, wrapped, instance, args, kwargs):
+    """Instruments and calls every function defined in TO_WRAP."""
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    name = to_wrap.get("span_name")
+    kind = to_wrap.get("kind") or TraceloopSpanKindValues.TASK.value
+    with tracer.start_as_current_span(f"{name}.task") as span:
+        span.set_attribute(
+            SpanAttributes.TRACELOOP_SPAN_KIND,
+            kind,
+        )
+        span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name)
+
+        return_value = wrapped(*args, **kwargs)
+
+    return return_value

packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/utils.py

@@ -0,0 +1,10 @@
+def _with_tracer_wrapper(func):
+    """Helper for providing tracer for wrapper functions."""
+
+    def _with_tracer(tracer, to_wrap):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
+
+        return wrapper
+
+    return _with_tracer

packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.12"

packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/workflow_wrapper.py

@@ -0,0 +1,34 @@
+from opentelemetry import context as context_api
+from opentelemetry.context import attach, set_value
+from inflection import underscore
+
+from opentelemetry.instrumentation.utils import (
+    _SUPPRESS_INSTRUMENTATION_KEY,
+)
+
+from opentelemetry.semconv.ai import SpanAttributes, TraceloopSpanKindValues
+
+from opentelemetry.instrumentation.llamaindex.utils import _with_tracer_wrapper
+
+
+@_with_tracer_wrapper
+def workflow_wrapper(tracer, to_wrap, wrapped, instance, args, kwargs):
+    """Instruments and calls every function defined in TO_WRAP."""
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    name = f"llama_index_{underscore(instance.__class__.__name__)}"
+    kind = to_wrap.get("kind") or TraceloopSpanKindValues.WORKFLOW.value
+
+    attach(set_value("workflow_name", name))
+
+    with tracer.start_as_current_span(f"{name}.workflow") as span:
+        span.set_attribute(
+            SpanAttributes.TRACELOOP_SPAN_KIND,
+            kind,
+        )
+        span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, name)
+
+        return_value = wrapped(*args, **kwargs)
+
+    return return_value