Query API

Search your knowledge graph using the Python API

The query API provides multiple search methods to query your indexed knowledge graph. Each search method has both a standard and streaming variant.

Search methods

GraphRAG provides four search methods:

• Global search - Query the entire knowledge graph using hierarchical community summaries
• Local search - Query specific entities and their local context
• DRIFT search - Dynamic reasoning with iterative feedback and traversal
• Basic search - Simple vector similarity search over text units

Global search

Perform a global search across the entire knowledge graph.

from graphrag.api import global_search
import pandas as pd

response, context = await global_search(
    config=config,
    entities=entities_df,
    communities=communities_df,
    community_reports=reports_df,
    community_level=2,
    dynamic_community_selection=False,
    response_type="multiple paragraphs",
    query="What are the main themes in the documents?"
)

Parameters

config · GraphRagConfig · required

GraphRAG configuration object loaded from settings.yaml or constructed programmatically.

entities · pd.DataFrame · required

DataFrame containing the final entities from entities.parquet.

communities · pd.DataFrame · required

DataFrame containing the final communities from communities.parquet.

community_reports · pd.DataFrame · required

DataFrame containing the final community reports from community_reports.parquet.

community_level · int | None · required

The community level to search at. Higher levels provide broader context, lower levels provide more detail. Use None to search all levels.

dynamic_community_selection · bool · required

Enable dynamic community selection instead of using all community reports at a fixed level. When True, the search engine intelligently selects relevant communities. You can still provide community_level to cap the maximum level.

response_type · str · required

The type of response to generate. Common options:

• "multiple paragraphs" - Detailed multi-paragraph response
• "single paragraph" - Concise single paragraph
• "single sentence" - Brief single sentence
• "list of 3-7 items" - Bullet point list
• "multi-page report" - Comprehensive report

query · str · required

The user query to search for.

callbacks · list[QueryCallbacks] | None · default: "None"

List of callback objects to receive query events and context data.

verbose · bool · default: "False"

Enable verbose logging output.

Returns

response · str | dict | list[dict]

The generated response to the query. Format depends on the response_type.

context · str | list[pd.DataFrame] | dict[str, pd.DataFrame]

The context data used to generate the response, including relevant community reports and entities.

Global search streaming

Stream the global search response as it’s generated.

from graphrag.api import global_search_streaming

async for chunk in global_search_streaming(
    config=config,
    entities=entities_df,
    communities=communities_df,
    community_reports=reports_df,
    community_level=2,
    dynamic_community_selection=False,
    response_type="multiple paragraphs",
    query="What are the main themes?"
):
    print(chunk, end="", flush=True)

Parameters are identical to global_search. Returns an AsyncGenerator that yields response chunks as strings.

Local search

Perform a local search focused on specific entities and their context.

from graphrag.api import local_search

response, context = await local_search(
    config=config,
    entities=entities_df,
    communities=communities_df,
    community_reports=reports_df,
    text_units=text_units_df,
    relationships=relationships_df,
    covariates=covariates_df,
    community_level=2,
    response_type="multiple paragraphs",
    query="What is the relationship between entity A and entity B?"
)

Parameters

config · GraphRagConfig · required

GraphRAG configuration object.

entities · pd.DataFrame · required

DataFrame containing the final entities from entities.parquet.

communities · pd.DataFrame · required

DataFrame containing the final communities from communities.parquet.

community_reports · pd.DataFrame · required

DataFrame containing the final community reports from community_reports.parquet.

text_units · pd.DataFrame · required

DataFrame containing the final text units from text_units.parquet.

relationships · pd.DataFrame · required

DataFrame containing the final relationships from relationships.parquet.

covariates · pd.DataFrame | None · required

DataFrame containing the final covariates from covariates.parquet, or None if covariates are not used.

community_level · int · required

The community level to search at.

response_type · str · required

The type of response to generate.

query · str · required

The user query to search for.

callbacks · list[QueryCallbacks] | None · default: "None"

List of callback objects to receive query events.

verbose · bool · default: "False"

Enable verbose logging output.

Returns

Returns a tuple of (response, context) similar to global search.

Local search streaming

Stream the local search response as it’s generated.

from graphrag.api import local_search_streaming

async for chunk in local_search_streaming(
    config=config,
    entities=entities_df,
    communities=communities_df,
    community_reports=reports_df,
    text_units=text_units_df,
    relationships=relationships_df,
    covariates=covariates_df,
    community_level=2,
    response_type="multiple paragraphs",
    query="What is the relationship between entity A and entity B?"
):
    print(chunk, end="", flush=True)

DRIFT search

Perform a DRIFT (Dynamic Reasoning with Iterative Feedback and Traversal) search.

from graphrag.api import drift_search

response, context = await drift_search(
    config=config,
    entities=entities_df,
    communities=communities_df,
    community_reports=reports_df,
    text_units=text_units_df,
    relationships=relationships_df,
    community_level=2,
    response_type="multiple paragraphs",
    query="Explain the connections between these topics"
)

Parameters

config · GraphRagConfig · required

GraphRAG configuration object.

entities · pd.DataFrame · required

DataFrame containing the final entities from entities.parquet.

communities · pd.DataFrame · required

DataFrame containing the final communities from communities.parquet.

community_reports · pd.DataFrame · required

DataFrame containing the final community reports from community_reports.parquet.

text_units · pd.DataFrame · required

DataFrame containing the final text units from text_units.parquet.

relationships · pd.DataFrame · required

DataFrame containing the final relationships from relationships.parquet.

community_level · int · required

The community level to search at.

response_type · str · required

The type of response to generate.

query · str · required

The user query to search for.

callbacks · list[QueryCallbacks] | None · default: "None"

List of callback objects to receive query events.

verbose · bool · default: "False"

Enable verbose logging output.

Returns

Returns a tuple of (response, context) similar to other search methods.

DRIFT search streaming

Stream the DRIFT search response as it’s generated.

from graphrag.api import drift_search_streaming

async for chunk in drift_search_streaming(
    config=config,
    entities=entities_df,
    communities=communities_df,
    community_reports=reports_df,
    text_units=text_units_df,
    relationships=relationships_df,
    community_level=2,
    response_type="multiple paragraphs",
    query="Explain the connections between these topics"
):
    print(chunk, end="", flush=True)

Basic search

Perform a basic vector similarity search over text units.

from graphrag.api import basic_search

response, context = await basic_search(
    config=config,
    text_units=text_units_df,
    response_type="multiple paragraphs",
    query="Find information about topic X"
)

Parameters

config · GraphRagConfig · required

GraphRAG configuration object.

text_units · pd.DataFrame · required

DataFrame containing the final text units from text_units.parquet.

response_type · str · required

The type of response to generate.

query · str · required

The user query to search for.

callbacks · list[QueryCallbacks] | None · default: "None"

List of callback objects to receive query events.

verbose · bool · default: "False"

Enable verbose logging output.

Returns

Returns a tuple of (response, context) similar to other search methods.

Basic search streaming

Stream the basic search response as it’s generated.

from graphrag.api import basic_search_streaming

async for chunk in basic_search_streaming(
    config=config,
    text_units=text_units_df,
    response_type="multiple paragraphs",
    query="Find information about topic X"
):
    print(chunk, end="", flush=True)

Complete example

Here’s a complete example showing how to load data and perform different types of searches:

import asyncio
import pandas as pd
from graphrag.config.models.graph_rag_config import GraphRagConfig
from graphrag.api import (
    global_search,
    local_search,
    drift_search,
    basic_search
)

async def main():
    # Load configuration
    config = GraphRagConfig.from_file("settings.yaml")
    
    # Load indexed data
    entities = pd.read_parquet("output/entities.parquet")
    communities = pd.read_parquet("output/communities.parquet")
    reports = pd.read_parquet("output/community_reports.parquet")
    text_units = pd.read_parquet("output/text_units.parquet")
    relationships = pd.read_parquet("output/relationships.parquet")
    
    # Optional: load covariates if available
    try:
        covariates = pd.read_parquet("output/covariates.parquet")
    except FileNotFoundError:
        covariates = None
    
    # Global search - broad questions
    print("Global search:")
    response, context = await global_search(
        config=config,
        entities=entities,
        communities=communities,
        community_reports=reports,
        community_level=2,
        dynamic_community_selection=False,
        response_type="multiple paragraphs",
        query="What are the main themes in the documents?"
    )
    print(response)
    print()
    
    # Local search - specific questions
    print("Local search:")
    response, context = await local_search(
        config=config,
        entities=entities,
        communities=communities,
        community_reports=reports,
        text_units=text_units,
        relationships=relationships,
        covariates=covariates,
        community_level=2,
        response_type="multiple paragraphs",
        query="What is the relationship between entity A and entity B?"
    )
    print(response)
    print()
    
    # DRIFT search - complex reasoning
    print("DRIFT search:")
    response, context = await drift_search(
        config=config,
        entities=entities,
        communities=communities,
        community_reports=reports,
        text_units=text_units,
        relationships=relationships,
        community_level=2,
        response_type="multiple paragraphs",
        query="Explain the connections between these topics"
    )
    print(response)
    print()
    
    # Basic search - simple vector similarity
    print("Basic search:")
    response, context = await basic_search(
        config=config,
        text_units=text_units,
        response_type="single paragraph",
        query="Find information about topic X"
    )
    print(response)

if __name__ == "__main__":
    asyncio.run(main())

Streaming example

import asyncio
import pandas as pd
from graphrag.config.models.graph_rag_config import GraphRagConfig
from graphrag.api import global_search_streaming

async def main():
    config = GraphRagConfig.from_file("settings.yaml")
    
    entities = pd.read_parquet("output/entities.parquet")
    communities = pd.read_parquet("output/communities.parquet")
    reports = pd.read_parquet("output/community_reports.parquet")
    
    print("Streaming response:")
    async for chunk in global_search_streaming(
        config=config,
        entities=entities,
        communities=communities,
        community_reports=reports,
        community_level=2,
        dynamic_community_selection=False,
        response_type="multiple paragraphs",
        query="What are the main themes?"
    ):
        print(chunk, end="", flush=True)
    print()

if __name__ == "__main__":
    asyncio.run(main())

Related

• Index API - Build knowledge graph indexes
• Prompt tune API - Generate custom prompts
• Configuration - Configure search settings