From 6875401cc92a08c2e3e90848b0bbe612ea60d673 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Sun, 7 Jan 2024 09:55:43 -0800
Subject: [PATCH] Lint

---
 permchain/langgraph/__init__.py | 17 ++++++-----------
 tests/test_pregel.py            |  6 +++---
 tests/test_pregel_async.py      |  6 +++---
 3 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/permchain/langgraph/__init__.py b/permchain/langgraph/__init__.py
index 02573bfc7..1b5fcde3a 100644
--- a/permchain/langgraph/__init__.py
+++ b/permchain/langgraph/__init__.py
@@ -12,11 +12,6 @@ from langchain_core.runnables.base import (
 from permchain.pregel import Channel, Pregel
 
 
-class Edge(NamedTuple):
-    start: str
-    end: str
-
-
 class Branch(NamedTuple):
     condition: Callable[..., str]
     ends: dict[str, str]
@@ -30,9 +25,9 @@ END = "__end__"
 
 
 class Graph:
-    def __init__(self):
+    def __init__(self) -> None:
         self.nodes: dict[str, Runnable] = {}
-        self.edges = set[Edge]()
+        self.edges = set[tuple[str, str]]()
         self.branches: defaultdict[str, list[Branch]] = defaultdict(list)
 
     def add_node(self, key: str, action: RunnableLike) -> None:
@@ -60,7 +55,7 @@ class Graph:
         start_key: str,
         condition: Callable[..., str],
         conditional_edge_mapping: Dict[str, str],
-    ):
+    ) -> None:
         if start_key not in self.nodes:
             raise ValueError(f"Need to add_node `{start_key}` first")
         if iscoroutinefunction(condition):
@@ -68,17 +63,17 @@ class Graph:
 
         self.branches[start_key].append(Branch(condition, conditional_edge_mapping))
 
-    def set_entry_point(self, key: str):
+    def set_entry_point(self, key: str) -> None:
         if key not in self.nodes:
             raise ValueError(f"Need to add_node `{key}` first")
         self.entry_point = key
 
-    def set_finish_point(self, key: str):
+    def set_finish_point(self, key: str) -> None:
         if key not in self.nodes:
             raise ValueError(f"Need to add_node `{key}` first")
         self.finish_point = key
 
-    def compile(self):
+    def compile(self) -> Pregel:
         ################################################
         #       STEP 1: VALIDATE GRAPH STRUCTURE       #
         ################################################
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index 642098321..30a8f891a 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -597,8 +597,8 @@ def test_conditional_graph() -> None:
     agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser)
 
     # Define tool execution logic
-    def execute_tools(data):
-        agent_action: AgentAction | AgentFinish = data.pop("agent_outcome")
+    def execute_tools(data: dict) -> dict:
+        agent_action: AgentAction = data.pop("agent_outcome")
         observation = {t.name: t for t in tools}[agent_action.tool].invoke(
             agent_action.tool_input
         )
@@ -608,7 +608,7 @@ def test_conditional_graph() -> None:
         return data
 
     # Define decision-making logic
-    def should_continue(data):
+    def should_continue(data: dict) -> str:
         # Logic to decide whether to continue in the loop or exit
         if isinstance(data["agent_outcome"], AgentFinish):
             return "exit"
diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py
index c205e0dac..082b8569a 100644
--- a/tests/test_pregel_async.py
+++ b/tests/test_pregel_async.py
@@ -625,8 +625,8 @@ async def test_conditional_graph() -> None:
     agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser)
 
     # Define tool execution logic
-    async def execute_tools(data):
-        agent_action: AgentAction | AgentFinish = data.pop("agent_outcome")
+    async def execute_tools(data: dict) -> dict:
+        agent_action: AgentAction = data.pop("agent_outcome")
         observation = await {t.name: t for t in tools}[agent_action.tool].ainvoke(
             agent_action.tool_input
         )
@@ -636,7 +636,7 @@ async def test_conditional_graph() -> None:
         return data
 
     # Define decision-making logic
-    def should_continue(data):
+    def should_continue(data: dict) -> str:
         # Logic to decide whether to continue in the loop or exit
         if isinstance(data["agent_outcome"], AgentFinish):
             return "exit"