Skip to main content

yoagent_state/
query.rs

1use crate::{Graph, Node, NodeId, Relation};
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
4
5#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
6pub struct Lineage {
7    pub root: Option<Node>,
8    pub incoming: Vec<Relation>,
9    pub outgoing: Vec<Relation>,
10    pub related_nodes: Vec<Node>,
11}
12
13impl Lineage {
14    pub fn from_graph(graph: &Graph, id: &NodeId) -> Self {
15        let incoming = graph.incoming(id, None);
16        let outgoing = graph.outgoing(id, None);
17        let mut related = BTreeMap::new();
18
19        for rel in incoming.iter().chain(outgoing.iter()) {
20            for node_id in [&rel.from, &rel.to] {
21                if node_id != id
22                    && let Some(node) = graph.nodes.get(node_id)
23                {
24                    related.insert(node_id.clone(), node.clone());
25                }
26            }
27        }
28
29        Self {
30            root: graph.nodes.get(id).cloned(),
31            incoming,
32            outgoing,
33            related_nodes: related.into_values().collect(),
34        }
35    }
36
37    pub fn to_markdown(&self) -> String {
38        let Some(root) = &self.root else {
39            return "Node not found.\n".to_string();
40        };
41
42        let title = root
43            .props
44            .get("title")
45            .and_then(|value| value.as_str())
46            .unwrap_or(root.id.as_str());
47        let status = root
48            .props
49            .get("status")
50            .and_then(|value| value.as_str())
51            .unwrap_or("unknown");
52
53        let mut lines = vec![
54            format!("# {}", title),
55            String::new(),
56            format!("- id: {}", root.id),
57            format!("- kind: {}", root.kind),
58            format!("- status: {}", status),
59        ];
60
61        if !root.artifacts.is_empty() {
62            lines.push(String::new());
63            lines.push("## Artifacts".to_string());
64            for artifact in &root.artifacts {
65                lines.push(format!("- {}: {}", artifact.kind, artifact.uri));
66            }
67        }
68
69        if !self.outgoing.is_empty() {
70            lines.push(String::new());
71            lines.push("## Outgoing".to_string());
72            for rel in &self.outgoing {
73                lines.push(format!("- {} -> {}", rel.rel, rel.to));
74            }
75        }
76
77        if !self.incoming.is_empty() {
78            lines.push(String::new());
79            lines.push("## Incoming".to_string());
80            for rel in &self.incoming {
81                lines.push(format!("- {} <- {}", rel.rel, rel.from));
82            }
83        }
84
85        lines.push(String::new());
86        lines.join("\n")
87    }
88}