Skip to main content

agentic_memory/v3/indexes/
procedural.rs

1//! Procedural index for ordered action sequences and workflow tracking.
2
3use super::{Index, IndexResult};
4use crate::v3::block::{Block, BlockContent, BlockHash, BlockType, BoundaryType};
5use std::collections::HashMap;
6
7/// Workflow step record
8#[derive(Debug, Clone)]
9pub struct WorkflowStep {
10    pub sequence: u64,
11    pub step_type: String,
12    pub description: String,
13}
14
15/// A tracked workflow
16#[derive(Debug, Clone)]
17pub struct Workflow {
18    pub id: String,
19    pub name: String,
20    pub start_sequence: u64,
21    pub end_sequence: Option<u64>,
22    pub steps: Vec<WorkflowStep>,
23}
24
25/// Procedural index for sessions and workflows.
26pub struct ProceduralIndex {
27    /// Session ID -> ordered list of block sequences
28    sessions: HashMap<String, Vec<u64>>,
29
30    /// Block -> session it belongs to
31    block_to_session: HashMap<u64, String>,
32
33    /// Current session ID
34    current_session: String,
35
36    /// Block hashes
37    hashes: HashMap<u64, BlockHash>,
38
39    /// Tracked workflows
40    workflows: Vec<Workflow>,
41}
42
43impl ProceduralIndex {
44    pub fn new() -> Self {
45        let session_id = uuid::Uuid::new_v4().to_string();
46        Self {
47            sessions: HashMap::new(),
48            block_to_session: HashMap::new(),
49            current_session: session_id,
50            hashes: HashMap::new(),
51            workflows: Vec::new(),
52        }
53    }
54
55    /// Get all blocks in a session, in order
56    pub fn get_session(&self, session_id: &str) -> Vec<IndexResult> {
57        self.sessions
58            .get(session_id)
59            .map(|blocks| {
60                blocks
61                    .iter()
62                    .filter_map(|&seq| {
63                        self.hashes.get(&seq).map(|&hash| IndexResult {
64                            block_sequence: seq,
65                            block_hash: hash,
66                            score: 1.0,
67                        })
68                    })
69                    .collect()
70            })
71            .unwrap_or_default()
72    }
73
74    /// Get current session blocks
75    pub fn get_current_session(&self) -> Vec<IndexResult> {
76        self.get_session(&self.current_session.clone())
77    }
78
79    /// Get current session ID
80    pub fn current_session_id(&self) -> &str {
81        &self.current_session
82    }
83
84    /// Get all session IDs
85    pub fn get_sessions(&self) -> Vec<String> {
86        self.sessions.keys().cloned().collect()
87    }
88
89    /// Get the last N blocks in current session
90    pub fn get_recent_steps(&self, n: usize) -> Vec<IndexResult> {
91        self.sessions
92            .get(&self.current_session)
93            .map(|blocks| {
94                blocks
95                    .iter()
96                    .rev()
97                    .take(n)
98                    .filter_map(|&seq| {
99                        self.hashes.get(&seq).map(|&hash| IndexResult {
100                            block_sequence: seq,
101                            block_hash: hash,
102                            score: 1.0,
103                        })
104                    })
105                    .collect()
106            })
107            .unwrap_or_default()
108    }
109
110    /// Start a new workflow
111    pub fn start_workflow(&mut self, name: &str, start_sequence: u64) -> String {
112        let id = uuid::Uuid::new_v4().to_string();
113        self.workflows.push(Workflow {
114            id: id.clone(),
115            name: name.to_string(),
116            start_sequence,
117            end_sequence: None,
118            steps: Vec::new(),
119        });
120        id
121    }
122
123    /// End a workflow
124    pub fn end_workflow(&mut self, workflow_id: &str, end_sequence: u64) {
125        if let Some(workflow) = self.workflows.iter_mut().find(|w| w.id == workflow_id) {
126            workflow.end_sequence = Some(end_sequence);
127        }
128    }
129
130    /// Add step to current workflow
131    pub fn add_workflow_step(&mut self, sequence: u64, step_type: &str, description: &str) {
132        if let Some(workflow) = self.workflows.last_mut() {
133            if workflow.end_sequence.is_none() {
134                workflow.steps.push(WorkflowStep {
135                    sequence,
136                    step_type: step_type.to_string(),
137                    description: description.to_string(),
138                });
139            }
140        }
141    }
142
143    /// Get workflow by ID
144    pub fn get_workflow(&self, workflow_id: &str) -> Option<&Workflow> {
145        self.workflows.iter().find(|w| w.id == workflow_id)
146    }
147
148    /// Get all workflows
149    pub fn get_all_workflows(&self) -> &[Workflow] {
150        &self.workflows
151    }
152}
153
154impl Default for ProceduralIndex {
155    fn default() -> Self {
156        Self::new()
157    }
158}
159
160impl Index for ProceduralIndex {
161    fn index(&mut self, block: &Block) {
162        self.hashes.insert(block.sequence, block.hash);
163
164        // Check for session boundary
165        if let BlockContent::Boundary {
166            boundary_type: BoundaryType::SessionStart | BoundaryType::Compaction,
167            ..
168        } = &block.content
169        {
170            self.current_session = uuid::Uuid::new_v4().to_string();
171        }
172
173        // Add to current session
174        self.sessions
175            .entry(self.current_session.clone())
176            .or_default()
177            .push(block.sequence);
178        self.block_to_session
179            .insert(block.sequence, self.current_session.clone());
180
181        // Auto-detect workflow steps
182        match block.block_type {
183            BlockType::ToolCall => {
184                if let BlockContent::Tool { tool_name, .. } = &block.content {
185                    self.add_workflow_step(block.sequence, "tool_call", tool_name);
186                }
187            }
188            BlockType::FileOperation => {
189                if let BlockContent::File {
190                    path, operation, ..
191                } = &block.content
192                {
193                    self.add_workflow_step(
194                        block.sequence,
195                        "file_op",
196                        &format!("{:?} {}", operation, path),
197                    );
198                }
199            }
200            BlockType::Decision => {
201                if let BlockContent::Decision { decision, .. } = &block.content {
202                    self.add_workflow_step(block.sequence, "decision", decision);
203                }
204            }
205            _ => {}
206        }
207    }
208
209    fn remove(&mut self, sequence: u64) {
210        self.hashes.remove(&sequence);
211        if let Some(session_id) = self.block_to_session.remove(&sequence) {
212            if let Some(blocks) = self.sessions.get_mut(&session_id) {
213                blocks.retain(|&s| s != sequence);
214            }
215        }
216    }
217
218    fn rebuild(&mut self, blocks: impl Iterator<Item = Block>) {
219        self.sessions.clear();
220        self.block_to_session.clear();
221        self.hashes.clear();
222        self.workflows.clear();
223        self.current_session = uuid::Uuid::new_v4().to_string();
224        for block in blocks {
225            self.index(&block);
226        }
227    }
228}