adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

knowledge_retrieval/bug/fix

Browse Source
This commit is contained in:
lixinyue
2026-02-04 15:44:26 +08:00
parent 33cc3c1c3f
commit 41550d4a41
1 changed files with 24 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
api/app/core/rag/nlp/search.py
View File

@@ -62,7 +62,15 @@ def knowledge_retrieval(
merge_strategy = config.get("merge_strategy", "weight") merge_strategy = config.get("merge_strategy", "weight")
reranker_id = config.get("reranker_id") reranker_id = config.get("reranker_id")
reranker_top_k = config.get("reranker_top_k", 1024) reranker_top_k = config.get("reranker_top_k", 1024)
use_graph = config.get("use_graph", "false").lower() == "true" # use_graph = config.get("use_graph", "false").lower() == "true"
use_graph_value = config.get("use_graph", False)
if isinstance(use_graph_value, bool):
use_graph = use_graph_value
elif isinstance(use_graph_value, str):
use_graph = use_graph_value.lower() in ("true", "1", "yes")
else:
use_graph = False
file_names_filter = [] file_names_filter = []
if user_ids: if user_ids:
@@ -159,13 +167,23 @@ def knowledge_retrieval(
# Use the specified reranker for re-ranking # Use the specified reranker for re-ranking
if reranker_id: if reranker_id:
return rerank(db=db, reranker_id=reranker_id, query=query, docs=all_results, top_k=reranker_top_k) try:
return rerank(db=db, reranker_id=reranker_id, query=query, docs=all_results, top_k=reranker_top_k)
except Exception as rerank_error:
# If reranker fails, log warning and continue with original results
print(f"Failed to rerank documents: {str(rerank_error)}")
print(f"Continuing with original retrieval results (count: {len(all_results)})")
# use graph # use graph
if use_graph: if use_graph:
from app.core.rag.common.settings import kg_retriever try:
doc = kg_retriever.retrieval(question=query, workspace_ids=workspace_ids, kb_ids=kb_ids, emb_mdl=embedding_model, llm=chat_model) from app.core.rag.common.settings import kg_retriever
if doc: doc = kg_retriever.retrieval(question=query, workspace_ids=workspace_ids, kb_ids=kb_ids, emb_mdl=embedding_model, llm=chat_model)
all_results.insert(0, doc) if doc:
all_results.insert(0, doc)
except Exception as graph_error:
print(f"Failed to retrieve from knowledge graph: {str(graph_error)}")
return all_results return all_results
except Exception as e: except Exception as e: