Merge pull request #23 from aurelio-labs/feat/str-nodes

feat: str nodes

Author: jamescalam

graphai/graph.py

@@ -6,7 +6,7 @@
 
 class Graph:
     def __init__(self, max_steps: int = 10):
-        self.nodes = []
+        self.nodes = {}
         self.edges = []
         self.start_node = None
         self.end_nodes = []
@@ -15,7 +15,9 @@ def __init__(self, max_steps: int = 10):
         self.max_steps = max_steps
 
     def add_node(self, node):
-        self.nodes.append(node)
+        if node.name in self.nodes:
+            raise Exception(f"Node with name '{node.name}' already exists.")
+        self.nodes[node.name] = node
         if node.is_start:
             if self.start_node is not None:
                 raise Exception(
@@ -27,10 +29,37 @@ def add_node(self, node):
         if node.is_end:
             self.end_nodes.append(node)
 
-    def add_edge(self, source: _Node, destination: _Node):
-        # TODO add logic to check that source and destination are nodes
-        # and they exist in the graph object already
-        edge = Edge(source, destination)
+    def add_edge(self, source: _Node | str, destination: _Node | str):
+        """Adds an edge between two nodes that already exist in the graph.
+        
+        Args:
+            source: The source node or its name.
+            destination: The destination node or its name.
+        """
+        source_node, destination_node = None, None
+        # get source node from graph
+        if isinstance(source, str):
+            source_node = self.nodes.get(source)
+        else:
+            # Check if it's a node-like object by looking for required attributes
+            if hasattr(source, 'name'):
+                source_node = self.nodes.get(source.name)
+        if source_node is None:
+            raise ValueError(
+                f"Node with name '{source.name if hasattr(source, 'name') else source}' not found."
+            )
+        # get destination node from graph
+        if isinstance(destination, str):
+            destination_node = self.nodes.get(destination)
+        else:
+            # Check if it's a node-like object by looking for required attributes
+            if hasattr(destination, 'name'):
+                destination_node = self.nodes.get(destination.name)
+        if destination_node is None:
+            raise ValueError(
+                f"Node with name '{destination.name if hasattr(destination, 'name') else destination}' not found."
+            )
+        edge = Edge(source_node, destination_node)
         self.edges.append(edge)
 
     def add_router(self, sources: list[_Node], router: _Node, destinations: List[_Node]):
@@ -139,7 +168,7 @@ def visualize(self):
 
         G = nx.DiGraph()
 
-        for node in self.nodes:
+        for node in self.nodes.values():
             G.add_node(node.name)
 
         for edge in self.edges:
@@ -173,10 +202,11 @@ def visualize(self):
                     pos[node] = (pos[node][0] - x_center, pos[node][1])
 
             # Scale the layout
-            max_x = max(abs(p[0]) for p in pos.values())
-            max_y = max(abs(p[1]) for p in pos.values())
-            scale = min(0.8 / max_x, 0.8 / max_y)
-            pos = {node: (x * scale, y * scale) for node, (x, y) in pos.items()}
+            max_x = max(abs(p[0]) for p in pos.values()) if pos else 1
+            max_y = max(abs(p[1]) for p in pos.values()) if pos else 1
+            if max_x > 0 and max_y > 0:
+                scale = min(0.8 / max_x, 0.8 / max_y)
+                pos = {node: (x * scale, y * scale) for node, (x, y) in pos.items()}
 
         else:
             print("Warning: The graph contains cycles. Visualization will use a spring layout.")

graphai/nodes/base.py

@@ -31,6 +31,7 @@ def _node(
         start: bool = False,
         end: bool = False,
         stream: bool = False,
+        name: str | None = None,
     ) -> Callable:
         """Decorator validating node structure.
         """
@@ -117,7 +118,7 @@ async def invoke(cls, input: Dict[str, Any], callback: Optional[Callback] = None
                 return out
 
         NodeClass.__name__ = func.__name__
-        NodeClass.name = func.__name__
+        NodeClass.name = name or func.__name__
         NodeClass.__doc__ = func.__doc__
         NodeClass.is_start = start
         NodeClass.is_end = end
@@ -132,14 +133,15 @@ def __call__(
         start: bool = False,
         end: bool = False,
         stream: bool = False,
+        name: str | None = None,
     ):
         # We must wrap the call to the decorator in a function for it to work
         # correctly with or without parenthesis
-        def wrap(func: Callable, start=start, end=end, stream=stream) -> Callable:
-            return self._node(func=func, start=start, end=end, stream=stream)
+        def wrap(func: Callable, start=start, end=end, stream=stream, name=name) -> Callable:
+            return self._node(func=func, start=start, end=end, stream=stream, name=name)
         if func:
             # Decorator is called without parenthesis
-            return wrap(func=func, start=start, end=end, stream=stream)
+            return wrap(func=func, start=start, end=end, stream=stream, name=name)
         # Decorator is called with parenthesis
         return wrap
 

tests/unit/test_graph.py

@@ -45,3 +45,45 @@ async def node_end(input: str):
 
         response = await graph.execute(input={"input": "ABC"})
         assert response == {"input": "ABCDEFG"}
+
+    @pytest.mark.asyncio
+    async def test_graph_with_name(self):
+        @node(start=True, name="start")
+        async def node_start(input: str):
+            """Start node"""
+            return {"input": input+"D"}
+        
+        @node(name="a")
+        async def node_a(input: str):
+            """Node A"""
+            return {"input": input+"E"}
+        
+        assert node_a.name == "a"
+
+        @node(name="b")
+        async def node_b(input: str):
+            """Node B"""
+            return {"input": input+"F"}
+
+        @node(end=True, name="end")
+        async def node_end(input: str):
+            """End node"""
+            return {"input": input+"G"}
+
+        graph = Graph()
+
+        nodes = [
+            (node_start, "start"),
+            (node_a, "a"),
+            (node_b, "b"),
+            (node_end, "end")
+        ]
+
+        for i, (node_fn, name) in enumerate(nodes):
+            graph.add_node(node_fn)
+            if i > 0:
+                # add each edge using the name only
+                graph.add_edge(nodes[i-1][1], name)
+
+        response = await graph.execute(input={"input": "ABC"})
+        assert response == {"input": "ABCDEFG"}