fix(llamaindex): instrument agents & tools (#533)

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

@@ -7,6 +7,9 @@
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 
+from opentelemetry.instrumentation.llamaindex.base_agent_instrumentor import (
+    BaseAgentInstrumentor,
+)
 from opentelemetry.instrumentation.llamaindex.retriever_query_engine_instrumentor import (
     RetrieverQueryEngineInstrumentor,
 )
@@ -16,6 +19,9 @@
 from opentelemetry.instrumentation.llamaindex.base_synthesizer_instrumentor import (
     BaseSynthesizerInstrumentor,
 )
+from opentelemetry.instrumentation.llamaindex.base_tool_instrumentor import (
+    BaseToolInstrumentor,
+)
 from opentelemetry.instrumentation.llamaindex.base_embedding_instrumentor import (
     BaseEmbeddingInstrumentor,
 )
@@ -48,6 +54,8 @@ def _instrument(self, **kwargs):
         BaseEmbeddingInstrumentor(tracer).instrument()
         CustomLLMInstrumentor(tracer).instrument()
         QueryPipelineInstrumentor(tracer).instrument()
+        BaseAgentInstrumentor(tracer).instrument()
+        BaseToolInstrumentor(tracer).instrument()
 
     def _uninstrument(self, **kwargs):
         pass

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

@@ -0,0 +1,72 @@
+from importlib.metadata import version as package_version, PackageNotFoundError
+
+from wrapt import wrap_function_wrapper
+
+from opentelemetry.instrumentation.llamaindex.utils import (
+    _with_tracer_wrapper,
+    start_as_current_span_async,
+)
+from opentelemetry.semconv.ai import SpanAttributes, TraceloopSpanKindValues
+
+
+TO_INSTRUMENT = [
+    {
+        "class": "AgentRunner",
+        "v9_module": "llama_index.agent.runner.base",
+        "v10_module": "llama_index.core.agent.runner.base",
+        "v10_legacy_module": "llama_index.legacy.agent.runner.base",
+    },
+    {
+        "class": "OpenAIAssistantAgent",
+        "v9_module": "llama_index.agent.openai_assistant_agent",
+        "v10_module": "llama_index.agent.openai.openai_assistant_agent",
+        "v10_legacy_module": "llama_index.legacy.agent.openai_assistant_agent",
+    },
+]
+
+
+class BaseAgentInstrumentor:
+    def __init__(self, tracer):
+        self._tracer = tracer
+
+    def instrument(self):
+        for module in TO_INSTRUMENT:
+            try:
+                package_version("llama-index-core")
+                self._instrument_module(module["v10_module"], module["class"])
+                self._instrument_module(module["v10_legacy_module"], module["class"])
+
+            except PackageNotFoundError:
+                self._instrument_module(module["v9_module"], module["class"])
+
+    def _instrument_module(self, module_name, class_name):
+        wrap_function_wrapper(
+            module_name, f"{class_name}.chat", query_wrapper(self._tracer)
+        )
+        wrap_function_wrapper(
+            module_name, f"{class_name}.achat", aquery_wrapper(self._tracer)
+        )
+
+
+@_with_tracer_wrapper
+def query_wrapper(tracer, wrapped, instance, args, kwargs):
+    with tracer.start_as_current_span(f"{instance.__class__.__name__}.agent") as span:
+        span.set_attribute(
+            SpanAttributes.TRACELOOP_SPAN_KIND,
+            TraceloopSpanKindValues.AGENT.value,
+        )
+
+        return wrapped(*args, **kwargs)
+
+
+@_with_tracer_wrapper
+async def aquery_wrapper(tracer, wrapped, instance, args, kwargs):
+    async with start_as_current_span_async(
+        tracer=tracer, name=f"{instance.__class__.__name__}.agent"
+    ) as span:
+        span.set_attribute(
+            SpanAttributes.TRACELOOP_SPAN_KIND,
+            TraceloopSpanKindValues.AGENT.value,
+        )
+
+        return await wrapped(*args, **kwargs)

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

@@ -0,0 +1,72 @@
+from importlib.metadata import version as package_version, PackageNotFoundError
+
+from wrapt import wrap_function_wrapper
+
+from opentelemetry.instrumentation.llamaindex.utils import (
+    _with_tracer_wrapper,
+    start_as_current_span_async,
+)
+from opentelemetry.semconv.ai import SpanAttributes, TraceloopSpanKindValues
+
+
+TO_INSTRUMENT = [
+    {
+        "class": "FunctionTool",
+        "v9_module": "llama_index.tools.function_tool",
+        "v10_module": "llama_index.core.tools.function_tool",
+        "v10_legacy_module": "llama_index.legacy.tools.function_tool",
+    },
+    {
+        "class": "QueryEngineTool",
+        "v9_module": "llama_index.tools.query_engine",
+        "v10_module": "llama_index.core.tools.query_engine",
+        "v10_legacy_module": "llama_index.legacy.tools.query_engine",
+    },
+]
+
+
+class BaseToolInstrumentor:
+    def __init__(self, tracer):
+        self._tracer = tracer
+
+    def instrument(self):
+        for module in TO_INSTRUMENT:
+            try:
+                package_version("llama-index-core")
+                self._instrument_module(module["v10_module"], module["class"])
+                self._instrument_module(module["v10_legacy_module"], module["class"])
+
+            except PackageNotFoundError:
+                self._instrument_module(module["v9_module"], module["class"])
+
+    def _instrument_module(self, module_name, class_name):
+        wrap_function_wrapper(
+            module_name, f"{class_name}.call", query_wrapper(self._tracer)
+        )
+        wrap_function_wrapper(
+            module_name, f"{class_name}.acall", aquery_wrapper(self._tracer)
+        )
+
+
+@_with_tracer_wrapper
+def query_wrapper(tracer, wrapped, instance, args, kwargs):
+    with tracer.start_as_current_span(f"{instance.__class__.__name__}.tool") as span:
+        span.set_attribute(
+            SpanAttributes.TRACELOOP_SPAN_KIND,
+            TraceloopSpanKindValues.TOOL.value,
+        )
+
+        return wrapped(*args, **kwargs)
+
+
+@_with_tracer_wrapper
+async def aquery_wrapper(tracer, wrapped, instance, args, kwargs):
+    async with start_as_current_span_async(
+        tracer=tracer, name=f"{instance.__class__.__name__}.tool"
+    ) as span:
+        span.set_attribute(
+            SpanAttributes.TRACELOOP_SPAN_KIND,
+            TraceloopSpanKindValues.TOOL.value,
+        )
+
+        return await wrapped(*args, **kwargs)

packages/opentelemetry-instrumentation-llamaindex/pyproject.toml

@@ -44,6 +44,8 @@ llama-index-postprocessor-cohere-rerank = "^0.1.2"
 opentelemetry-instrumentation-openai = {path="../opentelemetry-instrumentation-openai", develop=true}
 opentelemetry-instrumentation-cohere = {path="../opentelemetry-instrumentation-cohere", develop=true}
 opentelemetry-instrumentation-chromadb = {path="../opentelemetry-instrumentation-chromadb", develop=true}
+sqlalchemy = "^2.0.27"
+llama-index-agent-openai = "^0.1.5"
 
 
 [build-system]