253 lines
7.6 KiB
Python
253 lines
7.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from copy import deepcopy
|
|
from typing import Any
|
|
|
|
from behave import given, then, when
|
|
|
|
from cleveragents.agents.graphs.auto_debug import AutoDebugAgent, AutoDebugState
|
|
|
|
|
|
class _SequentialLLM:
|
|
"""Stub LLM that returns sequential responses with a content attribute."""
|
|
|
|
def __init__(self, responses: list[str]):
|
|
self._responses = list(responses)
|
|
self.invocations: list[list[Any]] = []
|
|
|
|
def invoke(self, messages):
|
|
self.invocations.append(messages)
|
|
if not self._responses:
|
|
raise AssertionError("No responses left in stub LLM")
|
|
content = self._responses.pop(0)
|
|
return type("Response", (), {"content": content})()
|
|
|
|
|
|
@given("a custom auto debug LLM yields JSON responses")
|
|
def step_custom_llm(context):
|
|
context.custom_llm = _SequentialLLM(
|
|
[
|
|
"Custom analysis output",
|
|
json.dumps(
|
|
{
|
|
"description": "JSON described fix",
|
|
"code": "# patched code",
|
|
"files_to_modify": ["file.py"],
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
"is_valid": True,
|
|
"reasoning": "Structured JSON shows success",
|
|
"issues": [],
|
|
}
|
|
),
|
|
]
|
|
)
|
|
|
|
|
|
@when("I run the auto debug agent graph with sample state")
|
|
def step_run_auto_debug_graph(context):
|
|
context.agent = AutoDebugAgent(llm=context.custom_llm, max_fix_attempts=2)
|
|
|
|
input_state: AutoDebugState = {
|
|
"messages": [],
|
|
"context": {},
|
|
"result": None,
|
|
"error": None,
|
|
"metadata": {},
|
|
"error_message": "NameError: missing value",
|
|
"code_context": "value = missing",
|
|
"attempted_fixes": [],
|
|
"current_fix": {},
|
|
"fix_validated": False,
|
|
}
|
|
|
|
context.final_state = context.agent.invoke(input_state)
|
|
|
|
|
|
@then("the provided LLM instance should be used")
|
|
def step_assert_llm_used(context):
|
|
assert context.agent.llm is context.custom_llm
|
|
assert len(context.custom_llm.invocations) == 3
|
|
|
|
|
|
@then("the fix JSON should populate the current fix")
|
|
def step_assert_fix_parsed(context):
|
|
fix = context.final_state["current_fix"]
|
|
assert fix["description"] == "JSON described fix"
|
|
assert fix["code"] == "# patched code"
|
|
assert fix["files_to_modify"] == ["file.py"]
|
|
|
|
|
|
@then("the validation JSON should mark the fix as valid")
|
|
def step_assert_validation_parsed(context):
|
|
assert context.final_state["fix_validated"] is True
|
|
result = context.final_state.get("result")
|
|
assert result is not None
|
|
assert result["success"] is True
|
|
assert result["fix"] == context.final_state["current_fix"]
|
|
assert result["attempts"] == 0
|
|
|
|
|
|
@given("a retrying auto debug LLM yields a failed then successful validation")
|
|
def step_retrying_llm(context):
|
|
context.retry_llm = _SequentialLLM(
|
|
[
|
|
"Initial analysis output",
|
|
json.dumps(
|
|
{
|
|
"description": "First fix",
|
|
"code": "# first attempt",
|
|
"files_to_modify": ["first.py"],
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
"is_valid": False,
|
|
"reasoning": "Validation failed",
|
|
"issues": ["still broken"],
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
"description": "Second fix",
|
|
"code": "# second attempt",
|
|
"files_to_modify": ["second.py"],
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
"is_valid": True,
|
|
"reasoning": "Validation succeeded",
|
|
"issues": [],
|
|
}
|
|
),
|
|
]
|
|
)
|
|
|
|
|
|
@when("I run the auto debug agent graph with a retry")
|
|
def step_run_auto_debug_with_retry(context):
|
|
agent = AutoDebugAgent(llm=context.retry_llm, max_fix_attempts=2)
|
|
|
|
input_state: AutoDebugState = {
|
|
"messages": [],
|
|
"context": {},
|
|
"result": None,
|
|
"error": None,
|
|
"metadata": {},
|
|
"error_message": "NameError: missing value",
|
|
"code_context": "value = missing",
|
|
"attempted_fixes": [],
|
|
"current_fix": {},
|
|
"fix_validated": False,
|
|
}
|
|
|
|
context.retry_state = agent.invoke(input_state)
|
|
|
|
|
|
@then("the agent should retry and capture the failed fix")
|
|
def step_assert_retry_behavior(context):
|
|
state = context.retry_state
|
|
assert state["fix_validated"] is True
|
|
assert len(state["attempted_fixes"]) == 1
|
|
assert state["attempted_fixes"][0]["description"] == "First fix"
|
|
assert state["current_fix"]["description"] == "Second fix"
|
|
result = state.get("result")
|
|
assert result is not None
|
|
assert result["success"] is True
|
|
assert result["attempts"] == 1
|
|
|
|
|
|
@given("a simple auto debug LLM for async calls")
|
|
def step_async_llm(context):
|
|
context.async_llm_responses = [
|
|
"Async analysis output",
|
|
json.dumps(
|
|
{
|
|
"description": "Async fix",
|
|
"code": "# async patched",
|
|
"files_to_modify": ["async.py"],
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
"is_valid": True,
|
|
"reasoning": "Async validation succeeded",
|
|
"issues": [],
|
|
}
|
|
),
|
|
]
|
|
|
|
|
|
@when("I run the auto debug async interfaces")
|
|
def step_run_async_interfaces(context):
|
|
base_state: AutoDebugState = {
|
|
"messages": [],
|
|
"context": {},
|
|
"result": None,
|
|
"error": None,
|
|
"metadata": {},
|
|
"error_message": "NameError: missing value",
|
|
"code_context": "value = missing",
|
|
"attempted_fixes": [],
|
|
"current_fix": {},
|
|
"fix_validated": False,
|
|
}
|
|
|
|
def make_agent() -> AutoDebugAgent:
|
|
return AutoDebugAgent(
|
|
llm=_SequentialLLM(list(context.async_llm_responses)), max_fix_attempts=1
|
|
)
|
|
|
|
async def collect_async_results():
|
|
agent_for_ainvoke = make_agent()
|
|
context.async_invoke_state = await agent_for_ainvoke.ainvoke(
|
|
deepcopy(base_state)
|
|
)
|
|
|
|
agent_for_stream = make_agent()
|
|
context.stream_events = list(agent_for_stream.stream(deepcopy(base_state)))
|
|
|
|
agent_for_astream = make_agent()
|
|
async_events: list[dict[str, Any]] = []
|
|
async for event in agent_for_astream.astream(deepcopy(base_state)):
|
|
async_events.append(event)
|
|
context.async_stream_events = async_events
|
|
|
|
asyncio.run(collect_async_results())
|
|
|
|
|
|
@then("the async auto debug results should complete")
|
|
def step_assert_async_results(context):
|
|
assert context.async_invoke_state["result"]["success"] is True
|
|
|
|
assert context.stream_events
|
|
assert isinstance(context.stream_events[-1], dict)
|
|
|
|
assert context.async_stream_events
|
|
assert isinstance(context.async_stream_events[-1], dict)
|
|
|
|
def extract_success(events):
|
|
for event in events:
|
|
if not isinstance(event, dict):
|
|
continue
|
|
for value in event.values():
|
|
if not isinstance(value, dict):
|
|
continue
|
|
result = value.get("result")
|
|
if isinstance(result, dict) and "success" in result:
|
|
return result.get("success")
|
|
return None
|
|
|
|
stream_success = extract_success(context.stream_events)
|
|
if stream_success is not None:
|
|
assert stream_success is True
|
|
|
|
async_success = extract_success(context.async_stream_events)
|
|
if async_success is not None:
|
|
assert async_success is True
|