kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat(knowledge-graph): 添加嵌入模型选择功能到图数据库实体上传（仅初次）

Browse Source 
允许用户在上传图数据库实体时选择嵌入模型
添加前端模型选择组件和后端支持
确保模型选择与现有数据兼容
This commit is contained in:
Wenjie Zhang 2026-01-01 22:19:48 +08:00
parent 295d019d16
commit 20f0f793dc
4 changed files with 65 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
server/routers/graph_router.py
View File

@ -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()}")

34
src/knowledge/services/upload_graph_service.py
View File

@ -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:

6
web/src/apis/graph_api.js
View File

@ -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)
},

31
web/src/views/GraphView.vue
View File

@ -115,6 +115,17 @@
<div class="note">
<p>上传的文件内容参考 test/data/A_Dream_of_Red_Mansions_tiny.jsonl 中的格式：</p>
</div>
<div style="margin-bottom: 16px;">
<span>嵌入模型:</span>
<EmbeddingModelSelector
v-model:value="state.embedModelName"
:disabled="!embedModelConfigurable"
style="width: 100%;"
/>
<div v-if="!embedModelConfigurable" style="font-size: 12px; margin-top: 4px;">
* 图数据库已有数据或已设定模型，不可更改
</div>
</div>
<a-upload-dragger
class="upload-dragger"
v-model:fileList="fileList"
@ -174,7 +185,7 @@ import { neo4jApi, unifiedApi } from '@/apis/graph_api';
import { useUserStore } from '@/stores/user';
import GraphCanvas from '@/components/GraphCanvas.vue';
import GraphDetailPanel from '@/components/GraphDetailPanel.vue';
import UploadModal from '@/components/FileUploadModal.vue';
import EmbeddingModelSelector from '@/components/EmbeddingModelSelector.vue';
import { useGraph } from '@/composables/useGraph';
const configStore = useConfigStore();
@ -203,12 +214,17 @@ const state = reactive({
selectedDbId: 'neo4j',
dbOptions: [],
lightragStats: null,
embedModelName: '',
})
const isNeo4j = computed(() => {
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);