Local Search
In [1]:
Copied!
# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License.
# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License.
In [2]:
Copied!
import os
import pandas as pd
import tiktoken
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.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.embedding import OpenAIEmbedding
from graphrag.query.llm.oai.typing import OpenaiApiType
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.vector_stores.lancedb import LanceDBVectorStore
import os
import pandas as pd
import tiktoken
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.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.embedding import OpenAIEmbedding
from graphrag.query.llm.oai.typing import OpenaiApiType
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.vector_stores.lancedb import LanceDBVectorStore
Local Search Example¶
Local search method generates answers by combining relevant data from the AI-extracted knowledge-graph with text chunks of the raw documents. This method is suitable for questions that require an understanding of specific entities mentioned in the documents (e.g. What are the healing properties of chamomile?).
Load text units and graph data tables as context for local search¶
- In this test we first load indexing outputs from parquet files to dataframes, then convert these dataframes into collections of data objects aligning with the knowledge model.
Load tables to dataframes¶
In [3]:
Copied!
INPUT_DIR = "./inputs/operation dulce"
LANCEDB_URI = f"{INPUT_DIR}/lancedb"
COMMUNITY_REPORT_TABLE = "create_final_community_reports"
ENTITY_TABLE = "create_final_nodes"
ENTITY_EMBEDDING_TABLE = "create_final_entities"
RELATIONSHIP_TABLE = "create_final_relationships"
COVARIATE_TABLE = "create_final_covariates"
TEXT_UNIT_TABLE = "create_final_text_units"
COMMUNITY_LEVEL = 2
INPUT_DIR = "./inputs/operation dulce"
LANCEDB_URI = f"{INPUT_DIR}/lancedb"
COMMUNITY_REPORT_TABLE = "create_final_community_reports"
ENTITY_TABLE = "create_final_nodes"
ENTITY_EMBEDDING_TABLE = "create_final_entities"
RELATIONSHIP_TABLE = "create_final_relationships"
COVARIATE_TABLE = "create_final_covariates"
TEXT_UNIT_TABLE = "create_final_text_units"
COMMUNITY_LEVEL = 2
Read entities¶
In [4]:
Copied!
# read nodes table to get community and degree data
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet")
entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)
# load description embeddings to an in-memory lancedb vectorstore
# to connect to a remote db, specify url and port values.
description_embedding_store = LanceDBVectorStore(
collection_name="default-entity-description",
)
description_embedding_store.connect(db_uri=LANCEDB_URI)
print(f"Entity count: {len(entity_df)}")
entity_df.head()
# read nodes table to get community and degree data
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet")
entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)
# load description embeddings to an in-memory lancedb vectorstore
# to connect to a remote db, specify url and port values.
description_embedding_store = LanceDBVectorStore(
collection_name="default-entity-description",
)
description_embedding_store.connect(db_uri=LANCEDB_URI)
print(f"Entity count: {len(entity_df)}")
entity_df.head()
Entity count: 888
Out[4]:
id | human_readable_id | title | community | level | degree | x | y | |
---|---|---|---|---|---|---|---|---|
0 | 9a062709-56dd-4bf2-8b41-926124b7a6f7 | 0 | JACOB COLLIER | 1 | 0 | 111 | 0.0 | 0.0 |
1 | 9a062709-56dd-4bf2-8b41-926124b7a6f7 | 0 | JACOB COLLIER | 11 | 1 | 111 | 0.0 | 0.0 |
2 | f8c54a6d-9d28-4332-a0d6-363796160430 | 1 | KEVIN SCOTT | 6 | 0 | 23 | 0.0 | 0.0 |
3 | f8c54a6d-9d28-4332-a0d6-363796160430 | 1 | KEVIN SCOTT | 35 | 1 | 23 | 0.0 | 0.0 |
4 | e97c00f8-6140-465f-beec-23c604331eae | 2 | MICROSOFT | 6 | 0 | 4 | 0.0 | 0.0 |
Read relationships¶
In [5]:
Copied!
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()
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()
Relationship count: 812
Out[5]:
id | human_readable_id | source | target | description | weight | combined_degree | text_unit_ids | |
---|---|---|---|---|---|---|---|---|
0 | 9af066c8-031b-4c52-b93b-b37763f6f0f7 | 0 | JACOB COLLIER | BEN BLOOMBERG | Jacob Collier and Ben Bloomberg collaborate on... | 81.0 | 170 | [aa55265004ced76e9050ed4b7a45c0496e10faa0eddb8... |
1 | 5b155801-2560-4449-914a-80addd22ed6d | 1 | JACOB COLLIER | GRAMMY AWARDS | Jacob Collier is a multi-Grammy Award-winning ... | 1.0 | 112 | [aa55265004ced76e9050ed4b7a45c0496e10faa0eddb8... |
2 | 07e7d26a-8823-4676-b210-e49a0e1a8788 | 2 | JACOB COLLIER | ALL I NEED | Jacob Collier's song 'All I Need' is popular i... | 7.0 | 112 | [aa55265004ced76e9050ed4b7a45c0496e10faa0eddb8... |
3 | 83f75ee8-c2ee-45c0-9886-d45f9aa5a105 | 3 | JACOB COLLIER | BEHIND THE TECH | Jacob Collier is a guest on the podcast 'Behin... | 1.0 | 116 | [aa55265004ced76e9050ed4b7a45c0496e10faa0eddb8... |
4 | 94ed887e-3d83-438b-a4a8-0b772ace3782 | 4 | JACOB COLLIER | LONDON | Jacob Collier is based in London | 8.0 | 112 | [aa55265004ced76e9050ed4b7a45c0496e10faa0eddb8... |
In [6]:
Copied!
# NOTE: covariates are turned off by default, because they generally need prompt tuning to be valuable
# Please see the GRAPHRAG_CLAIM_* settings
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}
# NOTE: covariates are turned off by default, because they generally need prompt tuning to be valuable
# Please see the GRAPHRAG_CLAIM_* settings
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}
Claim records: 178
Read community reports¶
In [7]:
Copied!
report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)
print(f"Report records: {len(report_df)}")
report_df.head()
report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)
print(f"Report records: {len(report_df)}")
report_df.head()
Report records: 72
Out[7]:
id | human_readable_id | community | parent | level | title | summary | full_content | rank | rank_explanation | findings | full_content_json | period | size | |
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
0 | 16949a5d17b740b2b4a6f787b0a637f1 | 43 | 43 | 10 | 2 | Ben Bloomberg and the Harmoniser Project | The community centers around Ben Bloomberg, a ... | # Ben Bloomberg and the Harmoniser Project\n\n... | 7.5 | The impact severity rating is high due to the ... | [{'explanation': 'Ben Bloomberg is a pivotal f... | {\n "title": "Ben Bloomberg and the Harmoni... | 2025-01-10 | 35 |
1 | 4ff756b7041f4dcab6612e016af2b14d | 44 | 44 | 10 | 2 | North Hampton and Influential Musicians | The community centers around North Hampton, a ... | # North Hampton and Influential Musicians\n\nT... | 6.5 | The impact severity rating is moderately high ... | [{'explanation': 'North Hampton serves as the ... | {\n "title": "North Hampton and Influential... | 2025-01-10 | 4 |
2 | 2d3df394272743a781606ad80ccb5312 | 45 | 45 | 10 | 2 | Prince of Monaco and Monaco | The community revolves around the Prince of Mo... | # Prince of Monaco and Monaco\n\nThe community... | 4.0 | The impact severity rating is moderate due to ... | [{'explanation': 'The Prince of Monaco is a ke... | {\n "title": "Prince of Monaco and Monaco",... | 2025-01-10 | 2 |
3 | becbd958973f42b0bd53cca9250feaf1 | 46 | 46 | 10 | 2 | Robot Opera and Broadway | The community revolves around the Robot Opera,... | # Robot Opera and Broadway\n\nThe community re... | 7.5 | The impact severity rating is high due to the ... | [{'explanation': 'The Robot Opera is a notable... | {\n "title": "Robot Opera and Broadway",\n ... | 2025-01-10 | 2 |
4 | f7d29921ae3e41a79ae7f88dae584892 | 47 | 47 | 13 | 2 | Ben and Jacob's Fusion of Art and Technology | The community centers around Ben and Jacob, wh... | # Ben and Jacob's Fusion of Art and Technology... | 7.5 | The impact severity rating is high due to the ... | [{'explanation': 'Ben and Jacob are key collab... | {\n "title": "Ben and Jacob's Fusion of Art... | 2025-01-10 | 5 |
Read text units¶
In [8]:
Copied!
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()
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()
Text unit records: 38
Out[8]:
id | human_readable_id | text | n_tokens | document_ids | entity_ids | relationship_ids | covariate_ids | |
---|---|---|---|---|---|---|---|---|
0 | aa55265004ced76e9050ed4b7a45c0496e10faa0eddb8a... | 1 | ../\nJACOB COLLIER: Honestly, I think mastery... | 1200 | [1e0886ae010728d10b2972f66b88608dc82b8645d3085... | [9a062709-56dd-4bf2-8b41-926124b7a6f7, f8c54a6... | [9af066c8-031b-4c52-b93b-b37763f6f0f7, 5b15580... | [f91209d1-0939-452e-b51b-be1763e2a27d, f2274c3... |
1 | 7f0fb1d3bf517dc76dffa984eec7a25e851e44ead0df82... | 2 | OMBERG: I grew up and started getting into al... | 1200 | [1e0886ae010728d10b2972f66b88608dc82b8645d3085... | [9a062709-56dd-4bf2-8b41-926124b7a6f7, f8c54a6... | [c2ac3612-3aaf-440c-babd-e21f474e0366, 9aab0b4... | [13c74c18-439b-4419-8427-4ba826503055, 7120179... |
2 | 27b739ceeddfa100f7be3cf002fd3a27aea2228f1a02c4... | 3 | , you know, and it’s a very linear pathway too... | 1200 | [1e0886ae010728d10b2972f66b88608dc82b8645d3085... | [9a062709-56dd-4bf2-8b41-926124b7a6f7, f8c54a6... | [9af066c8-031b-4c52-b93b-b37763f6f0f7, 5456bcb... | [a67f1c21-32ab-4eaa-b063-c815e7f3ea9d] |
3 | d97017305e234cc51554d653447d73b58441e1ff0f99e4... | 4 | you know, we started taking a lot of parts an... | 1200 | [1e0886ae010728d10b2972f66b88608dc82b8645d3085... | [9a062709-56dd-4bf2-8b41-926124b7a6f7, f8c54a6... | [bac3dd9b-f5c8-4966-9dc3-87f1f5976e36, 4f41be8... | [695ac017-7c10-44ad-a681-3b4c1ae86a87] |
4 | 2b6d29f8a74b16ea9a70423bce803a08a5b9ed4e6a946b... | 5 | only so much processing that my mind can do i... | 1200 | [1e0886ae010728d10b2972f66b88608dc82b8645d3085... | [9a062709-56dd-4bf2-8b41-926124b7a6f7, f8c54a6... | [d225ff7b-ca47-4fab-8d9a-4f86111526f8, 3fddb40... | [d8a7ad5c-a170-430e-9f9d-902047371ee2] |
In [9]:
Copied!
api_key = os.environ["GRAPHRAG_API_KEY"]
llm_model = os.environ["GRAPHRAG_LLM_MODEL"]
embedding_model = os.environ["GRAPHRAG_EMBEDDING_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")
text_embedder = OpenAIEmbedding(
api_key=api_key,
api_base=None,
api_type=OpenaiApiType.OpenAI,
model=embedding_model,
deployment_name=embedding_model,
max_retries=20,
)
api_key = os.environ["GRAPHRAG_API_KEY"]
llm_model = os.environ["GRAPHRAG_LLM_MODEL"]
embedding_model = os.environ["GRAPHRAG_EMBEDDING_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")
text_embedder = OpenAIEmbedding(
api_key=api_key,
api_base=None,
api_type=OpenaiApiType.OpenAI,
model=embedding_model,
deployment_name=embedding_model,
max_retries=20,
)
Create local search context builder¶
In [10]:
Copied!
context_builder = LocalSearchMixedContext(
community_reports=reports,
text_units=text_units,
entities=entities,
relationships=relationships,
# if you did not run covariates during indexing, set this to None
covariates=covariates,
entity_text_embeddings=description_embedding_store,
embedding_vectorstore_key=EntityVectorStoreKey.ID, # if the vectorstore uses entity title as ids, set this to EntityVectorStoreKey.TITLE
text_embedder=text_embedder,
token_encoder=token_encoder,
)
context_builder = LocalSearchMixedContext(
community_reports=reports,
text_units=text_units,
entities=entities,
relationships=relationships,
# if you did not run covariates during indexing, set this to None
covariates=covariates,
entity_text_embeddings=description_embedding_store,
embedding_vectorstore_key=EntityVectorStoreKey.ID, # if the vectorstore uses entity title as ids, set this to EntityVectorStoreKey.TITLE
text_embedder=text_embedder,
token_encoder=token_encoder,
)
Create local search engine¶
In [11]:
Copied!
# text_unit_prop: proportion of context window dedicated to related text units
# community_prop: proportion of context window dedicated to community reports.
# The remaining proportion is dedicated to entities and relationships. Sum of text_unit_prop and community_prop should be <= 1
# conversation_history_max_turns: maximum number of turns to include in the conversation history.
# conversation_history_user_turns_only: if True, only include user queries in the conversation history.
# top_k_mapped_entities: number of related entities to retrieve from the entity description embedding store.
# top_k_relationships: control the number of out-of-network relationships to pull into the context window.
# include_entity_rank: if True, include the entity rank in the entity table in the context window. Default entity rank = node degree.
# include_relationship_weight: if True, include the relationship weight in the context window.
# include_community_rank: if True, include the community rank in the context window.
# return_candidate_context: if True, return a set of dataframes containing all candidate entity/relationship/covariate records that
# could be relevant. Note that not all of these records will be included in the context window. The "in_context" column in these
# dataframes indicates whether the record is included in the context window.
# max_tokens: maximum number of tokens to use for the context window.
local_context_params = {
"text_unit_prop": 0.5,
"community_prop": 0.1,
"conversation_history_max_turns": 5,
"conversation_history_user_turns_only": True,
"top_k_mapped_entities": 10,
"top_k_relationships": 10,
"include_entity_rank": True,
"include_relationship_weight": True,
"include_community_rank": False,
"return_candidate_context": False,
"embedding_vectorstore_key": EntityVectorStoreKey.ID, # set this to EntityVectorStoreKey.TITLE if the vectorstore uses entity title as ids
"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)
}
llm_params = {
"max_tokens": 2_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 1000=1500)
"temperature": 0.0,
}
# text_unit_prop: proportion of context window dedicated to related text units
# community_prop: proportion of context window dedicated to community reports.
# The remaining proportion is dedicated to entities and relationships. Sum of text_unit_prop and community_prop should be <= 1
# conversation_history_max_turns: maximum number of turns to include in the conversation history.
# conversation_history_user_turns_only: if True, only include user queries in the conversation history.
# top_k_mapped_entities: number of related entities to retrieve from the entity description embedding store.
# top_k_relationships: control the number of out-of-network relationships to pull into the context window.
# include_entity_rank: if True, include the entity rank in the entity table in the context window. Default entity rank = node degree.
# include_relationship_weight: if True, include the relationship weight in the context window.
# include_community_rank: if True, include the community rank in the context window.
# return_candidate_context: if True, return a set of dataframes containing all candidate entity/relationship/covariate records that
# could be relevant. Note that not all of these records will be included in the context window. The "in_context" column in these
# dataframes indicates whether the record is included in the context window.
# max_tokens: maximum number of tokens to use for the context window.
local_context_params = {
"text_unit_prop": 0.5,
"community_prop": 0.1,
"conversation_history_max_turns": 5,
"conversation_history_user_turns_only": True,
"top_k_mapped_entities": 10,
"top_k_relationships": 10,
"include_entity_rank": True,
"include_relationship_weight": True,
"include_community_rank": False,
"return_candidate_context": False,
"embedding_vectorstore_key": EntityVectorStoreKey.ID, # set this to EntityVectorStoreKey.TITLE if the vectorstore uses entity title as ids
"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)
}
llm_params = {
"max_tokens": 2_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 1000=1500)
"temperature": 0.0,
}
In [12]:
Copied!
search_engine = LocalSearch(
llm=llm,
context_builder=context_builder,
token_encoder=token_encoder,
llm_params=llm_params,
context_builder_params=local_context_params,
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 = LocalSearch(
llm=llm,
context_builder=context_builder,
token_encoder=token_encoder,
llm_params=llm_params,
context_builder_params=local_context_params,
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
)
Run local search on sample queries¶
In [13]:
Copied!
result = await search_engine.asearch("Tell me about Agent Mercer")
print(result.response)
result = await search_engine.asearch("Tell me about Agent Mercer")
print(result.response)
Error embedding chunk {'OpenAIEmbedding': "Error code: 429 - {'error': {'message': 'You exceeded your current quota, please check your plan and billing details. For more information on this error, read the docs: https://platform.openai.com/docs/guides/error-codes/api-errors.', 'type': 'insufficient_quota', 'param': None, 'code': 'insufficient_quota'}}"}
--------------------------------------------------------------------------- ZeroDivisionError Traceback (most recent call last) Cell In[13], line 1 ----> 1 result = await search_engine.asearch("Tell me about Agent Mercer") 2 print(result.response) File ~/work/graphrag/graphrag/graphrag/query/structured_search/local_search/search.py:67, in LocalSearch.asearch(self, query, conversation_history, **kwargs) 65 search_prompt = "" 66 llm_calls, prompt_tokens, output_tokens = {}, {}, {} ---> 67 context_result = self.context_builder.build_context( 68 query=query, 69 conversation_history=conversation_history, 70 **kwargs, 71 **self.context_builder_params, 72 ) 73 llm_calls["build_context"] = context_result.llm_calls 74 prompt_tokens["build_context"] = context_result.prompt_tokens File ~/work/graphrag/graphrag/graphrag/query/structured_search/local_search/mixed_context.py:139, in LocalSearchMixedContext.build_context(self, query, conversation_history, include_entity_names, exclude_entity_names, conversation_history_max_turns, conversation_history_user_turns_only, max_tokens, text_unit_prop, community_prop, top_k_mapped_entities, top_k_relationships, include_community_rank, include_entity_rank, rank_description, include_relationship_weight, relationship_ranking_attribute, return_candidate_context, use_community_summary, min_community_rank, community_context_name, column_delimiter, **kwargs) 134 pre_user_questions = "\n".join( 135 conversation_history.get_user_turns(conversation_history_max_turns) 136 ) 137 query = f"{query}\n{pre_user_questions}" --> 139 selected_entities = map_query_to_entities( 140 query=query, 141 text_embedding_vectorstore=self.entity_text_embeddings, 142 text_embedder=self.text_embedder, 143 all_entities_dict=self.entities, 144 embedding_vectorstore_key=self.embedding_vectorstore_key, 145 include_entity_names=include_entity_names, 146 exclude_entity_names=exclude_entity_names, 147 k=top_k_mapped_entities, 148 oversample_scaler=2, 149 ) 151 # build context 152 final_context = list[str]() File ~/work/graphrag/graphrag/graphrag/query/context_builder/entity_extraction.py:58, in map_query_to_entities(query, text_embedding_vectorstore, text_embedder, all_entities_dict, embedding_vectorstore_key, include_entity_names, exclude_entity_names, k, oversample_scaler) 54 matched_entities = [] 55 if query != "": 56 # get entities with highest semantic similarity to query 57 # oversample to account for excluded entities ---> 58 search_results = text_embedding_vectorstore.similarity_search_by_text( 59 text=query, 60 text_embedder=lambda t: text_embedder.embed(t), 61 k=k * oversample_scaler, 62 ) 63 for result in search_results: 64 if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance( 65 result.document.id, str 66 ): File ~/work/graphrag/graphrag/graphrag/vector_stores/lancedb.py:134, in LanceDBVectorStore.similarity_search_by_text(self, text, text_embedder, k, **kwargs) 130 def similarity_search_by_text( 131 self, text: str, text_embedder: TextEmbedder, k: int = 10, **kwargs: Any 132 ) -> list[VectorStoreSearchResult]: 133 """Perform a similarity search using a given input text.""" --> 134 query_embedding = text_embedder(text) 135 if query_embedding: 136 return self.similarity_search_by_vector(query_embedding, k) File ~/work/graphrag/graphrag/graphrag/query/context_builder/entity_extraction.py:60, in map_query_to_entities.<locals>.<lambda>(t) 54 matched_entities = [] 55 if query != "": 56 # get entities with highest semantic similarity to query 57 # oversample to account for excluded entities 58 search_results = text_embedding_vectorstore.similarity_search_by_text( 59 text=query, ---> 60 text_embedder=lambda t: text_embedder.embed(t), 61 k=k * oversample_scaler, 62 ) 63 for result in search_results: 64 if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance( 65 result.document.id, str 66 ): File ~/work/graphrag/graphrag/graphrag/query/llm/oai/embedding.py:97, in OpenAIEmbedding.embed(self, text, **kwargs) 91 self._reporter.error( 92 message="Error embedding chunk", 93 details={self.__class__.__name__: str(e)}, 94 ) 96 continue ---> 97 chunk_embeddings = np.average(chunk_embeddings, axis=0, weights=chunk_lens) 98 chunk_embeddings = chunk_embeddings / np.linalg.norm(chunk_embeddings) 99 return chunk_embeddings.tolist() File ~/.cache/pypoetry/virtualenvs/graphrag-F2jvqev7-py3.11/lib/python3.11/site-packages/numpy/lib/function_base.py:550, in average(a, axis, weights, returned, keepdims) 548 scl = wgt.sum(axis=axis, dtype=result_dtype, **keepdims_kw) 549 if np.any(scl == 0.0): --> 550 raise ZeroDivisionError( 551 "Weights sum to zero, can't be normalized") 553 avg = avg_as_array = np.multiply(a, wgt, 554 dtype=result_dtype).sum(axis, **keepdims_kw) / scl 556 if returned: ZeroDivisionError: Weights sum to zero, can't be normalized
In [14]:
Copied!
question = "Tell me about Dr. Jordan Hayes"
result = await search_engine.asearch(question)
print(result.response)
question = "Tell me about Dr. Jordan Hayes"
result = await search_engine.asearch(question)
print(result.response)
Error embedding chunk {'OpenAIEmbedding': "Error code: 429 - {'error': {'message': 'You exceeded your current quota, please check your plan and billing details. For more information on this error, read the docs: https://platform.openai.com/docs/guides/error-codes/api-errors.', 'type': 'insufficient_quota', 'param': None, 'code': 'insufficient_quota'}}"}
--------------------------------------------------------------------------- ZeroDivisionError Traceback (most recent call last) Cell In[14], line 2 1 question = "Tell me about Dr. Jordan Hayes" ----> 2 result = await search_engine.asearch(question) 3 print(result.response) File ~/work/graphrag/graphrag/graphrag/query/structured_search/local_search/search.py:67, in LocalSearch.asearch(self, query, conversation_history, **kwargs) 65 search_prompt = "" 66 llm_calls, prompt_tokens, output_tokens = {}, {}, {} ---> 67 context_result = self.context_builder.build_context( 68 query=query, 69 conversation_history=conversation_history, 70 **kwargs, 71 **self.context_builder_params, 72 ) 73 llm_calls["build_context"] = context_result.llm_calls 74 prompt_tokens["build_context"] = context_result.prompt_tokens File ~/work/graphrag/graphrag/graphrag/query/structured_search/local_search/mixed_context.py:139, in LocalSearchMixedContext.build_context(self, query, conversation_history, include_entity_names, exclude_entity_names, conversation_history_max_turns, conversation_history_user_turns_only, max_tokens, text_unit_prop, community_prop, top_k_mapped_entities, top_k_relationships, include_community_rank, include_entity_rank, rank_description, include_relationship_weight, relationship_ranking_attribute, return_candidate_context, use_community_summary, min_community_rank, community_context_name, column_delimiter, **kwargs) 134 pre_user_questions = "\n".join( 135 conversation_history.get_user_turns(conversation_history_max_turns) 136 ) 137 query = f"{query}\n{pre_user_questions}" --> 139 selected_entities = map_query_to_entities( 140 query=query, 141 text_embedding_vectorstore=self.entity_text_embeddings, 142 text_embedder=self.text_embedder, 143 all_entities_dict=self.entities, 144 embedding_vectorstore_key=self.embedding_vectorstore_key, 145 include_entity_names=include_entity_names, 146 exclude_entity_names=exclude_entity_names, 147 k=top_k_mapped_entities, 148 oversample_scaler=2, 149 ) 151 # build context 152 final_context = list[str]() File ~/work/graphrag/graphrag/graphrag/query/context_builder/entity_extraction.py:58, in map_query_to_entities(query, text_embedding_vectorstore, text_embedder, all_entities_dict, embedding_vectorstore_key, include_entity_names, exclude_entity_names, k, oversample_scaler) 54 matched_entities = [] 55 if query != "": 56 # get entities with highest semantic similarity to query 57 # oversample to account for excluded entities ---> 58 search_results = text_embedding_vectorstore.similarity_search_by_text( 59 text=query, 60 text_embedder=lambda t: text_embedder.embed(t), 61 k=k * oversample_scaler, 62 ) 63 for result in search_results: 64 if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance( 65 result.document.id, str 66 ): File ~/work/graphrag/graphrag/graphrag/vector_stores/lancedb.py:134, in LanceDBVectorStore.similarity_search_by_text(self, text, text_embedder, k, **kwargs) 130 def similarity_search_by_text( 131 self, text: str, text_embedder: TextEmbedder, k: int = 10, **kwargs: Any 132 ) -> list[VectorStoreSearchResult]: 133 """Perform a similarity search using a given input text.""" --> 134 query_embedding = text_embedder(text) 135 if query_embedding: 136 return self.similarity_search_by_vector(query_embedding, k) File ~/work/graphrag/graphrag/graphrag/query/context_builder/entity_extraction.py:60, in map_query_to_entities.<locals>.<lambda>(t) 54 matched_entities = [] 55 if query != "": 56 # get entities with highest semantic similarity to query 57 # oversample to account for excluded entities 58 search_results = text_embedding_vectorstore.similarity_search_by_text( 59 text=query, ---> 60 text_embedder=lambda t: text_embedder.embed(t), 61 k=k * oversample_scaler, 62 ) 63 for result in search_results: 64 if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance( 65 result.document.id, str 66 ): File ~/work/graphrag/graphrag/graphrag/query/llm/oai/embedding.py:97, in OpenAIEmbedding.embed(self, text, **kwargs) 91 self._reporter.error( 92 message="Error embedding chunk", 93 details={self.__class__.__name__: str(e)}, 94 ) 96 continue ---> 97 chunk_embeddings = np.average(chunk_embeddings, axis=0, weights=chunk_lens) 98 chunk_embeddings = chunk_embeddings / np.linalg.norm(chunk_embeddings) 99 return chunk_embeddings.tolist() File ~/.cache/pypoetry/virtualenvs/graphrag-F2jvqev7-py3.11/lib/python3.11/site-packages/numpy/lib/function_base.py:550, in average(a, axis, weights, returned, keepdims) 548 scl = wgt.sum(axis=axis, dtype=result_dtype, **keepdims_kw) 549 if np.any(scl == 0.0): --> 550 raise ZeroDivisionError( 551 "Weights sum to zero, can't be normalized") 553 avg = avg_as_array = np.multiply(a, wgt, 554 dtype=result_dtype).sum(axis, **keepdims_kw) / scl 556 if returned: ZeroDivisionError: Weights sum to zero, can't be normalized
Inspecting the context data used to generate the response¶
In [15]:
Copied!
result.context_data["entities"].head()
result.context_data["entities"].head()
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[15], line 1 ----> 1 result.context_data["entities"].head() NameError: name 'result' is not defined
In [16]:
Copied!
result.context_data["relationships"].head()
result.context_data["relationships"].head()
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[16], line 1 ----> 1 result.context_data["relationships"].head() NameError: name 'result' is not defined
In [17]:
Copied!
if "reports" in result.context_data:
result.context_data["reports"].head()
if "reports" in result.context_data:
result.context_data["reports"].head()
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[17], line 1 ----> 1 if "reports" in result.context_data: 2 result.context_data["reports"].head() NameError: name 'result' is not defined
In [18]:
Copied!
result.context_data["sources"].head()
result.context_data["sources"].head()
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[18], line 1 ----> 1 result.context_data["sources"].head() NameError: name 'result' is not defined
In [19]:
Copied!
if "claims" in result.context_data:
print(result.context_data["claims"].head())
if "claims" in result.context_data:
print(result.context_data["claims"].head())
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[19], line 1 ----> 1 if "claims" in result.context_data: 2 print(result.context_data["claims"].head()) NameError: name 'result' is not defined
Question Generation¶
This function takes a list of user queries and generates the next candidate questions.
In [20]:
Copied!
question_generator = LocalQuestionGen(
llm=llm,
context_builder=context_builder,
token_encoder=token_encoder,
llm_params=llm_params,
context_builder_params=local_context_params,
)
question_generator = LocalQuestionGen(
llm=llm,
context_builder=context_builder,
token_encoder=token_encoder,
llm_params=llm_params,
context_builder_params=local_context_params,
)
In [21]:
Copied!
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)
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)
Error embedding chunk {'OpenAIEmbedding': "Error code: 429 - {'error': {'message': 'You exceeded your current quota, please check your plan and billing details. For more information on this error, read the docs: https://platform.openai.com/docs/guides/error-codes/api-errors.', 'type': 'insufficient_quota', 'param': None, 'code': 'insufficient_quota'}}"}
--------------------------------------------------------------------------- ZeroDivisionError Traceback (most recent call last) Cell In[21], line 5 1 question_history = [ 2 "Tell me about Agent Mercer", 3 "What happens in Dulce military base?", 4 ] ----> 5 candidate_questions = await question_generator.agenerate( 6 question_history=question_history, context_data=None, question_count=5 7 ) 8 print(candidate_questions.response) File ~/work/graphrag/graphrag/graphrag/query/question_gen/local_gen.py:79, in LocalQuestionGen.agenerate(self, question_history, context_data, question_count, **kwargs) 73 conversation_history = ConversationHistory.from_list(history) 75 if context_data is None: 76 # generate context data based on the question history 77 result = cast( 78 "ContextBuilderResult", ---> 79 self.context_builder.build_context( 80 query=question_text, 81 conversation_history=conversation_history, 82 **kwargs, 83 **self.context_builder_params, 84 ), 85 ) 86 context_data = cast("str", result.context_chunks) 87 context_records = result.context_records File ~/work/graphrag/graphrag/graphrag/query/structured_search/local_search/mixed_context.py:139, in LocalSearchMixedContext.build_context(self, query, conversation_history, include_entity_names, exclude_entity_names, conversation_history_max_turns, conversation_history_user_turns_only, max_tokens, text_unit_prop, community_prop, top_k_mapped_entities, top_k_relationships, include_community_rank, include_entity_rank, rank_description, include_relationship_weight, relationship_ranking_attribute, return_candidate_context, use_community_summary, min_community_rank, community_context_name, column_delimiter, **kwargs) 134 pre_user_questions = "\n".join( 135 conversation_history.get_user_turns(conversation_history_max_turns) 136 ) 137 query = f"{query}\n{pre_user_questions}" --> 139 selected_entities = map_query_to_entities( 140 query=query, 141 text_embedding_vectorstore=self.entity_text_embeddings, 142 text_embedder=self.text_embedder, 143 all_entities_dict=self.entities, 144 embedding_vectorstore_key=self.embedding_vectorstore_key, 145 include_entity_names=include_entity_names, 146 exclude_entity_names=exclude_entity_names, 147 k=top_k_mapped_entities, 148 oversample_scaler=2, 149 ) 151 # build context 152 final_context = list[str]() File ~/work/graphrag/graphrag/graphrag/query/context_builder/entity_extraction.py:58, in map_query_to_entities(query, text_embedding_vectorstore, text_embedder, all_entities_dict, embedding_vectorstore_key, include_entity_names, exclude_entity_names, k, oversample_scaler) 54 matched_entities = [] 55 if query != "": 56 # get entities with highest semantic similarity to query 57 # oversample to account for excluded entities ---> 58 search_results = text_embedding_vectorstore.similarity_search_by_text( 59 text=query, 60 text_embedder=lambda t: text_embedder.embed(t), 61 k=k * oversample_scaler, 62 ) 63 for result in search_results: 64 if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance( 65 result.document.id, str 66 ): File ~/work/graphrag/graphrag/graphrag/vector_stores/lancedb.py:134, in LanceDBVectorStore.similarity_search_by_text(self, text, text_embedder, k, **kwargs) 130 def similarity_search_by_text( 131 self, text: str, text_embedder: TextEmbedder, k: int = 10, **kwargs: Any 132 ) -> list[VectorStoreSearchResult]: 133 """Perform a similarity search using a given input text.""" --> 134 query_embedding = text_embedder(text) 135 if query_embedding: 136 return self.similarity_search_by_vector(query_embedding, k) File ~/work/graphrag/graphrag/graphrag/query/context_builder/entity_extraction.py:60, in map_query_to_entities.<locals>.<lambda>(t) 54 matched_entities = [] 55 if query != "": 56 # get entities with highest semantic similarity to query 57 # oversample to account for excluded entities 58 search_results = text_embedding_vectorstore.similarity_search_by_text( 59 text=query, ---> 60 text_embedder=lambda t: text_embedder.embed(t), 61 k=k * oversample_scaler, 62 ) 63 for result in search_results: 64 if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance( 65 result.document.id, str 66 ): File ~/work/graphrag/graphrag/graphrag/query/llm/oai/embedding.py:97, in OpenAIEmbedding.embed(self, text, **kwargs) 91 self._reporter.error( 92 message="Error embedding chunk", 93 details={self.__class__.__name__: str(e)}, 94 ) 96 continue ---> 97 chunk_embeddings = np.average(chunk_embeddings, axis=0, weights=chunk_lens) 98 chunk_embeddings = chunk_embeddings / np.linalg.norm(chunk_embeddings) 99 return chunk_embeddings.tolist() File ~/.cache/pypoetry/virtualenvs/graphrag-F2jvqev7-py3.11/lib/python3.11/site-packages/numpy/lib/function_base.py:550, in average(a, axis, weights, returned, keepdims) 548 scl = wgt.sum(axis=axis, dtype=result_dtype, **keepdims_kw) 549 if np.any(scl == 0.0): --> 550 raise ZeroDivisionError( 551 "Weights sum to zero, can't be normalized") 553 avg = avg_as_array = np.multiply(a, wgt, 554 dtype=result_dtype).sum(axis, **keepdims_kw) / scl 556 if returned: ZeroDivisionError: Weights sum to zero, can't be normalized