Local search notebook

Local search generates answers by combining relevant knowledge graph data with text chunks from the original documents. This method is ideal for questions about specific entities or detailed information.

This page references the local_search.ipynb notebook from the GraphRAG repository.

When to use local search

Local search excels at:

Entity-specific questions - “Who is Dr. Jordan Hayes?”
Relationship queries - “What is the connection between X and Y?”
Detailed information - “What are the properties of this concept?”
Fact retrieval - “When did this event occur?”

How local search works

Query embedding

Your question is converted to an embedding vector.

Entity retrieval

Semantically similar entities are retrieved from the vector store.

Context expansion

Related relationships, text units, and community reports are gathered.

Response generation

The LLM synthesizes an answer from the assembled context.

Setting up the notebook

Import required libraries

import os
import pandas as pd
from graphrag.query.context_builder.entity_extraction import EntityVectorStoreKey
from graphrag.query.indexer_adapters import (
    read_indexer_covariates,
    read_indexer_entities,
    read_indexer_relationships,
    read_indexer_reports,
    read_indexer_text_units,
)
from graphrag.query.question_gen.local_gen import LocalQuestionGen
from graphrag.query.structured_search.local_search.mixed_context import (
    LocalSearchMixedContext,
)
from graphrag.query.structured_search.local_search.search import LocalSearch
from graphrag_vectors import IndexSchema, LanceDBVectorStore

Configure paths

INPUT_DIR = "./inputs/operation dulce"
LANCEDB_URI = f"{INPUT_DIR}/lancedb"

COMMUNITY_REPORT_TABLE = "community_reports"
ENTITY_TABLE = "entities"
COMMUNITY_TABLE = "communities"
RELATIONSHIP_TABLE = "relationships"
COVARIATE_TABLE = "covariates"
TEXT_UNIT_TABLE = "text_units"
COMMUNITY_LEVEL = 2

Load data tables

Load entities

# Read entity and community data
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
community_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_TABLE}.parquet")

entities = read_indexer_entities(entity_df, community_df, COMMUNITY_LEVEL)

# Connect to entity description embeddings
description_embedding_store = LanceDBVectorStore(
    index_schema=IndexSchema(index_name="default-entity-description")
)
description_embedding_store.connect(db_uri=LANCEDB_URI)

print(f"Entity count: {len(entity_df)}")
entity_df.head()

Load relationships

relationship_df = pd.read_parquet(f"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet")
relationships = read_indexer_relationships(relationship_df)

print(f"Relationship count: {len(relationship_df)}")
relationship_df.head()

Load covariates (optional)

# Covariates (claims) are optional and disabled by default
covariate_df = pd.read_parquet(f"{INPUT_DIR}/{COVARIATE_TABLE}.parquet")
claims = read_indexer_covariates(covariate_df)

print(f"Claim records: {len(claims)}")
covariates = {"claims": claims}

Covariates extract claims from your documents and generally require prompt tuning to be valuable. See the GRAPHRAG_CLAIM_* settings.

Load community reports

report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
reports = read_indexer_reports(report_df, community_df, COMMUNITY_LEVEL)

print(f"Report records: {len(report_df)}")
report_df.head()

Load text units

text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet")
text_units = read_indexer_text_units(text_unit_df)

print(f"Text unit records: {len(text_unit_df)}")
text_unit_df.head()

Configure language models

from graphrag.config.enums import ModelType
from graphrag.config.models.language_model_config import LanguageModelConfig
from graphrag.language_model.manager import ModelManager
from graphrag.tokenizer.get_tokenizer import get_tokenizer

api_key = os.environ["GRAPHRAG_API_KEY"]

# Chat model configuration
chat_config = LanguageModelConfig(
    api_key=api_key,
    type=ModelType.Chat,
    model_provider="openai",
    model="gpt-4.1",
    max_retries=20,
)
chat_model = ModelManager().get_or_create_chat_model(
    name="local_search",
    model_type=ModelType.Chat,
    config=chat_config,
)

# Embedding model configuration
embedding_config = LanguageModelConfig(
    api_key=api_key,
    type=ModelType.Embedding,
    model_provider="openai",
    model="text-embedding-3-small",
    max_retries=20,
)
text_embedder = ModelManager().get_or_create_embedding_model(
    name="local_search_embedding",
    model_type=ModelType.Embedding,
    config=embedding_config,
)

tokenizer = get_tokenizer(chat_config)

Create local search context builder

context_builder = LocalSearchMixedContext(
    community_reports=reports,
    text_units=text_units,
    entities=entities,
    relationships=relationships,
    covariates=covariates,  # Set to None if not using claims
    entity_text_embeddings=description_embedding_store,
    embedding_vectorstore_key=EntityVectorStoreKey.ID,
    text_embedder=text_embedder,
    tokenizer=tokenizer,
)

If you use entity titles as vector store IDs, set embedding_vectorstore_key=EntityVectorStoreKey.TITLE.

Configure search parameters

Context parameters
Model parameters

local_context_params = {
    "text_unit_prop": 0.5,  # 50% of context for text chunks
    "community_prop": 0.1,  # 10% for community reports
    # Remaining 40% for entities and relationships
    
    "conversation_history_max_turns": 5,
    "conversation_history_user_turns_only": True,
    
    "top_k_mapped_entities": 10,  # Top entities from vector search
    "top_k_relationships": 10,  # Related relationships to include
    
    "include_entity_rank": True,  # Include node degree
    "include_relationship_weight": True,
    "include_community_rank": False,
    
    "return_candidate_context": False,  # Debug mode
    "embedding_vectorstore_key": EntityVectorStoreKey.ID,
    
    "max_tokens": 12_000,  # Total context window
}

model_params = {
    "max_tokens": 2_000,  # Response length
    "temperature": 0.0,  # Deterministic
}

Create search engine

search_engine = LocalSearch(
    model=chat_model,
    context_builder=context_builder,
    tokenizer=tokenizer,
    model_params=model_params,
    context_builder_params=local_context_params,
    response_type="multiple paragraphs",
)

Run queries

Basic queries

result = await search_engine.search("Tell me about Agent Mercer")
print(result.response)

Inspect context data

See which entities, relationships, and sources were used:

Entities
Relationships
Community reports
Source chunks
Claims

result.context_data["entities"].head()

result.context_data["relationships"].head()

if "reports" in result.context_data:
    result.context_data["reports"].head()

result.context_data["sources"].head()

if "claims" in result.context_data:
    print(result.context_data["claims"].head())

Question generation

Generate follow-up questions based on conversation history:

question_generator = LocalQuestionGen(
    model=chat_model,
    context_builder=context_builder,
    tokenizer=tokenizer,
    model_params=model_params,
    context_builder_params=local_context_params,
)

question_history = [
    "Tell me about Agent Mercer",
    "What happens in Dulce military base?",
]

candidate_questions = await question_generator.agenerate(
    question_history=question_history,
    context_data=None,
    question_count=5
)

print(candidate_questions.response)

Example queries

Entity information

queries = [
    "Who is Agent Mercer and what is their role?",
    "Tell me about Dr. Jordan Hayes",
    "What is the Dulce military base?",
]

for query in queries:
    result = await search_engine.search(query)
    print(f"Q: {query}")
    print(f"A: {result.response}\n")

Relationship queries

result = await search_engine.search(
    "What is the relationship between Agent Mercer and Dr. Hayes?"
)
print(result.response)

# Inspect supporting relationships
print(result.context_data["relationships"])

Multi-hop reasoning

result = await search_engine.search(
    "How are the different organizations connected through their members?"
)
print(result.response)

Tuning parameters

Context allocation

Adjust the proportion of context dedicated to different data types:

# More emphasis on original text
local_context_params["text_unit_prop"] = 0.7
local_context_params["community_prop"] = 0.1
# 20% for entities and relationships

# More emphasis on graph structure
local_context_params["text_unit_prop"] = 0.3
local_context_params["community_prop"] = 0.2
# 50% for entities and relationships

Entity retrieval

# Retrieve more entities for comprehensive coverage
local_context_params["top_k_mapped_entities"] = 20
local_context_params["top_k_relationships"] = 20

# Focus on fewer, most relevant entities
local_context_params["top_k_mapped_entities"] = 5
local_context_params["top_k_relationships"] = 5

Debug mode

Enable candidate context to see what was considered:

local_context_params["return_candidate_context"] = True

result = await search_engine.search("Your question")

# View all candidate entities (in_context column shows if used)
candidate_entities = result.context_data["entities"]
print(f"Total candidates: {len(candidate_entities)}")
print(f"Used in context: {candidate_entities['in_context'].sum()}")

Performance optimization

Reduce context size

Lower max_tokens for faster responses

max_tokens=8000

Limit entity retrieval

Reduce top_k_mapped_entities

top_k_mapped_entities=5

Skip community reports

Set community proportion to 0

community_prop=0.0

Use conversation history

Enable for context-aware follow-ups

conversation_history_max_turns=5

Advanced features

Conversation history

# First question
result1 = await search_engine.search("Who is Agent Mercer?")
print(result1.response)

# Follow-up leveraging history
result2 = await search_engine.search(
    "What are their main relationships?",
    conversation_history=[
        {"role": "user", "content": "Who is Agent Mercer?"},
        {"role": "assistant", "content": result1.response},
    ]
)
print(result2.response)

Custom response types

response_types = [
    "single paragraph",
    "bullet points",
    "detailed analysis",
    "timeline",
    "list of key points",
]

search_engine.response_type = "bullet points"
result = await search_engine.search("What are Agent Mercer's key activities?")

Comparison with global search

Aspect	Local Search	Global Search
Question Type	Specific, detailed	High-level, broad
Data Source	Entities + text chunks + reports	Community reports only
Cost	Lower	Higher
Response Time	Faster	Slower
Best For	Entity details, facts	Themes, summaries

Troubleshooting

Empty or irrelevant results

Solutions:

Increase top_k_mapped_entities for broader retrieval
Check if entities exist in your knowledge graph
Verify embedding model matches indexing model
Increase text_unit_prop for more context

Token limit exceeded

Solutions:

Reduce max_tokens in context parameters
Lower top_k_mapped_entities and top_k_relationships
Decrease text_unit_prop
Use a model with larger context window

Slow performance

Solutions:

Reduce context size parameters
Disable community reports if not needed
Use faster embedding model
Cache frequent queries

Next steps

Global search

Learn about global search for dataset-wide questions

DRIFT search

Explore dynamic iterative search

Search comparison

Compare all search methods side-by-side

Query guide

Complete local search documentation

Tutorials

Notebooks

Use Cases

​When to use local search

​How local search works

​Setting up the notebook

​Import required libraries

​Configure paths

​Load data tables

​Load entities

​Load relationships

​Load covariates (optional)

​Load community reports

​Load text units

​Configure language models

​Create local search context builder

​Configure search parameters

​Create search engine

​Run queries

​Basic queries

​Inspect context data

​Question generation

​Example queries

​Tuning parameters

​Context allocation

​Entity retrieval

​Debug mode

​Performance optimization

Reduce context size

Limit entity retrieval

Skip community reports

Use conversation history

​Advanced features

​Conversation history

​Custom response types

​Comparison with global search

​Troubleshooting

​Next steps

Global search

DRIFT search

Search comparison

Query guide

Build docs developers (and LLMs) love

When to use local search

How local search works

Setting up the notebook

Import required libraries

Configure paths

Load data tables

Load entities

Load relationships

Load covariates (optional)

Load community reports

Load text units

Configure language models

Create local search context builder

Configure search parameters

Create search engine

Run queries

Basic queries

Inspect context data

Question generation

Example queries

Tuning parameters

Context allocation

Entity retrieval

Debug mode

Performance optimization

Advanced features

Conversation history

Custom response types

Comparison with global search

Troubleshooting

Next steps