From 15c7916ad856ac4a4fe78979bbcd8c9b45e7abe5 Mon Sep 17 00:00:00 2001 From: "Tuan Anh Nguyen Dang (Tadashi_Cin)" Date: Fri, 11 Oct 2024 12:01:06 +0700 Subject: [PATCH] fix: improve GRAPHRAG key passing (#384) #none --- libs/ktem/ktem/index/file/graph/pipelines.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/libs/ktem/ktem/index/file/graph/pipelines.py b/libs/ktem/ktem/index/file/graph/pipelines.py index 9c1952b..c7d107b 100644 --- a/libs/ktem/ktem/index/file/graph/pipelines.py +++ b/libs/ktem/ktem/index/file/graph/pipelines.py @@ -221,9 +221,11 @@ class GraphRAGRetrieverPipeline(BaseFileIndexRetriever): text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet") 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( - api_key=os.getenv("OPENAI_API_KEY"), + api_key=os.getenv("GRAPHRAG_API_KEY"), api_base=None, api_type=OpenaiApiType.OpenAI, model=embedding_model,