langchain: add support for vectorize (#451)

Author: nicoloboschi

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

@@ -87,6 +87,16 @@ def openai_gpt4_llm():
     }
 
 
+@pytest.fixture
+def openai_gpt4o_llm():
+    model = "gpt-4o"
+
+    return {
+        "llm": _chat_openai(model=model, streaming=False),
+        "nemo_config": {"engine": "openai", "model": model},
+    }
+
+
 def _openai_embeddings(**kwargs) -> callable:
     return lambda: OpenAIEmbeddings(
         openai_api_key=get_required_env("OPENAI_API_KEY"), **kwargs
@@ -108,6 +118,22 @@ def openai_3large_embedding():
     return _openai_embeddings(model="text-embedding-3-large")
 
 
+@pytest.fixture
+def astra_vectorize_openai_small():
+    def call():
+        from astrapy.info import CollectionVectorServiceOptions
+
+        return {
+            "collection_vector_service_options": CollectionVectorServiceOptions(
+                provider="openai",
+                model_name="text-embedding-3-small",
+            ),
+            "collection_embedding_api_key": get_required_env("OPENAI_API_KEY"),
+        }
+
+    return call
+
+
 @pytest.fixture
 def azure_openai_gpt35turbo_llm():
     # model is configurable because it can be different from the deployment
@@ -221,31 +247,30 @@ def huggingface_hub_flant5xxl_llm():
 
 @pytest.fixture
 def huggingface_hub_minilml6v2_embedding():
-
     return lambda: HuggingFaceInferenceAPIEmbeddings(
         api_key=get_required_env("HUGGINGFACE_HUB_KEY"),
         model_name="sentence-transformers/all-MiniLM-l6-v2",
     )
 
 
 @pytest.fixture
-def nvidia_aifoundation_nvolveqa40k_embedding():
+def nvidia_aifoundation_embedqa4_embedding():
     def embedding():
         get_required_env("NVIDIA_API_KEY")
         from langchain_nvidia_ai_endpoints.embeddings import NVIDIAEmbeddings
 
-        return NVIDIAEmbeddings(model="playground_nvolveqa_40k")
+        return NVIDIAEmbeddings(model="ai-embed-qa-4")
 
     return embedding
 
 
 @pytest.fixture
-def nvidia_aifoundation_mixtral8x7b_llm():
+def nvidia_aifoundation_mistral_llm():
     def llm():
         get_required_env("NVIDIA_API_KEY")
         from langchain_nvidia_ai_endpoints import ChatNVIDIA
 
-        return ChatNVIDIA(model="playground_mixtral_8x7b")
+        return ChatNVIDIA(model="ai-mistral-large")
 
     return {"llm": llm, "nemo_config": None}
 
@@ -264,15 +289,16 @@ def 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"),
         # ("huggingface_hub_minilml6v2_embedding", "huggingface_hub_flant5xxl_llm"),
         (
-            "nvidia_aifoundation_nvolveqa40k_embedding",
-            "nvidia_aifoundation_mixtral8x7b_llm",
+            "nvidia_aifoundation_embedqa4_embedding",
+            "nvidia_aifoundation_mistral_llm",
         ),
     ],
 )
@@ -304,8 +330,15 @@ def _run_test(
     if test_case == "nemo_guardrails" and not resolved_llm["nemo_config"]:
         skip_test_due_to_implementation_not_supported("nemo_guardrails")
 
+    embedding = embedding_fn()
+    vector_store_kwargs = {}
+    if isinstance(embedding, dict):
+        vector_store_kwargs = embedding
+    else:
+        vector_store_kwargs["embedding"] = embedding
+
     vector_store = vector_store_context.new_langchain_vector_store(
-        embedding=embedding_fn()
+        **vector_store_kwargs
     )
     llm = resolved_llm["llm"]()  # llm is a callable
 

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

@@ -17,6 +17,7 @@
 from llama_index.core.vector_stores.types import (
     VectorStoreQuery,
 )
+from e2e_tests.test_utils import skip_test_due_to_implementation_not_supported
 
 
 from e2e_tests.test_utils import (
@@ -133,6 +134,8 @@ def __init__(self, handler: CassandraVectorStoreHandler):
     def new_langchain_vector_store(
         self, **kwargs
     ) -> EnhancedCassandraLangChainVectorStore:
+        if "embedding" not in kwargs:
+            skip_test_due_to_implementation_not_supported("astra vectorize")
         return EnhancedCassandraLangChainVectorStore(
             session=self.handler.cassandra_session,
             keyspace="default_keyspace",

libs/e2e-tests/pyproject.llamaindex.toml

@@ -44,7 +44,7 @@ llama-parse = { git = "https://github.com/run-llama/llama_parse.git", branch = "
 langchain = "0.1.17"
 langchain-core = "0.1.48"
 langchain-community = "0.0.36"
-langchain-astradb = "0.2.0"
+langchain-astradb = "0.3.3"
 langchain-openai = "0.1.3"
 langchain-google-genai = "0.0.11"
 langchain-google-vertexai = "1.0.1"

libs/langchain/pyproject.toml

@@ -20,7 +20,7 @@ ragstack-ai-colbert = "1.0.5"
 langchain = "0.1.19"
 langchain-core = "0.1.52"
 langchain-community = "0.0.38"
-langchain-astradb = "0.3.0"
+langchain-astradb = "0.3.3"
 langchain-openai = "0.1.3"
 langchain-google-genai = { version = "0.0.11", optional = true }
 langchain-google-vertexai = { version = "1.0.1", optional = true }