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

Merge branch 'feature/20251219_lxc' into develop

Browse Source
This commit is contained in:
lixiangcheng1
2025-12-30 15:22:27 +08:00
parent c2eff4f359 724eb4f801
commit eabf6bb1a9
1 changed files with 30 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -26,6 +26,8 @@ from app.core.rag.utils.doc_store_conn import DocStoreConnection, MatchDenseExpr
from app.core.rag.common.string_utils import remove_redundant_spaces from app.core.rag.common.string_utils import remove_redundant_spaces
from app.core.rag.common.float_utils import get_float from app.core.rag.common.float_utils import get_float
from app.core.rag.common.constants import PAGERANK_FLD, TAG_FLD from app.core.rag.common.constants import PAGERANK_FLD, TAG_FLD
from app.core.rag.llm.chat_model import Base
from app.core.rag.llm.embedding_model import OpenAIEmbed
def knowledge_retrieval( def knowledge_retrieval(
@@ -48,6 +50,7 @@ def knowledge_retrieval(
- merge_strategy: "weight" or other strategies - merge_strategy: "weight" or other strategies
- reranker_id: UUID of the reranker to use - reranker_id: UUID of the reranker to use
- reranker_top_k: int - reranker_top_k: int
- use_graph: bool, whether to use a graph
Returns: Returns:
Rearranged document block list (in descending order of relevance) Rearranged document block list (in descending order of relevance)
@@ -59,6 +62,7 @@ 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"
file_names_filter = [] file_names_filter = []
if user_ids: if user_ids:
@@ -67,6 +71,10 @@ def knowledge_retrieval(
if not knowledge_bases: if not knowledge_bases:
return [] return []
kb_ids = []
workspace_ids = []
chat_model = None
embedding_model = None
all_results = [] all_results = []
# Search each knowledge base # Search each knowledge base
for kb_config in knowledge_bases: for kb_config in knowledge_bases:
@@ -87,6 +95,22 @@ def knowledge_retrieval(
else: else:
continue continue
if str(db_knowledge.id) not in kb_ids:
kb_ids.append(str(db_knowledge.id))
if str(db_knowledge.workspace_id) not in workspace_ids:
workspace_ids.append(str(db_knowledge.workspace_id))
if not chat_model:
chat_model = Base(
key=db_knowledge.llm.api_keys[0].api_key,
model_name=db_knowledge.llm.api_keys[0].model_name,
base_url=db_knowledge.llm.api_keys[0].api_base
)
if not embedding_model:
embedding_model = OpenAIEmbed(
key=db_knowledge.embedding.api_keys[0].api_key,
model_name=db_knowledge.embedding.api_keys[0].model_name,
base_url=db_knowledge.embedding.api_keys[0].api_base
)
vector_service = ElasticSearchVectorFactory().init_vector(knowledge=db_knowledge) vector_service = ElasticSearchVectorFactory().init_vector(knowledge=db_knowledge)
# Retrieve according to the configured retrieval type # Retrieve according to the configured retrieval type
match kb_config["retrieve_type"]: match kb_config["retrieve_type"]:
@@ -136,6 +160,12 @@ 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) return rerank(db=db, reranker_id=reranker_id, query=query, docs=all_results, top_k=reranker_top_k)
# use graph
if use_graph:
from app.core.rag.common.settings import kg_retriever
doc = kg_retriever.retrieval(question=query, workspace_ids=workspace_ids, kb_ids=kb_ids, emb_mdl=embedding_model, llm=chat_model)
if doc:
all_results.insert(0, doc)
return all_results return all_results
except Exception as e: except Exception as e: