Multi-index query CLI support (#1675)

* Add vector store id reference to embeddings config.

* changed structure of output config section

* added cli integration for multi index global

* added cli integration for multi index local

* added cli integration for multi index drift and basic

* finished local testing of multi-index cli

* ruff fixes

* partially refactored test code to align with new output section

* more test changes for new output structure

* semversioner

* refactored to align with new multi index config proposal

* locally tested new multi-index output proposal

* cleaned up tests to align with new structure

---------

Co-authored-by: Derek Worthen <[email protected]>

Author: KennyZhang1

.semversioner/next-release/patch-20250204201317556265.json

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "multi index query cli support"
+}

graphrag/api/__init__.py

@@ -18,6 +18,10 @@
     global_search_streaming,
     local_search,
     local_search_streaming,
+    multi_index_basic_search,
+    multi_index_drift_search,
+    multi_index_global_search,
+    multi_index_local_search,
 )
 from graphrag.prompt_tune.types import DocSelectionType
 
@@ -33,6 +37,10 @@
     "drift_search_streaming",
     "basic_search",
     "basic_search_streaming",
+    "multi_index_basic_search",
+    "multi_index_drift_search",
+    "multi_index_global_search",
+    "multi_index_local_search",
     # prompt tuning API
     "DocSelectionType",
     "generate_indexing_prompts",

graphrag/api/query.py

@@ -220,13 +220,10 @@ async def multi_index_global_search(
     response_type: str,
     streaming: bool,
     query: str,
-) -> (
-    tuple[
-        str | dict[str, Any] | list[dict[str, Any]],
-        str | list[pd.DataFrame] | dict[str, pd.DataFrame],
-    ]
-    | AsyncGenerator
-):
+) -> tuple[
+    str | dict[str, Any] | list[dict[str, Any]],
+    str | list[pd.DataFrame] | dict[str, pd.DataFrame],
+]:
     """Perform a global search across multiple indexes and return the context data and response.
 
     Parameters
@@ -422,7 +419,6 @@ async def local_search(
     entities_ = read_indexer_entities(nodes, entities, community_level)
     covariates_ = read_indexer_covariates(covariates) if covariates is not None else []
     prompt = load_search_prompt(config.root_dir, config.local_search.prompt)
-
     search_engine = get_local_search_engine(
         config=config,
         reports=read_indexer_reports(community_reports, nodes, community_level),
@@ -531,13 +527,10 @@ async def multi_index_local_search(
     response_type: str,
     streaming: bool,
     query: str,
-) -> (
-    tuple[
-        str | dict[str, Any] | list[dict[str, Any]],
-        str | list[pd.DataFrame] | dict[str, pd.DataFrame],
-    ]
-    | AsyncGenerator
-):
+) -> tuple[
+    str | dict[str, Any] | list[dict[str, Any]],
+    str | list[pd.DataFrame] | dict[str, pd.DataFrame],
+]:
     """Perform a local search across multiple indexes and return the context data and response.
 
     Parameters
@@ -584,7 +577,6 @@ async def multi_index_local_search(
         "relationships": -1,
         "covariates": 0,
     }
-
     community_reports_dfs = []
     entities_dfs = []
     nodes_dfs = []
@@ -732,7 +724,6 @@ async def multi_index_local_search(
         covariates_combined = pd.concat(
             covariates_dfs, axis=0, ignore_index=True, sort=False
         )
-
     result = await local_search(
         config,
         nodes=nodes_combined,
@@ -927,13 +918,10 @@ async def multi_index_drift_search(
     response_type: str,
     streaming: bool,
     query: str,
-) -> (
-    tuple[
-        str | dict[str, Any] | list[dict[str, Any]],
-        str | list[pd.DataFrame] | dict[str, pd.DataFrame],
-    ]
-    | AsyncGenerator
-):
+) -> tuple[
+    str | dict[str, Any] | list[dict[str, Any]],
+    str | list[pd.DataFrame] | dict[str, pd.DataFrame],
+]:
     """Perform a DRIFT search across multiple indexes and return the context data and response.
 
     Parameters
@@ -1240,13 +1228,10 @@ async def multi_index_basic_search(
     index_names: list[str],
     streaming: bool,
     query: str,
-) -> (
-    tuple[
-        str | dict[str, Any] | list[dict[str, Any]],
-        str | list[pd.DataFrame] | dict[str, pd.DataFrame],
-    ]
-    | AsyncGenerator
-):
+) -> tuple[
+    str | dict[str, Any] | list[dict[str, Any]],
+    str | list[pd.DataFrame] | dict[str, pd.DataFrame],
+]:
     """Perform a basic search across multiple indexes and return the context data and response.
 
     Parameters

graphrag/cli/query.py

@@ -6,8 +6,7 @@
 import asyncio
 import sys
 from pathlib import Path
-
-import pandas as pd
+from typing import TYPE_CHECKING, Any
 
 import graphrag.api as api
 from graphrag.config.load_config import load_config
@@ -16,6 +15,9 @@
 from graphrag.storage.factory import StorageFactory
 from graphrag.utils.storage import load_table_from_storage, storage_has_table
 
+if TYPE_CHECKING:
+    import pandas as pd
+
 logger = PrintProgressLogger("")
 
 
@@ -49,14 +51,43 @@ def run_global_search(
         ],
         optional_list=[],
     )
+
+    # Call the Multi-Index Global Search API
+    if dataframe_dict["multi-index"]:
+        final_nodes_list = dataframe_dict["create_final_nodes"]
+        final_entities_list = dataframe_dict["create_final_entities"]
+        final_communities_list = dataframe_dict["create_final_communities"]
+        final_community_reports_list = dataframe_dict["create_final_community_reports"]
+        index_names = dataframe_dict["index_names"]
+
+        response, context_data = asyncio.run(
+            api.multi_index_global_search(
+                config=config,
+                nodes_list=final_nodes_list,
+                entities_list=final_entities_list,
+                communities_list=final_communities_list,
+                community_reports_list=final_community_reports_list,
+                index_names=index_names,
+                community_level=community_level,
+                dynamic_community_selection=dynamic_community_selection,
+                response_type=response_type,
+                streaming=streaming,
+                query=query,
+            )
+        )
+        logger.success(f"Global Search Response:\n{response}")
+        # NOTE: we return the response and context data here purely as a complete demonstration of the API.
+        # External users should use the API directly to get the response and context data.
+        return response, context_data
+
+    # Otherwise, call the Single-Index Global Search API
     final_nodes: pd.DataFrame = dataframe_dict["create_final_nodes"]
     final_entities: pd.DataFrame = dataframe_dict["create_final_entities"]
     final_communities: pd.DataFrame = dataframe_dict["create_final_communities"]
     final_community_reports: pd.DataFrame = dataframe_dict[
         "create_final_community_reports"
     ]
 
-    # call the Query API
     if streaming:
 
         async def run_streaming_search():
@@ -137,6 +168,46 @@ def run_local_search(
             "create_final_covariates",
         ],
     )
+    # Call the Multi-Index Local Search API
+    if dataframe_dict["multi-index"]:
+        final_nodes_list = dataframe_dict["create_final_nodes"]
+        final_entities_list = dataframe_dict["create_final_entities"]
+        final_community_reports_list = dataframe_dict["create_final_community_reports"]
+        final_text_units_list = dataframe_dict["create_final_text_units"]
+        final_relationships_list = dataframe_dict["create_final_relationships"]
+        index_names = dataframe_dict["index_names"]
+
+        # If any covariates tables are missing from any index, set the covariates list to None
+        if (
+            len(dataframe_dict["create_final_covariates"])
+            != dataframe_dict["num_indexes"]
+        ):
+            final_covariates_list = None
+        else:
+            final_covariates_list = dataframe_dict["create_final_covariates"]
+
+        response, context_data = asyncio.run(
+            api.multi_index_local_search(
+                config=config,
+                nodes_list=final_nodes_list,
+                entities_list=final_entities_list,
+                community_reports_list=final_community_reports_list,
+                text_units_list=final_text_units_list,
+                relationships_list=final_relationships_list,
+                covariates_list=final_covariates_list,
+                index_names=index_names,
+                community_level=community_level,
+                response_type=response_type,
+                streaming=streaming,
+                query=query,
+            )
+        )
+        logger.success(f"Local Search Response:\n{response}")
+        # NOTE: we return the response and context data here purely as a complete demonstration of the API.
+        # External users should use the API directly to get the response and context data.
+        return response, context_data
+
+    # Otherwise, call the Single-Index Local Search API
     final_nodes: pd.DataFrame = dataframe_dict["create_final_nodes"]
     final_community_reports: pd.DataFrame = dataframe_dict[
         "create_final_community_reports"
@@ -146,7 +217,6 @@ def run_local_search(
     final_entities: pd.DataFrame = dataframe_dict["create_final_entities"]
     final_covariates: pd.DataFrame | None = dataframe_dict["create_final_covariates"]
 
-    # call the Query API
     if streaming:
 
         async def run_streaming_search():
@@ -226,6 +296,37 @@ def run_drift_search(
             "create_final_entities",
         ],
     )
+
+    # Call the Multi-Index Drift Search API
+    if dataframe_dict["multi-index"]:
+        final_nodes_list = dataframe_dict["create_final_nodes"]
+        final_entities_list = dataframe_dict["create_final_entities"]
+        final_community_reports_list = dataframe_dict["create_final_community_reports"]
+        final_text_units_list = dataframe_dict["create_final_text_units"]
+        final_relationships_list = dataframe_dict["create_final_relationships"]
+        index_names = dataframe_dict["index_names"]
+
+        response, context_data = asyncio.run(
+            api.multi_index_drift_search(
+                config=config,
+                nodes_list=final_nodes_list,
+                entities_list=final_entities_list,
+                community_reports_list=final_community_reports_list,
+                text_units_list=final_text_units_list,
+                relationships_list=final_relationships_list,
+                index_names=index_names,
+                community_level=community_level,
+                response_type=response_type,
+                streaming=streaming,
+                query=query,
+            )
+        )
+        logger.success(f"DRIFT Search Response:\n{response}")
+        # NOTE: we return the response and context data here purely as a complete demonstration of the API.
+        # External users should use the API directly to get the response and context data.
+        return response, context_data
+
+    # Otherwise, call the Single-Index Drift Search API
     final_nodes: pd.DataFrame = dataframe_dict["create_final_nodes"]
     final_community_reports: pd.DataFrame = dataframe_dict[
         "create_final_community_reports"
@@ -234,7 +335,6 @@ def run_drift_search(
     final_relationships: pd.DataFrame = dataframe_dict["create_final_relationships"]
     final_entities: pd.DataFrame = dataframe_dict["create_final_entities"]
 
-    # call the Query API
     if streaming:
 
         async def run_streaming_search():
@@ -308,9 +408,29 @@ def run_basic_search(
             "create_final_text_units",
         ],
     )
+
+    # Call the Multi-Index Basic Search API
+    if dataframe_dict["multi-index"]:
+        final_text_units_list = dataframe_dict["create_final_text_units"]
+        index_names = dataframe_dict["index_names"]
+
+        response, context_data = asyncio.run(
+            api.multi_index_basic_search(
+                config=config,
+                text_units_list=final_text_units_list,
+                index_names=index_names,
+                streaming=streaming,
+                query=query,
+            )
+        )
+        logger.success(f"Basic Search Response:\n{response}")
+        # NOTE: we return the response and context data here purely as a complete demonstration of the API.
+        # External users should use the API directly to get the response and context data.
+        return response, context_data
+
+    # Otherwise, call the Single-Index Basic Search API
     final_text_units: pd.DataFrame = dataframe_dict["create_final_text_units"]
 
-    # # call the Query API
     if streaming:
 
         async def run_streaming_search():
@@ -351,9 +471,46 @@ def _resolve_output_files(
     config: GraphRagConfig,
     output_list: list[str],
     optional_list: list[str] | None = None,
-) -> dict[str, pd.DataFrame]:
+) -> dict[str, Any]:
     """Read indexing output files to a dataframe dict."""
     dataframe_dict = {}
+
+    # Loading output files for multi-index search
+    if config.outputs:
+        dataframe_dict["multi-index"] = True
+        dataframe_dict["num_indexes"] = len(config.outputs)
+        dataframe_dict["index_names"] = config.outputs.keys()
+        for output in config.outputs.values():
+            output_config = output.model_dump()
+            storage_obj = StorageFactory().create_storage(
+                storage_type=output_config["type"], kwargs=output_config
+            )
+            for name in output_list:
+                if name not in dataframe_dict:
+                    dataframe_dict[name] = []
+                df_value = asyncio.run(
+                    load_table_from_storage(name=name, storage=storage_obj)
+                )
+                dataframe_dict[name].append(df_value)
+
+            # for optional output files, do not append if the dataframe does not exist
+            if optional_list:
+                for optional_file in optional_list:
+                    if optional_file not in dataframe_dict:
+                        dataframe_dict[optional_file] = []
+                    file_exists = asyncio.run(
+                        storage_has_table(optional_file, storage_obj)
+                    )
+                    if file_exists:
+                        df_value = asyncio.run(
+                            load_table_from_storage(
+                                name=optional_file, storage=storage_obj
+                            )
+                        )
+                        dataframe_dict[optional_file].append(df_value)
+        return dataframe_dict
+    # Loading output files for single-index search
+    dataframe_dict["multi-index"] = False
     output_config = config.output.model_dump()  # type: ignore
     storage_obj = StorageFactory().create_storage(
         storage_type=output_config["type"], kwargs=output_config
@@ -373,5 +530,4 @@ def _resolve_output_files(
                 dataframe_dict[optional_file] = df_value
             else:
                 dataframe_dict[optional_file] = None
-
     return dataframe_dict

graphrag/config/defaults.py

@@ -132,6 +132,7 @@
 SNAPSHOTS_EMBEDDINGS = False
 SNAPSHOTS_TRANSIENT = False
 OUTPUT_BASE_DIR = "output"
+OUTPUT_DEFAULT_ID = "default_output"
 OUTPUT_TYPE = OutputType.file
 SUMMARIZE_DESCRIPTIONS_MAX_LENGTH = 500
 SUMMARIZE_MODEL_ID = DEFAULT_CHAT_MODEL_ID