Fix: Mem0 graph data handling (#248)

* Update handling logic

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

* Update lint

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

---------

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

Author: andy-k-improving

src/strands_tools/mem0_memory.py

@@ -454,6 +454,47 @@ def format_retrieve_response(memories: List[Dict]) -> Panel:
     return Panel(table, title="[bold green]Search Results", border_style="green")
 
 
+def format_retrieve_graph_response(memories: List[Dict]) -> Panel:
+    """Format retrieve response for graph data"""
+    if not memories:
+        return Panel("No graph memories found matching the query.",
+                     title="[bold yellow]No Matches", border_style="yellow")
+
+    table = Table(title="Search Results", show_header=True, header_style="bold magenta")
+    table.add_column("Source", style="cyan")
+    table.add_column("Relationship", style="yellow", width=50)
+    table.add_column("Destination", style="green")
+
+    for memory in memories:
+        source = memory.get("source", "N/A")
+        relationship = memory.get("relationship", "N/A")
+        destination = memory.get("destination", "N/A")
+
+        table.add_row(source, relationship, destination)
+
+    return Panel(table, title="[bold green]Search Results (Graph)", border_style="green")
+
+
+def format_list_graph_response(memories: List[Dict]) -> Panel:
+    """Format list response for graph data"""
+    if not memories:
+        return Panel("No graph memories found.", title="[bold yellow]No Memories", border_style="yellow")
+
+    table = Table(title="Graph Memories", 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.get("source", "N/A")
+        relationship = memory.get("relationship", "N/A")
+        destination = memory.get("target", "N/A")
+
+        table.add_row(source, relationship, destination)
+
+    return Panel(table, title="[bold green]Memories List (Graph)", border_style="green")
+
+
 def format_history_response(history: List[Dict]) -> Panel:
     """Format memory history response."""
     if not history:
@@ -637,6 +678,14 @@ def mem0_memory(tool: ToolUse, **kwargs: Any) -> ToolResult:
             results_list = memories if isinstance(memories, list) else memories.get("results", [])
             panel = format_list_response(results_list)
             console.print(panel)
+
+            # Process graph relations (If any)
+            if "relations" in memories:
+                relationships_list = memories.get("relations", [])
+                results_list.extend(relationships_list)
+                panel_graph = format_list_graph_response(relationships_list)
+                console.print(panel_graph)
+
             return ToolResult(
                 toolUseId=tool_use_id,
                 status="success",
@@ -656,6 +705,14 @@ def mem0_memory(tool: ToolUse, **kwargs: Any) -> ToolResult:
             results_list = memories if isinstance(memories, list) else memories.get("results", [])
             panel = format_retrieve_response(results_list)
             console.print(panel)
+
+            # Process graph relations (If any)
+            if "relations" in memories:
+                relationships_list = memories.get("relations", [])
+                results_list.extend(relationships_list)
+                panel_graph = format_retrieve_graph_response(relationships_list)
+                console.print(panel_graph)
+
             return ToolResult(
                 toolUseId=tool_use_id,
                 status="success",