agent.py

from langgraph.graph import END, StateGraph
import os
from typing import TypedDict, List
import subprocess
from nodes import unit_agent, coder_agent, refactor_agent

class AgentCoder(TypedDict):
    requirement: str
    unittest: str
    code: str
    refactored_code: str
    state_of_code: List[str]

def unittester(state):
    if not os.path.exists('result'):
        os.makedirs('result')
    open('result/test_function.py', 'w').close()
    open('result/function.py', 'w').close()
    requirement = state['requirement']
    tests = unit_agent.invoke({'requirement': requirement})[0]
    file = open('result/test_function.py', 'w')
    file.write(tests)
    file.close()
    return {'unittest': tests}
    
def programmer(state):
    tests = state['unittest']
    error = state['state_of_code'][0]
    code = coder_agent.invoke({'unittest': tests, 'error': error})[0]
    file = open('result/function.py', 'w')
    file.write(code)
    file.close()
    return {'code': code}

def refactorer(state):
    tests = state['unittest']
    error = state['state_of_code'][0]
    code = state['code']
    refactored_code = refactor_agent.invoke({'code': code})[0]
    file = open('result/function.py', 'w')
    file.write(refactored_code)
    file.close()
    return {'refactored_code': refactored_code}

def executer(state):
    if 'state_of_code' not in state:
      code_state = []
    else:
      code_state = state['state_of_code']
    result = subprocess.run(['pytest', 'result/test_function.py'], stdout=subprocess.PIPE)
    code_state.append(result.stdout.decode())
    return {'state_of_code': code_state}

workflow = StateGraph(AgentCoder)

# Define the nodes
workflow.add_node("programmer", programmer)
workflow.add_node("refactorer", refactorer)
workflow.add_node("red_executer", executer)
workflow.add_node("green_executer", executer)
workflow.add_node("blue_executer", executer)
workflow.add_node("unittester", unittester)

# Build graph
workflow.set_entry_point("unittester")
workflow.add_edge("unittester", "red_executer")
workflow.add_edge("red_executer", "programmer")
workflow.add_edge("programmer", "green_executer")
workflow.add_edge("green_executer", "refactorer")
workflow.add_edge("refactorer", "blue_executer")
workflow.add_edge("blue_executer", END)

# Compile
graph = workflow.compile()