feat: support both new and legacy llama_index versions (#420)

traceloop · Feb 13, 2024 · 5e5bc88 · 5e5bc88
1 parent 18a5537
commit 5e5bc88
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/...umentation-llamaindex/opentelemetry/instrumentation/llamaindex/custom_llm_instrumentor.py b/...umentation-llamaindex/opentelemetry/instrumentation/llamaindex/custom_llm_instrumentor.py
@@ -11,9 +11,12 @@
     _with_tracer_wrapper, start_as_current_span_async, should_send_prompts
 )
 
-from llama_index.llms.custom import CustomLLM
-
-MODULE_NAME = "llama_index.llms"
+try:
+    from llama_index.core.llms import CustomLLM
+    MODULE_NAME = "llama_index.core.llms"
+except ModuleNotFoundError:
+    from llama_index.legacy.llms.custom import CustomLLM
+    MODULE_NAME = "llama_index.legacy.llms"
 
 
 class CustomLLMInstrumentor: