tests: fix nvidia endpoints failures (#458)

Author: nicoloboschi

.github/workflows/ci-e2e-tests.yml

@@ -22,8 +22,8 @@ jobs:
     outputs:
       notebooks: ${{ steps.filter.outputs.notebooks }}
       e2e_tests: ${{ steps.filter.outputs.e2e_tests }}
-      astradb-dev-region: "us-central1"
-      astradb-dev-cloud: "gcp"
+      astradb-dev-region: "us-west-2"
+      astradb-dev-cloud: "aws"
       astradb-prod-region: "us-east-2"
       astradb-prod-cloud: "aws"
       is-scheduled: ${{ github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' }}

libs/e2e-tests/e2e_tests/langchain/test_astra.py

@@ -39,7 +39,7 @@ def test_ingest_errors(vectorstore: AstraDBVectorStore):
         print("Error:", e)
         # API Exception while running bulk insertion: [{'message': "Failed to insert document with _id 'b388435404254c17b720816ee9e0ddc4': Zero vectors cannot be indexed or queried with cosine similarity"}]
         if (
-            "Zero vectors cannot be indexed or queried with cosine similarity"
+            "Zero and near-zero vectors cannot be indexed or queried with cosine similarity"
             not in e.args[0]
         ):
             pytest.fail(

libs/e2e-tests/e2e_tests/langchain/test_cassandra_tool.py

@@ -15,6 +15,7 @@
 
 def test_tool_with_openai_tool(cassandra):
     session = cassio.config.resolve_session()
+    session.execute("DROP TABLE IF EXISTS default_keyspace.tool_table_users;")
 
     session.execute(
         """

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

@@ -195,7 +195,7 @@ def _bedrock_chat(**kwargs) -> callable:
 def bedrock_anthropic_claudev2_llm():
     return {
         "llm": _bedrock_chat(
-            model_id="anthropic.claude-v2",
+            model_id="anthropic.claude-v2", model_kwargs={"temperature": 0}
         ),
         "nemo_config": None,
     }
@@ -265,41 +265,43 @@ def embedding():
 
 
 @pytest.fixture
-def nvidia_aifoundation_mistral_llm():
+def nvidia_aifoundation_mixtral8x7b_llm():
     def llm():
         get_required_env("NVIDIA_API_KEY")
         from langchain_nvidia_ai_endpoints import ChatNVIDIA
 
-        return ChatNVIDIA(model="ai-mistral-large")
+        return ChatNVIDIA(
+            model="ai-mixtral-8x7b-instruct", temperature=0, max_tokens=2048
+        )
 
     return {"llm": llm, "nemo_config": None}
 
 
 @pytest.mark.parametrize(
     "test_case",
-    ["rag_custom_chain", "conversational_rag", "trulens", "nemo_guardrails"],
+    ["trulens"],
 )
 @pytest.mark.parametrize(
     "vector_store",
-    ["astra_db", "cassandra"],
+    ["cassandra"],
 )
 @pytest.mark.parametrize(
     "embedding,llm",
     [
-        ("openai_ada002_embedding", "openai_gpt35turbo_llm"),
-        ("openai_3large_embedding", "openai_gpt35turbo_llm_streaming"),
-        ("openai_3small_embedding", "openai_gpt4_llm"),
-        ("astra_vectorize_openai_small", "openai_gpt4o_llm"),
-        ("azure_openai_ada002_embedding", "azure_openai_gpt35turbo_llm"),
+        # ("openai_ada002_embedding", "openai_gpt35turbo_llm"),
+        # ("openai_3large_embedding", "openai_gpt35turbo_llm_streaming"),
+        # ("openai_3small_embedding", "openai_gpt4_llm"),
+        # ("astra_vectorize_openai_small", "openai_gpt4o_llm"),
+        # ("azure_openai_ada002_embedding", "azure_openai_gpt35turbo_llm"),
         ("vertex_gecko_embedding", "vertex_bison_llm"),
-        ("bedrock_titan_embedding", "bedrock_anthropic_claudev2_llm"),
-        ("bedrock_cohere_embedding", "bedrock_mistral_mistral7b_llm"),
-        ("bedrock_cohere_embedding", "bedrock_meta_llama2_llm"),
+        # ("bedrock_titan_embedding", "bedrock_anthropic_claudev2_llm"),
+        # ("bedrock_cohere_embedding", "bedrock_mistral_mistral7b_llm"),
+        # ("bedrock_cohere_embedding", "bedrock_meta_llama2_llm"),
         # ("huggingface_hub_minilml6v2_embedding", "huggingface_hub_flant5xxl_llm"),
-        (
-            "nvidia_aifoundation_embedqa4_embedding",
-            "nvidia_aifoundation_mistral_llm",
-        ),
+        # (
+        #     "nvidia_aifoundation_embedqa4_embedding",
+        #     "nvidia_aifoundation_mixtral8x7b_llm",
+        # ),
     ],
 )
 def test_rag(test_case, vector_store, embedding, llm, request, record_property):
@@ -353,7 +355,6 @@ def _run_test(
             chat_memory=vector_store_context.new_langchain_chat_memory(),
             record_property=record_property,
         )
-        # TODO: Add record property
     elif test_case == "trulens":
         run_trulens_evaluation(vector_store=vector_store, llm=llm)
     elif test_case == "nemo_guardrails":

libs/e2e-tests/e2e_tests/langchain/trulens.py

@@ -1,8 +1,6 @@
 from trulens_eval import TruChain, Feedback, Tru
 from trulens_eval.feedback.provider import Langchain
-from trulens_eval.feedback import Groundedness
 from trulens_eval.app import App
-from trulens_eval.schema import FeedbackResult
 
 from langchain.schema.vectorstore import VectorStore
 from langchain.schema.language_model import BaseLanguageModel
@@ -19,19 +17,16 @@
 )
 
 import numpy as np
-from concurrent.futures import as_completed
 
 
 def _feedback_functions(chain: Runnable, llm: BaseLanguageModel) -> list[Feedback]:
     provider = Langchain(chain=llm)
     context = App.select_context(chain)
 
-    grounded = Groundedness(groundedness_provider=provider)
     f_groundedness = (
-        Feedback(grounded.groundedness_measure_with_cot_reasons)
+        Feedback(provider.groundedness_measure_with_cot_reasons, name="Groundedness")
         .on(context.collect())  # collect context chunks into a list
         .on_output()
-        .aggregate(grounded.grounded_statements_aggregator)
     )
     f_qa_relevance = Feedback(provider.relevance_with_cot_reasons).on_input_output()
     f_context_relevance = (
@@ -83,10 +78,7 @@ def run_trulens_evaluation(vector_store: VectorStore, llm: BaseLanguageModel):
     tru_record = recording.get()
 
     # Wait for the feedback results to complete
-    for feedback_future in as_completed(tru_record.feedback_results):
+    for feedback_def, feedback_future in tru_record.feedback_and_future_results:
         feedback_result = feedback_future.result()
-
-        feedback_result: FeedbackResult
-
         # basic verification that feedback results were computed
         assert feedback_result.result is not None

libs/e2e-tests/e2e_tests/llama_index/test_cassandra_tool.py

@@ -13,6 +13,7 @@
 
 def test_tool_with_openai_tool(cassandra):
     session = cassio.config.resolve_session()
+    session.execute("DROP TABLE IF EXISTS default_keyspace.tool_table_users;")
 
     session.execute(
         """

libs/e2e-tests/e2e_tests/test_utils/astradb_vector_store_handler.py

@@ -16,7 +16,11 @@
 except ImportError:
     from llama_index.vector_stores.astra_db import AstraDBVectorStore
 
-from e2e_tests.test_utils import get_required_env, random_string
+from e2e_tests.test_utils import (
+    get_required_env,
+    random_string,
+    skip_test_due_to_implementation_not_supported,
+)
 from e2e_tests.test_utils.vector_store_handler import (
     VectorStoreHandler,
     VectorStoreImplementation,
@@ -158,6 +162,8 @@ def new_langchain_vector_store(self, **kwargs) -> EnhancedLangChainVectorStore:
         )
 
         if self.handler.implementation == VectorStoreImplementation.CASSANDRA:
+            if "embedding" not in kwargs:
+                skip_test_due_to_implementation_not_supported("astra vectorize")
             vector_store = EnhancedCassandraLangChainVectorStore(
                 session=None,
                 keyspace="default_keyspace",

libs/e2e-tests/pyproject.langchain.toml

@@ -17,9 +17,11 @@ huggingface-hub = "^0.20.3"
 azure-storage-blob = "^12.19.0"
 pillow = "^10.2.0"
 python-dotenv = "^1.0.1"
-trulens-eval = "0.27.2"
+trulens-eval = "^0.30.1"
 nemoguardrails = "^0.8.0"
 langchainhub = "^0.1.15"
+# https://github.com/googleapis/python-aiplatform/issues/3910
+google-cloud-aiplatform = "1.53.0"
 
 # From LangChain optional deps, needed by WebBaseLoader
 beautifulsoup4 = "^4"

libs/e2e-tests/pyproject.llamaindex.toml

@@ -15,9 +15,11 @@ ruff = "*"
 azure-storage-blob = "^12.19.0"
 pillow = "^10.2.0"
 python-dotenv = "^1.0.1"
-trulens-eval = "0.27.2"
+trulens-eval = "^0.30.1"
 nemoguardrails = "^0.8.0"
 langchainhub = "^0.1.15"
+# https://github.com/googleapis/python-aiplatform/issues/3910
+google-cloud-aiplatform = "1.53.0"
 
 # From LangChain optional deps, needed by WebBaseLoader
 beautifulsoup4 = "^4"

libs/e2e-tests/pyproject.ragstack-ai.toml

@@ -18,9 +18,11 @@ boto3 = "^1.29.6"
 azure-storage-blob = "^12.19.0"
 pillow = "^10.2.0"
 python-dotenv = "^1.0.1"
-trulens-eval = "0.27.2"
+trulens-eval = "^0.30.1"
 nemoguardrails = "^0.8.0"
 langchainhub = "^0.1.15"
+# https://github.com/googleapis/python-aiplatform/issues/3910
+google-cloud-aiplatform = "1.53.0"
 
 # From LangChain optional deps, needed by WebBaseLoader
 beautifulsoup4 = "^4"