mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-09 22:31:05 +00:00
86 lines
3.4 KiB
Python
86 lines
3.4 KiB
Python
"""Return the digest graph, rooted by its three memory categories."""
|
|
|
|
from ..base_step import BaseStep
|
|
from ...components import R
|
|
from ...enumeration import DreamBucketEnum
|
|
from ...schema import GraphSnapshot, GraphSnapshotEdge, GraphSnapshotNode
|
|
|
|
_CATEGORY_BUCKETS = (
|
|
DreamBucketEnum.WIKI,
|
|
DreamBucketEnum.PERSONAL,
|
|
DreamBucketEnum.PROCEDURE,
|
|
)
|
|
|
|
|
|
@R.register("graph_snapshot_step")
|
|
class GraphSnapshotStep(BaseStep):
|
|
"""Build the frontend digest graph through the file-store contract.
|
|
|
|
Category nodes connect to every indexed Markdown file in their bucket.
|
|
Digest files retain wikilinks to other digest files and to daily notes. Daily
|
|
notes are leaves: their own outgoing links are intentionally not returned.
|
|
"""
|
|
|
|
async def execute(self):
|
|
assert self.context is not None
|
|
indexed_nodes = await self.file_store.get_nodes()
|
|
node_by_path = {node.path: node for node in indexed_nodes if node.path.lower().endswith(".md")}
|
|
|
|
digest_dir = str(self.config_value("digest_dir")).strip("/")
|
|
daily_dir = str(self.config_value("daily_dir")).strip("/")
|
|
category_paths = {
|
|
bucket: f"{digest_dir}/{bucket.value}" if digest_dir else bucket.value for bucket in _CATEGORY_BUCKETS
|
|
}
|
|
daily_prefix = f"{daily_dir}/" if daily_dir else ""
|
|
|
|
digest_paths_by_category = {
|
|
label: sorted(path for path in node_by_path if path.startswith(f"{category_path}/"))
|
|
for label, category_path in category_paths.items()
|
|
}
|
|
digest_paths = {path for paths in digest_paths_by_category.values() for path in paths}
|
|
|
|
edge_keys: set[tuple[str, str, str | None]] = set()
|
|
for source in digest_paths:
|
|
for link in node_by_path[source].links:
|
|
target = link.target_path
|
|
if target in digest_paths or (target in node_by_path and target.startswith(daily_prefix)):
|
|
edge_keys.add((source, target, link.target_anchor))
|
|
|
|
daily_paths = {target for _source, target, _anchor in edge_keys if target.startswith(daily_prefix)}
|
|
|
|
nodes = [
|
|
GraphSnapshotNode(
|
|
id=f"virtual:{bucket.value}",
|
|
path=category_paths[bucket],
|
|
name=bucket.value,
|
|
indexed=False,
|
|
virtual=True,
|
|
)
|
|
for bucket in _CATEGORY_BUCKETS
|
|
]
|
|
for bucket in _CATEGORY_BUCKETS:
|
|
for path in digest_paths_by_category[bucket]:
|
|
edge_keys.add((f"virtual:{bucket.value}", path, None))
|
|
|
|
for path in sorted(digest_paths | daily_paths):
|
|
node = node_by_path[path]
|
|
nodes.append(
|
|
GraphSnapshotNode(
|
|
id=path,
|
|
path=path,
|
|
name=node.front_matter.name,
|
|
description=node.front_matter.description,
|
|
indexed=True,
|
|
),
|
|
)
|
|
|
|
edges = [
|
|
GraphSnapshotEdge(source=source, target=target, target_anchor=anchor)
|
|
for source, target, anchor in sorted(edge_keys, key=lambda edge: (edge[0], edge[1], edge[2] or ""))
|
|
]
|
|
graph = GraphSnapshot(nodes=nodes, edges=edges)
|
|
|
|
self.context.response.success = True
|
|
self.context.response.answer = graph.model_dump()
|
|
self.logger.info(f"[{self.name}] nodes={len(nodes)} edges={len(edges)}")
|
|
return self.context.response
|