better graph displays in jupyter (#32)

* better graph displays in jupyter

* sharing .show among all local memory implementations

* fixed docstring

hybridagi/memory/integration/local/local_document_memory.py

@@ -5,7 +5,10 @@
 from hybridagi.core.datatypes import Document, DocumentList
 import networkx as nx
 
-class LocalDocumentMemory(DocumentMemory):
+from .local_memory import LocalMemory
+
+
+class LocalDocumentMemory(LocalMemory, DocumentMemory):
     """
     A class used to manage and store documents locally.
 
@@ -145,18 +148,3 @@ def clear(self):
         self._documents = {}
         self._embeddings = {}
         self._graph = nx.DiGraph()
-
-    def show(self, notebook=False):
-        """
-        Visualize the local document memory as a network graph.
-
-        Parameters:
-            notebook (bool): Whether to display the graph in a Jupyter notebook or not.
-        """
-        from pyvis.network import Network
-        net = Network(notebook=notebook, directed=True)
-        net.from_nx(self._graph)
-        net.toggle_physics(True)
-        net.show(f'{self.index_name}_document_memory.html', notebook=notebook)
-
-

hybridagi/memory/integration/local/local_fact_memory.py

@@ -1,18 +1,21 @@
-from abc import ABC, abstractmethod
-from typing import Union, List, Dict, Optional
 from collections import OrderedDict
+from typing import Union, List, Dict, Optional
 from uuid import UUID
 from hybridagi.memory.fact_memory import FactMemory
 from hybridagi.core.datatypes import Entity, EntityList
 from hybridagi.core.datatypes import Fact, FactList
 import networkx as nx
 import random
 
-# Utilitary function to generate a random color
+from .local_memory import LocalMemory
+
+
 def random_color():
+    """Utilitary function to generate a random color"""
     return "#%06x" % random.randint(0, 0xFFFFFF)
 
-class LocalFactMemory(FactMemory):
+
+class LocalFactMemory(LocalMemory, FactMemory):
     """
     A class used to manage and store facts locally.
 
@@ -209,16 +212,3 @@ def clear(self):
         self._entities_embeddings = OrderedDict()
         self._facts_embeddings = OrderedDict()
         self._labels_colors = {}
-
-    def show(self, notebook=False):
-        """
-        Visualize the local fact memory as a network graph.
-
-        Parameters:
-            notebook (bool): Whether to display the graph in a Jupyter notebook or not.
-        """
-        from pyvis.network import Network
-        net = Network(notebook=notebook, directed=True)
-        net.from_nx(self._graph)
-        net.toggle_physics(True)
-        net.show(f'{self.index_name}_fact_memory.html', notebook=notebook)

hybridagi/memory/integration/local/local_memory.py

@@ -0,0 +1,49 @@
+from urllib.parse import quote
+from uuid import uuid4
+
+
+def isolate(html_code: str) -> str:
+    """
+    When drawing the same pyvis graph from multiple cells in a jupyter notebook
+    They end up sharing the same id for the javascript <script> and thus messing each other
+    display(display_id=...) or making unique IDs in HTML() doesn't work unfortunately
+    Here's the issue that proposes this workaround https://github.com/jupyter/notebook/issues/6598
+    """
+    content = quote(html_code, safe='')
+    return """
+    <iframe
+        width="100%"
+        height="610px"
+        style="border: none"
+        sandbox="allow-scripts allow-modals"
+        referrerpolicy="no-referrer"
+        src="data:text/html;charset=UTF-8,{content}"
+    ></iframe>
+    """.format(content=content)
+
+
+class LocalMemory:       
+    def show(self, notebook: bool = False, cdn_resources: str = 'in_line') -> None:
+        """
+        Visualize the local memory as a network graph.
+
+        Parameters:
+            notebook (bool): Whether to display the graph in a Jupyter notebook or not.
+            cdn_resources (str): Where to pull resources for css and js files. Defaults to local.
+                Options ['local','in_line','remote'].
+                local: pull resources from local lib folder.
+                in_line: insert lib resources as inline script tags.
+                remote: pull resources from hash checked cdns.
+        """
+        from pyvis.network import Network
+        net = Network(notebook=notebook, directed=True, cdn_resources=cdn_resources)
+        net.from_nx(self._graph)
+        net.toggle_physics(True)
+
+        if notebook:
+            from IPython.display import display, HTML
+            unique_id = f'{uuid4().hex}.html'
+            html = net.generate_html(unique_id, notebook=True)
+            display(HTML(isolate(html)), display_id=unique_id)
+        else:
+            net.show(f'{self.index_name}.html', notebook=notebook)

hybridagi/memory/integration/local/local_program_memory.py

@@ -6,7 +6,10 @@
 from hybridagi.core.datatypes import GraphProgramList
 import networkx as nx
 
-class LocalProgramMemory(ProgramMemory):
+from .local_memory import LocalMemory
+
+
+class LocalProgramMemory(LocalMemory, ProgramMemory):
     """
     A class used to manage and store programs locally.
 
@@ -161,16 +164,3 @@ def clear(self):
         self._programs = {}
         self._embeddings = {}
         self._graph = nx.DiGraph()
-
-    def show(self, notebook: bool=False):
-        """
-        Visualize the program memory as a network graph.
-
-        Parameters:
-            notebook (bool): Whether to display the graph in a Jupyter notebook or not.
-        """
-        from pyvis.network import Network
-        net = Network(notebook=notebook, directed=True)
-        net.from_nx(self._graph)
-        net.toggle_physics(True)
-        net.show(f'{self.index_name}_program_memory.html', notebook=notebook)

hybridagi/memory/integration/local/local_trace_memory.py

@@ -5,7 +5,10 @@
 from hybridagi.core.datatypes import AgentStep, AgentStepList, AgentStepType
 import networkx as nx
 
-class LocalTraceMemory(TraceMemory):
+from .local_memory import LocalMemory
+
+
+class LocalTraceMemory(LocalMemory, TraceMemory):
     """
     A class used to manage and store agent steps locally.
 
@@ -136,16 +139,3 @@ def clear(self):
         self._steps = {}
         self._embeddings = OrderedDict()
         self._graph = nx.DiGraph()
-
-    def show(self, notebook: bool=False):
-        """
-        Visualize the trace memory as a network graph.
-
-        Parameters:
-            notebook (bool): Whether to display the graph in a Jupyter notebook or not.
-        """
-        from pyvis.network import Network
-        net = Network(notebook=notebook, directed=True)
-        net.from_nx(self._graph)
-        net.toggle_physics(True)
-        net.show(f'{self.index_name}_trace_memory.html', notebook=notebook)