feat(llama-index): migrate to v0.10.x (#331)

Ref: #330

examples/extensions/llama_index/llama_index_llm.py

@@ -1,8 +1,7 @@
 """Use a model through LLamaIndex"""
 
 from dotenv import load_dotenv
-from llama_index.llms.base import ChatMessage
-from llama_index.llms.types import MessageRole
+from llama_index.core.llms import ChatMessage, MessageRole
 
 from genai import Client
 from genai.credentials import Credentials

pyproject.toml

@@ -40,7 +40,7 @@ langchain-core = { version = "^0.1.0", optional = true }
 pyyaml = { version = "^6.0.0", optional = true }
 datasets = { version = "^2.13.0",  optional = true }
 transformers = { version = "^4.33.3",  optional = true, extras=["agents"] }
-llama-index = { version = "^0.9.15",  optional = true }
+llama-index-core = { version = "^0.10.0", optional = true }
 uvicorn = { version = "^0.22.0",  optional = true }
 fastapi = { version = "^0.100.0",  optional = true }
 
@@ -110,7 +110,7 @@ sqlalchemy = "^2.0.25"
 [tool.poetry.extras]
 langchain = ["langchain-core", "pyyaml"]
 huggingface = ["datasets", "transformers"]
-llama-index = ["llama-index"]
+llama-index = ["llama-index-core"]
 localserver = ["uvicorn", "fastapi"]
 
 [tool.pytest.ini_options]

src/genai/extensions/llama_index/embeddings.py

@@ -7,7 +7,7 @@
 from genai.text.embedding.embedding_service import CreateExecutionOptions
 
 try:
-    from llama_index.embeddings.base import BaseEmbedding, Embedding
+    from llama_index.core.base.embeddings.base import BaseEmbedding, Embedding
 except ImportError:
     raise ImportError("Could not import llamaindex: Please install ibm-generative-ai[llama-index] extension.")  # noqa: B904
 

src/genai/extensions/llama_index/llm.py

@@ -24,8 +24,9 @@
 
 
 try:
-    from llama_index.callbacks import CallbackManager
-    from llama_index.llms.base import (
+    from llama_index.core.callbacks import CallbackManager
+    from llama_index.core.llms import (
+        LLM,
         ChatMessage,
         ChatResponse,
         ChatResponseAsyncGen,
@@ -34,11 +35,9 @@
         CompletionResponseAsyncGen,
         CompletionResponseGen,
         LLMMetadata,
-        llm_chat_callback,
-        llm_completion_callback,
+        MessageRole,
     )
-    from llama_index.llms.llm import LLM
-    from llama_index.llms.types import MessageRole
+    from llama_index.core.llms.callbacks import llm_chat_callback, llm_completion_callback
 
 
 except ImportError:

tests/integration/extensions/test_llama_index.py

@@ -1,7 +1,12 @@
 import pytest
-from llama_index.llms.base import ChatMessage, CompletionResponse
-from llama_index.llms.base import ChatResponse as LlamaIndexChatResponse
-from llama_index.llms.types import MessageRole
+from llama_index.core.llms import (
+    ChatMessage,
+    CompletionResponse,
+    MessageRole,
+)
+from llama_index.core.llms import (
+    ChatResponse as LlamaIndexChatResponse,
+)
 
 from genai.extensions._common.utils import create_generation_info_from_response
 from genai.extensions.llama_index import IBMGenAILlamaIndex