opentelemetry-instrumentation-llamaindex 0.12.0__tar.gz → 0.12.2__tar.gz

{opentelemetry_instrumentation_llamaindex-0.12.0 → opentelemetry_instrumentation_llamaindex-0.12.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: opentelemetry-instrumentation-llamaindex
-Version: 0.12.0
+Version: 0.12.2
 Summary: OpenTelemetry LlamaIndex instrumentation
 License: Apache-2.0
 Author: Gal Kleinman

{opentelemetry_instrumentation_llamaindex-0.12.0 → opentelemetry_instrumentation_llamaindex-0.12.2}/opentelemetry/instrumentation/llamaindex/__init__.py

@@ -22,6 +22,9 @@ from opentelemetry.instrumentation.llamaindex.base_embedding_instrumentor import
 from opentelemetry.instrumentation.llamaindex.custom_llm_instrumentor import (
     CustomLLMInstrumentor,
 )
+from opentelemetry.instrumentation.llamaindex.query_pipeline_instrumentor import (
+    QueryPipelineInstrumentor,
+)
 from opentelemetry.instrumentation.llamaindex.version import __version__
 
 logger = logging.getLogger(__name__)
@@ -44,6 +47,7 @@ class LlamaIndexInstrumentor(BaseInstrumentor):
         BaseSynthesizerInstrumentor(tracer).instrument()
         BaseEmbeddingInstrumentor(tracer).instrument()
         CustomLLMInstrumentor(tracer).instrument()
+        QueryPipelineInstrumentor(tracer).instrument()
 
     def _uninstrument(self, **kwargs):
         pass

opentelemetry_instrumentation_llamaindex-0.12.2/opentelemetry/instrumentation/llamaindex/query_pipeline_instrumentor.py

@@ -0,0 +1,71 @@
+from importlib.metadata import version as package_version, PackageNotFoundError
+
+from wrapt import wrap_function_wrapper
+from opentelemetry.context import attach, set_value
+
+from opentelemetry.instrumentation.llamaindex.utils import (
+    _with_tracer_wrapper,
+    start_as_current_span_async,
+)
+from opentelemetry.semconv.ai import SpanAttributes, TraceloopSpanKindValues
+
+V9_MODULE_NAME = "llama_index.query_pipeline.query"
+V10_MODULE_NAME = "llama_index.core.query_pipeline.query"
+V10_LEGACY_MODULE_NAME = "llama_index.legacy.query_pipeline.query"
+
+CLASS_NAME = "QueryPipeline"
+WORKFLOW_NAME = "llama_index_query_pipeline"
+
+
+class QueryPipelineInstrumentor:
+    def __init__(self, tracer):
+        self._tracer = tracer
+
+    def instrument(self):
+        try:
+            package_version("llama-index-core")
+            self._instrument_module(V10_MODULE_NAME)
+            self._instrument_module(V10_LEGACY_MODULE_NAME)
+
+        except PackageNotFoundError:
+            self._instrument_module(V9_MODULE_NAME)
+
+    def _instrument_module(self, module_name):
+        wrap_function_wrapper(
+            module_name, f"{CLASS_NAME}.run", run_wrapper(self._tracer)
+        )
+        wrap_function_wrapper(
+            module_name, f"{CLASS_NAME}.arun", arun_wrapper(self._tracer)
+        )
+
+
+def set_workflow_context():
+    attach(set_value("workflow_name", WORKFLOW_NAME))
+
+
+@_with_tracer_wrapper
+def run_wrapper(tracer, wrapped, instance, args, kwargs):
+    set_workflow_context()
+
+    with tracer.start_as_current_span(f"{WORKFLOW_NAME}.workflow") as span:
+        span.set_attribute(
+            SpanAttributes.TRACELOOP_SPAN_KIND,
+            TraceloopSpanKindValues.WORKFLOW.value,
+        )
+
+        return wrapped(*args, **kwargs)
+
+
+@_with_tracer_wrapper
+async def arun_wrapper(tracer, wrapped, instance, args, kwargs):
+    set_workflow_context()
+
+    async with start_as_current_span_async(
+        tracer=tracer, name=f"{WORKFLOW_NAME}.workflow"
+    ) as span:
+        span.set_attribute(
+            SpanAttributes.TRACELOOP_SPAN_KIND,
+            TraceloopSpanKindValues.WORKFLOW.value,
+        )
+
+        return await wrapped(*args, **kwargs)

opentelemetry_instrumentation_llamaindex-0.12.2/opentelemetry/instrumentation/llamaindex/version.py

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

{opentelemetry_instrumentation_llamaindex-0.12.0 → opentelemetry_instrumentation_llamaindex-0.12.2}/pyproject.toml

@@ -8,7 +8,7 @@ show_missing = true
 
 [tool.poetry]
 name = "opentelemetry-instrumentation-llamaindex"
-version = "0.12.0"
+version = "0.12.2"
 description = "OpenTelemetry LlamaIndex instrumentation"
 authors = [
   "Gal Kleinman <gal@traceloop.com>",
@@ -40,8 +40,10 @@ chromadb = "^0.4.22"
 openai = "^1.12.0"
 opentelemetry-sdk = "^1.22.0"
 llama-index = "^0.10.10"
-opentelemetry-instrumentation-openai = "^0.12.0"
-opentelemetry-instrumentation-chromadb = "^0.12.0"
+llama-index-postprocessor-cohere-rerank = "^0.1.2"
+opentelemetry-instrumentation-openai = "^0.12.2"
+opentelemetry-instrumentation-cohere = "^0.12.2"
+opentelemetry-instrumentation-chromadb = "^0.12.2"
 
 
 [build-system]

opentelemetry_instrumentation_llamaindex-0.12.0/opentelemetry/instrumentation/llamaindex/version.py

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