feat: add llama-parse as dependency and bump llama-index to 0.9.48 (#279)

Author: NadirJ

.github/workflows/_run_e2e_tests.yml

@@ -88,6 +88,7 @@ jobs:
           ASTRA_DB_ID: "${{ steps.astra-db.outputs.db_id }}"
           OPENAI_API_KEY: "${{ secrets.E2E_TESTS_OPEN_AI_KEY }}"
           LANGCHAIN_API_KEY: "${{ secrets.E2E_TESTS_LANGCHAIN_API_KEY }}"
+          LLAMA_CLOUD_API_KEY: "${{ secrets.E2E_TESTS_LLAMA_CLOUD_API_KEY }}"
           GCLOUD_ACCOUNT_KEY_JSON: "${{ secrets.E2E_TESTS_GCLOUD_ACCOUNT_KEY_JSON }}"
         run: |
           source scripts/ci-common-env.sh
@@ -114,6 +115,7 @@ jobs:
           HUGGINGFACE_HUB_KEY: "${{ secrets.E2E_TESTS_HUGGINGFACE_HUB_KEY }}"
           NVIDIA_API_KEY: "${{ secrets.E2E_TESTS_NVIDIA_API_KEY }}"
           LANGCHAIN_API_KEY: "${{ secrets.E2E_TESTS_LANGCHAIN_API_KEY }}"
+          LLAMA_CLOUD_API_KEY: "${{ secrets.E2E_TESTS_LLAMA_CLOUD_API_KEY }}"
         run: |
           source scripts/ci-common-env.sh
           if [ "${{ inputs.suite-name == 'ragstack' }}" == "true" ]; then

pyproject.toml

@@ -7,23 +7,22 @@ authors = ["DataStax"]
 readme = "PACKAGE_README.md"
 repository = "https://github.com/datastax/ragstack-ai"
 documentation = "https://docs.datastax.com/en/ragstack"
-packages = [
-    {include = "ragstack"}
-]
+packages = [{ include = "ragstack" }]
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"
 astrapy = "~0.7.0"
 cassio = "~0.1.3"
 unstructured = "^0.10"
-llama-index = { version = "0.9.34", extras = ["langchain"] }
-langchain = {version = "0.1.4"}
+llama-index = { version = "0.9.48", extras = ["langchain"] }
+llama-parse = { version = "0.1.4" }
+langchain = { version = "0.1.4" }
 langchain-core = "0.1.16"
 langchain-community = "0.0.16"
-langchain-openai = {version = "0.0.3"}
-langchain-google-genai = {version = "0.0.6", optional = true}
-langchain-google-vertexai = {version = "0.0.3", optional = true}
-langchain-nvidia-ai-endpoints = {version = "0.0.1", optional = true}
+langchain-openai = { version = "0.0.3" }
+langchain-google-genai = { version = "0.0.6", optional = true }
+langchain-google-vertexai = { version = "0.0.3", optional = true }
+langchain-nvidia-ai-endpoints = { version = "0.0.1", optional = true }
 
 [tool.poetry.extras]
 langchain-google = ["langchain-google-genai", "langchain-google-vertexai"]
@@ -39,5 +38,4 @@ requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry.dev-dependencies]
-yamllint = "^1.34.0" 
-
+yamllint = "^1.34.0"

ragstack-e2e-tests/.env.template

@@ -29,4 +29,7 @@ VECTOR_DATABASE_TYPE=astradb
 # HUGGINGFACE_HUB_KEY=
 
 # Nvidia
-# NVIDIA_API_KEY=
+# NVIDIA_API_KEY=
+
+# LlamaIndex
+# LLAMA_CLOUD_API_KEY=

ragstack-e2e-tests/e2e_tests/langchain/test_compatibility_rag.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from typing import List
 
 import pytest
@@ -15,6 +14,7 @@
     run_conversational_rag,
 )
 from e2e_tests.langchain.trulens import run_trulens_evaluation
+from e2e_tests.test_utils import get_local_resource_path
 
 from langchain.chat_models import ChatOpenAI, AzureChatOpenAI, ChatVertexAI, BedrockChat
 from langchain.embeddings import (
@@ -341,12 +341,6 @@ def embed_query(self, text: str) -> List[float]:
         assert "Coffee Machine Ultra Cool" in response.content
 
 
-def get_local_resource_path(filename: str):
-    dirname = os.path.dirname(__file__)
-    e2e_tests_dir = os.path.dirname(dirname)
-    return os.path.join(e2e_tests_dir, "resources", filename)
-
-
 @pytest.mark.parametrize("chat", ["vertex_gemini_pro_llm", "gemini_pro_llm"])
 def test_chat(chat, request, record_property):
     set_current_test_info(

ragstack-e2e-tests/e2e_tests/llama_index/conftest.py

@@ -0,0 +1,25 @@
+import pytest
+
+from e2e_tests.conftest import (
+    get_vector_store_handler,
+)
+
+from e2e_tests.test_utils.vector_store_handler import (
+    VectorStoreImplementation,
+)
+
+
+@pytest.fixture
+def astra_db():
+    handler = get_vector_store_handler(VectorStoreImplementation.ASTRADB)
+    context = handler.before_test()
+    yield context
+    handler.after_test()
+
+
+@pytest.fixture
+def cassandra():
+    handler = get_vector_store_handler(VectorStoreImplementation.CASSANDRA)
+    context = handler.before_test()
+    yield context
+    handler.after_test()

ragstack-e2e-tests/e2e_tests/llama_index/test_compatibility_rag.py

@@ -1,5 +1,4 @@
 import logging
-import os
 
 import pytest
 from langchain.embeddings import VertexAIEmbeddings, HuggingFaceInferenceAPIEmbeddings
@@ -29,32 +28,15 @@
 from e2e_tests.conftest import (
     set_current_test_info,
     get_required_env,
-    get_vector_store_handler,
 )
 from vertexai.vision_models import MultiModalEmbeddingModel, Image
 
+from e2e_tests.test_utils import get_local_resource_path
 from e2e_tests.test_utils.vector_store_handler import (
-    VectorStoreImplementation,
     VectorStoreTestContext,
 )
 
 
-@pytest.fixture
-def astra_db():
-    handler = get_vector_store_handler(VectorStoreImplementation.ASTRADB)
-    context = handler.before_test()
-    yield context
-    handler.after_test()
-
-
-@pytest.fixture
-def cassandra():
-    handler = get_vector_store_handler(VectorStoreImplementation.CASSANDRA)
-    context = handler.before_test()
-    yield context
-    handler.after_test()
-
-
 @pytest.fixture
 def openai_llm():
     return "openai", OpenAI(api_key=get_required_env("OPEN_AI_KEY"))
@@ -110,7 +92,7 @@ def bedrock_anthropic_llm():
         model="anthropic.claude-v2",
         aws_access_key_id=get_required_env("AWS_ACCESS_KEY_ID"),
         aws_secret_access_key=get_required_env("AWS_SECRET_ACCESS_KEY"),
-        aws_region_name=get_required_env("BEDROCK_AWS_REGION"),
+        region_name=get_required_env("BEDROCK_AWS_REGION"),
     )
 
 
@@ -120,7 +102,7 @@ def bedrock_meta_llm():
         model="meta.llama2-13b-chat-v1",
         aws_access_key_id=get_required_env("AWS_ACCESS_KEY_ID"),
         aws_secret_access_key=get_required_env("AWS_SECRET_ACCESS_KEY"),
-        aws_region_name=get_required_env("BEDROCK_AWS_REGION"),
+        region_name=get_required_env("BEDROCK_AWS_REGION"),
     )
 
 
@@ -138,12 +120,16 @@ def bedrock_titan_embedding():
 
 @pytest.fixture
 def bedrock_cohere_embedding():
+    import boto3
+
     return (
         "bedrock-cohere",
         1024,
-        BedrockEmbedding.from_credentials(
-            model_name="cohere.embed-english-v3",
-            aws_region=get_required_env("BEDROCK_AWS_REGION"),
+        BedrockEmbedding(
+            client=boto3.Session(
+                region_name=get_required_env("BEDROCK_AWS_REGION")
+            ).client("bedrock-runtime"),
+            model="cohere.embed-english-v3",
         ),
     )
 
@@ -339,12 +325,6 @@ def test_multimodal(vector_store, embedding, llm, request):
     assert "Coffee Machine Ultra Cool" in response
 
 
-def get_local_resource_path(filename: str):
-    dirname = os.path.dirname(__file__)
-    e2e_tests_dir = os.path.dirname(dirname)
-    return os.path.join(e2e_tests_dir, "resources", filename)
-
-
 @pytest.mark.parametrize(
     "chat",
     ["gemini_pro_llm", "vertex_gemini_pro_llm"],

ragstack-e2e-tests/e2e_tests/llama_index/test_llama_parse.py

@@ -0,0 +1,66 @@
+import pytest
+
+try:
+    from llama_parse import LlamaParse
+except ImportError:
+    pytest.skip("llama_parse is not supported, skipping tests", allow_module_level=True)
+
+from llama_index import (
+    VectorStoreIndex,
+    StorageContext,
+    ServiceContext,
+)
+
+from llama_index.embeddings import OpenAIEmbedding
+from llama_index.llms import OpenAI
+
+from e2e_tests.conftest import (
+    set_current_test_info,
+    get_required_env,
+)
+from e2e_tests.test_utils import get_local_resource_path
+from e2e_tests.test_utils.vector_store_handler import (
+    VectorStoreTestContext,
+)
+
+
+@pytest.fixture
+def llama_parse_text():
+    return "text", LlamaParse(result_type="text")
+
+
+@pytest.fixture
+def llama_parse_markdown():
+    return "markdown", LlamaParse(result_type="markdown")
+
+
+@pytest.mark.parametrize("vector_store", ["cassandra", "astra_db"])
+@pytest.mark.parametrize(
+    "llama_parse_instance",
+    ["llama_parse_text", "llama_parse_markdown"],
+)
+def test_llama_parse(vector_store, llama_parse_instance, request):
+    vector_store_context: VectorStoreTestContext = request.getfixturevalue(vector_store)
+    lp_type, lp = request.getfixturevalue(llama_parse_instance)
+    llm = OpenAI(api_key=get_required_env("OPEN_AI_KEY"))
+    embedding = OpenAIEmbedding(api_key=get_required_env("OPEN_AI_KEY"))
+
+    set_current_test_info(
+        "llama_index::llama_parse",
+        f"{lp_type},{vector_store}",
+    )
+    vector_store = vector_store_context.new_llamaindex_vector_store(
+        embedding_dimension=1536
+    )
+    storage_context = StorageContext.from_defaults(vector_store=vector_store)
+    service_context = ServiceContext.from_defaults(llm=llm, embed_model=embedding)
+
+    file_path = get_local_resource_path("tree.pdf")
+    documents = lp.load_data(file_path)
+
+    index = VectorStoreIndex.from_documents(
+        documents, storage_context=storage_context, service_context=service_context
+    )
+
+    retriever = index.as_retriever()
+    assert len(retriever.retrieve("What was Eldenroot?")) > 0

ragstack-e2e-tests/e2e_tests/resources/tree.pdf

@@ -18,6 +18,12 @@ def get_required_env(name) -> str:
     return value
 
 
+def get_local_resource_path(filename: str):
+    dirname = os.path.dirname(__file__)
+    e2e_tests_dir = os.path.dirname(dirname)
+    return os.path.join(e2e_tests_dir, "resources", filename)
+
+
 def random_string() -> str:
     return str(uuid.uuid4()).split("-")[0]
 

ragstack-e2e-tests/e2e_tests/test_utils/__init__.py

@@ -114,18 +114,8 @@ def search_documents(self, vector: List[float], limit: int) -> List[str]:
         return docs
 
 
-def metaclass_resolver(*classes):
-    metaclass = tuple(set(type(cls) for cls in classes))
-    metaclass = (
-        metaclass[0]
-        if len(metaclass) == 1
-        else type("_".join(mcls.__name__ for mcls in metaclass), metaclass, {})
-    )  # class M_C
-    return metaclass("_".join(cls.__name__ for cls in classes), classes, {})
-
-
 class EnhancedAstraDBLlamaIndexVectorStore(
-    metaclass_resolver(EnhancedLlamaIndexVectorStore, AstraDBVectorStore)
+    AstraDBVectorStore, EnhancedLlamaIndexVectorStore
 ):
 
     def put_document(