Global Search

In [1]:

# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License.

# Copyright (c) 2024 Microsoft Corporation. # Licensed under the MIT License.

In [2]:

import os

import pandas as pd
import tiktoken

from graphrag.query.indexer_adapters import read_indexer_entities, read_indexer_reports
from graphrag.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.typing import OpenaiApiType
from graphrag.query.structured_search.global_search.community_context import (
    GlobalCommunityContext,
)
from graphrag.query.structured_search.global_search.search import GlobalSearch

import os import pandas as pd import tiktoken from graphrag.query.indexer_adapters import read_indexer_entities, read_indexer_reports from graphrag.query.llm.oai.chat_openai import ChatOpenAI from graphrag.query.llm.oai.typing import OpenaiApiType from graphrag.query.structured_search.global_search.community_context import ( GlobalCommunityContext, ) from graphrag.query.structured_search.global_search.search import GlobalSearch

Global Search example¶

Global search method generates answers by searching over all AI-generated community reports in a map-reduce fashion. This is a resource-intensive method, but often gives good responses for questions that require an understanding of the dataset as a whole (e.g. What are the most significant values of the herbs mentioned in this notebook?).

LLM setup¶

In [3]:

api_key = os.environ["GRAPHRAG_API_KEY"]
llm_model = os.environ["GRAPHRAG_LLM_MODEL"]

llm = ChatOpenAI(
    api_key=api_key,
    model=llm_model,
    api_type=OpenaiApiType.OpenAI,  # OpenaiApiType.OpenAI or OpenaiApiType.AzureOpenAI
    max_retries=20,
)

token_encoder = tiktoken.get_encoding("cl100k_base")

api_key = os.environ["GRAPHRAG_API_KEY"] llm_model = os.environ["GRAPHRAG_LLM_MODEL"] llm = ChatOpenAI( api_key=api_key, model=llm_model, api_type=OpenaiApiType.OpenAI, # OpenaiApiType.OpenAI or OpenaiApiType.AzureOpenAI max_retries=20, ) token_encoder = tiktoken.get_encoding("cl100k_base")

Load community reports as context for global search¶

• Load all community reports in the create_final_community_reports table from the ire-indexing engine, to be used as context data for global search.
• Load entities from the create_final_nodes and create_final_entities tables from the ire-indexing engine, to be used for calculating community weights for context ranking. Note that this is optional (if no entities are provided, we will not calculate community weights and only use the rank attribute in the community reports table for context ranking)

In [4]:

# parquet files generated from indexing pipeline
INPUT_DIR = "./inputs/operation dulce"
COMMUNITY_REPORT_TABLE = "create_final_community_reports"
ENTITY_TABLE = "create_final_nodes"
ENTITY_EMBEDDING_TABLE = "create_final_entities"

# community level in the Leiden community hierarchy from which we will load the community reports
# higher value means we use reports from more fine-grained communities (at the cost of higher computation cost)
COMMUNITY_LEVEL = 2

# parquet files generated from indexing pipeline INPUT_DIR = "./inputs/operation dulce" COMMUNITY_REPORT_TABLE = "create_final_community_reports" ENTITY_TABLE = "create_final_nodes" ENTITY_EMBEDDING_TABLE = "create_final_entities" # community level in the Leiden community hierarchy from which we will load the community reports # higher value means we use reports from more fine-grained communities (at the cost of higher computation cost) COMMUNITY_LEVEL = 2

In [5]:

entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet")

reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)
entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)
print(f"Total report count: {len(report_df)}")
print(
    f"Report count after filtering by community level {COMMUNITY_LEVEL}: {len(reports)}"
)
report_df.head()

entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet") report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet") entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet") reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL) entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL) print(f"Total report count: {len(report_df)}") print( f"Report count after filtering by community level {COMMUNITY_LEVEL}: {len(reports)}" ) report_df.head()

Total report count: 6
Report count after filtering by community level 2: 5

Out[5]:

	community	full_content	level	rank	title	rank_explanation	summary	findings	full_content_json	id
0	4	# Dulce Base and the Paranormal Military Squad...	1	8.5	Dulce Base and the Paranormal Military Squad: ...	The impact severity rating is high due to the ...	The community is centered around Dulce Base, a...	[{'explanation': 'Dulce Base is the primary lo...	{\n "title": "Dulce Base and the Paranormal...	6f8ba6b6-506e-46c1-83ce-982d59622554
1	5	# Sam Rivera and the Paranormal Military Squad...	1	7.5	Sam Rivera and the Paranormal Military Squad a...	The impact severity rating is high due to the ...	The community is centered around Sam Rivera, a...	[{'explanation': 'Sam Rivera is recognized for...	{\n "title": "Sam Rivera and the Paranormal...	418f4536-d673-4212-8a7c-ca1aac547d0f
2	0	# Dulce Base and the Paranormal Military Squad...	0	8.5	Dulce Base and the Paranormal Military Squad Team	The impact severity rating is high due to the ...	Dulce Base serves as the operational hub for t...	[{'explanation': 'The Paranormal Military Squa...	{\n "title": "Dulce Base and the Paranormal...	251df57e-fd49-49a7-b262-ccaff95d7a51
3	1	# Dulce Base and the Paranormal Military Squad...	0	8.5	Dulce Base and the Paranormal Military Squad: ...	The impact severity rating is high due to the ...	The community is centered around Dulce Base, a...	[{'explanation': 'Sam Rivera is recognized for...	{\n "title": "Dulce Base and the Paranormal...	6e536385-8056-4a82-8670-c0ccaf007fb4
4	2	# Dulce Base: Extraterrestrial Research and Co...	0	8.5	Dulce Base: Extraterrestrial Research and Comm...	The impact severity rating is high due to the ...	Dulce Base is a highly classified facility in ...	[{'explanation': 'Dulce Base is the epicenter ...	{\n "title": "Dulce Base: Extraterrestrial ...	660c8fba-b65f-4fcd-9c98-203600cf1981

Build global context based on community reports¶

In [6]:

context_builder = GlobalCommunityContext(
    community_reports=reports,
    entities=entities,  # default to None if you don't want to use community weights for ranking
    token_encoder=token_encoder,
)

context_builder = GlobalCommunityContext( community_reports=reports, entities=entities, # default to None if you don't want to use community weights for ranking token_encoder=token_encoder, )

Perform global search¶

In [7]:

context_builder_params = {
    "use_community_summary": False,  # False means using full community reports. True means using community short summaries.
    "shuffle_data": True,
    "include_community_rank": True,
    "min_community_rank": 0,
    "community_rank_name": "rank",
    "include_community_weight": True,
    "community_weight_name": "occurrence weight",
    "normalize_community_weight": True,
    "max_tokens": 12_000,  # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 5000)
    "context_name": "Reports",
}

map_llm_params = {
    "max_tokens": 1000,
    "temperature": 0.0,
    "response_format": {"type": "json_object"},
}

reduce_llm_params = {
    "max_tokens": 2000,  # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 1000-1500)
    "temperature": 0.0,
}

context_builder_params = { "use_community_summary": False, # False means using full community reports. True means using community short summaries. "shuffle_data": True, "include_community_rank": True, "min_community_rank": 0, "community_rank_name": "rank", "include_community_weight": True, "community_weight_name": "occurrence weight", "normalize_community_weight": True, "max_tokens": 12_000, # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 5000) "context_name": "Reports", } map_llm_params = { "max_tokens": 1000, "temperature": 0.0, "response_format": {"type": "json_object"}, } reduce_llm_params = { "max_tokens": 2000, # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 1000-1500) "temperature": 0.0, }

In [8]:

search_engine = GlobalSearch(
    llm=llm,
    context_builder=context_builder,
    token_encoder=token_encoder,
    max_data_tokens=12_000,  # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 5000)
    map_llm_params=map_llm_params,
    reduce_llm_params=reduce_llm_params,
    allow_general_knowledge=False,  # set this to True will add instruction to encourage the LLM to incorporate general knowledge in the response, which may increase hallucinations, but could be useful in some use cases.
    json_mode=True,  # set this to False if your LLM model does not support JSON mode.
    context_builder_params=context_builder_params,
    concurrent_coroutines=32,
    response_type="multiple paragraphs",  # free form text describing the response type and format, can be anything, e.g. prioritized list, single paragraph, multiple paragraphs, multiple-page report
)

search_engine = GlobalSearch( llm=llm, context_builder=context_builder, token_encoder=token_encoder, max_data_tokens=12_000, # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 5000) map_llm_params=map_llm_params, reduce_llm_params=reduce_llm_params, allow_general_knowledge=False, # set this to True will add instruction to encourage the LLM to incorporate general knowledge in the response, which may increase hallucinations, but could be useful in some use cases. json_mode=True, # set this to False if your LLM model does not support JSON mode. context_builder_params=context_builder_params, concurrent_coroutines=32, response_type="multiple paragraphs", # free form text describing the response type and format, can be anything, e.g. prioritized list, single paragraph, multiple paragraphs, multiple-page report )

In [9]:

result = await search_engine.asearch(
    "What is the major conflict in this story and who are the protagonist and antagonist?"
)

print(result.response)

result = await search_engine.asearch( "What is the major conflict in this story and who are the protagonist and antagonist?" ) print(result.response)

### Major Conflict

The central conflict in the story revolves around the Paranormal Military Squad's mission to establish contact with extraterrestrial intelligence. This mission involves deciphering alien signals and managing the potential implications of first contact. The conflict is marked by the secrecy and high stakes associated with the mission, as well as the challenges posed by the unknown nature of the extraterrestrial entities [Data: Reports (4, 5, 2, 3)].

### Protagonists

The protagonists are the key members of the Paranormal Military Squad, including Taylor Cruz, Dr. Jordan Hayes, Alex Mercer, and Sam Rivera. Each of these individuals plays a crucial role in the mission, bringing their expertise in leadership, signal decryption, diplomatic engagement, and technical innovation to the forefront [Data: Reports (4, 5, 2, 3)].

### Antagonist

The antagonist in the story is not a single entity or character. Instead, it may be considered the unknown and potentially threatening nature of the extraterrestrial signals and the challenges they present to the Paranormal Military Squad's mission [Data: Reports (4, 5, 2, 3)].

In [10]:

# inspect the data used to build the context for the LLM responses
result.context_data["reports"]

# inspect the data used to build the context for the LLM responses result.context_data["reports"]

Out[10]:

	id	title	occurrence weight	content	rank
0	4	Dulce Base and the Paranormal Military Squad: ...	1.000000	# Dulce Base and the Paranormal Military Squad...	8.5
1	5	Sam Rivera and the Paranormal Military Squad a...	1.000000	# Sam Rivera and the Paranormal Military Squad...	7.5
2	2	Dulce Base: Extraterrestrial Research and Comm...	0.750000	# Dulce Base: Extraterrestrial Research and Co...	8.5
3	3	Operation: Dulce and the Paranormal Military S...	0.583333	# Operation: Dulce and the Paranormal Military...	8.5
4	0	Dulce Base and the Paranormal Military Squad Team	0.416667	# Dulce Base and the Paranormal Military Squad...	8.5

In [11]:

# inspect number of LLM calls and tokens
print(f"LLM calls: {result.llm_calls}. LLM tokens: {result.prompt_tokens}")

# inspect number of LLM calls and tokens print(f"LLM calls: {result.llm_calls}. LLM tokens: {result.prompt_tokens}")

LLM calls: 2. LLM tokens: 5261