feat: add test and fix final node

.github/workflows/test.yml

@@ -0,0 +1,49 @@
+name: Test
+
+on:
+  pull_request:
+
+
+env:
+  POETRY_VERSION: "1.8.4"
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version:
+        - "3.10"
+        - "3.11"
+        - "3.12"
+        - "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
+      with:
+        python-version: ${{ matrix.python-version }}
+        cache: poetry
+    - name: Install dependencies
+      run: |
+        poetry install --all-extras
+    - name: Pytest All
+      run: |
+        poetry 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

graphai/graph.py

@@ -83,6 +83,11 @@ async def execute(self, input):
             else:
                 output = await current_node.invoke(input=state)
                 self._validate_output(output=output, node_name=current_node.name)
+            # add output to state
+            state = {**state, **output}
+            if current_node.is_end:
+                # finish loop if this was an end node
+                break
             if current_node.is_router:
                 # if we have a router node we let the router decide the next node
                 next_node_name = str(output["choice"])
@@ -91,10 +96,6 @@ async def execute(self, input):
             else:
                 # otherwise, we have linear path
                 current_node = self._get_next_node(current_node=current_node)
-            # add output to state
-            state = {**state, **output}
-            if current_node.is_end:
-                break
             steps += 1
             if steps >= self.max_steps:
                 raise Exception(