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 index 8f6ba10162f53..594b1d577a111 100644 --- 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 index 71e8197b41482..38cd852fd32a0 100644 --- 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(