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
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,
    VectorStoreConfig,
    VectorStoreType,
    create_vector_store,
)
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,
    VectorStoreConfig,
    VectorStoreType,
    create_vector_store,
)

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 = "community_reports"
ENTITY_TABLE = "entities"
COMMUNITY_TABLE = "communities"
RELATIONSHIP_TABLE = "relationships"
COVARIATE_TABLE = "covariates"
TEXT_UNIT_TABLE = "text_units"
COMMUNITY_LEVEL = 2
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

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")
community_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_TABLE}.parquet")

entities = read_indexer_entities(entity_df, community_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 = create_vector_store(
    config=VectorStoreConfig(type=VectorStoreType.LanceDB, uri=LANCEDB_URI),
    index_schema=IndexSchema(index_name="default-entity-description"),
)
description_embedding_store.connect()

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")
community_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_TABLE}.parquet")

entities = read_indexer_entities(entity_df, community_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 = create_vector_store(
    config=VectorStoreConfig(type=VectorStoreType.LanceDB, uri=LANCEDB_URI),
    index_schema=IndexSchema(index_name="default-entity-description"),
)
description_embedding_store.connect()

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

Entity count: 39

Out[4]:

	id	human_readable_id	title	type	description	text_unit_ids	frequency	degree
0	bbb2ad85-3d0b-4c26-b6f9-71ed76ee64bd	0	ALEX MERCER	PERSON	Agent Alex Mercer is a determined and thoughtf...	[ff65e3e3764d360b35c55a83aaa35b0138af32551d0bc...	5	17
1	66112f07-d45e-4394-92ad-66e87c8579ff	1	TAYLOR CRUZ	PERSON	Taylor Cruz is a confident, driven, and ambiti...	[ff65e3e3764d360b35c55a83aaa35b0138af32551d0bc...	4	17
2	e1274d98-d87d-4307-9e9d-851146e85637	2	JORDAN HAYES	PERSON	Dr. Jordan Hayes is a highly analytical and sk...	[ff65e3e3764d360b35c55a83aaa35b0138af32551d0bc...	4	13
3	d1135bbc-3bb8-450f-acf9-f3199e9c45cf	3	SAM RIVERA	PERSON	Sam Rivera is a highly skilled and adaptable m...	[ff65e3e3764d360b35c55a83aaa35b0138af32551d0bc...	4	17
4	6e9b5c3b-2ec4-4e48-b111-cebfbd02ee0f	4	PARANORMAL MILITARY SQUAD	ORGANIZATION	The Paranormal Military Squad is an elite mili...	[ff65e3e3764d360b35c55a83aaa35b0138af32551d0bc...	2	12

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: 107

Out[5]:

	id	human_readable_id	source	target	description	weight	combined_degree	text_unit_ids
0	7c9039f8-43d7-4685-a317-f7bbbfac8e3a	0	ALEX MERCER	PARANORMAL MILITARY SQUAD	Alex Mercer is a prominent agent and operative...	17.0	29	[ff65e3e3764d360b35c55a83aaa35b0138af32551d0bc...
1	f0310841-7ad1-4286-8e1d-1641c180b3ab	1	TAYLOR CRUZ	PARANORMAL MILITARY SQUAD	Taylor Cruz is a leading agent in the Paranorm...	9.0	29	[ff65e3e3764d360b35c55a83aaa35b0138af32551d0bc...
2	84441d65-6f4e-40eb-92ef-743ac6d1229c	2	JORDAN HAYES	PARANORMAL MILITARY SQUAD	Dr. Jordan Hayes is a key member of the Parano...	17.0	25	[ff65e3e3764d360b35c55a83aaa35b0138af32551d0bc...
3	ec2c0109-3a2f-479e-89d9-7a64abe07a63	3	SAM RIVERA	PARANORMAL MILITARY SQUAD	Sam Rivera is a technical expert and agent wit...	9.0	29	[ff65e3e3764d360b35c55a83aaa35b0138af32551d0bc...
4	19d249b1-487e-4db9-85db-410cfb6ab986	4	PARANORMAL MILITARY SQUAD	OPERATION: DULCE	The Paranormal Military Squad is the team assi...	10.0	20	[ff65e3e3764d360b35c55a83aaa35b0138af32551d0bc...

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: 41

Read community reports¶

In [7]:

Copied!

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()
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()

Report records: 10

Out[7]:

	id	human_readable_id	community	level	parent	children	title	summary	full_content	rank	rating_explanation	findings	full_content_json	period	size
0	9785e64b429dacd70cce8238051f3438b387ebcd6920ef...	7	7	1	0	[]	Paranormal Military Squad and Operation Dulce	This community centers on the Paranormal Milit...	# Paranormal Military Squad and Operation Dulc...	8.5	The impact severity rating is high due to the ...	[{'explanation': 'The Paranormal Military Squa...	{\n "title": "Paranormal Military Squad and...	2026-01-13	4
1	212c8ff1796bff731abe3d54f3216511ab15d5fb0d97eb...	8	8	1	0	[]	Paranormal Military Squad: Leadership and Scie...	This community centers on the Paranormal Milit...	# Paranormal Military Squad: Leadership and Sc...	8.0	The community poses a high impact due to its i...	[{'explanation': 'Taylor Cruz, often referred ...	{\n "title": "Paranormal Military Squad: Le...	2026-01-13	3
2	dbff99eaacade8ac4337a9bde0f533313445f0643ab3d0...	9	9	1	0	[]	Paranormal Military Squad and Operation: Dulce...	This community centers around the Paranormal M...	# Paranormal Military Squad and Operation: Dul...	8.5	The community poses a high impact severity due...	[{'explanation': 'The military complex is the ...	{\n "title": "Paranormal Military Squad and...	2026-01-13	3
3	b6e8ab1ac8ecc605bbec09ad78819b95ac08d62c5a9991...	0	0	0	-1	[7, 8, 9]	Paranormal Military Squad and Operation: Dulce	This community centers on the Paranormal Milit...	# Paranormal Military Squad and Operation: Dul...	8.5	The community poses a high impact due to its e...	[{'explanation': 'Taylor Cruz is the central a...	{\n "title": "Paranormal Military Squad and...	2026-01-13	10
4	97038bee58306d13b709b615fb2b816cae0c11273a9e73...	1	1	0	-1	[]	Team of Agents Investigating Dulce Base	This community centers on a specialized team o...	# Team of Agents Investigating Dulce Base\n\nT...	7.5	The impact severity rating is high due to the ...	[{'explanation': 'The core of this community i...	{\n "title": "Team of Agents Investigating ...	2026-01-13	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: 5

Out[8]:

	id	human_readable_id	text	n_tokens	document_id	entity_ids	relationship_ids	covariate_ids
0	ff65e3e3764d360b35c55a83aaa35b0138af32551d0bc0...	0	# Operation: Dulce\n\n## Chapter 1\n\nThe thru...	1200	6e81f882f89dd5596e1925dd3ae8a4f0a0edcb55b35a80...	[bbb2ad85-3d0b-4c26-b6f9-71ed76ee64bd, 66112f0...	[7c9039f8-43d7-4685-a317-f7bbbfac8e3a, f031084...	[4c6e019b-29c5-456c-951c-0cd3df308634, b2d2c91...
1	56388c08c6e58a2aeaa235be38fe973bbdcbfb5ca72fd2...	1	Taylor offered a brief nod, perhaps too curt,...	1200	6e81f882f89dd5596e1925dd3ae8a4f0a0edcb55b35a80...	[bbb2ad85-3d0b-4c26-b6f9-71ed76ee64bd, 66112f0...	[e15afdbb-6757-4f10-897d-54b9bb3dddc7, 319a7dc...	[9cdfde65-90c1-46a2-bfa2-9808ea4c821a, 83a704b...
2	6012990fbcfad49511dce9a360dcc28d0c9bd37dc7b337...	2	team was descending into the earth, carrying ...	1200	6e81f882f89dd5596e1925dd3ae8a4f0a0edcb55b35a80...	[bbb2ad85-3d0b-4c26-b6f9-71ed76ee64bd, 66112f0...	[4f42f1c3-b769-48f6-a01d-6a4a89b1c084, 319a7dc...	[d3d99801-1d52-4398-9e36-3dd3d8a06a56, 1a2f60a...
3	402ae07f78c0154e8c24a93c00a90ee7c85e7478a755da...	3	Cruz squared their shoulders, clearly ruffled ...	1200	6e81f882f89dd5596e1925dd3ae8a4f0a0edcb55b35a80...	[bbb2ad85-3d0b-4c26-b6f9-71ed76ee64bd, e1274d9...	[7c9039f8-43d7-4685-a317-f7bbbfac8e3a, 84441d6...	[b39dd61f-5a5a-4894-9047-23bee3e03f73, 23029bd...
4	999a787eccebd488865de53a048f60889847f42057d5ad...	4	to him, and he couldn't help but feel the wei...	382	6e81f882f89dd5596e1925dd3ae8a4f0a0edcb55b35a80...	[bbb2ad85-3d0b-4c26-b6f9-71ed76ee64bd, 66112f0...	[d3ea1355-14f3-494e-94cf-d43ac1077ab4, 72e568c...	[7c71874a-1443-40a6-bbad-6bc78ef23ee4, dac744e...

In [9]:

Copied!





from graphrag.tokenizer.get_tokenizer import get_tokenizer
from graphrag_llm.completion import create_completion
from graphrag_llm.config import ModelConfig
from graphrag_llm.embedding import create_embedding

api_key = os.environ["GRAPHRAG_API_KEY"]

chat_config = ModelConfig(
    type="litellm",
    model_provider="openai",
    model="gpt-4.1",
    api_key=api_key,
)
chat_model = create_completion(chat_config)

embedding_config = ModelConfig(
    type="litellm",
    model_provider="openai",
    model="text-embedding-3-small",
    api_key=api_key,
)

text_embedder = create_embedding(embedding_config)

tokenizer = get_tokenizer(chat_config)
from graphrag.tokenizer.get_tokenizer import get_tokenizer
from graphrag_llm.completion import create_completion
from graphrag_llm.config import ModelConfig
from graphrag_llm.embedding import create_embedding

api_key = os.environ["GRAPHRAG_API_KEY"]

chat_config = ModelConfig(
    type="litellm",
    model_provider="openai",
    model="gpt-4.1",
    api_key=api_key,
)
chat_model = create_completion(chat_config)

embedding_config = ModelConfig(
    type="litellm",
    model_provider="openai",
    model="text-embedding-3-small",
    api_key=api_key,
)

text_embedder = create_embedding(embedding_config)

tokenizer = get_tokenizer(chat_config)

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,
    tokenizer=tokenizer,
)
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,
    tokenizer=tokenizer,
)

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)
}

model_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)
}

model_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(
    model=chat_model,
    context_builder=context_builder,
    tokenizer=tokenizer,
    model_params=model_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(
    model=chat_model,
    context_builder=context_builder,
    tokenizer=tokenizer,
    model_params=model_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.search("Tell me about Agent Mercer")
print(result.response)
result = await search_engine.search("Tell me about Agent Mercer")
print(result.response)

Request failed with exception=litellm.AuthenticationError: AuthenticationError: OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}
Traceback (most recent call last):
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1310, in embedding
    headers, sync_embedding_response = self.make_sync_openai_embedding_request(
                                       ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        openai_client=openai_client,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ...<2 lines>...
        logging_obj=logging_obj,
        ^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/logging_utils.py", line 347, in sync_wrapper
    result: Final = func(*args, **kwargs)
                    ~~~~^^^^^^^^^^^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1185, in make_sync_openai_embedding_request
    raise e
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1179, in make_sync_openai_embedding_request
    raw_response = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout)
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_legacy_response.py", line 369, in wrapped
    return cast(LegacyAPIResponse[R], func(*args, **kwargs))
                                      ~~~~^^^^^^^^^^^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/resources/embeddings.py", line 136, in create
    return self._post(
           ~~~~~~~~~~^
        "/embeddings",
        ^^^^^^^^^^^^^^
    ...<9 lines>...
        cast_to=CreateEmbeddingResponse,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py", line 1375, in post
    return cast(ResponseT, self.request(cast_to, opts, stream=stream, stream_cls=stream_cls))
                           ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py", line 1148, in request
    raise self._make_status_error_from_response(err.response) from None
openai.AuthenticationError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py", line 6308, in embedding
    response = openai_chat_completions.embedding(
        model=model,
    ...<10 lines>...
        shared_session=shared_session,
    )
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1341, in embedding
    raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
litellm.llms.openai.common_utils.OpenAIError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_logging.py", line 49, in _request_count_middleware
    return sync_middleware(**kwargs)
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_request_count.py", line 53, in _request_count_middleware
    result = sync_middleware(**kwargs)
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_metrics.py", line 56, in _metrics_middleware
    response = sync_middleware(**kwargs)
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/embedding/lite_llm_embedding.py", line 188, in _base_embedding
    response = litellm.embedding(**new_args)
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py", line 1745, in wrapper
    raise e
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py", line 1591, in wrapper
    result = original_function(*args, **kwargs)
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py", line 7110, in embedding
    raise exception_type(
          ~~~~~~~~~~~~~~^
        model=model,
        ^^^^^^^^^^^^
    ...<2 lines>...
        extra_kwargs=kwargs,
        ^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py", line 2675, in exception_type
    raise e  # it's already mapped
    ^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py", line 2436, in exception_type
    _map_openai_exception(
    ~~~~~~~~~~~~~~~~~~~~~^
        model=model,
        ^^^^^^^^^^^^
    ...<5 lines>...
        extra_information=extra_information,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py", line 409, in _map_openai_exception
    raise AuthenticationError(
    ...<5 lines>...
    )
litellm.exceptions.AuthenticationError: litellm.AuthenticationError: AuthenticationError: OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

---------------------------------------------------------------------------
AuthenticationError                       Traceback (most recent call last)
File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1310, in OpenAIChatCompletion.embedding(self, model, input, timeout, logging_obj, model_response, optional_params, api_key, api_base, client, aembedding, max_retries, shared_session)
   1309 headers: dict | None = None
-> 1310 headers, sync_embedding_response = self.make_sync_openai_embedding_request(
   1311     openai_client=openai_client,
   1312     data=data,
   1313     timeout=timeout,
   1314     logging_obj=logging_obj,
   1315 )
   1317 ## LOGGING

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/logging_utils.py:347, in track_llm_api_timing.<locals>.decorator.<locals>.sync_wrapper(*args, **kwargs)
    346 try:
--> 347     result: Final = func(*args, **kwargs)
    348     return result

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1185, in OpenAIChatCompletion.make_sync_openai_embedding_request(self, openai_client, data, timeout, logging_obj)
   1184 except Exception as e:
-> 1185     raise e

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1179, in OpenAIChatCompletion.make_sync_openai_embedding_request(self, openai_client, data, timeout, logging_obj)
   1178 try:
-> 1179     raw_response = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout)
   1181     headers: Final = dict(raw_response.headers)

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_legacy_response.py:369, in to_raw_response_wrapper.<locals>.wrapped(*args, **kwargs)
    367 kwargs["extra_headers"] = extra_headers
--> 369 return cast(LegacyAPIResponse[R], func(*args, **kwargs))

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/resources/embeddings.py:136, in Embeddings.create(self, input, model, dimensions, encoding_format, user, extra_headers, extra_query, extra_body, timeout)
    134     return obj
--> 136 return self._post(
    137     "/embeddings",
    138     body=maybe_transform(params, embedding_create_params.EmbeddingCreateParams),
    139     options=make_request_options(
    140         extra_headers=extra_headers,
    141         extra_query=extra_query,
    142         extra_body=extra_body,
    143         timeout=timeout,
    144         post_parser=parser,
    145         security={"bearer_auth": True},
    146     ),
    147     cast_to=CreateEmbeddingResponse,
    148 )

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py:1375, in SyncAPIClient.post(self, path, cast_to, body, content, options, files, stream, stream_cls)
   1372 opts = FinalRequestOptions.construct(
   1373     method="post", url=path, json_data=body, content=content, files=to_httpx_files(files), **options
   1374 )
-> 1375 return cast(ResponseT, self.request(cast_to, opts, stream=stream, stream_cls=stream_cls))

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py:1148, in SyncAPIClient.request(self, cast_to, options, stream, stream_cls)
   1147     log.debug("Re-raising status error")
-> 1148     raise self._make_status_error_from_response(err.response) from None
   1150 break

AuthenticationError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

OpenAIError                               Traceback (most recent call last)
File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py:6308, in embedding(model, input, dimensions, encoding_format, timeout, api_base, api_version, api_key, api_type, caching, user, custom_llm_provider, litellm_call_id, logger_fn, **kwargs)
   6307     ## EMBEDDING CALL
-> 6308     response = openai_chat_completions.embedding(
   6309         model=model,
   6310         input=input,
   6311         api_base=api_base,
   6312         api_key=api_key,
   6313         logging_obj=logging,
   6314         timeout=timeout,
   6315         model_response=EmbeddingResponse(),
   6316         optional_params=optional_params,
   6317         client=client,
   6318         aembedding=aembedding,
   6319         max_retries=max_retries,
   6320         shared_session=shared_session,
   6321     )
   6322 elif custom_llm_provider == "databricks":

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1341, in OpenAIChatCompletion.embedding(self, model, input, timeout, logging_obj, model_response, optional_params, api_key, api_base, client, aembedding, max_retries, shared_session)
   1340     error_headers = getattr(error_response, "headers", None)
-> 1341 raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)

OpenAIError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

AuthenticationError                       Traceback (most recent call last)
Cell In[13], line 1
----> 1 result = await search_engine.search("Tell me about Agent Mercer")
      2 print(result.response)

File ~/work/graphrag/graphrag/packages/graphrag/graphrag/query/structured_search/local_search/search.py:66, in LocalSearch.search(self, query, conversation_history, **kwargs)
     64 search_prompt = ""
     65 llm_calls, prompt_tokens, output_tokens = {}, {}, {}
---> 66 context_result = self.context_builder.build_context(
     67     query=query,
     68     conversation_history=conversation_history,
     69     **kwargs,
     70     **self.context_builder_params,
     71 )
     72 llm_calls["build_context"] = context_result.llm_calls
     73 prompt_tokens["build_context"] = context_result.prompt_tokens

File ~/work/graphrag/graphrag/packages/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_context_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/packages/graphrag/graphrag/query/context_builder/entity_extraction.py:62, 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)
     58 matched_entities = []
     59 if query != "":
     60     # get entities with highest semantic similarity to query
     61     # oversample to account for excluded entities
---> 62     search_results = text_embedding_vectorstore.similarity_search_by_text(
     63         text=query,
     64         text_embedder=lambda t: text_embedder.embedding(input=[t]).first_embedding,
     65         k=k * oversample_scaler,
     66     )
     67     for result in search_results:
     68         if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance(
     69             result.document.id, str
     70         ):

File ~/work/graphrag/graphrag/packages/graphrag-vectors/graphrag_vectors/vector_store.py:186, in VectorStore.similarity_search_by_text(self, text, text_embedder, k, select, filters, include_vectors)
    176 def similarity_search_by_text(
    177     self,
    178     text: str,
   (...)    183     include_vectors: bool = True,
    184 ) -> list[VectorStoreSearchResult]:
    185     """Perform a text-based similarity search."""
--> 186     query_embedding = text_embedder(text)
    187     if query_embedding:
    188         return self.similarity_search_by_vector(
    189             query_embedding=query_embedding,
    190             k=k,
   (...)    193             include_vectors=include_vectors,
    194         )

File ~/work/graphrag/graphrag/packages/graphrag/graphrag/query/context_builder/entity_extraction.py:64, in map_query_to_entities.<locals>.<lambda>(t)
     58 matched_entities = []
     59 if query != "":
     60     # get entities with highest semantic similarity to query
     61     # oversample to account for excluded entities
     62     search_results = text_embedding_vectorstore.similarity_search_by_text(
     63         text=query,
---> 64         text_embedder=lambda t: text_embedder.embedding(input=[t]).first_embedding,
     65         k=k * oversample_scaler,
     66     )
     67     for result in search_results:
     68         if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance(
     69             result.document.id, str
     70         ):

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/embedding/lite_llm_embedding.py:130, in LiteLLMEmbedding.embedding(self, **kwargs)
    127     request_metrics = None
    129 try:
--> 130     return self._embedding(metrics=request_metrics, **kwargs)
    131 finally:
    132     if request_metrics:

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_logging.py:49, in with_logging.<locals>._request_count_middleware(**kwargs)
     47 metrics: Metrics | None = kwargs.get("metrics")
     48 try:
---> 49     return sync_middleware(**kwargs)
     50 except Exception as e:
     51     retries = metrics.get("retries", None) if metrics else None

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_request_count.py:53, in with_request_count.<locals>._request_count_middleware(**kwargs)
     51     metrics["failed_response_count"] = 0
     52 try:
---> 53     result = sync_middleware(**kwargs)
     54     if metrics is not None:
     55         metrics["successful_response_count"] = 1

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_metrics.py:56, in with_metrics.<locals>._metrics_middleware(**kwargs)
     54 metrics: Metrics | None = kwargs.get("metrics")
     55 start_time = time.time()
---> 56 response = sync_middleware(**kwargs)
     57 end_time = time.time()
     59 if metrics is not None:

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/embedding/lite_llm_embedding.py:188, in _create_base_embeddings.<locals>._base_embedding(**kwargs)
    185 kwargs.pop("metrics", None)  # Remove metrics if present
    186 new_args: dict[str, Any] = {**base_args, **kwargs}
--> 188 response = litellm.embedding(**new_args)
    189 return LLMEmbeddingResponse(**response.model_dump())

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py:1745, in client.<locals>.wrapper(*args, **kwargs)
   1741 if logging_obj:
   1742     logging_obj.failure_handler(
   1743         e, traceback_exception, start_time, end_time
   1744     )  # DO NOT MAKE THREADED - router retry fallback relies on this!
-> 1745 raise e

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py:1591, in client.<locals>.wrapper(*args, **kwargs)
   1589         print_verbose(f"Error while checking max token limit: {e}")
   1590 # MODEL CALL
-> 1591 result = original_function(*args, **kwargs)
   1592 end_time = datetime.datetime.now()
   1593 if _is_streaming_request(
   1594     kwargs=kwargs,
   1595     call_type=call_type,
   1596 ):

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py:7110, in embedding(model, input, dimensions, encoding_format, timeout, api_base, api_version, api_key, api_type, caching, user, custom_llm_provider, litellm_call_id, logger_fn, **kwargs)
   7104 litellm_logging_obj.post_call(
   7105     input=input,
   7106     api_key=api_key,
   7107     original_response=str(e),
   7108 )
   7109 ## Map to OpenAI Exception
-> 7110 raise exception_type(
   7111     model=model,
   7112     original_exception=e,
   7113     custom_llm_provider=custom_llm_provider,
   7114     extra_kwargs=kwargs,
   7115 )

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py:2675, in exception_type(model, original_exception, custom_llm_provider, completion_kwargs, extra_kwargs)
   2673     if isinstance(e, error_type):
   2674         setattr(e, "litellm_response_headers", litellm_response_headers)
-> 2675         raise e  # it's already mapped
   2676 raised_exc: Final = APIConnectionError(
   2677     message=f"{original_exception}\n{_redact_string(traceback.format_exc())}",
   2678     llm_provider="",
   2679     model="",
   2680 )
   2681 setattr(raised_exc, "litellm_response_headers", litellm_response_headers)

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py:2436, in exception_type(model, original_exception, custom_llm_provider, completion_kwargs, extra_kwargs)
   2422     extract_and_raise_litellm_exception(
   2423         response=getattr(original_exception, "response", None),
   2424         error_str=error_str,
   2425         model=model,
   2426         custom_llm_provider=custom_llm_provider,
   2427     )
   2428 if (
   2429     custom_llm_provider == "openai"
   2430     or custom_llm_provider == "text-completion-openai"
   (...)   2434     or custom_llm_provider == "runwayml"
   2435 ):
-> 2436     _map_openai_exception(
   2437         model=model,
   2438         original_exception=mappable_exception,
   2439         custom_llm_provider=custom_llm_provider,
   2440         error_str=error_str,
   2441         exception_type=exception_type,
   2442         exception_provider=exception_provider,
   2443         extra_information=extra_information,
   2444     )
   2445 elif custom_llm_provider == "anthropic" or custom_llm_provider == "anthropic_text":  # one of the anthropics
   2446     _map_anthropic_exception(
   2447         model=model,
   2448         original_exception=mappable_exception,
   (...)   2453         extra_information=extra_information,
   2454     )

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py:409, in _map_openai_exception(model, original_exception, custom_llm_provider, error_str, exception_type, exception_provider, extra_information)
    401     raise BadRequestError(
    402         message=f"{exception_provider} - {message}",
    403         llm_provider=custom_llm_provider,
   (...)    406         litellm_debug_info=extra_information,
    407     )
    408 elif original_exception.status_code == 401:
--> 409     raise AuthenticationError(
    410         message=f"AuthenticationError: {exception_provider} - {message}",
    411         llm_provider=custom_llm_provider,
    412         model=model,
    413         response=getattr(original_exception, "response", None),
    414         litellm_debug_info=extra_information,
    415     )
    416 elif original_exception.status_code == 404:
    417     raise NotFoundError(
    418         message=f"NotFoundError: {exception_provider} - {message}",
    419         model=model,
   (...)    422         litellm_debug_info=extra_information,
    423     )

AuthenticationError: litellm.AuthenticationError: AuthenticationError: OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

In [14]:

Copied!

question = "Tell me about Dr. Jordan Hayes"
result = await search_engine.search(question)
print(result.response)
question = "Tell me about Dr. Jordan Hayes"
result = await search_engine.search(question)
print(result.response)

Request failed with exception=litellm.AuthenticationError: AuthenticationError: OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}
Traceback (most recent call last):
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1310, in embedding
    headers, sync_embedding_response = self.make_sync_openai_embedding_request(
                                       ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        openai_client=openai_client,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ...<2 lines>...
        logging_obj=logging_obj,
        ^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/logging_utils.py", line 347, in sync_wrapper
    result: Final = func(*args, **kwargs)
                    ~~~~^^^^^^^^^^^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1185, in make_sync_openai_embedding_request
    raise e
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1179, in make_sync_openai_embedding_request
    raw_response = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout)
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_legacy_response.py", line 369, in wrapped
    return cast(LegacyAPIResponse[R], func(*args, **kwargs))
                                      ~~~~^^^^^^^^^^^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/resources/embeddings.py", line 136, in create
    return self._post(
           ~~~~~~~~~~^
        "/embeddings",
        ^^^^^^^^^^^^^^
    ...<9 lines>...
        cast_to=CreateEmbeddingResponse,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py", line 1375, in post
    return cast(ResponseT, self.request(cast_to, opts, stream=stream, stream_cls=stream_cls))
                           ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py", line 1148, in request
    raise self._make_status_error_from_response(err.response) from None
openai.AuthenticationError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py", line 6308, in embedding
    response = openai_chat_completions.embedding(
        model=model,
    ...<10 lines>...
        shared_session=shared_session,
    )
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1341, in embedding
    raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
litellm.llms.openai.common_utils.OpenAIError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_logging.py", line 49, in _request_count_middleware
    return sync_middleware(**kwargs)
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_request_count.py", line 53, in _request_count_middleware
    result = sync_middleware(**kwargs)
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_metrics.py", line 56, in _metrics_middleware
    response = sync_middleware(**kwargs)
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/embedding/lite_llm_embedding.py", line 188, in _base_embedding
    response = litellm.embedding(**new_args)
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py", line 1745, in wrapper
    raise e
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py", line 1591, in wrapper
    result = original_function(*args, **kwargs)
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py", line 7110, in embedding
    raise exception_type(
          ~~~~~~~~~~~~~~^
        model=model,
        ^^^^^^^^^^^^
    ...<2 lines>...
        extra_kwargs=kwargs,
        ^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py", line 2675, in exception_type
    raise e  # it's already mapped
    ^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py", line 2436, in exception_type
    _map_openai_exception(
    ~~~~~~~~~~~~~~~~~~~~~^
        model=model,
        ^^^^^^^^^^^^
    ...<5 lines>...
        extra_information=extra_information,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py", line 409, in _map_openai_exception
    raise AuthenticationError(
    ...<5 lines>...
    )
litellm.exceptions.AuthenticationError: litellm.AuthenticationError: AuthenticationError: OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

---------------------------------------------------------------------------
AuthenticationError                       Traceback (most recent call last)
File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1310, in OpenAIChatCompletion.embedding(self, model, input, timeout, logging_obj, model_response, optional_params, api_key, api_base, client, aembedding, max_retries, shared_session)
   1309 headers: dict | None = None
-> 1310 headers, sync_embedding_response = self.make_sync_openai_embedding_request(
   1311     openai_client=openai_client,
   1312     data=data,
   1313     timeout=timeout,
   1314     logging_obj=logging_obj,
   1315 )
   1317 ## LOGGING

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/logging_utils.py:347, in track_llm_api_timing.<locals>.decorator.<locals>.sync_wrapper(*args, **kwargs)
    346 try:
--> 347     result: Final = func(*args, **kwargs)
    348     return result

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1185, in OpenAIChatCompletion.make_sync_openai_embedding_request(self, openai_client, data, timeout, logging_obj)
   1184 except Exception as e:
-> 1185     raise e

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1179, in OpenAIChatCompletion.make_sync_openai_embedding_request(self, openai_client, data, timeout, logging_obj)
   1178 try:
-> 1179     raw_response = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout)
   1181     headers: Final = dict(raw_response.headers)

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_legacy_response.py:369, in to_raw_response_wrapper.<locals>.wrapped(*args, **kwargs)
    367 kwargs["extra_headers"] = extra_headers
--> 369 return cast(LegacyAPIResponse[R], func(*args, **kwargs))

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/resources/embeddings.py:136, in Embeddings.create(self, input, model, dimensions, encoding_format, user, extra_headers, extra_query, extra_body, timeout)
    134     return obj
--> 136 return self._post(
    137     "/embeddings",
    138     body=maybe_transform(params, embedding_create_params.EmbeddingCreateParams),
    139     options=make_request_options(
    140         extra_headers=extra_headers,
    141         extra_query=extra_query,
    142         extra_body=extra_body,
    143         timeout=timeout,
    144         post_parser=parser,
    145         security={"bearer_auth": True},
    146     ),
    147     cast_to=CreateEmbeddingResponse,
    148 )

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py:1375, in SyncAPIClient.post(self, path, cast_to, body, content, options, files, stream, stream_cls)
   1372 opts = FinalRequestOptions.construct(
   1373     method="post", url=path, json_data=body, content=content, files=to_httpx_files(files), **options
   1374 )
-> 1375 return cast(ResponseT, self.request(cast_to, opts, stream=stream, stream_cls=stream_cls))

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py:1148, in SyncAPIClient.request(self, cast_to, options, stream, stream_cls)
   1147     log.debug("Re-raising status error")
-> 1148     raise self._make_status_error_from_response(err.response) from None
   1150 break

AuthenticationError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

OpenAIError                               Traceback (most recent call last)
File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py:6308, in embedding(model, input, dimensions, encoding_format, timeout, api_base, api_version, api_key, api_type, caching, user, custom_llm_provider, litellm_call_id, logger_fn, **kwargs)
   6307     ## EMBEDDING CALL
-> 6308     response = openai_chat_completions.embedding(
   6309         model=model,
   6310         input=input,
   6311         api_base=api_base,
   6312         api_key=api_key,
   6313         logging_obj=logging,
   6314         timeout=timeout,
   6315         model_response=EmbeddingResponse(),
   6316         optional_params=optional_params,
   6317         client=client,
   6318         aembedding=aembedding,
   6319         max_retries=max_retries,
   6320         shared_session=shared_session,
   6321     )
   6322 elif custom_llm_provider == "databricks":

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1341, in OpenAIChatCompletion.embedding(self, model, input, timeout, logging_obj, model_response, optional_params, api_key, api_base, client, aembedding, max_retries, shared_session)
   1340     error_headers = getattr(error_response, "headers", None)
-> 1341 raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)

OpenAIError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

AuthenticationError                       Traceback (most recent call last)
Cell In[14], line 2
      1 question = "Tell me about Dr. Jordan Hayes"
----> 2 result = await search_engine.search(question)
      3 print(result.response)

File ~/work/graphrag/graphrag/packages/graphrag/graphrag/query/structured_search/local_search/search.py:66, in LocalSearch.search(self, query, conversation_history, **kwargs)
     64 search_prompt = ""
     65 llm_calls, prompt_tokens, output_tokens = {}, {}, {}
---> 66 context_result = self.context_builder.build_context(
     67     query=query,
     68     conversation_history=conversation_history,
     69     **kwargs,
     70     **self.context_builder_params,
     71 )
     72 llm_calls["build_context"] = context_result.llm_calls
     73 prompt_tokens["build_context"] = context_result.prompt_tokens

File ~/work/graphrag/graphrag/packages/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_context_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/packages/graphrag/graphrag/query/context_builder/entity_extraction.py:62, 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)
     58 matched_entities = []
     59 if query != "":
     60     # get entities with highest semantic similarity to query
     61     # oversample to account for excluded entities
---> 62     search_results = text_embedding_vectorstore.similarity_search_by_text(
     63         text=query,
     64         text_embedder=lambda t: text_embedder.embedding(input=[t]).first_embedding,
     65         k=k * oversample_scaler,
     66     )
     67     for result in search_results:
     68         if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance(
     69             result.document.id, str
     70         ):

File ~/work/graphrag/graphrag/packages/graphrag-vectors/graphrag_vectors/vector_store.py:186, in VectorStore.similarity_search_by_text(self, text, text_embedder, k, select, filters, include_vectors)
    176 def similarity_search_by_text(
    177     self,
    178     text: str,
   (...)    183     include_vectors: bool = True,
    184 ) -> list[VectorStoreSearchResult]:
    185     """Perform a text-based similarity search."""
--> 186     query_embedding = text_embedder(text)
    187     if query_embedding:
    188         return self.similarity_search_by_vector(
    189             query_embedding=query_embedding,
    190             k=k,
   (...)    193             include_vectors=include_vectors,
    194         )

File ~/work/graphrag/graphrag/packages/graphrag/graphrag/query/context_builder/entity_extraction.py:64, in map_query_to_entities.<locals>.<lambda>(t)
     58 matched_entities = []
     59 if query != "":
     60     # get entities with highest semantic similarity to query
     61     # oversample to account for excluded entities
     62     search_results = text_embedding_vectorstore.similarity_search_by_text(
     63         text=query,
---> 64         text_embedder=lambda t: text_embedder.embedding(input=[t]).first_embedding,
     65         k=k * oversample_scaler,
     66     )
     67     for result in search_results:
     68         if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance(
     69             result.document.id, str
     70         ):

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/embedding/lite_llm_embedding.py:130, in LiteLLMEmbedding.embedding(self, **kwargs)
    127     request_metrics = None
    129 try:
--> 130     return self._embedding(metrics=request_metrics, **kwargs)
    131 finally:
    132     if request_metrics:

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_logging.py:49, in with_logging.<locals>._request_count_middleware(**kwargs)
     47 metrics: Metrics | None = kwargs.get("metrics")
     48 try:
---> 49     return sync_middleware(**kwargs)
     50 except Exception as e:
     51     retries = metrics.get("retries", None) if metrics else None

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_request_count.py:53, in with_request_count.<locals>._request_count_middleware(**kwargs)
     51     metrics["failed_response_count"] = 0
     52 try:
---> 53     result = sync_middleware(**kwargs)
     54     if metrics is not None:
     55         metrics["successful_response_count"] = 1

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_metrics.py:56, in with_metrics.<locals>._metrics_middleware(**kwargs)
     54 metrics: Metrics | None = kwargs.get("metrics")
     55 start_time = time.time()
---> 56 response = sync_middleware(**kwargs)
     57 end_time = time.time()
     59 if metrics is not None:

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/embedding/lite_llm_embedding.py:188, in _create_base_embeddings.<locals>._base_embedding(**kwargs)
    185 kwargs.pop("metrics", None)  # Remove metrics if present
    186 new_args: dict[str, Any] = {**base_args, **kwargs}
--> 188 response = litellm.embedding(**new_args)
    189 return LLMEmbeddingResponse(**response.model_dump())

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py:1745, in client.<locals>.wrapper(*args, **kwargs)
   1741 if logging_obj:
   1742     logging_obj.failure_handler(
   1743         e, traceback_exception, start_time, end_time
   1744     )  # DO NOT MAKE THREADED - router retry fallback relies on this!
-> 1745 raise e

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py:1591, in client.<locals>.wrapper(*args, **kwargs)
   1589         print_verbose(f"Error while checking max token limit: {e}")
   1590 # MODEL CALL
-> 1591 result = original_function(*args, **kwargs)
   1592 end_time = datetime.datetime.now()
   1593 if _is_streaming_request(
   1594     kwargs=kwargs,
   1595     call_type=call_type,
   1596 ):

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py:7110, in embedding(model, input, dimensions, encoding_format, timeout, api_base, api_version, api_key, api_type, caching, user, custom_llm_provider, litellm_call_id, logger_fn, **kwargs)
   7104 litellm_logging_obj.post_call(
   7105     input=input,
   7106     api_key=api_key,
   7107     original_response=str(e),
   7108 )
   7109 ## Map to OpenAI Exception
-> 7110 raise exception_type(
   7111     model=model,
   7112     original_exception=e,
   7113     custom_llm_provider=custom_llm_provider,
   7114     extra_kwargs=kwargs,
   7115 )

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py:2675, in exception_type(model, original_exception, custom_llm_provider, completion_kwargs, extra_kwargs)
   2673     if isinstance(e, error_type):
   2674         setattr(e, "litellm_response_headers", litellm_response_headers)
-> 2675         raise e  # it's already mapped
   2676 raised_exc: Final = APIConnectionError(
   2677     message=f"{original_exception}\n{_redact_string(traceback.format_exc())}",
   2678     llm_provider="",
   2679     model="",
   2680 )
   2681 setattr(raised_exc, "litellm_response_headers", litellm_response_headers)

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py:2436, in exception_type(model, original_exception, custom_llm_provider, completion_kwargs, extra_kwargs)
   2422     extract_and_raise_litellm_exception(
   2423         response=getattr(original_exception, "response", None),
   2424         error_str=error_str,
   2425         model=model,
   2426         custom_llm_provider=custom_llm_provider,
   2427     )
   2428 if (
   2429     custom_llm_provider == "openai"
   2430     or custom_llm_provider == "text-completion-openai"
   (...)   2434     or custom_llm_provider == "runwayml"
   2435 ):
-> 2436     _map_openai_exception(
   2437         model=model,
   2438         original_exception=mappable_exception,
   2439         custom_llm_provider=custom_llm_provider,
   2440         error_str=error_str,
   2441         exception_type=exception_type,
   2442         exception_provider=exception_provider,
   2443         extra_information=extra_information,
   2444     )
   2445 elif custom_llm_provider == "anthropic" or custom_llm_provider == "anthropic_text":  # one of the anthropics
   2446     _map_anthropic_exception(
   2447         model=model,
   2448         original_exception=mappable_exception,
   (...)   2453         extra_information=extra_information,
   2454     )

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py:409, in _map_openai_exception(model, original_exception, custom_llm_provider, error_str, exception_type, exception_provider, extra_information)
    401     raise BadRequestError(
    402         message=f"{exception_provider} - {message}",
    403         llm_provider=custom_llm_provider,
   (...)    406         litellm_debug_info=extra_information,
    407     )
    408 elif original_exception.status_code == 401:
--> 409     raise AuthenticationError(
    410         message=f"AuthenticationError: {exception_provider} - {message}",
    411         llm_provider=custom_llm_provider,
    412         model=model,
    413         response=getattr(original_exception, "response", None),
    414         litellm_debug_info=extra_information,
    415     )
    416 elif original_exception.status_code == 404:
    417     raise NotFoundError(
    418         message=f"NotFoundError: {exception_provider} - {message}",
    419         model=model,
   (...)    422         litellm_debug_info=extra_information,
    423     )

AuthenticationError: litellm.AuthenticationError: AuthenticationError: OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

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(
    model=chat_model,
    context_builder=context_builder,
    tokenizer=tokenizer,
    model_params=model_params,
    context_builder_params=local_context_params,
)
question_generator = LocalQuestionGen(
    model=chat_model,
    context_builder=context_builder,
    tokenizer=tokenizer,
    model_params=model_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)

Request failed with exception=litellm.AuthenticationError: AuthenticationError: OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}
Traceback (most recent call last):
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1310, in embedding
    headers, sync_embedding_response = self.make_sync_openai_embedding_request(
                                       ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        openai_client=openai_client,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ...<2 lines>...
        logging_obj=logging_obj,
        ^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/logging_utils.py", line 347, in sync_wrapper
    result: Final = func(*args, **kwargs)
                    ~~~~^^^^^^^^^^^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1185, in make_sync_openai_embedding_request
    raise e
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1179, in make_sync_openai_embedding_request
    raw_response = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout)
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_legacy_response.py", line 369, in wrapped
    return cast(LegacyAPIResponse[R], func(*args, **kwargs))
                                      ~~~~^^^^^^^^^^^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/resources/embeddings.py", line 136, in create
    return self._post(
           ~~~~~~~~~~^
        "/embeddings",
        ^^^^^^^^^^^^^^
    ...<9 lines>...
        cast_to=CreateEmbeddingResponse,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py", line 1375, in post
    return cast(ResponseT, self.request(cast_to, opts, stream=stream, stream_cls=stream_cls))
                           ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py", line 1148, in request
    raise self._make_status_error_from_response(err.response) from None
openai.AuthenticationError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py", line 6308, in embedding
    response = openai_chat_completions.embedding(
        model=model,
    ...<10 lines>...
        shared_session=shared_session,
    )
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py", line 1341, in embedding
    raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
litellm.llms.openai.common_utils.OpenAIError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_logging.py", line 49, in _request_count_middleware
    return sync_middleware(**kwargs)
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_request_count.py", line 53, in _request_count_middleware
    result = sync_middleware(**kwargs)
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_metrics.py", line 56, in _metrics_middleware
    response = sync_middleware(**kwargs)
  File "/home/runner/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/embedding/lite_llm_embedding.py", line 188, in _base_embedding
    response = litellm.embedding(**new_args)
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py", line 1745, in wrapper
    raise e
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py", line 1591, in wrapper
    result = original_function(*args, **kwargs)
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py", line 7110, in embedding
    raise exception_type(
          ~~~~~~~~~~~~~~^
        model=model,
        ^^^^^^^^^^^^
    ...<2 lines>...
        extra_kwargs=kwargs,
        ^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py", line 2675, in exception_type
    raise e  # it's already mapped
    ^^^^^^^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py", line 2436, in exception_type
    _map_openai_exception(
    ~~~~~~~~~~~~~~~~~~~~~^
        model=model,
        ^^^^^^^^^^^^
    ...<5 lines>...
        extra_information=extra_information,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/runner/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py", line 409, in _map_openai_exception
    raise AuthenticationError(
    ...<5 lines>...
    )
litellm.exceptions.AuthenticationError: litellm.AuthenticationError: AuthenticationError: OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

---------------------------------------------------------------------------
AuthenticationError                       Traceback (most recent call last)
File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1310, in OpenAIChatCompletion.embedding(self, model, input, timeout, logging_obj, model_response, optional_params, api_key, api_base, client, aembedding, max_retries, shared_session)
   1309 headers: dict | None = None
-> 1310 headers, sync_embedding_response = self.make_sync_openai_embedding_request(
   1311     openai_client=openai_client,
   1312     data=data,
   1313     timeout=timeout,
   1314     logging_obj=logging_obj,
   1315 )
   1317 ## LOGGING

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/logging_utils.py:347, in track_llm_api_timing.<locals>.decorator.<locals>.sync_wrapper(*args, **kwargs)
    346 try:
--> 347     result: Final = func(*args, **kwargs)
    348     return result

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1185, in OpenAIChatCompletion.make_sync_openai_embedding_request(self, openai_client, data, timeout, logging_obj)
   1184 except Exception as e:
-> 1185     raise e

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1179, in OpenAIChatCompletion.make_sync_openai_embedding_request(self, openai_client, data, timeout, logging_obj)
   1178 try:
-> 1179     raw_response = openai_client.embeddings.with_raw_response.create(**data, timeout=timeout)
   1181     headers: Final = dict(raw_response.headers)

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_legacy_response.py:369, in to_raw_response_wrapper.<locals>.wrapped(*args, **kwargs)
    367 kwargs["extra_headers"] = extra_headers
--> 369 return cast(LegacyAPIResponse[R], func(*args, **kwargs))

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/resources/embeddings.py:136, in Embeddings.create(self, input, model, dimensions, encoding_format, user, extra_headers, extra_query, extra_body, timeout)
    134     return obj
--> 136 return self._post(
    137     "/embeddings",
    138     body=maybe_transform(params, embedding_create_params.EmbeddingCreateParams),
    139     options=make_request_options(
    140         extra_headers=extra_headers,
    141         extra_query=extra_query,
    142         extra_body=extra_body,
    143         timeout=timeout,
    144         post_parser=parser,
    145         security={"bearer_auth": True},
    146     ),
    147     cast_to=CreateEmbeddingResponse,
    148 )

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py:1375, in SyncAPIClient.post(self, path, cast_to, body, content, options, files, stream, stream_cls)
   1372 opts = FinalRequestOptions.construct(
   1373     method="post", url=path, json_data=body, content=content, files=to_httpx_files(files), **options
   1374 )
-> 1375 return cast(ResponseT, self.request(cast_to, opts, stream=stream, stream_cls=stream_cls))

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/openai/_base_client.py:1148, in SyncAPIClient.request(self, cast_to, options, stream, stream_cls)
   1147     log.debug("Re-raising status error")
-> 1148     raise self._make_status_error_from_response(err.response) from None
   1150 break

AuthenticationError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

OpenAIError                               Traceback (most recent call last)
File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py:6308, in embedding(model, input, dimensions, encoding_format, timeout, api_base, api_version, api_key, api_type, caching, user, custom_llm_provider, litellm_call_id, logger_fn, **kwargs)
   6307     ## EMBEDDING CALL
-> 6308     response = openai_chat_completions.embedding(
   6309         model=model,
   6310         input=input,
   6311         api_base=api_base,
   6312         api_key=api_key,
   6313         logging_obj=logging,
   6314         timeout=timeout,
   6315         model_response=EmbeddingResponse(),
   6316         optional_params=optional_params,
   6317         client=client,
   6318         aembedding=aembedding,
   6319         max_retries=max_retries,
   6320         shared_session=shared_session,
   6321     )
   6322 elif custom_llm_provider == "databricks":

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/llms/openai/openai.py:1341, in OpenAIChatCompletion.embedding(self, model, input, timeout, logging_obj, model_response, optional_params, api_key, api_base, client, aembedding, max_retries, shared_session)
   1340     error_headers = getattr(error_response, "headers", None)
-> 1341 raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)

OpenAIError: Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}

During handling of the above exception, another exception occurred:

AuthenticationError                       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/packages/graphrag/graphrag/query/question_gen/local_gen.py:85, in LocalQuestionGen.agenerate(self, question_history, context_data, question_count, **kwargs)
     79     conversation_history = ConversationHistory.from_list(history)
     81 if context_data is None:
     82     # generate context data based on the question history
     83     result = cast(
     84         "ContextBuilderResult",
---> 85         self.context_builder.build_context(
     86             query=question_text,
     87             conversation_history=conversation_history,
     88             **kwargs,
     89             **self.context_builder_params,
     90         ),
     91     )
     92     context_data = cast("str", result.context_chunks)
     93     context_records = result.context_records

File ~/work/graphrag/graphrag/packages/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_context_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/packages/graphrag/graphrag/query/context_builder/entity_extraction.py:62, 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)
     58 matched_entities = []
     59 if query != "":
     60     # get entities with highest semantic similarity to query
     61     # oversample to account for excluded entities
---> 62     search_results = text_embedding_vectorstore.similarity_search_by_text(
     63         text=query,
     64         text_embedder=lambda t: text_embedder.embedding(input=[t]).first_embedding,
     65         k=k * oversample_scaler,
     66     )
     67     for result in search_results:
     68         if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance(
     69             result.document.id, str
     70         ):

File ~/work/graphrag/graphrag/packages/graphrag-vectors/graphrag_vectors/vector_store.py:186, in VectorStore.similarity_search_by_text(self, text, text_embedder, k, select, filters, include_vectors)
    176 def similarity_search_by_text(
    177     self,
    178     text: str,
   (...)    183     include_vectors: bool = True,
    184 ) -> list[VectorStoreSearchResult]:
    185     """Perform a text-based similarity search."""
--> 186     query_embedding = text_embedder(text)
    187     if query_embedding:
    188         return self.similarity_search_by_vector(
    189             query_embedding=query_embedding,
    190             k=k,
   (...)    193             include_vectors=include_vectors,
    194         )

File ~/work/graphrag/graphrag/packages/graphrag/graphrag/query/context_builder/entity_extraction.py:64, in map_query_to_entities.<locals>.<lambda>(t)
     58 matched_entities = []
     59 if query != "":
     60     # get entities with highest semantic similarity to query
     61     # oversample to account for excluded entities
     62     search_results = text_embedding_vectorstore.similarity_search_by_text(
     63         text=query,
---> 64         text_embedder=lambda t: text_embedder.embedding(input=[t]).first_embedding,
     65         k=k * oversample_scaler,
     66     )
     67     for result in search_results:
     68         if embedding_vectorstore_key == EntityVectorStoreKey.ID and isinstance(
     69             result.document.id, str
     70         ):

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/embedding/lite_llm_embedding.py:130, in LiteLLMEmbedding.embedding(self, **kwargs)
    127     request_metrics = None
    129 try:
--> 130     return self._embedding(metrics=request_metrics, **kwargs)
    131 finally:
    132     if request_metrics:

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_logging.py:49, in with_logging.<locals>._request_count_middleware(**kwargs)
     47 metrics: Metrics | None = kwargs.get("metrics")
     48 try:
---> 49     return sync_middleware(**kwargs)
     50 except Exception as e:
     51     retries = metrics.get("retries", None) if metrics else None

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_request_count.py:53, in with_request_count.<locals>._request_count_middleware(**kwargs)
     51     metrics["failed_response_count"] = 0
     52 try:
---> 53     result = sync_middleware(**kwargs)
     54     if metrics is not None:
     55         metrics["successful_response_count"] = 1

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/middleware/with_metrics.py:56, in with_metrics.<locals>._metrics_middleware(**kwargs)
     54 metrics: Metrics | None = kwargs.get("metrics")
     55 start_time = time.time()
---> 56 response = sync_middleware(**kwargs)
     57 end_time = time.time()
     59 if metrics is not None:

File ~/work/graphrag/graphrag/packages/graphrag-llm/graphrag_llm/embedding/lite_llm_embedding.py:188, in _create_base_embeddings.<locals>._base_embedding(**kwargs)
    185 kwargs.pop("metrics", None)  # Remove metrics if present
    186 new_args: dict[str, Any] = {**base_args, **kwargs}
--> 188 response = litellm.embedding(**new_args)
    189 return LLMEmbeddingResponse(**response.model_dump())

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py:1745, in client.<locals>.wrapper(*args, **kwargs)
   1741 if logging_obj:
   1742     logging_obj.failure_handler(
   1743         e, traceback_exception, start_time, end_time
   1744     )  # DO NOT MAKE THREADED - router retry fallback relies on this!
-> 1745 raise e

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/utils.py:1591, in client.<locals>.wrapper(*args, **kwargs)
   1589         print_verbose(f"Error while checking max token limit: {e}")
   1590 # MODEL CALL
-> 1591 result = original_function(*args, **kwargs)
   1592 end_time = datetime.datetime.now()
   1593 if _is_streaming_request(
   1594     kwargs=kwargs,
   1595     call_type=call_type,
   1596 ):

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/main.py:7110, in embedding(model, input, dimensions, encoding_format, timeout, api_base, api_version, api_key, api_type, caching, user, custom_llm_provider, litellm_call_id, logger_fn, **kwargs)
   7104 litellm_logging_obj.post_call(
   7105     input=input,
   7106     api_key=api_key,
   7107     original_response=str(e),
   7108 )
   7109 ## Map to OpenAI Exception
-> 7110 raise exception_type(
   7111     model=model,
   7112     original_exception=e,
   7113     custom_llm_provider=custom_llm_provider,
   7114     extra_kwargs=kwargs,
   7115 )

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py:2675, in exception_type(model, original_exception, custom_llm_provider, completion_kwargs, extra_kwargs)
   2673     if isinstance(e, error_type):
   2674         setattr(e, "litellm_response_headers", litellm_response_headers)
-> 2675         raise e  # it's already mapped
   2676 raised_exc: Final = APIConnectionError(
   2677     message=f"{original_exception}\n{_redact_string(traceback.format_exc())}",
   2678     llm_provider="",
   2679     model="",
   2680 )
   2681 setattr(raised_exc, "litellm_response_headers", litellm_response_headers)

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py:2436, in exception_type(model, original_exception, custom_llm_provider, completion_kwargs, extra_kwargs)
   2422     extract_and_raise_litellm_exception(
   2423         response=getattr(original_exception, "response", None),
   2424         error_str=error_str,
   2425         model=model,
   2426         custom_llm_provider=custom_llm_provider,
   2427     )
   2428 if (
   2429     custom_llm_provider == "openai"
   2430     or custom_llm_provider == "text-completion-openai"
   (...)   2434     or custom_llm_provider == "runwayml"
   2435 ):
-> 2436     _map_openai_exception(
   2437         model=model,
   2438         original_exception=mappable_exception,
   2439         custom_llm_provider=custom_llm_provider,
   2440         error_str=error_str,
   2441         exception_type=exception_type,
   2442         exception_provider=exception_provider,
   2443         extra_information=extra_information,
   2444     )
   2445 elif custom_llm_provider == "anthropic" or custom_llm_provider == "anthropic_text":  # one of the anthropics
   2446     _map_anthropic_exception(
   2447         model=model,
   2448         original_exception=mappable_exception,
   (...)   2453         extra_information=extra_information,
   2454     )

File ~/work/graphrag/graphrag/.venv/lib/python3.13/site-packages/litellm/litellm_core_utils/exception_mapping_utils.py:409, in _map_openai_exception(model, original_exception, custom_llm_provider, error_str, exception_type, exception_provider, extra_information)
    401     raise BadRequestError(
    402         message=f"{exception_provider} - {message}",
    403         llm_provider=custom_llm_provider,
   (...)    406         litellm_debug_info=extra_information,
    407     )
    408 elif original_exception.status_code == 401:
--> 409     raise AuthenticationError(
    410         message=f"AuthenticationError: {exception_provider} - {message}",
    411         llm_provider=custom_llm_provider,
    412         model=model,
    413         response=getattr(original_exception, "response", None),
    414         litellm_debug_info=extra_information,
    415     )
    416 elif original_exception.status_code == 404:
    417     raise NotFoundError(
    418         message=f"NotFoundError: {exception_provider} - {message}",
    419         model=model,
   (...)    422         litellm_debug_info=extra_information,
    423     )

AuthenticationError: litellm.AuthenticationError: AuthenticationError: OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: ArzIHXTm************************************************************************dHdW. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}