From 20f0f793dc6eec2b418bbb1043ec4668c6e3363f Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Thu, 1 Jan 2026 22:19:48 +0800 Subject: [PATCH] =?UTF-8?q?feat(knowledge-graph):=20=E6=B7=BB=E5=8A=A0?= =?UTF-8?q?=E5=B5=8C=E5=85=A5=E6=A8=A1=E5=9E=8B=E9=80=89=E6=8B=A9=E5=8A=9F?= =?UTF-8?q?=E8=83=BD=E5=88=B0=E5=9B=BE=E6=95=B0=E6=8D=AE=E5=BA=93=E5=AE=9E?= =?UTF-8?q?=E4=BD=93=E4=B8=8A=E4=BC=A0=EF=BC=88=E4=BB=85=E5=88=9D=E6=AC=A1?= =?UTF-8?q?=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 允许用户在上传图数据库实体时选择嵌入模型 添加前端模型选择组件和后端支持 确保模型选择与现有数据兼容 --- server/routers/graph_router.py | 7 ++-- .../services/upload_graph_service.py | 34 +++++++++++++++---- web/src/apis/graph_api.js | 6 ++-- web/src/views/GraphView.vue | 31 +++++++++++++++-- 4 files changed, 65 insertions(+), 13 deletions(-) diff --git a/server/routers/graph_router.py b/server/routers/graph_router.py index 8f10a562..2cbafb6e 100644 --- a/server/routers/graph_router.py +++ b/server/routers/graph_router.py @@ -266,7 +266,10 @@ async def index_neo4j_entities(data: dict = Body(default={}), current_user: User @graph.post("/neo4j/add-entities") async def add_neo4j_entities( - file_path: str = Body(...), kgdb_name: str | None = Body(None), current_user: User = Depends(get_admin_user) + file_path: str = Body(...), + kgdb_name: str | None = Body(None), + embed_model_name: str | None = Body(None), + current_user: User = Depends(get_admin_user), ): """通过JSONL文件添加图谱实体到Neo4j""" try: @@ -281,7 +284,7 @@ async def add_neo4j_entities( if not file_path.endswith(".jsonl"): return {"success": False, "message": "文件格式错误，请上传jsonl文件", "status": "failed"} - await graph_base.jsonl_file_add_entity(file_path, kgdb_name) + await graph_base.jsonl_file_add_entity(file_path, kgdb_name, embed_model_name) return {"success": True, "message": "实体添加成功", "status": "success"} except Exception as e: logger.error(f"添加实体失败: {e}, {traceback.format_exc()}") diff --git a/src/knowledge/services/upload_graph_service.py b/src/knowledge/services/upload_graph_service.py index 26c3ddc0..c2abd30b 100644 --- a/src/knowledge/services/upload_graph_service.py +++ b/src/knowledge/services/upload_graph_service.py @@ -29,6 +29,7 @@ class UploadGraphService: self.embed_model = select_embedding_model(self.embed_model_name) self.work_dir = os.path.join(config.save_dir, "knowledge_graph", self.kgdb_name) os.makedirs(self.work_dir, exist_ok=True) + self.is_initialized_from_file = False # 尝试加载已保存的图数据库信息 if not self.load_graph_info(): @@ -51,7 +52,6 @@ class UploadGraphService: self.connection._connect() logger.info(f"Connected to Neo4j: {self.get_graph_info(self.kgdb_name)}") # 连接成功后保存图数据库信息 - self.save_graph_info(self.kgdb_name) def close(self): """关闭数据库连接""" @@ -84,7 +84,7 @@ class UploadGraphService: if self.status == "closed": self.start() - async def jsonl_file_add_entity(self, file_path, kgdb_name="neo4j"): + async def jsonl_file_add_entity(self, file_path, kgdb_name="neo4j", embed_model_name=None): """从JSONL文件添加实体三元组到Neo4j""" assert self.driver is not None, "Database is not connected" self.connection.status = "processing" @@ -134,7 +134,7 @@ class UploadGraphService: triples = list(read_triples(actual_file_path)) - await self.txt_add_vector_entity(triples, kgdb_name) + await self.txt_add_vector_entity(triples, kgdb_name, embed_model_name) except Exception as e: logger.error(f"处理文件失败: {e}") @@ -153,7 +153,7 @@ class UploadGraphService: self.save_graph_info() return kgdb_name - async def txt_add_vector_entity(self, triples, kgdb_name="neo4j"): + async def txt_add_vector_entity(self, triples, kgdb_name="neo4j", embed_model_name=None): """添加实体三元组""" assert self.driver is not None, "Database is not connected" self.use_database(kgdb_name) @@ -258,12 +258,24 @@ class UploadGraphService: embedding=embedding, ) + # 检查是否允许更新模型 + if embed_model_name and not self.is_initialized_from_file: + if embed_model_name != self.embed_model_name: + logger.info(f"Changing embedding model from {self.embed_model_name} to {embed_model_name}") + self.embed_model_name = embed_model_name + self.embed_model = select_embedding_model(self.embed_model_name) + # 判断模型名称是否匹配 - self.embed_model_name = self.embed_model_name or config.embed_model + if not self.embed_model_name: + self.embed_model_name = config.embed_model + cur_embed_info = config.embed_model_names.get(self.embed_model_name) logger.warning(f"embed_model_name={self.embed_model_name}, {cur_embed_info=}") - assert self.embed_model_name == config.embed_model or self.embed_model_name is None, ( - f"embed_model_name={self.embed_model_name}, {config.embed_model=}" + + # 允许 self.embed_model_name 与 config.embed_model 不同（用户自定义选择的情况） + # 但必须在支持的模型列表中 + assert self.embed_model_name in config.embed_model_names, ( + f"Unsupported embed model: {self.embed_model_name}" ) with self.driver.session() as session: @@ -412,6 +424,7 @@ class UploadGraphService: "labels": labels, "status": self.status, "embed_model_name": self.embed_model_name, + "embed_model_configurable": not self.is_initialized_from_file, "unindexed_node_count": len(self.query_nodes_without_embedding(graph_name)), } @@ -448,6 +461,8 @@ class UploadGraphService: json.dump(graph_info, f, ensure_ascii=False, indent=2) # logger.info(f"图数据库信息已保存到：{info_file_path}") + # 保存成功后，标记为从文件初始化（锁定配置） + self.is_initialized_from_file = True return True except Exception as e: logger.error(f"保存图数据库信息失败：{e}") @@ -471,9 +486,14 @@ class UploadGraphService: if graph_info.get("embed_model_name"): self.embed_model_name = graph_info["embed_model_name"] + # 重新选择embedding model + if self.embed_model_name: + self.embed_model = select_embedding_model(self.embed_model_name) + # 如果需要，可以加载更多信息 # 注意：这里不更新self.kgdb_name，因为它是在初始化时设置的 + self.is_initialized_from_file = True logger.info(f"已加载图数据库信息，最后更新时间：{graph_info.get('last_updated')}") return True except Exception as e: diff --git a/web/src/apis/graph_api.js b/web/src/apis/graph_api.js index be5b67d1..6c7d53d9 100644 --- a/web/src/apis/graph_api.js +++ b/web/src/apis/graph_api.js @@ -122,12 +122,14 @@ export const neo4jApi = { * 通过JSONL文件添加图谱实体到Neo4j * @param {string} file_path - JSONL文件路径 * @param {string} kgdb_name - Neo4j数据库名称（默认为'neo4j'） + * @param {string} embed_model_name - 嵌入模型名称 (可选) * @returns {Promise} - 添加结果 */ - addEntities: async (file_path, kgdb_name = 'neo4j') => { + addEntities: async (file_path, kgdb_name = 'neo4j', embed_model_name = null) => { return await apiPost('/api/graph/neo4j/add-entities', { file_path: file_path, - kgdb_name: kgdb_name + kgdb_name: kgdb_name, + embed_model_name: embed_model_name }, {}, true) }, diff --git a/web/src/views/GraphView.vue b/web/src/views/GraphView.vue index 17cacaec..209bcd8c 100644 --- a/web/src/views/GraphView.vue +++ b/web/src/views/GraphView.vue @@ -115,6 +115,17 @@

上传的文件内容参考 test/data/A_Dream_of_Red_Mansions_tiny.jsonl 中的格式：

+
+ 嵌入模型: + +
+ * 图数据库已有数据或已设定模型，不可更改 +
+
{ return state.selectedDbId === 'neo4j'; }); +const embedModelConfigurable = computed(() => { + return graphInfo.value?.embed_model_configurable ?? true +}) + // 检查是否有有效的已上传文件 const hasValidFile = computed(() => { return fileList.value.some(file => file.status === 'done' && file.response?.file_path); @@ -287,6 +303,12 @@ const loadGraphInfo = () => { .then(data => { console.log(data) graphInfo.value = data.data + if (graphInfo.value?.embed_model_name) { + state.embedModelName = graphInfo.value.embed_model_name + } else { + // Default if not set (though backend usually sends default) + state.embedModelName = cur_embed_model.value + } state.loadingGraphInfo = false }) .catch(error => { @@ -303,6 +325,11 @@ const addDocumentByFile = () => { return } + if (!state.embedModelName) { + message.error('请选择嵌入模型') + return + } + state.processing = true // 获取已上传的文件路径 @@ -316,7 +343,7 @@ const addDocumentByFile = () => { return } - neo4jApi.addEntities(filePath) + neo4jApi.addEntities(filePath, 'neo4j', state.embedModelName) .then((data) => { if (data.status === 'success') { message.success(data.message);