chore: Support the batch mode in Llama Index Query Pipeline.

PiperOrigin-RevId: 742454899

Author: shawn-yang-google

tests/unit/vertex_llama_index/test_reasoning_engine_templates_llama_index.py

@@ -108,6 +108,15 @@ def llama_index_instrumentor_none_mock():
         yield llama_index_instrumentor_mock
 
 
+@pytest.fixture
+def nest_asyncio_apply_mock():
+    with mock.patch.object(
+        _utils,
+        "_import_nest_asyncio_or_warn",
+    ) as nest_asyncio_apply_mock:
+        yield nest_asyncio_apply_mock
+
+
 @pytest.mark.usefixtures("google_auth_mock")
 class TestLlamaIndexQueryPipelineAgent:
     def setup_method(self):
@@ -199,6 +208,20 @@ def test_query_with_kwargs_and_input_dict(self, json_loads_mock):
         agent.query(input={"input": "test query"})
         mocks.assert_has_calls([mock.call.run.run(input="test query")])
 
+    def test_query_with_batch_input(self, json_loads_mock, nest_asyncio_apply_mock):
+        agent = llama_index.LlamaIndexQueryPipelineAgent(
+            model=_TEST_MODEL,
+            prompt=self.prompt,
+        )
+        agent._runnable = mock.Mock()
+        mocks = mock.Mock()
+        mocks.attach_mock(mock=agent._runnable, attribute="run")
+        agent.query(input={"input": ["test query 1", "test query 2"]}, batch=True)
+        mocks.assert_has_calls(
+            [mock.call.run.run(input=["test query 1", "test query 2"], batch=True)]
+        )
+        nest_asyncio_apply_mock.assert_called_once()
+
     @pytest.mark.usefixtures("caplog")
     def test_enable_tracing(
         self,

vertexai/preview/reasoning_engines/templates/llama_index.py

@@ -544,6 +544,10 @@ def query(
         if not self._runnable:
             self.set_up()
 
+        if kwargs.get("batch"):
+            nest_asyncio = _utils._import_nest_asyncio_or_warn()
+            nest_asyncio.apply()
+
         return _utils.to_json_serializable_llama_index_object(
             self._runnable.run(**input, **kwargs)
         )

vertexai/reasoning_engines/_utils.py

@@ -472,3 +472,16 @@ def _import_autogen_tools_or_warn() -> Optional[types.ModuleType]:
             "autogen.tools is not installed. Please call: `pip install ag2[tools]`"
         )
     return None
+
+
+def _import_nest_asyncio_or_warn() -> Optional[types.ModuleType]:
+    """Tries to import the nest_asyncio module."""
+    try:
+        import nest_asyncio
+
+        return nest_asyncio
+    except ImportError:
+        _LOGGER.warning(
+            "nest_asyncio is not installed. Please call: `pip install nest-asyncio`"
+        )
+    return None