Skip to main content

ainl_memory/
trajectory_persist.rs

1//! Persist graph [`TrajectoryNode`] + `ainl_trajectories` row in one shot (OpenFang + ainl-runtime).
2
3use crate::node::{AinlMemoryNode, TrajectoryNode};
4use crate::trajectory_table::TrajectoryDetailRecord;
5use crate::GraphMemory;
6
7use ainl_contracts::{TrajectoryOutcome, TrajectoryStep};
8use uuid::Uuid;
9
10/// When **unset** or any non-falsy value, trajectory rows are written after each successful episode
11/// (same opt-out semantics as `AINL_EXTRACTOR_ENABLED` in OpenFang: `0`, `false`, `no`, `off`).
12#[must_use]
13pub fn trajectory_env_enabled() -> bool {
14    match std::env::var("AINL_TRAJECTORY_ENABLED") {
15        Ok(s) => {
16            let v = s.trim().to_ascii_lowercase();
17            !(v == "0" || v == "false" || v == "no" || v == "off")
18        }
19        Err(_) => true,
20    }
21}
22
23fn coarse_steps_from_tools(tools: &[String]) -> Vec<TrajectoryStep> {
24    let base_ms = chrono::Utc::now().timestamp_millis();
25    tools
26        .iter()
27        .enumerate()
28        .map(|(i, name)| TrajectoryStep {
29            step_id: format!("step_{i}"),
30            timestamp_ms: base_ms + i as i64,
31            adapter: "builtin".into(),
32            operation: name.clone(),
33            inputs_preview: None,
34            outputs_preview: None,
35            duration_ms: 0,
36            success: true,
37            error: None,
38            vitals: None,
39            freshness_at_step: None,
40            frame_vars: None,
41            tool_telemetry: None,
42        })
43        .collect()
44}
45
46/// Write trajectory graph node, `trajectory_of` edge, and `ainl_trajectories` detail row.
47///
48/// Returns `(graph_trajectory_node_id, detail_table_row_id)`.
49#[allow(clippy::too_many_arguments)] // wide signature mirrors the trajectory schema columns
50pub fn persist_trajectory_for_episode(
51    memory: &GraphMemory,
52    agent_id: &str,
53    episode_graph_id: Uuid,
54    steps: Vec<TrajectoryStep>,
55    outcome: TrajectoryOutcome,
56    session_id: &str,
57    project_id: Option<&str>,
58    ainl_source_hash: Option<&str>,
59    duration_ms: u64,
60    frame_vars: Option<serde_json::Value>,
61    fitness_delta: Option<f32>,
62) -> Result<(Uuid, Uuid), String> {
63    let recorded_at = chrono::Utc::now().timestamp();
64    let traj_body = TrajectoryNode {
65        episode_id: episode_graph_id,
66        recorded_at,
67        session_id: session_id.to_string(),
68        project_id: project_id.map(str::to_string),
69        ainl_source_hash: ainl_source_hash.map(str::to_string),
70        outcome,
71        steps: steps.clone(),
72        duration_ms,
73        frame_vars: frame_vars.clone(),
74        fitness_delta,
75    };
76    let mut node = AinlMemoryNode::new_trajectory(traj_body, agent_id);
77    if let Some(p) = project_id.map(str::trim).filter(|s| !s.is_empty()) {
78        node.project_id = Some(p.to_string());
79    }
80    let graph_traj_id = node.id;
81    memory.write_node(&node)?;
82    memory.insert_graph_edge_checked(graph_traj_id, episode_graph_id, "trajectory_of")?;
83
84    let detail_id = Uuid::new_v4();
85    let row = TrajectoryDetailRecord {
86        id: detail_id,
87        episode_id: episode_graph_id,
88        graph_trajectory_node_id: Some(graph_traj_id),
89        agent_id: agent_id.to_string(),
90        session_id: session_id.to_string(),
91        project_id: project_id.map(str::to_string),
92        recorded_at,
93        outcome,
94        ainl_source_hash: ainl_source_hash.map(str::to_string),
95        duration_ms,
96        steps,
97        frame_vars,
98        fitness_delta,
99    };
100    memory.insert_trajectory_detail(&row)?;
101    Ok((graph_traj_id, detail_id))
102}
103
104/// Convenience when only coarse tool names are known (no per-call timings).
105#[inline]
106#[allow(clippy::too_many_arguments)] // forwards every column to `persist_trajectory_for_episode`
107pub fn persist_trajectory_coarse_tools(
108    memory: &GraphMemory,
109    agent_id: &str,
110    episode_graph_id: Uuid,
111    tools: &[String],
112    outcome: TrajectoryOutcome,
113    session_id: &str,
114    project_id: Option<&str>,
115    ainl_source_hash: Option<&str>,
116) -> Result<(Uuid, Uuid), String> {
117    let steps = coarse_steps_from_tools(tools);
118    persist_trajectory_for_episode(
119        memory,
120        agent_id,
121        episode_graph_id,
122        steps,
123        outcome,
124        session_id,
125        project_id,
126        ainl_source_hash,
127        0,
128        None,
129        None,
130    )
131}