hwchase17 · February 7, 2024 03:14 · Feb 7, 2024
diff --git a/coding.py b/coding.py
@@ -0,0 +1,41 @@
+class AgentState(TypedDict):
+    code: str
+    tests: str
+    errors: Optional[str]
+
+def initial_writer(state):
+    ...
+    return {"code": ..., "tests": ...}
+
+
+def run_tests(state):
+    ...
+    # error = None if no errors
+    return {"errors": error}
+
+
+def reprogram(state):
+    ...
+    return {"code": ...}
+
+
+workflow = StateGraph(AgentState)
+
+workflow.add_node("initial_writer", initial_writer)
+workflow.add_node("execute_tests", run_tests)
+workflow.add_node("reprogram", reprogram)
+
+workflow.set_entry_point("initial_writer")
+
+workflow.add_conditional_edges(
+    "execute_tests",
+    should_continue,
+    {
+        "continue": "reprogram",
+        "end": END
+    }
+)
+
+workflow.add_edge('initial_writer', 'execute_tests')
+workflow.add_edge('reprogram', 'execute_tests')
+graph = workflow.compile()