Merge pull request #24 from aurelio-labs/feat/graph-state

feat: v1 of graph state

Author: jamescalam

.github/workflows/release.yml

@@ -9,34 +9,35 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python
-      uses: actions/setup-python@v2
+    - uses: actions/checkout@v4
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
       with:
-        python-version: '3.13'
-    - name: Install Poetry
-      run: |
-        curl -sSL https://install.python-poetry.org | python - -y --version 1.8.4
+        enable-cache: true
+        cache-dependency-glob: "uv.lock"
+        python-version: "3.10"
     - name: Install dependencies
-      run: poetry install
+      run: uv sync --extra dev
     - name: Build
-      run: poetry build
+      run: uv build
 
   publish:
     needs: build
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python
-      uses: actions/setup-python@v2
+    - uses: actions/checkout@v4
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
       with:
-        python-version: '3.13'
-    - name: Install Poetry
-      run: |
-        curl -sSL https://install.python-poetry.org | python - -y --version 1.8.4
+        enable-cache: true
+        cache-dependency-glob: "uv.lock"
+        python-version: "3.10"
+    - name: Install dependencies
+      run: uv sync --extra dev
+    - name: Build
+      run: uv build
     - name: Publish to PyPI
       run: |
-        poetry config repositories.remote https://upload.pypi.org/legacy/
-        poetry --no-interaction -v publish --build --repository remote --username "__token__" --password "$PYPI_API_TOKEN"
+        uv publish --username "__token__" --password "$PYPI_API_TOKEN"
       env:
         PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}

.github/workflows/test.yml

@@ -3,13 +3,14 @@ name: Test
 on:
   pull_request:
 
-
 env:
-  POETRY_VERSION: "1.8.4"
+  UV_VERSION: "0.1.19"
 
 jobs:
   build:
     runs-on: ubuntu-latest
+    env:
+      CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
     strategy:
       matrix:
         python-version:
@@ -19,31 +20,19 @@ jobs:
         - "3.13"
     steps:
     - uses: actions/checkout@v4
-    - name: Cache Poetry
-      uses: actions/cache@v4
-      with:
-        path: ~/.poetry
-        key: ${{ runner.os }}-poetry-${{ hashFiles('**/poetry.lock') }}
-        restore-keys: |
-          ${{ runner.os }}-poetry-
-    - name: Install poetry
-      run: |
-        pipx install poetry==$POETRY_VERSION
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
       with:
+        enable-cache: true
+        cache-dependency-glob: "uv.lock"
         python-version: ${{ matrix.python-version }}
-        cache: poetry
-    - name: Install dependencies
-      run: |
-        poetry install --all-extras
+    - name: Install Dependencies
+      run: uv sync --extra dev
     - name: Pytest All
       run: |
-        poetry run pytest -vv -n 20 tests
+        uv run pytest -vv -n 20 tests
     - name: Upload coverage to Codecov
       uses: codecov/codecov-action@v2
-      env:
-        CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
       with:
         file: ./coverage.xml
         fail_ci_if_error: false

.python-version

@@ -0,0 +1 @@
+3.12

README.md

@@ -3,4 +3,10 @@
 1. Async-first
 2. Minimize abstractions
 3. One way to do one thing
-4. Graph-based AI
+4. Graph-based AI
+
+## Installation
+
+```
+pip install -qU graphai-lib
+```

graphai/graph.py

@@ -1,18 +1,36 @@
-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):
         if node.name in self.nodes:
@@ -100,17 +118,20 @@ async def execute(self, input):
             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}
@@ -142,10 +163,21 @@ def get_callback(self):
         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:

graphai/nodes/base.py

@@ -105,14 +105,18 @@ def get_signature(cls):
                 return "\n".join(signature_components)
 
             @classmethod
-            async def invoke(cls, input: Dict[str, Any], callback: Optional[Callback] = None):
+            async def invoke(cls, input: Dict[str, Any], callback: Optional[Callback] = None, state: Optional[Dict[str, Any]] = None):
                 if callback:
                     if stream:
                         input["callback"] = callback
                     else:
                         raise ValueError(
                             f"Error in node {func.__name__}. When callback provided, stream must be True."
                         )
+                # Add state to the input if present and the parameter exists in the function signature
+                if state is not None and "state" in cls._func_signature.parameters:
+                    input["state"] = state
+                
                 instance = cls()
                 out = await instance.execute(**input)
                 return out