diff --git a/src/config/base.yaml b/src/config/base.yaml index e493cdb9..f7e34f54 100644 --- a/src/config/base.yaml +++ b/src/config/base.yaml @@ -3,7 +3,7 @@ name: base ## model ### model_provider, option in deepseek, zhipu -model_provider: qianfan +model_provider: zhipu ## model dir 可以写相对路径和绝对路径 ### 相对路径是相对于环境变量中 MODEL_ROOT_DIR 的路径 diff --git a/src/core/retriever.py b/src/core/retriever.py index ba818577..cdd37188 100644 --- a/src/core/retriever.py +++ b/src/core/retriever.py @@ -53,8 +53,33 @@ class Retriever: _, entities = self.rewrite_query(query, history) for entitie in entities: result = dbm.graph_base.query_entity_like(entitie) - results.append(result) if result else None - return results + results.extend(result) if result else None + return {"results": self.format_query_results(results)} + + def format_query_results(self, results): + formatted_results = {"nodes": [], "edges": []} + for row in results: + n, relations, m = row + formatted_results["nodes"].append({ + "id": n.id, + "name": n._properties["name"], + "properties": n._properties + }) + formatted_results["nodes"].append({ + "id": m.id, + "name": m._properties["name"], + "properties": m._properties + }) + for rel in relations: + formatted_results["edges"].append({ + "id": rel.id, + "type": rel.type, + "source": rel.start_node.id, + "target": rel.end_node.id, + "source_name": rel.start_node._properties["name"], + "target_name": rel.end_node._properties["name"], + }) + return formatted_results def query_knowledgebase(self, query, history, meta):