feat: support llama-index new instrumentation paradigm under feature …

…flag `use_experimental_instrumentation` (#462)

python/instrumentation/openinference-instrumentation-llama-index/examples/agent_calculator_tools.py

@@ -0,0 +1,34 @@
+from llama_index.agent.openai import OpenAIAgent
+from llama_index.core import Settings
+from llama_index.core.tools import FunctionTool
+from llama_index.llms.openai import OpenAI
+from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+endpoint = "http://127.0.0.1:6006/v1/traces"
+tracer_provider = trace_sdk.TracerProvider()
+tracer_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint)))
+
+LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
+
+
+def multiply(a: int, b: int) -> int:
+    """Multiple two integers and returns the result integer"""
+    return a * b
+
+
+def add(a: int, b: int) -> int:
+    """Add two integers and returns the result integer"""
+    return a + b
+
+
+multiply_tool = FunctionTool.from_defaults(fn=multiply)
+add_tool = FunctionTool.from_defaults(fn=add)
+agent = OpenAIAgent.from_tools([multiply_tool, add_tool])
+Settings.llm = OpenAI(model="gpt-3.5-turbo")
+
+if __name__ == "__main__":
+    response = agent.query("What is (121 * 3) + 42?")
+    print(response)

python/instrumentation/openinference-instrumentation-llama-index/examples/chroma_query_engine.py

@@ -1,29 +1,36 @@
+import tempfile
+from urllib.request import urlretrieve
+
 import chromadb
-from llama_index.core import SimpleDirectoryReader, StorageContext, VectorStoreIndex
+from llama_index.core import Settings, SimpleDirectoryReader, StorageContext, VectorStoreIndex
+from llama_index.llms.openai import OpenAI
 from llama_index.vector_stores.chroma import ChromaVectorStore
 from openinference.instrumentation import using_attributes
 from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
-from opentelemetry import trace as trace_api
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk import trace as trace_sdk
 from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor
 
 endpoint = "http://127.0.0.1:6006/v1/traces"
 tracer_provider = trace_sdk.TracerProvider()
-trace_api.set_tracer_provider(tracer_provider)
 tracer_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint)))
 tracer_provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
 
-LlamaIndexInstrumentor().instrument()
+LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
+
 
 chroma_client = chromadb.EphemeralClient()
 chroma_collection = chroma_client.create_collection("essays")
-
-documents = SimpleDirectoryReader("./data/paul_graham/").load_data()
 vector_store = ChromaVectorStore(chroma_collection)
 storage_context = StorageContext.from_defaults(vector_store=vector_store)
+with tempfile.NamedTemporaryFile() as tf:
+    urlretrieve(
+        "https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt",
+        tf.name,
+    )
+    documents = SimpleDirectoryReader(input_files=[tf.name]).load_data()
 index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
-
+Settings.llm = OpenAI(model="gpt-3.5-turbo")
 
 if __name__ == "__main__":
     query_engine = index.as_query_engine()

python/instrumentation/openinference-instrumentation-llama-index/examples/query_engine_tools.py

@@ -0,0 +1,57 @@
+import pandas as pd
+import wikipedia
+from llama_index.core import Document, Settings
+from llama_index.core.indices import VectorStoreIndex
+from llama_index.core.query_engine import NLSQLTableQueryEngine, RouterQueryEngine
+from llama_index.core.selectors import LLMSingleSelector
+from llama_index.core.tools import QueryEngineTool
+from llama_index.core.utilities.sql_wrapper import SQLDatabase
+from llama_index.llms.openai import OpenAI
+from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from sqlalchemy import create_engine
+
+endpoint = "http://127.0.0.1:6006/v1/traces"
+tracer_provider = trace_sdk.TracerProvider()
+tracer_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint)))
+
+LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
+
+engine = create_engine("sqlite:///:memory:")
+pd.read_parquet(
+    "https://storage.googleapis.com/arize-phoenix-assets/datasets/structured/camera-info/cameras.parquet"
+).to_sql("cameras", engine, index=False)
+sql_tool = QueryEngineTool.from_defaults(
+    query_engine=NLSQLTableQueryEngine(
+        sql_database=SQLDatabase(engine, include_tables=["cameras"]),
+        tables=["cameras"],
+    ),
+    description=(
+        "Useful for translating a natural language query into a SQL query over"
+        " a table containing technical details about specific digital camera models: Model,"
+        " Release date, Max resolution, Low resolution, Effective pixels, Zoom wide (W),"
+        " Zoom tele (T), Normal focus range, Macro focus range, Storage included,"
+        " Weight (inc. batteries), Dimensions, Price"
+    ),
+)
+
+page = wikipedia.page(pageid=52797)
+vector_tool = QueryEngineTool.from_defaults(
+    query_engine=VectorStoreIndex.from_documents(
+        [Document(id_=page.pageid, text=page.content)]
+    ).as_query_engine(),
+    description="Useful for answering generic questions about digital cameras.",
+)
+query_engine = RouterQueryEngine(
+    selector=LLMSingleSelector.from_defaults(),
+    query_engine_tools=[sql_tool, vector_tool],
+)
+Settings.llm = OpenAI(model="gpt-3.5-turbo")
+
+if __name__ == "__main__":
+    response = query_engine.query("What is the most expensive digital camera?")
+    print(str(response))
+    response = query_engine.query("Tell me about the history of digital camera sensors.")
+    print(str(response))

python/instrumentation/openinference-instrumentation-llama-index/examples/requirements.txt

@@ -1,6 +1,10 @@
+chromadb
 llama-index >= 0.10.0
+llama-index-agent-openai
+llama-index-llms-openai
 llama-index-vector-stores-chroma
 openinference-instrumentation-llama-index
-opentelemetry-sdk
 opentelemetry-exporter-otlp
-chromadb
+opentelemetry-sdk
+sqlalchemy
+wikipedia

python/instrumentation/openinference-instrumentation-llama-index/examples/streaming.py

@@ -0,0 +1,31 @@
+import tempfile
+from urllib.request import urlretrieve
+
+from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
+from llama_index.llms.openai import OpenAI
+from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+endpoint = "http://127.0.0.1:6006/v1/traces"
+tracer_provider = trace_sdk.TracerProvider()
+tracer_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint)))
+
+LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
+
+
+with tempfile.NamedTemporaryFile() as tf:
+    urlretrieve(
+        "https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt",
+        tf.name,
+    )
+    documents = SimpleDirectoryReader(input_files=[tf.name]).load_data()
+
+index = VectorStoreIndex.from_documents(documents)
+Settings.llm = OpenAI(model="gpt-3.5-turbo")
+
+if __name__ == "__main__":
+    query_engine = index.as_query_engine(streaming=True, similarity_top_k=1)
+    response_stream = query_engine.query("What did the author do growing up?")
+    response_stream.print_response_stream()

python/instrumentation/openinference-instrumentation-llama-index/pyproject.toml

@@ -36,11 +36,11 @@ dependencies = [
 
 [project.optional-dependencies]
 instruments = [
-  "llama-index >= 0.10.5",
+  "llama-index-core >= 0.10.5",
 ]
 test = [
-  "llama-index == 0.10.5",
-  "llama-index-core == 0.10.10",
+  "llama-index == 0.10.36",
+  "llama-index-core == 0.10.36",
   "llama-index-llms-openai",
   "opentelemetry-sdk",
   "respx",

python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/__init__.py

@@ -1,3 +1,4 @@
+import importlib
 import logging
 from typing import Any, Collection
 
@@ -6,37 +7,101 @@
 from opentelemetry import trace as trace_api
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore
 
-_MODULE = "llama_index.core"
-
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
 
+_ELIGIBLE_VERSION_FOR_NEW_INSTRUMENTATION = (0, 10, 37)
+
 
 class LlamaIndexInstrumentor(BaseInstrumentor):  # type: ignore
     """
     An instrumentor for LlamaIndex
     """
 
-    __slots__ = ("_original_global_handler",)
+    __slots__ = (
+        "_original_global_handler",
+        "_use_experimental_instrumentation",  # feature flag
+        "_event_handler",
+    )
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
     def _instrument(self, **kwargs: Any) -> None:
         if not (tracer_provider := kwargs.get("tracer_provider")):
             tracer_provider = trace_api.get_tracer_provider()
+        self._use_experimental_instrumentation = kwargs.get("use_experimental_instrumentation")
+        if self._use_experimental_instrumentation:
+            if not _legacy_llama_index():
+                print(
+                    "`use_experimental_instrumentation` feature flag is set. Spans "
+                    "will be generated using the new instrumentation system "
+                    "For more information about the new instrumentation system, visit "
+                    "https://docs.llamaindex.ai/en/stable/module_guides/observability/instrumentation/"  # noqa E501
+                )
+            else:
+                print(
+                    f"`use_experimental_instrumentation` feature flag is set. But "
+                    f"the version of `llama-index-core` is not "
+                    f">={'.'.join(map(str, _ELIGIBLE_VERSION_FOR_NEW_INSTRUMENTATION))}, "
+                    f"so the flag is ignored."
+                )
         tracer = trace_api.get_tracer(__name__, __version__, tracer_provider)
         from openinference.instrumentation.llama_index._callback import (
             OpenInferenceTraceCallbackHandler,
         )
 
-        import llama_index.core
+        if (
+            _legacy_llama_index()
+            or not self._use_experimental_instrumentation
+            or not _legacy_llama_index()
+            and self._use_experimental_instrumentation == "both"
+        ):
+            import llama_index.core
+
+            self._original_global_handler = llama_index.core.global_handler
+            llama_index.core.global_handler = OpenInferenceTraceCallbackHandler(tracer=tracer)
+
+        self._event_handler = None
+        if not _legacy_llama_index() and self._use_experimental_instrumentation:
+            from llama_index.core.instrumentation import get_dispatcher
+
+            from ._handler import EventHandler
 
-        self._original_global_handler = llama_index.core.global_handler
-        llama_index.core.global_handler = OpenInferenceTraceCallbackHandler(tracer=tracer)
+            self._event_handler = EventHandler(tracer=tracer)
+            dispatcher = get_dispatcher()
+            dispatcher.add_event_handler(self._event_handler)
+            dispatcher.add_span_handler(self._event_handler.span_handler)
 
     def _uninstrument(self, **kwargs: Any) -> None:
-        import llama_index.core
+        if (
+            _legacy_llama_index()
+            or not self._use_experimental_instrumentation
+            or not _legacy_llama_index()
+            and self._use_experimental_instrumentation == "both"
+        ):
+            import llama_index.core
+
+            llama_index.core.global_handler = self._original_global_handler
+            self._original_global_handler = None
+
+        if not _legacy_llama_index() and self._use_experimental_instrumentation:
+            if self._event_handler is None:
+                return
+            from llama_index.core.instrumentation import get_dispatcher
+
+            dispatcher = get_dispatcher()
+            dispatcher.span_handlers[:] = filter(
+                lambda h: h is not self._event_handler.span_handler,  # type: ignore
+                dispatcher.span_handlers,
+            )
+            dispatcher.event_handlers[:] = filter(
+                lambda h: h is not self._event_handler,
+                dispatcher.event_handlers,
+            )
+            self._event_handler = None
+
 
-        llama_index.core.global_handler = self._original_global_handler
-        self._original_global_handler = None
+def _legacy_llama_index() -> bool:
+    v = importlib.metadata.version("llama-index-core")
+    return tuple(map(int, v.split(".")[:3])) < _ELIGIBLE_VERSION_FOR_NEW_INSTRUMENTATION