chore(genai,vertexai,community): Making integration tests async #960

libs/community/Makefile

@@ -15,7 +15,7 @@ test tests:
 	uv run --group test pytest --disable-socket --allow-unix-socket $(TEST_FILE)
 
 integration_test integration_tests:
-	uv run --group test --group test_integration pytest --retries 3 --retry-delay 1 $(TEST_FILE)
+	uv run --group test --group test_integration pytest -n auto --retries 3 --retry-delay 1 $(TEST_FILE)
 
 # Run unit tests and generate a coverage report.
 coverage:

libs/community/pyproject.toml

@@ -53,6 +53,7 @@ featurestore = ["google-cloud-bigquery-storage>=2.6.0,<3.0.0", "pandas>=1.0.0; p
 [dependency-groups]
 test = [
     "pytest>=7.3.0,<8.0.0",
+    "pytest-xdist>=3.8.0,<4.0.0",
     "freezegun>=1.2.2,<2.0.0",
     "pytest-mock>=3.10.0,<4.0.0",
     "syrupy>=4.0.2,<5.0.0",

libs/community/tests/integration_tests/test_check_grounding.py

@@ -93,14 +93,14 @@ def output_parser(
 
 
 @pytest.mark.extended
-def test_integration_parse(
+async def test_integration_parse(
     output_parser: VertexAICheckGroundingWrapper,
     input_documents: List[Document],
 ) -> None:
     answer_candidate = "Ulm, in the Kingdom of Württemberg in the German Empire"
-    response = output_parser.with_config(
+    response = await output_parser.with_config(
         configurable={"documents": input_documents}
-    ).invoke(answer_candidate)
+    ).ainvoke(answer_candidate)
 
     assert isinstance(response, VertexAICheckGroundingWrapper.CheckGroundingResponse)
     assert response.support_score >= 0 and response.support_score <= 1

libs/community/tests/integration_tests/test_rank.py

@@ -159,20 +159,20 @@ def ranker(
 
 
 @pytest.mark.extended
-def test_compression_retriever(
+async def test_compression_retriever(
     mock_vector_store_retriever: MockVectorStoreRetriever, ranker: VertexAIRank
 ) -> None:
     compression_retriever = CustomRankingRetriever(
         base_retriever=mock_vector_store_retriever, ranker=ranker
     )
     query = "What was the name of einstein's mother? Was she a scientist too?"
-    compressed_docs = compression_retriever.invoke(query)
+    compressed_docs = await compression_retriever.ainvoke(query)
 
     expected_docs = [
         Document(
             page_content=(
-                "Life and career\nChildhood, youth and education\n"
-                "See also: Einstein family\nEinstein in 1882, age\xa03\n"
+                "Life and career\\nChildhood, youth and education\\n"
+                "See also: Einstein family\\nEinstein in 1882, age\xa03\\n"
                 "Albert Einstein was born in Ulm,[19] in the Kingdom of "
                 "Württemberg in the German Empire, on 14 March 1879.[20][21] "
                 "His parents, secular Ashkenazi Jews, were Hermann Einstein, "
@@ -181,7 +181,7 @@ def test_compression_retriever(
                 "Isarvorstadt, where Einstein's father and his uncle Jakob "
                 "founded Elektrotechnische Fabrik J. Einstein & Cie, a "
                 "company that manufactured electrical equipment based on "
-                "direct current.[19]\nAlbert attended a Catholic elementary "
+                "direct current.[19]\\nAlbert attended a Catholic elementary "
                 "school in Munich from the age of five. When he was eight, "
                 "he was transferred to the Luitpold Gymnasium, where he "
                 "received advanced primary and then secondary school "
@@ -195,17 +195,17 @@ def test_compression_retriever(
         ),
         Document(
             page_content=(
-                "Marriages, relationships and children\n"
-                "Albert Einstein and Mileva Marić Einstein, 1912\n"
-                "Albert Einstein and Elsa Einstein, 1930\n"
+                "Marriages, relationships and children\\n"
+                "Albert Einstein and Mileva Marić Einstein, 1912\\n"
+                "Albert Einstein and Elsa Einstein, 1930\\n"
                 "Correspondence between Einstein and Marić, discovered and "
                 "published in 1987, revealed that in early 1902, while "
                 "Marić was visiting her parents in Novi Sad, she gave birth "
                 "to a daughter, Lieserl. When Marić returned to Switzerland "
                 "it was without the child, whose fate is uncertain. A "
                 "letter of Einstein's that he wrote in September 1903 "
                 "suggests that the girl was either given up for adoption or "
-                "died of scarlet fever in infancy.[45][46]\n"
+                "died of scarlet fever in infancy.[45][46]\\n"
                 "Einstein and Marić married in January 1903. In May 1904, "
                 "their son Hans Albert was born in Bern, Switzerland. Their "
                 "son Eduard was born in Zürich in July 1910. In letters "
@@ -237,7 +237,7 @@ def test_compression_retriever(
                 "by some to be a Russian spy; her husband, the Russian "
                 "sculptor Sergei Konenkov, created the bronze bust of "
                 "Einstein at the Institute for Advanced Study at "
-                "Princeton.[65][66][failed verification]\n"
+                "Princeton.[65][66][failed verification]\\n"
                 "Following an episode of acute mental illness at about the "
                 "age of twenty, Einstein's son Eduard was diagnosed with "
                 "schizophrenia.[67] He spent the remainder of his life "

libs/genai/Makefile

@@ -15,7 +15,7 @@ test tests:
 	uv run --group test pytest --disable-socket --allow-unix-socket $(TEST_FILE)
 
 integration_test integration_tests:
-	uv run --group test --group test_integration pytest --retries 3 --retry-delay 1 $(TEST_FILE)
+	uv run --group test --group test_integration pytest -n auto --retries 3 --retry-delay 1 $(TEST_FILE)
 
 check_imports: $(shell find langchain_google_genai -name '*.py')
 	uv run --all-groups python ./scripts/check_imports.py $^

libs/genai/pyproject.toml

@@ -40,6 +40,7 @@ typing = [
 ]
 test = [
     "pytest>=8.4.0,<9.0.0",
+    "pytest-xdist>=3.8.0,<4.0.0",
     "freezegun>=1.5.0,<2.0.0",
     "pytest-mock>=3.14.0,<4.0.0",
     "syrupy>=4.9.0,<5.0.0",

libs/genai/tests/integration_tests/test_callbacks.py

@@ -27,12 +27,12 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
     "model_name",
     MODEL_NAMES,
 )
-def test_streaming_callback(model_name: str) -> None:
+async def test_streaming_callback(model_name: str) -> None:
     prompt_template = "Tell me details about the Company {name} with 2 bullet point?"
     cb = StreamingLLMCallbackHandler()
     llm = ChatGoogleGenerativeAI(model=model_name, callbacks=[cb])
     llm_chain = PromptTemplate.from_template(prompt_template) | llm
-    for _t in llm_chain.stream({"name": "Google"}):
+    async for _t in llm_chain.astream({"name": "Google"}):
         pass
     assert len(cb.tokens) > 1
     assert len(cb.generations) == 1

libs/genai/tests/integration_tests/test_chat_models.py

@@ -173,7 +173,7 @@
 
 
 @pytest.mark.flaky(retries=3, delay=1)
-def test_chat_google_genai_invoke_with_image() -> None:
+async def test_chat_google_genai_invoke_with_image() -> None:
     """Test generating an image and then text.
 
     Using `generation_config` to specify response modalities.
@@ -184,7 +184,7 @@
 
     for _ in range(3):
         # We break as soon as we get an image back, as it's not guaranteed
-        result = llm.invoke(
+        result = await llm.ainvoke(
             "Say 'meow!' and then Generate an image of a cat.",
             config={"tags": ["meow"]},
             generation_config={
@@ -210,7 +210,7 @@
 
     # Test we can pass back in
     next_message = {"role": "user", "content": "Thanks!"}
-    _ = llm.invoke([result, next_message])
+    _ = await llm.ainvoke([result, next_message])
 
     # Test content_blocks property
     content_blocks = result.content_blocks
@@ -225,13 +225,13 @@
     _ = llm.invoke(["What's this?", {"role": "assistant", "content": content_blocks}])
 
 
-def test_chat_google_genai_invoke_with_audio() -> None:
+async def test_chat_google_genai_invoke_with_audio() -> None:
     """Test generating audio."""
     llm = ChatGoogleGenerativeAI(
         model=_AUDIO_OUTPUT_MODEL, response_modalities=[Modality.AUDIO]
     )
 
-    result = llm.invoke(
+    result = await llm.ainvoke(
         "Please say The quick brown fox jumps over the lazy dog",
     )
     assert isinstance(result, AIMessage)
@@ -259,7 +259,7 @@
         (100, "explicit thinking budget"),
     ],
 )
-def test_chat_google_genai_invoke_thinking(
+async def test_chat_google_genai_invoke_thinking(
     thinking_budget: int | None, test_description: str
 ) -> None:
     """Test invoke a thinking model with different thinking budget configurations."""
@@ -269,7 +269,7 @@
 
     llm = ChatGoogleGenerativeAI(**llm_kwargs)
 
-    result = llm.invoke(
+    result = await llm.ainvoke(
         "How many O's are in Google? Please tell me how you double checked the result",
     )
 
@@ -728,14 +728,14 @@
     )
 
 
-def test_chat_google_genai_single_call_with_history() -> None:
+async def test_chat_google_genai_single_call_with_history() -> None:
     model = ChatGoogleGenerativeAI(model=_MODEL)
     text_question1, text_answer1 = "How much is 2+2?", "4"
     text_question2 = "How much is 3+3?"
     message1 = HumanMessage(content=text_question1)
     message2 = AIMessage(content=text_answer1)
     message3 = HumanMessage(content=text_question2)
-    response = model.invoke([message1, message2, message3])
+    response = await model.ainvoke([message1, message2, message3])
     assert isinstance(response, AIMessage)
     if isinstance(response.content, list):
         text_content = "".join(
@@ -752,7 +752,7 @@
     "model_name",
     [_MODEL, _PRO_MODEL],
 )
-def test_chat_google_genai_system_message(
+async def test_chat_google_genai_system_message(
     model_name: str,
 ) -> None:
     """Test system message handling.
@@ -769,7 +769,7 @@
     message1 = HumanMessage(content=text_question1)
     message2 = AIMessage(content=text_answer1)
     message3 = HumanMessage(content=text_question2)
-    response = model.invoke([system_message, message1, message2, message3])
+    response = await model.ainvoke([system_message, message1, message2, message3])
     assert isinstance(response, AIMessage)
     if isinstance(response.content, list):
         text_content = "".join(
@@ -815,7 +815,7 @@
         assert len(output.content) > 0
 
 
-def test_chat_function_calling_with_multiple_parts() -> None:
+async def test_chat_function_calling_with_multiple_parts() -> None:
     @tool
     def search(
         question: str,
@@ -849,7 +849,7 @@
             "sparrow, hawk, crow by using search tool."
         )
     )
-    response = llm_with_search_force.invoke([request])
+    response = await llm_with_search_force.ainvoke([request])
 
     assert isinstance(response, AIMessage)
     assert len(response.tool_calls) > 0
@@ -873,7 +873,7 @@
             "colors?"
         )
     )
-    result = llm_with_search.invoke([request, response, *tool_messages, follow_up])
+    result = await llm_with_search.ainvoke([request, response, *tool_messages, follow_up])
 
     assert isinstance(result, AIMessage)
 

libs/genai/tests/integration_tests/test_embeddings.py

@@ -19,11 +19,11 @@
         " model against the pickle rick?",
     ],
 )
-def test_embed_query_different_lengths(query: str) -> None:
+async def test_embed_query_different_lengths(query: str) -> None:
     """Test embedding queries of different lengths."""
     model = GoogleGenerativeAIEmbeddings(model=_MODEL)
     # Note: embed_query() is a sync method, but initialization needs the loop
-    result = model.embed_query(query, output_dimensionality=_OUTPUT_DIMENSIONALITY)
+    result = await model.aembed_query(query, output_dimensionality=_OUTPUT_DIMENSIONALITY)
     assert len(result) == 768
     assert isinstance(result, list)
 
@@ -47,12 +47,12 @@
     assert isinstance(result, list)
 
 
-def test_embed_documents() -> None:
+async def test_embed_documents() -> None:
     """Test embedding a query."""
     model = GoogleGenerativeAIEmbeddings(
         model=_MODEL,
     )
-    result = model.embed_documents(
+    result = await model.aembed_documents(
         ["Hello world", "Good day, world"], output_dimensionality=_OUTPUT_DIMENSIONALITY
     )
     assert len(result) == 2
@@ -94,24 +94,24 @@
         ).embed_query("Hello world", output_dimensionality=_OUTPUT_DIMENSIONALITY)
 
 
-def test_embed_documents_consistency() -> None:
+async def test_embed_documents_consistency() -> None:
     """Test embedding consistency for the same document."""
     model = GoogleGenerativeAIEmbeddings(model=_MODEL)
     doc = "Consistent document for testing"
-    result1 = model.embed_documents([doc], output_dimensionality=_OUTPUT_DIMENSIONALITY)
-    result2 = model.embed_documents([doc], output_dimensionality=_OUTPUT_DIMENSIONALITY)
+    result1 = await model.aembed_documents([doc], output_dimensionality=_OUTPUT_DIMENSIONALITY)
+    result2 = await model.aembed_documents([doc], output_dimensionality=_OUTPUT_DIMENSIONALITY)
     assert result1 == result2
 
 
-def test_embed_documents_quality() -> None:
+async def test_embed_documents_quality() -> None:
     """Smoke test embedding quality by comparing similar and dissimilar documents."""
     model = GoogleGenerativeAIEmbeddings(model=_MODEL)
     similar_docs = ["Document A", "Similar Document A"]
     dissimilar_docs = ["Document A", "Completely Different Zebra"]
-    similar_embeddings = model.embed_documents(
+    similar_embeddings = await model.aembed_documents(
         similar_docs, output_dimensionality=_OUTPUT_DIMENSIONALITY
     )
-    dissimilar_embeddings = model.embed_documents(
+    dissimilar_embeddings = await model.aembed_documents(
         dissimilar_docs, output_dimensionality=_OUTPUT_DIMENSIONALITY
     )
     similar_distance = np.linalg.norm(
@@ -123,22 +123,22 @@
     assert similar_distance < dissimilar_distance
 
 
-def test_embed_query_task_type() -> None:
+async def test_embed_query_task_type() -> None:
     """Test for task_type."""
     embeddings = GoogleGenerativeAIEmbeddings(model=_MODEL, task_type="clustering")
-    emb = embeddings.embed_query(
+    emb = await embeddings.aembed_query(
         "How does alphafold work?", output_dimensionality=_OUTPUT_DIMENSIONALITY
     )
 
     embeddings2 = GoogleGenerativeAIEmbeddings(model=_MODEL)
-    emb2 = embeddings2.embed_query(
+    emb2 = await embeddings2.aembed_query(
         "How does alphafold work?",
         task_type="clustering",
         output_dimensionality=_OUTPUT_DIMENSIONALITY,
     )
 
     embeddings3 = GoogleGenerativeAIEmbeddings(model=_MODEL)
-    emb3 = embeddings3.embed_query(
+    emb3 = await embeddings3.aembed_query(
         "How does alphafold work?", output_dimensionality=_OUTPUT_DIMENSIONALITY
     )