From 5e5bc88443ae80cc38a4aef3bf1e5744be355f9e Mon Sep 17 00:00:00 2001
From: Alex Feel <71711753+alex-feel@users.noreply.github.com>
Date: Tue, 13 Feb 2024 20:38:08 +0200
Subject: [PATCH] feat: support both new and legacy llama_index versions (#420)

---
 .../llamaindex/custom_llm_instrumentor.py                | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/custom_llm_instrumentor.py b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/custom_llm_instrumentor.py
index 3b22df2b9..0ce4b2df5 100644
--- a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/custom_llm_instrumentor.py
+++ b/packages/opentelemetry-instrumentation-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: