From 809dc3b1cc1de738b91f9508ab26a88c62b318bf Mon Sep 17 00:00:00 2001
From: Andrei Fajardo <92402603+nerdai@users.noreply.github.com>
Date: Wed, 14 Feb 2024 17:43:22 -0500
Subject: [PATCH] [BUG-FIX] retriever_mode param missing when constructing
KGTableRetriever (#10725)
add retriever_mode param when constructing KGTableRetriever
---
.../llama_index/core/indices/knowledge_graph/base.py | 1 +
.../llama_index/core/indices/knowledge_graph/retrievers.py | 7 ++++++-
2 files changed, 7 insertions(+), 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
index 8f6ba1016..594b1d577 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 @@ class KnowledgeGraphIndex(BaseIndex[KG]):
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 71e8197b4..38cd852fd 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 @@ class KGTableRetriever(BaseRetriever):
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(
--
GitLab