graphai-lib 0.0.5__tar.gz → 0.0.6__tar.gz

{graphai_lib-0.0.5 → graphai_lib-0.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphai-lib
-Version: 0.0.5
+Version: 0.0.6
 Summary: Not an AI framework
 Requires-Python: <3.14,>=3.10
 Description-Content-Type: text/markdown

{graphai_lib-0.0.5 → graphai_lib-0.0.6}/graphai/graph.py

@@ -1,19 +1,33 @@
-from typing import List, Dict, Any, Optional
-from graphai.nodes.base import _Node
+from typing import List, Dict, Any, Optional, Protocol, Type
 from graphai.callback import Callback
 from graphai.utils import logger
 
 
+class NodeProtocol(Protocol):
+    """Protocol defining the interface of a decorated node."""
+    name: str
+    is_start: bool
+    is_end: bool
+    is_router: bool
+    stream: bool
+    
+    async def invoke(
+        self, 
+        input: Dict[str, Any], 
+        callback: Optional[Callback] = None, 
+        state: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]: ...
+
+
 class Graph:
     def __init__(
         self, max_steps: int = 10, initial_state: Optional[Dict[str, Any]] = None
     ):
-        self.nodes: Dict[str, _Node] = {}
+        self.nodes: Dict[str, NodeProtocol] = {}
         self.edges: List[Any] = []
-        self.start_node: Optional[_Node] = None
-        self.end_nodes: List[_Node] = []
-        self.Callback = Callback
-        self.callback = None
+        self.start_node: Optional[NodeProtocol] = None
+        self.end_nodes: List[NodeProtocol] = []
+        self.Callback: Type[Callback] = Callback
         self.max_steps = max_steps
         self.state = initial_state or {}
 
@@ -34,7 +48,7 @@ class Graph:
         """Reset the graph state to an empty dict."""
         self.state = {}
 
-    def add_node(self, node):
+    def add_node(self, node: NodeProtocol):
         if node.name in self.nodes:
             raise Exception(f"Node with name '{node.name}' already exists.")
         self.nodes[node.name] = node
@@ -49,7 +63,7 @@ class Graph:
         if node.is_end:
             self.end_nodes.append(node)
 
-    def add_edge(self, source: _Node | str, destination: _Node | str):
+    def add_edge(self, source: NodeProtocol | str, destination: NodeProtocol | str):
         """Adds an edge between two nodes that already exist in the graph.
 
         Args:
@@ -83,7 +97,7 @@ class Graph:
         self.edges.append(edge)
 
     def add_router(
-        self, sources: list[_Node], router: _Node, destinations: List[_Node]
+        self, sources: list[NodeProtocol], router: NodeProtocol, destinations: List[NodeProtocol]
     ):
         if not router.is_router:
             raise TypeError("A router object must be passed to the router parameter.")
@@ -91,10 +105,10 @@ class Graph:
         for destination in destinations:
             self.add_edge(router, destination)
 
-    def set_start_node(self, node: _Node):
+    def set_start_node(self, node: NodeProtocol):
         self.start_node = node
 
-    def set_end_node(self, node: _Node):
+    def set_end_node(self, node: NodeProtocol):
         self.end_node = node
 
     def compile(self):
@@ -116,11 +130,15 @@ class Graph:
                 f"Instead, got {type(output)} from '{output}'."
             )
 
-    async def execute(self, input):
+    async def execute(self, input, callback: Callback | None = None):
         # TODO JB: may need to add init callback here to init the queue on every new execution
-        if self.callback is None:
-            self.callback = self.get_callback()
+        if callback is None:
+            callback = self.get_callback()
+        
+        # Type assertion to tell the type checker that start_node is not None after compile()
+        assert self.start_node is not None, "Graph must be compiled before execution"
         current_node = self.start_node
+        
         state = input
         # Don't reset the graph state if it was initialized with initial_state
         steps = 0
@@ -128,13 +146,13 @@ class Graph:
             # we invoke the node here
             if current_node.stream:
                 # add callback tokens and param here if we are streaming
-                await self.callback.start_node(node_name=current_node.name)
+                await callback.start_node(node_name=current_node.name)
                 # Include graph's internal state in the node execution context
                 output = await current_node.invoke(
-                    input=state, callback=self.callback, state=self.state
+                    input=state, callback=callback, state=self.state
                 )
                 self._validate_output(output=output, node_name=current_node.name)
-                await self.callback.end_node(node_name=current_node.name)
+                await callback.end_node(node_name=current_node.name)
             else:
                 # Include graph's internal state in the node execution context
                 output = await current_node.invoke(input=state, state=self.state)
@@ -159,16 +177,30 @@ class Graph:
                     "by setting `max_steps` when initializing the Graph object."
                 )
         # TODO JB: may need to add end callback here to close the queue for every execution
-        if self.callback and "callback" in state:
-            await self.callback.close()
+        if callback and "callback" in state:
+            await callback.close()
             del state["callback"]
         return state
 
     def get_callback(self):
-        self.callback = self.Callback()
-        return self.callback
+        """Get a new instance of the callback class.
+
+        :return: A new instance of the callback class.
+        :rtype: Callback
+        """
+        callback = self.Callback()
+        return callback
+
+    def set_callback(self, callback_class: Type[Callback]):
+        """Set the callback class that is returned by the `get_callback` method and used
+        as the default callback when no callback is passed to the `execute` method.
+
+        :param callback_class: The callback class to use as the default callback.
+        :type callback_class: Type[Callback]
+        """
+        self.Callback = callback_class
 
-    def _get_node_by_name(self, node_name: str) -> _Node:
+    def _get_node_by_name(self, node_name: str) -> NodeProtocol:
         """Get a node by its name.
 
         Args:

{graphai_lib-0.0.5 → graphai_lib-0.0.6}/graphai_lib.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphai-lib
-Version: 0.0.5
+Version: 0.0.6
 Summary: Not an AI framework
 Requires-Python: <3.14,>=3.10
 Description-Content-Type: text/markdown

{graphai_lib-0.0.5 → graphai_lib-0.0.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "graphai-lib"
-version = "0.0.5"
+version = "0.0.6"
 description = "Not an AI framework"
 readme = "README.md"
 requires-python = ">=3.10,<3.14"