Fix accidentally hardcoded edge collection name in searching upstream nodes

This commit is contained in:
Ken Chen 2025-06-28 16:25:44 +08:00
parent 6574dfb7ea
commit d0f4eee404

View file

@ -722,11 +722,31 @@ class MongoGraphStorage(BaseGraphStorage):
not in [
"_id",
"connected_edges",
"source_ids",
"edge_count",
]
},
)
def _construct_graph_edge(self, edge_id: str, edge: dict[str, str]):
return KnowledgeGraphEdge(
id=edge_id,
type=edge.get("relationship", ""),
source=edge["source_node_id"],
target=edge["target_node_id"],
properties={
k: v
for k, v in edge.items()
if k
not in [
"_id",
"source_node_id",
"target_node_id",
"relationship",
]
},
)
async def get_knowledge_graph(
self,
node_label: str,
@ -810,15 +830,16 @@ class MongoGraphStorage(BaseGraphStorage):
"coll": "chunk_entity_relation",
"pipeline": [
{"$match": {"_id": label}},
{"$project": project_doc},
{
"$graphLookup": {
"from": "chunk_entity_relation_edges",
"from": self._edge_collection_name,
"startWith": "$_id",
"connectFromField": "source_node_id",
"connectToField": "target_node_id",
"as": "connected_edges",
"maxDepth": max_depth,
"depthField": "depth",
"as": "connected_edges",
}
},
],
@ -891,25 +912,7 @@ class MongoGraphStorage(BaseGraphStorage):
edge_id = f"{edge['source_node_id']}-{edge['target_node_id']}"
if edge_id not in seen_edges:
result.edges.append(
KnowledgeGraphEdge(
id=edge_id,
type=edge.get("relationship", ""),
source=edge["source_node_id"],
target=edge["target_node_id"],
properties={
k: v
for k, v in edge.items()
if k
not in [
"_id",
"source_node_id",
"target_node_id",
"relationship",
]
},
)
)
result.edges.append(self._construct_graph_edge(edge_id, edge))
seen_edges.add(edge_id)
logger.info(