coding.py

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()