fix: improve GRAPHRAG key passing (#384) #none

This commit is contained in:
Tuan Anh Nguyen Dang (Tadashi_Cin) 2024-10-11 12:01:06 +07:00 committed by GitHub
parent 497cb40aaf
commit 15c7916ad8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -221,9 +221,11 @@ class GraphRAGRetrieverPipeline(BaseFileIndexRetriever):
text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet") text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet")
text_units = read_indexer_text_units(text_unit_df) text_units = read_indexer_text_units(text_unit_df)
embedding_model = os.getenv("GRAPHRAG_EMBEDDING_MODEL") embedding_model = os.getenv(
"GRAPHRAG_EMBEDDING_MODEL", "text-embedding-3-small"
)
text_embedder = OpenAIEmbedding( text_embedder = OpenAIEmbedding(
api_key=os.getenv("OPENAI_API_KEY"), api_key=os.getenv("GRAPHRAG_API_KEY"),
api_base=None, api_base=None,
api_type=OpenaiApiType.OpenAI, api_type=OpenaiApiType.OpenAI,
model=embedding_model, model=embedding_model,