[BUG-FIX] retriever_mode param missing when constructing KGTableRetri…

…ever (#10725) add retriever_mode param when constructing KGTableRetriever
run-llama · Feb 14, 2024 · 809dc3b · 809dc3b
1 parent c6b43b2
commit 809dc3b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/llama-index-core/llama_index/core/indices/knowledge_graph/base.py b/llama-index-core/llama_index/core/indices/knowledge_graph/base.py
@@ -138,6 +138,7 @@ def as_retriever(
             object_map=self._object_map,
             llm=self._llm,
             embed_model=embed_model or self._embed_model,
+            retriever_mode=retriever_mode,
             **kwargs,
         )
 

diff --git a/llama-index-core/llama_index/core/indices/knowledge_graph/retrievers.py b/llama-index-core/llama_index/core/indices/knowledge_graph/retrievers.py
@@ -1,4 +1,5 @@
 """KG Retrievers."""
+
 import logging
 from collections import defaultdict
 from enum import Enum
@@ -120,7 +121,11 @@ def __init__(
         self.query_keyword_extract_template = query_keyword_extract_template or DQKET
         self.similarity_top_k = similarity_top_k
         self._include_text = include_text
-        self._retriever_mode = KGRetrieverMode(retriever_mode)
+        self._retriever_mode = (
+            KGRetrieverMode(retriever_mode)
+            if retriever_mode
+            else KGRetrieverMode.KEYWORD
+        )
 
         self._llm = llm or llm_from_settings_or_context(Settings, index.service_context)
         self._embed_model = embed_model or embed_model_from_settings_or_context(