Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Swekit refactor and integration testing #657

Merged
merged 21 commits into from
Oct 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,9 @@ refactor_debug.ipynb
python/dockerfiles/e2b-dev/composio
python/dockerfiles/composio
python/swe/dockerfiles/generated
python/swe/dockerfiles/logs
python/swe/dockerfiles/logs*
python/swe/logs
indexed
python/swe/tests/logs_dir
python/swe/dockerfiles/indexed
python/swe/agentic
python/swe/dockerfiles/generated*
56 changes: 16 additions & 40 deletions python/swe/agent/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,17 +33,12 @@ def pop_thought_from_request(request: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
return request


def get_agent_graph(repo_name: str, workspace_id: str, test_command: str):
def get_agent_graph(repo_name: str, workspace_id: str):
import random
import string

random_string = "".join(random.choices(string.digits, k=6))

def save_test_response(response: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
with open(f"test_response_{random_string}.txt", "w") as f:
f.write(response["data"]["stdout"] + "\n" + response["data"]["stderr"])

return response
run_file = f"messages_{random_string}.txt"

bedrock_client = BedrockChat(
credentials_profile_name="default",
Expand All @@ -57,11 +52,7 @@ def save_test_response(response: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
metadata={
App.CODE_ANALYSIS_TOOL: {
"dir_to_index_path": f"/home/user/{repo_name}",
},
App.SHELLTOOL: {
"project_path": f"/home/user/{repo_name}",
"test_command": test_command,
},
}
},
processors={
"pre": {
Expand All @@ -74,9 +65,6 @@ def save_test_response(response: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
App.CODE_ANALYSIS_TOOL: add_thought_to_request,
App.SHELLTOOL: add_thought_to_request,
},
"post": {
Action.SHELLTOOL_TEST_COMMAND: save_test_response,
},
},
)
composio_toolset.set_workspace_id(workspace_id)
Expand Down Expand Up @@ -118,19 +106,10 @@ def save_test_response(response: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
),
]

test_tools = [
*composio_toolset.get_actions(
actions=[
Action.SHELLTOOL_TEST_COMMAND,
]
)
]

# Create two separate tool nodes
code_analysis_tool_node = ToolNode(code_analysis_tools)
file_tool_node = ToolNode(file_tools)
swe_tool_node = ToolNode(swe_tools)
test_tool_node = ToolNode(test_tools)

# Define AgentState
class AgentState(TypedDict):
Expand Down Expand Up @@ -181,6 +160,12 @@ def invoke_with_retry(agent, state):
},
name=name,
)
with open(run_file, "w") as handle:
message_str = ""
for message in state["messages"]:
message_type = type(message).__name__
message_str += f"{message_type}: {str(message.content)}\n"
handle.write(message_str)
return {"messages": [result], "sender": name}

return agent_node
Expand Down Expand Up @@ -209,7 +194,7 @@ def create_agent(system_prompt, tools):
code_analyzer_agent = create_agent(CODE_ANALYZER_PROMPT, code_analysis_tools)
code_analyzer_node = create_agent_node(code_analyzer_agent, code_analyzer_name)

editing_agent = create_agent(EDITING_AGENT_PROMPT, file_tools + test_tools)
editing_agent = create_agent(EDITING_AGENT_PROMPT, file_tools)
editing_node = create_agent_node(editing_agent, editor_name)

# Update router function
Expand All @@ -222,6 +207,7 @@ def router(
"continue",
"analyze_code",
"edit_file",
"swe_tool",
]:
messages = state["messages"]
for message in reversed(messages):
Expand Down Expand Up @@ -252,7 +238,6 @@ def router(
workflow.add_node("code_edit_tool", file_tool_node)
workflow.add_node("code_analysis_tool", code_analysis_tool_node)
workflow.add_node("swe_tool", swe_tool_node)
workflow.add_node("test_tool", test_tool_node)
# Add start and end
workflow.add_edge(START, software_engineer_name)

Expand All @@ -272,11 +257,6 @@ def router(
lambda x: x["sender"],
{software_engineer_name: software_engineer_name},
)
workflow.add_conditional_edges(
"test_tool",
lambda x: x["sender"],
{editor_name: editor_name},
)

# Update conditional edges for the coding agent
workflow.add_conditional_edges(
Expand Down Expand Up @@ -304,6 +284,8 @@ def code_analyzer_router(state):
return "code_analysis_tool"
if "ANALYSIS COMPLETE" in last_ai_message.content:
return "done"
if "EDIT FILE" in last_ai_message.content:
return "edit_file"
return "continue"

# Add conditional edges for the code analyzer
Expand All @@ -313,6 +295,7 @@ def code_analyzer_router(state):
{
"continue": code_analyzer_name,
"done": software_engineer_name,
"edit_file": editor_name,
"code_analysis_tool": "code_analysis_tool",
},
)
Expand All @@ -327,13 +310,7 @@ def code_editor_router(state):
last_ai_message = messages[-1]

if last_ai_message.tool_calls:
tool_name = last_ai_message.tool_calls[0]["name"]
if tool_name in [x.name for x in file_tools]:
return "code_edit_tool"
else:
return "test_tool"
if "EDITING AND TESTING COMPLETED" in last_ai_message.content:
return "done"
return "code_edit_tool"
if "EDITING COMPLETED" in last_ai_message.content:
return "done"
return "continue"
Expand All @@ -345,9 +322,8 @@ def code_editor_router(state):
"continue": editor_name,
"done": software_engineer_name,
"code_edit_tool": "code_edit_tool",
"test_tool": "test_tool",
},
)

graph = workflow.compile()
return graph, composio_toolset, f"test_response_{random_string}.txt"
return graph, composio_toolset, run_file
Binary file removed python/swe/agent/architecture.png
Binary file not shown.
Loading
Loading