graphai-lib 0.0.2__py3-none-any.whl → 0.0.3__py3-none-any.whl

graphai/graph.py

@@ -1,21 +1,41 @@
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional
 from graphai.nodes.base import _Node
 from graphai.callback import Callback
 from semantic_router.utils.logger import logger
 
 
 class Graph:
-    def __init__(self, max_steps: int = 10):
-        self.nodes = []
-        self.edges = []
-        self.start_node = None
-        self.end_nodes = []
+    def __init__(self, max_steps: int = 10, initial_state: Optional[Dict[str, Any]] = None):
+        self.nodes: Dict[str, _Node] = {}
+        self.edges: List[Any] = []
+        self.start_node: Optional[_Node] = None
+        self.end_nodes: List[_Node] = []
         self.Callback = Callback
         self.callback = None
         self.max_steps = max_steps
+        self.state = initial_state or {}
+
+    # Allow getting and setting the graph's internal state
+    def get_state(self) -> Dict[str, Any]:
+        """Get the current graph state."""
+        return self.state
+
+    def set_state(self, state: Dict[str, Any]):
+        """Set the graph state."""
+        self.state = state
+
+    def update_state(self, values: Dict[str, Any]):
+        """Update the graph state with new values."""
+        self.state.update(values)
+
+    def reset_state(self):
+        """Reset the graph state to an empty dict."""
+        self.state = {}
 
     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 +47,37 @@ class Graph:
         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]):
@@ -71,17 +118,20 @@ class Graph:
             self.callback = self.get_callback()
         current_node = self.start_node
         state = input
+        # Don't reset the graph state if it was initialized with initial_state
         steps = 0
         while True:
             # 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)
-                output = await current_node.invoke(input=state, callback=self.callback)
+                # Include graph's internal state in the node execution context
+                output = await current_node.invoke(input=state, callback=self.callback, state=self.state)
                 self._validate_output(output=output, node_name=current_node.name)
                 await self.callback.end_node(node_name=current_node.name)
             else:
-                output = await current_node.invoke(input=state)
+                # Include graph's internal state in the node execution context
+                output = await current_node.invoke(input=state, state=self.state)
                 self._validate_output(output=output, node_name=current_node.name)
             # add output to state
             state = {**state, **output}
@@ -113,10 +163,21 @@ class Graph:
         return self.callback
 
     def _get_node_by_name(self, node_name: str) -> _Node:
-        for node in self.nodes:
-            if node.name == node_name:
-                return node
-        raise Exception(f"Node with name {node_name} not found.")
+        """Get a node by its name.
+        
+        Args:
+            node_name: The name of the node to find.
+            
+        Returns:
+            The node with the given name.
+            
+        Raises:
+            Exception: If no node with the given name is found.
+        """
+        node = self.nodes.get(node_name)
+        if node is None:
+            raise Exception(f"Node with name {node_name} not found.")
+        return node
 
     def _get_next_node(self, current_node):
         for edge in self.edges:
@@ -139,7 +200,7 @@ class Graph:
 
         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 +234,11 @@ class Graph:
                     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_lib-0.0.3.dist-info/METADATA

@@ -0,0 +1,32 @@
+Metadata-Version: 2.4
+Name: graphai-lib
+Version: 0.0.3
+Summary: Not an AI framework
+Requires-Python: <3.14,>=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: semantic-router>=0.1.5
+Requires-Dist: networkx>=3.4.2
+Requires-Dist: matplotlib>=3.10.0
+Provides-Extra: dev
+Requires-Dist: ipykernel>=6.25.0; extra == "dev"
+Requires-Dist: ruff>=0.1.5; extra == "dev"
+Requires-Dist: pytest>=8.2.0; extra == "dev"
+Requires-Dist: pytest-mock>=3.12.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
+Requires-Dist: pytest-xdist>=3.5.0; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.24.0; extra == "dev"
+Requires-Dist: mypy>=1.7.1; extra == "dev"
+Requires-Dist: black[jupyter]<24.5.0,>=23.12.1; extra == "dev"
+
+# Philosophy
+
+1. Async-first
+2. Minimize abstractions
+3. One way to do one thing
+4. Graph-based AI
+
+## Installation
+
+```
+pip install -qU graphai-lib
+```

graphai_lib-0.0.3.dist-info/RECORD

@@ -0,0 +1,8 @@
+graphai/__init__.py,sha256=EHigFOWewDXLZXbdfjZH9kdLPhw6NT0ChS77lNAVAA8,109
+graphai/callback.py,sha256=K-h44pyL2VLXwJzIB_bcVYp5R6xv8zNca5FmN6994Uk,7598
+graphai/graph.py,sha256=EALHEhbXAaJmTvm7cL3Tdh0moRIw7lIyJDCNnCts2QA,10335
+graphai/utils.py,sha256=zrgpk82rIn7lwh631KhN-OgMAJMdbm0k5GPL1eMf2sQ,4522
+graphai_lib-0.0.3.dist-info/METADATA,sha256=fOjGeptDi6YD-fYZD9NtFv4_lzLWjMtvBnfwoQ9zLfo,879
+graphai_lib-0.0.3.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
+graphai_lib-0.0.3.dist-info/top_level.txt,sha256=TXlqmhLViX-3xGH2g5w6cavRd-QMf229Hl88jdMOGt8,8
+graphai_lib-0.0.3.dist-info/RECORD,,

{graphai_lib-0.0.2.dist-info → graphai_lib-0.0.3.dist-info}/WHEEL

@@ -1,4 +1,5 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.9.1
+Generator: setuptools (77.0.3)
 Root-Is-Purelib: true
 Tag: py3-none-any
+

graphai_lib-0.0.3.dist-info/top_level.txt

@@ -0,0 +1 @@
+graphai

graphai/nodes/__init__.py

@@ -1,3 +0,0 @@
-from graphai.nodes.base import node, router
-
-__all__ = ["node", "router"]

graphai/nodes/base.py

@@ -1,148 +0,0 @@
-import inspect
-from typing import Any, Callable, Dict, Optional
-
-from graphai.callback import Callback
-from graphai.utils import FunctionSchema
-
-
-class NodeMeta(type):
-    @staticmethod
-    def positional_to_kwargs(cls_type, args) -> Dict[str, Any]:
-        init_signature = inspect.signature(cls_type.__init__)
-        init_params = {name: arg for name, arg in init_signature.parameters.items() if name != "self"}
-        return init_params
-
-    def __call__(cls, *args, **kwargs):
-        named_positional_args = NodeMeta.positional_to_kwargs(cls, args)
-        kwargs.update(named_positional_args)
-        return super().__call__(**kwargs)
-
-
-class _Node:
-    def __init__(
-        self,
-        is_router: bool = False,
-    ):
-        self.is_router = is_router
-
-    def _node(
-        self,
-        func: Callable,
-        start: bool = False,
-        end: bool = False,
-        stream: bool = False,
-    ) -> Callable:
-        """Decorator validating node structure.
-        """
-        if not callable(func):
-            raise ValueError("Node must be a callable function.")
-        
-        func_signature = inspect.signature(func)
-        schema = FunctionSchema(func)
-
-        class NodeClass:
-            _func_signature = func_signature
-            is_router = None
-            _stream = stream
-
-            def __init__(self):
-                self._expected_params = set(self._func_signature.parameters.keys())
-
-            async def execute(self, *args, **kwargs):
-                # Prepare arguments, including callback if stream is True
-                params_dict = await self._parse_params(*args, **kwargs)
-                return await func(**params_dict)  # Pass only the necessary arguments
-
-            async def _parse_params(self, *args, **kwargs) -> Dict[str, Any]:
-                # filter out unexpected keyword args
-                expected_kwargs = {k: v for k, v in kwargs.items() if k in self._expected_params}
-                # Convert args to kwargs based on the function signature
-                args_names = list(self._func_signature.parameters.keys())[1:len(args)+1]  # skip 'self'
-                expected_args_kwargs = dict(zip(args_names, args))
-                # Combine filtered args and kwargs
-                combined_params = {**expected_args_kwargs, **expected_kwargs}
-
-                # Bind the current instance attributes to the function signature
-                if "callback" in self._expected_params and not stream:
-                    raise ValueError(
-                        f"Node {func.__name__}: requires stream=True when callback is defined."
-                    )
-                bound_params = self._func_signature.bind_partial(**combined_params)
-                # get the default parameters (if any)
-                bound_params.apply_defaults()
-                params_dict = bound_params.arguments.copy()
-                # Filter arguments to match the next node's parameters
-                filtered_params = {
-                    k: v for k, v in params_dict.items() if k in self._expected_params
-                }
-                # confirm all required parameters are present
-                missing_params = [
-                    p for p in self._expected_params if p not in filtered_params
-                ]
-                # if anything is missing we raise an error
-                if missing_params:
-                    raise ValueError(
-                        f"Missing required parameters for the {func.__name__} node: {', '.join(missing_params)}"
-                    )
-                return filtered_params
-
-
-            @classmethod
-            def get_signature(cls):
-                """Returns the signature of the decorated function as LLM readable
-                string.
-                """
-                signature_components = []
-                if NodeClass._func_signature:
-                    for param in NodeClass._func_signature.parameters.values():
-                        if param.default is param.empty:
-                            signature_components.append(f"{param.name}: {param.annotation}")
-                        else:
-                            signature_components.append(f"{param.name}: {param.annotation} = {param.default}")
-                else:
-                    return "No signature"
-                return "\n".join(signature_components)
-
-            @classmethod
-            async def invoke(cls, input: Dict[str, Any], callback: Optional[Callback] = None):
-                if callback:
-                    if stream:
-                        input["callback"] = callback
-                    else:
-                        raise ValueError(
-                            f"Error in node {func.__name__}. When callback provided, stream must be True."
-                        )
-                instance = cls()
-                out = await instance.execute(**input)
-                return out
-
-        NodeClass.__name__ = func.__name__
-        NodeClass.name = func.__name__
-        NodeClass.__doc__ = func.__doc__
-        NodeClass.is_start = start
-        NodeClass.is_end = end
-        NodeClass.is_router = self.is_router
-        NodeClass.stream = stream
-        NodeClass.schema = schema
-        return NodeClass
-
-    def __call__(
-        self,
-        func: Optional[Callable] = None,
-        start: bool = False,
-        end: bool = False,
-        stream: bool = False,
-    ):
-        # 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)
-        if func:
-            # Decorator is called without parenthesis
-            return wrap(func=func, start=start, end=end, stream=stream)
-        # Decorator is called with parenthesis
-        return wrap
-
-
-node = _Node()
-router = _Node(is_router=True)

graphai_lib-0.0.2.dist-info/METADATA

@@ -1,25 +0,0 @@
-Metadata-Version: 2.1
-Name: graphai-lib
-Version: 0.0.2
-Summary: 
-License: MIT
-Author: Aurelio AI
-Author-email: hello@aurelio.ai
-Requires-Python: >=3.10,<3.14
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.13
-Requires-Dist: matplotlib (>=3.10.0,<4.0.0)
-Requires-Dist: networkx (>=3.4.2,<4.0.0)
-Requires-Dist: semantic-router (>=0.1.0.dev4)
-Description-Content-Type: text/markdown
-
-# Philosophy
-
-1. Async-first
-2. Minimize abstractions
-3. One way to do one thing
-4. Graph-based AI

graphai_lib-0.0.2.dist-info/RECORD

@@ -1,9 +0,0 @@
-graphai/__init__.py,sha256=EHigFOWewDXLZXbdfjZH9kdLPhw6NT0ChS77lNAVAA8,109
-graphai/callback.py,sha256=K-h44pyL2VLXwJzIB_bcVYp5R6xv8zNca5FmN6994Uk,7598
-graphai/graph.py,sha256=Gn_bl0iAtcH9QT54G2WHJDbxv0TPkVwuPOgAfu4bHT4,7687
-graphai/nodes/__init__.py,sha256=4826Ubk5yUfbVH7F8DmoTKQyax624Q2QJHsGxqgQ_ng,73
-graphai/nodes/base.py,sha256=VHceReX-FcRTrq8eWF6vfDyRfc4F8Zdcya5UXSPFbmw,5937
-graphai/utils.py,sha256=zrgpk82rIn7lwh631KhN-OgMAJMdbm0k5GPL1eMf2sQ,4522
-graphai_lib-0.0.2.dist-info/METADATA,sha256=Asgg3vqJWJRAq0VmLnzUC8dsfUCcKoH4D1pLSUopUVg,732
-graphai_lib-0.0.2.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-graphai_lib-0.0.2.dist-info/RECORD,,