chore: tests

traceloop · Feb 27, 2024 · 13c42b9 · 13c42b9
1 parent f1032bf
commit 13c42b9
Show file tree

Hide file tree

Showing 8 changed files with 1,896 additions and 30 deletions.
diff --git a/...umentation-llamaindex/opentelemetry/instrumentation/llamaindex/base_agent_instrumentor.py b/...umentation-llamaindex/opentelemetry/instrumentation/llamaindex/base_agent_instrumentor.py
@@ -9,32 +9,42 @@
 from opentelemetry.semconv.ai import SpanAttributes, TraceloopSpanKindValues
 
 
-V9_MODULE_NAME = "llama_index.agent.types"
-V10_MODULE_NAME = "llama_index.core.agent.types"
-V10_LEGACY_MODULE_NAME = "llama_index.legacy.agent.types"
-
-CLASS_NAME = "BaseAgent"
+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):
-        try:
-            package_version("llama-index-core")
-            self._instrument_module(V10_MODULE_NAME)
-            self._instrument_module(V10_LEGACY_MODULE_NAME)
+        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(V9_MODULE_NAME)
+            except PackageNotFoundError:
+                self._instrument_module(module["v9_module"], module["class"])
 
-    def _instrument_module(self, module_name):
+    def _instrument_module(self, module_name, class_name):
         wrap_function_wrapper(
-            module_name, f"{CLASS_NAME}._query", query_wrapper(self._tracer)
+            module_name, f"{class_name}.chat", query_wrapper(self._tracer)
         )
         wrap_function_wrapper(
-            module_name, f"{CLASS_NAME}._aquery", aquery_wrapper(self._tracer)
+            module_name, f"{class_name}.achat", aquery_wrapper(self._tracer)
         )
 
 

diff --git a/...rumentation-llamaindex/opentelemetry/instrumentation/llamaindex/base_tool_instrumentor.py b/...rumentation-llamaindex/opentelemetry/instrumentation/llamaindex/base_tool_instrumentor.py
@@ -9,32 +9,42 @@
 from opentelemetry.semconv.ai import SpanAttributes, TraceloopSpanKindValues
 
 
-V9_MODULE_NAME = "llama_index.tools.types"
-V10_MODULE_NAME = "llama_index.core.tools.types"
-V10_LEGACY_MODULE_NAME = "llama_index.legacy.tools.types"
-
-CLASS_NAME = "AsyncBaseTool"
+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):
-        try:
-            package_version("llama-index-core")
-            self._instrument_module(V10_MODULE_NAME)
-            self._instrument_module(V10_LEGACY_MODULE_NAME)
+        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(V9_MODULE_NAME)
+            except PackageNotFoundError:
+                self._instrument_module(module["v9_module"], module["class"])
 
-    def _instrument_module(self, module_name):
+    def _instrument_module(self, module_name, class_name):
         wrap_function_wrapper(
-            module_name, f"{CLASS_NAME}.call", query_wrapper(self._tracer)
+            module_name, f"{class_name}.call", query_wrapper(self._tracer)
         )
         wrap_function_wrapper(
-            module_name, f"{CLASS_NAME}.acall", aquery_wrapper(self._tracer)
+            module_name, f"{class_name}.acall", aquery_wrapper(self._tracer)
         )
 
 

diff --git a/packages/opentelemetry-instrumentation-llamaindex/poetry.lock b/packages/opentelemetry-instrumentation-llamaindex/poetry.lock
diff --git a/packages/opentelemetry-instrumentation-llamaindex/pyproject.toml b/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]