@llm-dev-ops/agentics-cli 1.4.29 → 1.4.34

package/dist/bundled-agents/memory-graph-agents/src/engine/mod.rs

@@ -0,0 +1,1163 @@
+//! Core engine for the memory graph
+
+mod async_memory_graph;
+
+pub use async_memory_graph::AsyncMemoryGraph;
+
+use crate::error::{Error, Result};
+use crate::storage::{SledBackend, StorageBackend};
+use crate::types::{
+    AgentNode, Config, ConversationSession, Edge, EdgeType, Node, NodeId, PromptMetadata,
+    PromptNode, PromptTemplate, ResponseMetadata, ResponseNode, SessionId, TemplateId, TokenUsage,
+    ToolInvocation,
+};
+use parking_lot::RwLock;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+/// Main interface for interacting with the memory graph
+///
+/// `MemoryGraph` provides a thread-safe, high-level API for managing conversation
+/// sessions, prompts, responses, and their relationships in a graph structure.
+///
+/// # Examples
+///
+/// ```no_run
+/// use llm_memory_graph::{MemoryGraph, Config};
+///
+/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+/// let config = Config::new("./data/my_graph.db");
+/// let graph = MemoryGraph::open(config)?;
+///
+/// let session = graph.create_session()?;
+/// let prompt_id = graph.add_prompt(session.id, "What is Rust?".to_string(), None)?;
+/// # Ok(())
+/// # }
+/// ```
+pub struct MemoryGraph {
+    backend: Arc<dyn StorageBackend>,
+    sessions: Arc<RwLock<HashMap<SessionId, ConversationSession>>>,
+}
+
+impl MemoryGraph {
+    /// Open or create a memory graph with the given configuration
+    ///
+    /// This will create the database directory if it doesn't exist and initialize
+    /// all necessary storage trees.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - The database path is invalid or inaccessible
+    /// - Storage initialization fails
+    /// - Existing data is corrupted
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// use llm_memory_graph::{MemoryGraph, Config};
+    ///
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// let config = Config::new("./data/graph.db");
+    /// let graph = MemoryGraph::open(config)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn open(config: Config) -> Result<Self> {
+        let backend = SledBackend::open(&config.path)?;
+
+        Ok(Self {
+            backend: Arc::new(backend),
+            sessions: Arc::new(RwLock::new(HashMap::new())),
+        })
+    }
+
+    /// Create a new conversation session
+    ///
+    /// Sessions are used to group related prompts and responses together.
+    /// Each session has a unique ID and can store custom metadata.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the session cannot be persisted to storage.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// let session = graph.create_session()?;
+    /// println!("Created session: {}", session.id);
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn create_session(&self) -> Result<ConversationSession> {
+        let session = ConversationSession::new();
+        self.backend.store_node(&Node::Session(session.clone()))?;
+
+        // Cache the session
+        self.sessions.write().insert(session.id, session.clone());
+
+        Ok(session)
+    }
+
+    /// Create a session with custom metadata
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the session cannot be persisted to storage.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config};
+    /// # use std::collections::HashMap;
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// let mut metadata = HashMap::new();
+    /// metadata.insert("user_id".to_string(), "123".to_string());
+    /// let session = graph.create_session_with_metadata(metadata)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn create_session_with_metadata(
+        &self,
+        metadata: HashMap<String, String>,
+    ) -> Result<ConversationSession> {
+        let session = ConversationSession::with_metadata(metadata);
+        self.backend.store_node(&Node::Session(session.clone()))?;
+
+        // Cache the session
+        self.sessions.write().insert(session.id, session.clone());
+
+        Ok(session)
+    }
+
+    /// Get a session by ID
+    ///
+    /// This will first check the in-memory cache, then fall back to storage.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - The session doesn't exist
+    /// - Storage retrieval fails
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let created_session = graph.create_session()?;
+    /// let session = graph.get_session(created_session.id)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_session(&self, session_id: SessionId) -> Result<ConversationSession> {
+        // Check cache first
+        if let Some(session) = self.sessions.read().get(&session_id) {
+            return Ok(session.clone());
+        }
+
+        // Fall back to storage
+        let nodes = self.backend.get_session_nodes(&session_id)?;
+        for node in nodes {
+            if let Node::Session(session) = node {
+                if session.id == session_id {
+                    // Update cache
+                    self.sessions.write().insert(session_id, session.clone());
+                    return Ok(session);
+                }
+            }
+        }
+
+        Err(Error::SessionNotFound(session_id.to_string()))
+    }
+
+    /// Add a prompt to a session
+    ///
+    /// This creates a new prompt node and automatically creates edges linking it
+    /// to the session and to the previous prompt if one exists.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - The session doesn't exist
+    /// - Storage operations fail
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// let prompt_id = graph.add_prompt(
+    ///     session.id,
+    ///     "Explain quantum entanglement".to_string(),
+    ///     None,
+    /// )?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn add_prompt(
+        &self,
+        session_id: SessionId,
+        content: String,
+        metadata: Option<PromptMetadata>,
+    ) -> Result<NodeId> {
+        // Verify session exists
+        self.get_session(session_id)?;
+
+        let prompt = if let Some(meta) = metadata {
+            PromptNode::with_metadata(session_id, content, meta)
+        } else {
+            PromptNode::new(session_id, content)
+        };
+
+        let prompt_id = prompt.id;
+        self.backend.store_node(&Node::Prompt(prompt.clone()))?;
+
+        // Create edge from prompt to session
+        let session_nodes = self.backend.get_session_nodes(&session_id)?;
+        if let Some(session_node) = session_nodes.iter().find(|n| matches!(n, Node::Session(_))) {
+            let edge = Edge::new(prompt_id, session_node.id(), EdgeType::PartOf);
+            self.backend.store_edge(&edge)?;
+        }
+
+        // Find the previous prompt in this session and create a Follows edge
+        let session_prompts: Vec<_> = session_nodes
+            .into_iter()
+            .filter_map(|n| {
+                if let Node::Prompt(p) = n {
+                    Some(p)
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        if !session_prompts.is_empty() {
+            // Get the most recent prompt (excluding the one we just added)
+            let mut previous_prompts: Vec<_> = session_prompts
+                .into_iter()
+                .filter(|p| p.id != prompt_id)
+                .collect();
+            previous_prompts.sort_by(|a, b| b.timestamp.cmp(&a.timestamp));
+
+            if let Some(prev_prompt) = previous_prompts.first() {
+                let edge = Edge::new(prompt_id, prev_prompt.id, EdgeType::Follows);
+                self.backend.store_edge(&edge)?;
+            }
+        }
+
+        Ok(prompt_id)
+    }
+
+    /// Add a response to a prompt
+    ///
+    /// This creates a response node and a RespondsTo edge linking it to the prompt.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - The prompt doesn't exist
+    /// - Storage operations fail
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, TokenUsage};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt_id = graph.add_prompt(session.id, "Test".to_string(), None)?;
+    /// let usage = TokenUsage::new(10, 20);
+    /// let response_id = graph.add_response(
+    ///     prompt_id,
+    ///     "Quantum entanglement is...".to_string(),
+    ///     usage,
+    ///     None,
+    /// )?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn add_response(
+        &self,
+        prompt_id: NodeId,
+        content: String,
+        usage: TokenUsage,
+        metadata: Option<ResponseMetadata>,
+    ) -> Result<NodeId> {
+        // Verify prompt exists
+        self.get_node(prompt_id)?;
+
+        let response = if let Some(meta) = metadata {
+            ResponseNode::with_metadata(prompt_id, content, usage, meta)
+        } else {
+            ResponseNode::new(prompt_id, content, usage)
+        };
+
+        let response_id = response.id;
+        self.backend.store_node(&Node::Response(response.clone()))?;
+
+        // Create edge from response to prompt
+        let edge = Edge::new(response_id, prompt_id, EdgeType::RespondsTo);
+        self.backend.store_edge(&edge)?;
+
+        Ok(response_id)
+    }
+
+    /// Add a tool invocation node to the graph
+    ///
+    /// This creates a tool invocation record and automatically creates an INVOKES edge
+    /// from the response to the tool invocation.
+    ///
+    /// # Arguments
+    ///
+    /// * `tool` - The tool invocation to add
+    ///
+    /// # Returns
+    ///
+    /// The node ID of the created tool invocation
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::*;
+    /// # fn main() -> Result<()> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt_id = graph.add_prompt(session.id, "Test".to_string(), None)?;
+    /// # let response_id = graph.add_response(prompt_id, "Response".to_string(), TokenUsage::new(10, 20), None)?;
+    /// let params = serde_json::json!({"operation": "add", "a": 2, "b": 3});
+    /// let tool = ToolInvocation::new(response_id, "calculator".to_string(), params);
+    /// let tool_id = graph.add_tool_invocation(tool)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn add_tool_invocation(&self, tool: ToolInvocation) -> Result<NodeId> {
+        let tool_id = tool.id;
+        let response_id = tool.response_id;
+
+        // Store the tool invocation node
+        self.backend.store_node(&Node::ToolInvocation(tool))?;
+
+        // Create INVOKES edge from response to tool
+        let edge = Edge::new(response_id, tool_id, EdgeType::Invokes);
+        self.backend.store_edge(&edge)?;
+
+        Ok(tool_id)
+    }
+
+    /// Update an existing tool invocation with results
+    ///
+    /// This method updates a tool invocation's status, result, and duration after execution.
+    ///
+    /// # Arguments
+    ///
+    /// * `tool_id` - The ID of the tool invocation to update
+    /// * `success` - Whether the tool execution was successful
+    /// * `result_or_error` - Either the result (if successful) or error message (if failed)
+    /// * `duration_ms` - Execution duration in milliseconds
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::*;
+    /// # fn main() -> Result<()> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt_id = graph.add_prompt(session.id, "Test".to_string(), None)?;
+    /// # let response_id = graph.add_response(prompt_id, "Response".to_string(), TokenUsage::new(10, 20), None)?;
+    /// # let params = serde_json::json!({"operation": "add", "a": 2, "b": 3});
+    /// # let tool = ToolInvocation::new(response_id, "calculator".to_string(), params);
+    /// # let tool_id = graph.add_tool_invocation(tool)?;
+    /// // Mark tool invocation as successful
+    /// let result = serde_json::json!({"result": 5});
+    /// graph.update_tool_invocation(tool_id, true, serde_json::to_string(&result)?, 150)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn update_tool_invocation(
+        &self,
+        tool_id: NodeId,
+        success: bool,
+        result_or_error: String,
+        duration_ms: u64,
+    ) -> Result<()> {
+        // Get the tool invocation node
+        let node = self
+            .backend
+            .get_node(&tool_id)?
+            .ok_or_else(|| Error::NodeNotFound(tool_id.to_string()))?;
+
+        if let Node::ToolInvocation(mut tool) = node {
+            if success {
+                let result: serde_json::Value = serde_json::from_str(&result_or_error)?;
+                tool.mark_success(result, duration_ms);
+            } else {
+                tool.mark_failed(result_or_error, duration_ms);
+            }
+
+            // Update the node in storage
+            self.backend.store_node(&Node::ToolInvocation(tool))?;
+            Ok(())
+        } else {
+            Err(Error::InvalidNodeType {
+                expected: "ToolInvocation".to_string(),
+                actual: format!("{:?}", node.node_type()),
+            })
+        }
+    }
+
+    /// Get all tool invocations for a specific response
+    ///
+    /// # Arguments
+    ///
+    /// * `response_id` - The response node ID
+    ///
+    /// # Returns
+    ///
+    /// A vector of tool invocation nodes
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::*;
+    /// # fn main() -> Result<()> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt_id = graph.add_prompt(session.id, "Test".to_string(), None)?;
+    /// # let response_id = graph.add_response(prompt_id, "Response".to_string(), TokenUsage::new(10, 20), None)?;
+    /// let tools = graph.get_response_tools(response_id)?;
+    /// println!("Response invoked {} tools", tools.len());
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_response_tools(&self, response_id: NodeId) -> Result<Vec<ToolInvocation>> {
+        let edges = self.backend.get_outgoing_edges(&response_id)?;
+
+        let mut tools = Vec::new();
+        for edge in edges {
+            if edge.edge_type == EdgeType::Invokes {
+                if let Some(Node::ToolInvocation(tool)) = self.backend.get_node(&edge.to)? {
+                    tools.push(tool);
+                }
+            }
+        }
+
+        Ok(tools)
+    }
+
+    /// Create and register an agent in the graph
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage fails
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, AgentNode};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// let agent = AgentNode::new(
+    ///     "Researcher".to_string(),
+    ///     "research".to_string(),
+    ///     vec!["web_search".to_string(), "summarize".to_string()],
+    /// );
+    /// let agent_id = graph.add_agent(agent)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn add_agent(&self, agent: AgentNode) -> Result<NodeId> {
+        let node_id = agent.node_id;
+        self.backend.store_node(&Node::Agent(agent))?;
+        Ok(node_id)
+    }
+
+    /// Update an existing agent's data
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - The agent doesn't exist
+    /// - Storage update fails
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, AgentNode};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let agent = AgentNode::new("Test".to_string(), "test".to_string(), vec![]);
+    /// # let node_id = graph.add_agent(agent)?;
+    /// let node = graph.get_node(node_id)?;
+    /// if let llm_memory_graph::types::Node::Agent(mut agent) = node {
+    ///     agent.update_metrics(true, 250, 150);
+    ///     graph.update_agent(agent)?;
+    /// }
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn update_agent(&self, agent: AgentNode) -> Result<()> {
+        self.backend.store_node(&Node::Agent(agent))?;
+        Ok(())
+    }
+
+    /// Assign an agent to handle a prompt
+    ///
+    /// Creates a HandledBy edge from the prompt to the agent.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage fails
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, AgentNode};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt_id = graph.add_prompt(session.id, "Test".to_string(), None)?;
+    /// # let agent = AgentNode::new("Test".to_string(), "test".to_string(), vec![]);
+    /// # let agent_node_id = graph.add_agent(agent)?;
+    /// graph.assign_agent_to_prompt(prompt_id, agent_node_id)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn assign_agent_to_prompt(&self, prompt_id: NodeId, agent_node_id: NodeId) -> Result<()> {
+        let edge = Edge::new(prompt_id, agent_node_id, EdgeType::HandledBy);
+        self.backend.store_edge(&edge)?;
+        Ok(())
+    }
+
+    /// Create a transfer from a response to an agent
+    ///
+    /// Creates a TransfersTo edge indicating agent handoff.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage fails
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, AgentNode, TokenUsage};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt_id = graph.add_prompt(session.id, "Test".to_string(), None)?;
+    /// # let response_id = graph.add_response(prompt_id, "Test".to_string(), TokenUsage::new(10, 10), None)?;
+    /// # let agent = AgentNode::new("Test".to_string(), "test".to_string(), vec![]);
+    /// # let agent_node_id = graph.add_agent(agent)?;
+    /// graph.transfer_to_agent(response_id, agent_node_id)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn transfer_to_agent(&self, response_id: NodeId, agent_node_id: NodeId) -> Result<()> {
+        let edge = Edge::new(response_id, agent_node_id, EdgeType::TransfersTo);
+        self.backend.store_edge(&edge)?;
+        Ok(())
+    }
+
+    /// Get the agent assigned to handle a prompt
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - No agent is assigned
+    /// - Storage retrieval fails
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, AgentNode};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt_id = graph.add_prompt(session.id, "Test".to_string(), None)?;
+    /// # let agent = AgentNode::new("Test".to_string(), "test".to_string(), vec![]);
+    /// # let agent_id = graph.add_agent(agent)?;
+    /// # graph.assign_agent_to_prompt(prompt_id, agent_id)?;
+    /// let agent = graph.get_prompt_agent(prompt_id)?;
+    /// println!("Handled by: {}", agent.name);
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_prompt_agent(&self, prompt_id: NodeId) -> Result<AgentNode> {
+        let edges = self.backend.get_outgoing_edges(&prompt_id)?;
+        for edge in edges {
+            if edge.edge_type == EdgeType::HandledBy {
+                if let Some(Node::Agent(agent)) = self.backend.get_node(&edge.to)? {
+                    return Ok(agent);
+                }
+            }
+        }
+        Err(Error::TraversalError(
+            "No agent assigned to this prompt".to_string(),
+        ))
+    }
+
+    /// Get all agents a response was transferred to
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage retrieval fails
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, AgentNode, TokenUsage};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt_id = graph.add_prompt(session.id, "Test".to_string(), None)?;
+    /// # let response_id = graph.add_response(prompt_id, "Test".to_string(), TokenUsage::new(10, 10), None)?;
+    /// let agents = graph.get_agent_handoffs(response_id)?;
+    /// for agent in agents {
+    ///     println!("Transferred to: {}", agent.name);
+    /// }
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_agent_handoffs(&self, response_id: NodeId) -> Result<Vec<AgentNode>> {
+        let edges = self.backend.get_outgoing_edges(&response_id)?;
+        let mut agents = Vec::new();
+        for edge in edges {
+            if edge.edge_type == EdgeType::TransfersTo {
+                if let Some(Node::Agent(agent)) = self.backend.get_node(&edge.to)? {
+                    agents.push(agent);
+                }
+            }
+        }
+        Ok(agents)
+    }
+
+    /// Get a node by its ID
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - The node doesn't exist
+    /// - Storage retrieval fails
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt_id = graph.add_prompt(session.id, "Test".to_string(), None)?;
+    /// let node = graph.get_node(prompt_id)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_node(&self, node_id: NodeId) -> Result<Node> {
+        self.backend
+            .get_node(&node_id)?
+            .ok_or_else(|| Error::NodeNotFound(node_id.to_string()))
+    }
+
+    /// Add a custom edge between two nodes
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage operations fail.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, EdgeType};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt1 = graph.add_prompt(session.id, "Test1".to_string(), None)?;
+    /// # let prompt2 = graph.add_prompt(session.id, "Test2".to_string(), None)?;
+    /// graph.add_edge(prompt1, prompt2, EdgeType::Follows)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn add_edge(&self, from: NodeId, to: NodeId, edge_type: EdgeType) -> Result<()> {
+        let edge = Edge::new(from, to, edge_type);
+        self.backend.store_edge(&edge)?;
+        Ok(())
+    }
+
+    /// Get all edges originating from a node
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage retrieval fails.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt_id = graph.add_prompt(session.id, "Test".to_string(), None)?;
+    /// let edges = graph.get_outgoing_edges(prompt_id)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_outgoing_edges(&self, node_id: NodeId) -> Result<Vec<Edge>> {
+        self.backend.get_outgoing_edges(&node_id)
+    }
+
+    /// Get all edges pointing to a node
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage retrieval fails.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let prompt_id = graph.add_prompt(session.id, "Test".to_string(), None)?;
+    /// let edges = graph.get_incoming_edges(prompt_id)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_incoming_edges(&self, node_id: NodeId) -> Result<Vec<Edge>> {
+        self.backend.get_incoming_edges(&node_id)
+    }
+
+    /// Get all nodes in a session
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage retrieval fails.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// let nodes = graph.get_session_nodes(session.id)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_session_nodes(&self, session_id: SessionId) -> Result<Vec<Node>> {
+        self.backend.get_session_nodes(&session_id)
+    }
+
+    /// Flush all pending writes to disk
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the flush operation fails.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// graph.flush()?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn flush(&self) -> Result<()> {
+        self.backend.flush()
+    }
+
+    /// Get storage statistics
+    ///
+    /// Returns information about node count, edge count, storage size, etc.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if statistics cannot be retrieved.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// let stats = graph.stats()?;
+    /// println!("Nodes: {}, Edges: {}", stats.node_count, stats.edge_count);
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn stats(&self) -> Result<crate::storage::StorageStats> {
+        self.backend.stats()
+    }
+
+    // ===== Template Management Methods =====
+
+    /// Create and store a new prompt template
+    ///
+    /// Templates are versioned prompt structures that can be instantiated with variables.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage fails.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, PromptTemplate, VariableSpec};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// let variables = vec![
+    ///     VariableSpec::new(
+    ///         "user_input".to_string(),
+    ///         "String".to_string(),
+    ///         true,
+    ///         "User's question".to_string(),
+    ///     ),
+    /// ];
+    /// let template = PromptTemplate::new(
+    ///     "Question Answering".to_string(),
+    ///     "Answer this question: {{user_input}}".to_string(),
+    ///     variables,
+    /// );
+    /// let template_id = graph.create_template(template)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn create_template(&self, template: PromptTemplate) -> Result<TemplateId> {
+        let template_id = template.id;
+        self.backend.store_node(&Node::Template(template))?;
+        Ok(template_id)
+    }
+
+    /// Get a template by its ID
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - The template doesn't exist
+    /// - Storage retrieval fails
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, PromptTemplate, VariableSpec};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let template = PromptTemplate::new("Test".to_string(), "{{x}}".to_string(), vec![]);
+    /// # let template_id = graph.create_template(template)?;
+    /// let template = graph.get_template(template_id)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_template(&self, _template_id: TemplateId) -> Result<PromptTemplate> {
+        // Templates store template_id as their primary ID, but we need the node_id
+        // We'll need to search for it - for now, let's try a direct approach
+        // In practice, we might want to add a template index to the storage backend
+
+        // For now, search all nodes (this is inefficient - TODO: add template index)
+        let stats = self.backend.stats()?;
+        for _ in 0..stats.node_count {
+            // This is a placeholder - we need a way to iterate all nodes
+            // or maintain a template index
+        }
+
+        Err(Error::NodeNotFound(
+            "Template lookup by TemplateId not yet fully implemented - use get_template_by_node_id instead".to_string()
+        ))
+    }
+
+    /// Get a template by its node ID
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - The node doesn't exist or is not a template
+    /// - Storage retrieval fails
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, PromptTemplate};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let template = PromptTemplate::new("Test".to_string(), "{{x}}".to_string(), vec![]);
+    /// # let node_id = template.node_id;
+    /// # graph.create_template(template)?;
+    /// let template = graph.get_template_by_node_id(node_id)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn get_template_by_node_id(&self, node_id: NodeId) -> Result<PromptTemplate> {
+        let node = self
+            .backend
+            .get_node(&node_id)?
+            .ok_or_else(|| Error::NodeNotFound(format!("Node {} not found", node_id)))?;
+
+        match node {
+            Node::Template(template) => Ok(template),
+            _ => Err(Error::ValidationError(format!(
+                "Node {} is not a template",
+                node_id
+            ))),
+        }
+    }
+
+    /// Update an existing template
+    ///
+    /// This will store the updated template data. Note that the template's
+    /// version should be bumped appropriately before calling this method.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage update fails.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, PromptTemplate, VersionLevel};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let template = PromptTemplate::new("Test".to_string(), "{{x}}".to_string(), vec![]);
+    /// # let node_id = template.node_id;
+    /// # graph.create_template(template)?;
+    /// let mut template = graph.get_template_by_node_id(node_id)?;
+    /// template.record_usage();
+    /// template.bump_version(VersionLevel::Patch);
+    /// graph.update_template(template)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn update_template(&self, template: PromptTemplate) -> Result<()> {
+        self.backend.store_node(&Node::Template(template))?;
+        Ok(())
+    }
+
+    /// Link a prompt to the template it was instantiated from
+    ///
+    /// Creates an Instantiates edge from the prompt to the template.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage fails.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, PromptTemplate, VariableSpec};
+    /// # use std::collections::HashMap;
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let session = graph.create_session()?;
+    /// # let template = PromptTemplate::new("Test".to_string(), "Hello {{name}}".to_string(), vec![]);
+    /// # let template_node_id = template.node_id;
+    /// # graph.create_template(template.clone())?;
+    /// let mut values = HashMap::new();
+    /// values.insert("name".to_string(), "World".to_string());
+    /// let prompt_text = template.instantiate(&values)?;
+    /// let prompt_id = graph.add_prompt(session.id, prompt_text, None)?;
+    /// graph.link_prompt_to_template(prompt_id, template_node_id)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn link_prompt_to_template(
+        &self,
+        prompt_id: NodeId,
+        template_node_id: NodeId,
+    ) -> Result<()> {
+        let edge = Edge::new(prompt_id, template_node_id, EdgeType::Instantiates);
+        self.backend.store_edge(&edge)?;
+        Ok(())
+    }
+
+    /// Create a new template that inherits from a parent template
+    ///
+    /// This creates the new template and automatically establishes an Inherits edge.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if storage fails.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use llm_memory_graph::{MemoryGraph, Config, PromptTemplate, VariableSpec};
+    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+    /// # let graph = MemoryGraph::open(Config::default())?;
+    /// # let parent = PromptTemplate::new("Parent".to_string(), "Base: {{x}}".to_string(), vec![]);
+    /// # let parent_id = parent.id;
+    /// # let parent_node_id = parent.node_id;
+    /// # graph.create_template(parent)?;
+    /// let child = PromptTemplate::from_parent(
+    ///     parent_id,
+    ///     "Child Template".to_string(),
+    ///     "Extended: {{x}} with {{y}}".to_string(),
+    ///     vec![],
+    /// );
+    /// let child_id = graph.create_template_from_parent(child, parent_node_id)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn create_template_from_parent(
+        &self,
+        template: PromptTemplate,
+        parent_node_id: NodeId,
+    ) -> Result<TemplateId> {
+        let template_id = template.id;
+        let template_node_id = template.node_id;
+
+        // Store the new template
+        self.backend.store_node(&Node::Template(template))?;
+
+        // Create Inherits edge from child to parent
+        let edge = Edge::new(template_node_id, parent_node_id, EdgeType::Inherits);
+        self.backend.store_edge(&edge)?;
+
+        Ok(template_id)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tempfile::tempdir;
+
+    #[test]
+    fn test_create_graph() {
+        let dir = tempdir().unwrap();
+        let config = Config::new(dir.path());
+        let graph = MemoryGraph::open(config).unwrap();
+
+        let stats = graph.stats().unwrap();
+        assert_eq!(stats.node_count, 0);
+    }
+
+    #[test]
+    fn test_create_session() {
+        let dir = tempdir().unwrap();
+        let config = Config::new(dir.path());
+        let graph = MemoryGraph::open(config).unwrap();
+
+        let session = graph.create_session().unwrap();
+        let retrieved = graph.get_session(session.id).unwrap();
+
+        assert_eq!(session.id, retrieved.id);
+    }
+
+    #[test]
+    fn test_add_prompt() {
+        let dir = tempdir().unwrap();
+        let config = Config::new(dir.path());
+        let graph = MemoryGraph::open(config).unwrap();
+
+        let session = graph.create_session().unwrap();
+        let prompt_id = graph
+            .add_prompt(session.id, "Test prompt".to_string(), None)
+            .unwrap();
+
+        let node = graph.get_node(prompt_id).unwrap();
+        assert!(matches!(node, Node::Prompt(_)));
+    }
+
+    #[test]
+    fn test_add_response() {
+        let dir = tempdir().unwrap();
+        let config = Config::new(dir.path());
+        let graph = MemoryGraph::open(config).unwrap();
+
+        let session = graph.create_session().unwrap();
+        let prompt_id = graph
+            .add_prompt(session.id, "Test prompt".to_string(), None)
+            .unwrap();
+
+        let usage = TokenUsage::new(10, 20);
+        let response_id = graph
+            .add_response(prompt_id, "Test response".to_string(), usage, None)
+            .unwrap();
+
+        let node = graph.get_node(response_id).unwrap();
+        assert!(matches!(node, Node::Response(_)));
+    }
+
+    #[test]
+    fn test_conversation_chain() {
+        let dir = tempdir().unwrap();
+        let config = Config::new(dir.path());
+        let graph = MemoryGraph::open(config).unwrap();
+
+        let session = graph.create_session().unwrap();
+
+        // Add first prompt
+        let prompt1 = graph
+            .add_prompt(session.id, "First prompt".to_string(), None)
+            .unwrap();
+        let usage1 = TokenUsage::new(5, 10);
+        let _response1 = graph
+            .add_response(prompt1, "First response".to_string(), usage1, None)
+            .unwrap();
+
+        // Add second prompt
+        let prompt2 = graph
+            .add_prompt(session.id, "Second prompt".to_string(), None)
+            .unwrap();
+        let usage2 = TokenUsage::new(6, 12);
+        let _response2 = graph
+            .add_response(prompt2, "Second response".to_string(), usage2, None)
+            .unwrap();
+
+        // Verify session has all nodes
+        let nodes = graph.get_session_nodes(session.id).unwrap();
+        assert!(nodes.len() >= 4); // session + 2 prompts + 2 responses
+
+        // Verify edges exist
+        let outgoing = graph.get_outgoing_edges(prompt2).unwrap();
+        assert!(!outgoing.is_empty());
+    }
+
+    #[test]
+    fn test_session_not_found() {
+        let dir = tempdir().unwrap();
+        let config = Config::new(dir.path());
+        let graph = MemoryGraph::open(config).unwrap();
+
+        let fake_session = SessionId::new();
+        let result = graph.get_session(fake_session);
+
+        assert!(result.is_err());
+        assert!(matches!(result.unwrap_err(), Error::SessionNotFound(_)));
+    }
+
+    #[test]
+    fn test_node_not_found() {
+        let dir = tempdir().unwrap();
+        let config = Config::new(dir.path());
+        let graph = MemoryGraph::open(config).unwrap();
+
+        let fake_node = NodeId::new();
+        let result = graph.get_node(fake_node);
+
+        assert!(result.is_err());
+        assert!(matches!(result.unwrap_err(), Error::NodeNotFound(_)));
+    }
+
+    #[test]
+    fn test_session_with_metadata() {
+        let dir = tempdir().unwrap();
+        let config = Config::new(dir.path());
+        let graph = MemoryGraph::open(config).unwrap();
+
+        let mut metadata = HashMap::new();
+        metadata.insert("user".to_string(), "alice".to_string());
+
+        let session = graph.create_session_with_metadata(metadata).unwrap();
+        let retrieved = graph.get_session(session.id).unwrap();
+
+        assert_eq!(retrieved.metadata.get("user"), Some(&"alice".to_string()));
+    }
+}