Skip to content

Commit 94f0973

Browse files
committed
feat: implement external retrieval service for knowledge base with configuration and API endpoints
1 parent 68049cd commit 94f0973

19 files changed

Lines changed: 1091 additions & 38 deletions

File tree

‎apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -506,9 +506,9 @@ def _handle_mcp(
506506

507507
# 工具(workflow/custom) + 智能体(子应用) → LangChain tools;
508508
# MCP(自定义/库内) → mcp_servers 配置;技能 → 交给引擎侧 init_skills 初始化
509-
tools = get_tool_tools(source_type, source_id, tool_ids, workspace_id) + get_application_tools(
510-
source_type, source_id, application_ids, workspace_id, self.get_workflow_parameters()
511-
)
509+
tools = get_tool_tools(
510+
source_type, source_id, tool_ids, workspace_id, self.get_workflow_parameters()
511+
) + get_application_tools(source_type, source_id, application_ids, workspace_id, self.get_workflow_parameters())
512512
mcp_servers_config = get_mcp_servers(mcp_source, mcp_servers, mcp_tool_id, mcp_tool_ids, self._handle_variables)
513513
ToolExecutor().validate_mcp_transport(json.dumps(mcp_servers_config))
514514

‎apps/application/workflow/nodes/ai_chat_node/tools/tool/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,14 @@
1313
__all__ = ["get_tool_tools", "get_workflow_tools", "get_custom_tools"]
1414

1515

16-
def get_tool_tools(source_type, source_id, tool_ids, workspace_id):
16+
def get_tool_tools(source_type, source_id, tool_ids, workspace_id, workflow_params=None):
1717
"""
1818
构建工具(Tool)类工具:内部按 tool_type 拆分 workflow / custom,合并返回 LangChain tools。
1919
2020
节点只需传入混合的 tool_ids,各构建器各自按 tool_type 过滤。
2121
"""
2222
if not tool_ids:
2323
return []
24-
return get_workflow_tools(source_type, source_id, tool_ids, workspace_id) + get_custom_tools(
24+
return get_workflow_tools(source_type, source_id, tool_ids, workspace_id, workflow_params) + get_custom_tools(
2525
source_type, source_id, tool_ids, workspace_id
2626
)

‎apps/application/workflow/nodes/ai_chat_node/tools/tool/workflow.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,9 @@ def _save_workflow_tool_record(
8181
).save()
8282

8383

84-
def get_workflow_func(source_type, source_id, tool, qv, workspace_id):
84+
def get_workflow_func(source_type, source_id, tool, qv, workspace_id, workflow_params=None):
85+
from knowledge.services.retrieval_access import inherited_retrieval_context
86+
8587
tool_id = tool.id
8688

8789
def inner(**kwargs):
@@ -100,6 +102,7 @@ def inner(**kwargs):
100102
"workspace_id": workspace_id,
101103
"default_model_setting": qv.default_model_setting or {},
102104
**kwargs,
105+
**inherited_retrieval_context({"workspace_id": workspace_id, **(workflow_params or {})}),
103106
}
104107

105108
# WorkflowManage.run() 在后台线程异步执行节点,完成时机由 on_complete 回调驱动,
@@ -157,7 +160,7 @@ def get_start_node_fn(wf, wm):
157160
return inner
158161

159162

160-
def get_workflow_tools(source_type, source_id, tool_workflow_ids, workspace_id):
163+
def get_workflow_tools(source_type, source_id, tool_workflow_ids, workspace_id, workflow_params=None):
161164
tools = QuerySet(Tool).filter(
162165
id__in=tool_workflow_ids, is_active=True, tool_type=ToolType.WORKFLOW, workspace_id=workspace_id
163166
)
@@ -170,7 +173,7 @@ def get_workflow_tools(source_type, source_id, tool_workflow_ids, workspace_id):
170173
results = []
171174
for tool in tools:
172175
qv = qd.get(tool.id)
173-
func = get_workflow_func(source_type, source_id, tool, qv, workspace_id)
176+
func = get_workflow_func(source_type, source_id, tool, qv, workspace_id, workflow_params)
174177
args = get_workflow_args(tool, qv)
175178
tool = StructuredTool.from_function(
176179
func=func,

‎apps/application/workflow/nodes/search_document_node/search_document_node.py‎

Lines changed: 22 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,9 @@
1313

1414
from application.workflow.common import WorkflowType
1515
from application.workflow.i_node import INode
16-
from common.auth.constants.role_constants import RoleConstants
17-
from common.database_model_manage.database_model_manage import DatabaseModelManage
1816
from common.utils.shared_resource_auth import filter_authorized_ids
1917
from knowledge.models import Document, DocumentTag, Knowledge
18+
from knowledge.services.retrieval_access import filter_workflow_knowledge
2019

2120

2221
class SearchDocumentNodeSerializer(serializers.Serializer):
@@ -213,21 +212,20 @@ def execute(self):
213212
QuerySet(Document).filter(knowledge_id__in=ref_knowledge_ids).values_list("id", flat=True)
214213
)
215214

216-
get_knowledge_list_of_authorized = DatabaseModelManage.get_model("get_knowledge_list_of_authorized")
217-
chat_user_type = workflow_params.get("chat_user_type")
218-
219-
if get_knowledge_list_of_authorized is not None and RoleConstants.CHAT_USER.value.name == chat_user_type:
220-
actual_knowledge_ids = list(
221-
QuerySet(Document).filter(id__in=document_id_list).values_list("knowledge_id", flat=True).distinct()
222-
)
223-
authorized_knowledge_ids = get_knowledge_list_of_authorized(
224-
workflow_params.get("chat_user_id"), [str(k_id) for k_id in actual_knowledge_ids]
225-
)
226-
document_id_list = list(
227-
QuerySet(Document)
228-
.filter(id__in=document_id_list, knowledge_id__in=authorized_knowledge_ids)
229-
.values_list("id", flat=True)
230-
)
215+
actual_knowledge_ids = list(
216+
QuerySet(Document)
217+
.filter(id__in=document_id_list or [], is_active=True)
218+
.values_list("knowledge_id", flat=True)
219+
.distinct()
220+
)
221+
authorized_knowledge_ids = filter_workflow_knowledge(
222+
filter_authorized_ids("knowledge", actual_knowledge_ids, workspace_id), workflow_params
223+
)
224+
document_id_list = list(
225+
QuerySet(Document)
226+
.filter(id__in=document_id_list or [], knowledge_id__in=authorized_knowledge_ids, is_active=True)
227+
.values_list("id", flat=True)
228+
)
231229

232230
if search_mode == "auto":
233231
matched_doc_ids = _handle_auto_tags(self.workflow_manage, document_id_list, question_reference)
@@ -239,7 +237,13 @@ def execute(self):
239237
final_document_ids = list(matched_document_ids)
240238

241239
final_document_ids = [str(doc_id) for doc_id in final_document_ids]
242-
document_items = list(QuerySet(Document).filter(id__in=final_document_ids).values())
240+
authorized_knowledge_ids = filter_workflow_knowledge(authorized_knowledge_ids, workflow_params)
241+
document_items = list(
242+
QuerySet(Document)
243+
.filter(id__in=final_document_ids, knowledge_id__in=authorized_knowledge_ids, is_active=True)
244+
.values()
245+
)
246+
final_document_ids = [str(doc["id"]) for doc in document_items]
243247
final_knowledge_ids = list(set(str(doc["knowledge_id"]) for doc in document_items))
244248
knowledge_items = list(QuerySet(Knowledge).filter(id__in=final_knowledge_ids).values())
245249

‎apps/application/workflow/nodes/search_knowledge_node/search_knowledge_node.py‎

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -19,15 +19,14 @@
1919

2020
from application.workflow.common import WorkflowType
2121
from application.workflow.i_node import INode
22-
from common.auth.constants.role_constants import RoleConstants
2322
from common.config.embedding_config import VectorStore
24-
from common.database_model_manage.database_model_manage import DatabaseModelManage
2523
from common.db.search import native_search
2624
from common.utils.common import flat_map, get_file_content
2725
from common.utils.shared_resource_auth import filter_authorized_ids
2826
from knowledge.models import Document, Paragraph, Knowledge, SearchMode, SourceType
2927
from knowledge.services.multimodal_retrieval import get_hit_asset_map
3028
from knowledge.services.retrieval_stats import get_recall_tracker, record_recall_safely
29+
from knowledge.services.retrieval_access import filter_workflow_knowledge
3130
from maxkb.conf import PROJECT_DIR
3231
from models_provider.tools import get_model_instance_by_model_workspace_id
3332

@@ -178,18 +177,27 @@ def _record_recalled_items(embedding_list: List[Dict], paragraph_list: List[Dict
178177
record_recall_safely(recalled_items, tracker=get_recall_tracker(workflow_manage))
179178

180179

181-
def _list_paragraph(embedding_list: List, vector):
180+
def _list_paragraph(embedding_list: List, vector, knowledge_ids=None, document_ids=None):
182181
paragraph_id_list = [row.get("paragraph_id") for row in embedding_list]
183182
if paragraph_id_list is None or len(paragraph_id_list) == 0:
184183
return []
184+
query = QuerySet(Paragraph).filter(id__in=paragraph_id_list)
185+
if knowledge_ids is not None:
186+
query = query.filter(
187+
knowledge_id__in=knowledge_ids,
188+
is_active=True,
189+
document_id__in=QuerySet(Document).filter(knowledge_id__in=knowledge_ids, is_active=True).values("id"),
190+
)
191+
if document_ids is not None:
192+
query = query.filter(document_id__in=document_ids)
185193
paragraph_list = native_search(
186-
QuerySet(Paragraph).filter(id__in=paragraph_id_list),
194+
query,
187195
get_file_content(
188196
os.path.join(PROJECT_DIR, "apps", "application", "sql", "list_knowledge_paragraph_by_paragraph_id.sql")
189197
),
190198
with_table_name=True,
191199
)
192-
if len(paragraph_list) != len(paragraph_id_list):
200+
if knowledge_ids is None and len(paragraph_list) != len(paragraph_id_list):
193201
exist_paragraph_list = [row.get("id") for row in paragraph_list]
194202
for paragraph_id in paragraph_id_list:
195203
if paragraph_id not in exist_paragraph_list:
@@ -250,15 +258,11 @@ def execute(self):
250258
.distinct()
251259
]
252260

253-
get_knowledge_list_of_authorized = DatabaseModelManage.get_model("get_knowledge_list_of_authorized")
254-
chat_user_type = workflow_params.get("chat_user_type")
255-
if get_knowledge_list_of_authorized is not None and RoleConstants.CHAT_USER.value.name == chat_user_type:
256-
knowledge_id_list = get_knowledge_list_of_authorized(workflow_params.get("chat_user_id"), knowledge_id_list)
257-
258261
workspace_id = workflow_params.get("workspace_id")
259262
knowledge_id_list = filter_authorized_ids("knowledge", knowledge_id_list, workspace_id)
263+
knowledge_id_list = filter_workflow_knowledge(knowledge_id_list, workflow_params)
260264

261-
if len(knowledge_id_list) == 0:
265+
if len(knowledge_id_list) == 0 or document_id_list == []:
262266
self._write_empty_result(question)
263267
return
264268

@@ -293,7 +297,8 @@ def execute(self):
293297
self._write_empty_result(question)
294298
return
295299

296-
paragraph_list = _list_paragraph(embedding_list, vector)
300+
knowledge_id_list = filter_workflow_knowledge(knowledge_id_list, workflow_params)
301+
paragraph_list = _list_paragraph(embedding_list, vector, knowledge_id_list, document_id_list)
297302
hit_asset_map = get_hit_asset_map(embedding_list)
298303
result = [
299304
reset_paragraph

‎apps/application/workflow/nodes/tool_workflow_lib_node/tool_workflow_lib_node.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from common.exception.app_exception import ChatException, AppApiException
2323
from common.field.common import ObjectField
2424
from tools.models import Tool, ToolType, ToolWorkflowVersion
25+
from knowledge.services.retrieval_access import inherited_retrieval_context
2526

2627

2728
class InputField(serializers.Serializer):
@@ -108,6 +109,7 @@ def execute(self):
108109
"form_data": workflow_params.get("form_data"),
109110
"default_model_setting": tool_workflow_version.default_model_setting or {},
110111
**parameters,
112+
**inherited_retrieval_context(workflow_params),
111113
}
112114

113115
node_id = self.get_node_id()

‎apps/chat/mcp/knowledge.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
"""Knowledge equivalent of the application's MCPToolHandler."""
2+
3+
import json
4+
5+
from rest_framework.exceptions import ValidationError
6+
7+
from knowledge.services.external_retrieval import retrieve
8+
from knowledge.services.retrieval_access import RetrievalError
9+
10+
PROTOCOL_VERSIONS = ("2025-03-26", "2025-06-18", "2025-11-25")
11+
12+
13+
class KnowledgeMCPToolHandler:
14+
def __init__(self, knowledge, identity):
15+
self.knowledge, self.identity = knowledge, identity
16+
self.tool_name = f"knowledge_{knowledge.id}"
17+
18+
def initialize(self, params):
19+
version = params.get("protocolVersion")
20+
if (
21+
not isinstance(version, str)
22+
or not isinstance(params.get("capabilities"), dict)
23+
or not isinstance(params.get("clientInfo"), dict)
24+
):
25+
raise ValidationError("Invalid initialization parameters.")
26+
return {
27+
"protocolVersion": version if version in PROTOCOL_VERSIONS else "2025-06-18",
28+
"serverInfo": {"name": "maxkb-knowledge-mcp", "version": "1.0.0"},
29+
"capabilities": {"tools": {}},
30+
}
31+
32+
def list_tools(self):
33+
return {
34+
"tools": [
35+
{
36+
"name": self.tool_name,
37+
"description": f"检索知识库:{self.knowledge.name}",
38+
"inputSchema": {
39+
"type": "object",
40+
"additionalProperties": False,
41+
"required": ["query_text"],
42+
"properties": {
43+
"query_text": {"type": "string", "minLength": 1, "maxLength": 8000},
44+
"top_number": {"type": "integer", "minimum": 1, "maximum": 50, "default": 5},
45+
"similarity": {"type": "number", "minimum": 0, "maximum": 1, "default": 0},
46+
"search_mode": {
47+
"type": "string",
48+
"enum": ["embedding", "keywords", "blend"],
49+
"default": "embedding",
50+
},
51+
},
52+
},
53+
}
54+
]
55+
}
56+
57+
def call_tool(self, params):
58+
if params.get("name") != self.tool_name:
59+
raise ValidationError("Unknown tool.")
60+
try:
61+
output = retrieve(self.knowledge.id, self.identity, params.get("arguments", {}))
62+
except (RetrievalError, ValidationError):
63+
raise
64+
except Exception:
65+
return {"isError": True, "content": [{"type": "text", "text": "Knowledge retrieval failed."}]}
66+
return {"content": [{"type": "text", "text": json.dumps(output, ensure_ascii=False)}]}

‎apps/chat/serializers/chat.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
from application.workflow.nodes import get_start_node
4646
from application.workflow.workflow_manage import WorkflowManage, CallBack
4747
from application.workflow.workflow_run_registry import WorkflowRunRegistry
48+
from knowledge.services.retrieval_access import identity_from_server
4849
from chat.template.agent_simple import build_workflow
4950
from common import result
5051
from common.exception.app_exception import AppApiException, AppChatNumOutOfBoundsFailed, ChatException
@@ -285,6 +286,7 @@ def chat_work_flow(self, application, instance: dict, base_to_response):
285286
ChatRecord(id=chat_record_id_str, **self.get_defaults_record(message_dict)).save(force_insert=True)
286287

287288
parameters = {
289+
"retrieval_identity": identity_from_server(chat_user_id, chat_user_type, debug),
288290
"history_chat_record": history_chat_record,
289291
"question": message,
290292
"chat_id": chat_id,

‎apps/chat/urls.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,15 @@
22

33
from application.views import ChatRecordDetailView, ChatRecordLinkView
44
from chat.views import v2 as v2_views, v3 as v3_views
5+
from chat.views.v3.knowledge import knowledge_mcp_view, retrieve_view
56

67
app_name = "chat"
78
# @formatter:off
89
# fmt: off
910

1011
v3=[
12+
path('knowledge/<uuid:knowledge_id>/retrieve', retrieve_view),
13+
path('knowledge/<uuid:knowledge_id>/mcp', knowledge_mcp_view),
1114
# ---- application 作用域:application_id 从 path 获取 ----
1215
path('application/<str:application_id>/', include([
1316
path("profile",v3_views.ApplicationProfile.as_view(), name='v3_profile'),

0 commit comments

Comments
 (0)