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

Cinnamon · Oct 11, 2024 · 15c7916 · 15c7916
1 parent 497cb40
commit 15c7916
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/libs/ktem/ktem/index/file/graph/pipelines.py b/libs/ktem/ktem/index/file/graph/pipelines.py
@@ -221,9 +221,11 @@ def _build_graph_search(self):
         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,