Update display of stored message

Signed-off-by: Andy Kwok <[email protected]>

Author: andy-k-improving

src/strands_tools/mem0_memory.py

@@ -579,6 +579,26 @@ def format_store_response(results: List[Dict]) -> Panel:
     return Panel(table, title="[bold green]Memory Stored", border_style="green")
 
 
+def format_store_graph_response(memories: List[Dict]) -> Panel:
+    """Format store response for graph data"""
+    if not memories:
+        return Panel("No graph memories stored.", title="[bold yellow]No Memories Stored", border_style="yellow")
+
+    table = Table(title="Graph Memories Stored", show_header=True, header_style="bold magenta")
+    table.add_column("Source", style="cyan")
+    table.add_column("Relationship", style="yellow", width=50)
+    table.add_column("Target", style="green")
+
+    for memory in memories:
+        source = memory[0].get("source", "N/A")
+        relationship = memory[0].get("relationship", "N/A")
+        destination = memory[0].get("target", "N/A")
+
+        table.add_row(source, relationship, destination)
+
+    return Panel(table, title="[bold green]Memories Stored (Graph)", border_style="green")
+
+
 def mem0_memory(tool: ToolUse, **kwargs: Any) -> ToolResult:
     """
     Memory management tool for storing, retrieving, and managing memories in Mem0.
@@ -690,6 +710,14 @@ def mem0_memory(tool: ToolUse, **kwargs: Any) -> ToolResult:
             if results_list:
                 panel = format_store_response(results_list)
                 console.print(panel)
+
+            # Process graph relations (If any)
+            if "relations" in results:
+                relationships_list = results.get("relations", [])["added_entities"]
+                results_list.extend(relationships_list)
+                panel_graph = format_store_graph_response(relationships_list)
+                console.print(panel_graph)
+
             return ToolResult(
                 toolUseId=tool_use_id,
                 status="success",