first commit
This commit is contained in:
@@ -0,0 +1,200 @@
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
if PROJECT_ROOT not in sys.path:
|
||||
sys.path.insert(0, PROJECT_ROOT)
|
||||
WORKSPACE_ROOT = os.path.join(PROJECT_ROOT, "workspace")
|
||||
|
||||
|
||||
from SimpleLLMFunc.hooks.events import (
|
||||
ReActEventType,
|
||||
ToolCallEndEvent,
|
||||
ToolCallErrorEvent,
|
||||
ToolCallStartEvent,
|
||||
)
|
||||
from SimpleLLMFunc.hooks.stream import EventYield, ResponseYield
|
||||
|
||||
import tools.code_tools as code_tools_module
|
||||
from tools.code_tools import cad_code_generator
|
||||
|
||||
|
||||
class _FakeEmitter:
|
||||
def __init__(self):
|
||||
self.events: list[tuple[str, dict]] = []
|
||||
|
||||
async def emit(self, event_name: str, data):
|
||||
self.events.append((event_name, data))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cad_code_generator_bridges_nested_specialist_events(monkeypatch):
|
||||
monkeypatch.chdir(WORKSPACE_ROOT)
|
||||
|
||||
async def fake_specialist(**kwargs):
|
||||
assert "Current working directory:" in kwargs["message"]
|
||||
assert "Skill root: use the preferred skill root below." in kwargs["message"]
|
||||
assert "validation_command: uv run python part/model.py" in kwargs["message"]
|
||||
assert "references/docs/api/README.md" in kwargs["message"]
|
||||
yield EventYield(
|
||||
event=ToolCallStartEvent(
|
||||
event_type=ReActEventType.TOOL_CALL_START,
|
||||
timestamp=datetime.now(timezone.utc),
|
||||
trace_id="trace-1",
|
||||
func_name="cad_code_generator_specialist",
|
||||
iteration=0,
|
||||
tool_name="read_file",
|
||||
tool_call_id="nested-1",
|
||||
arguments={"file_path": "part/model.py"},
|
||||
tool_call=cast(
|
||||
Any,
|
||||
{
|
||||
"id": "nested-1",
|
||||
"type": "function",
|
||||
"function": {"name": "read_file", "arguments": "{}"},
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
yield EventYield(
|
||||
event=ToolCallEndEvent(
|
||||
event_type=ReActEventType.TOOL_CALL_END,
|
||||
timestamp=datetime.now(timezone.utc),
|
||||
trace_id="trace-1",
|
||||
func_name="cad_code_generator_specialist",
|
||||
iteration=0,
|
||||
tool_name="read_file",
|
||||
tool_call_id="nested-1",
|
||||
arguments={"file_path": "part/model.py"},
|
||||
result="print('old')",
|
||||
execution_time=0.05,
|
||||
success=True,
|
||||
)
|
||||
)
|
||||
yield EventYield(
|
||||
event=ToolCallErrorEvent(
|
||||
event_type=ReActEventType.TOOL_CALL_ERROR,
|
||||
timestamp=datetime.now(timezone.utc),
|
||||
trace_id="trace-1",
|
||||
func_name="cad_code_generator_specialist",
|
||||
iteration=0,
|
||||
tool_name="execute_command",
|
||||
tool_call_id="nested-2",
|
||||
arguments={"command": "python model.py"},
|
||||
error=RuntimeError("boom"),
|
||||
error_message="boom",
|
||||
error_type="RuntimeError",
|
||||
execution_time=0.12,
|
||||
)
|
||||
)
|
||||
yield ResponseYield(response="STATUS: SUCCESS\nSUMMARY: ok", messages=[])
|
||||
|
||||
monkeypatch.setattr(
|
||||
code_tools_module,
|
||||
"cad_code_generator_specialist",
|
||||
fake_specialist,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
code_tools_module, "_read_latest_code", lambda path: "print('ok')\n"
|
||||
)
|
||||
emitter = _FakeEmitter()
|
||||
result = await cad_code_generator(
|
||||
task="Create a cube as a new file. This is a create-new-file task.",
|
||||
target_file_path="part/model.py",
|
||||
event_emitter=emitter,
|
||||
)
|
||||
|
||||
event_names = [event_name for event_name, _ in emitter.events]
|
||||
assert event_names == [
|
||||
"subagent_status",
|
||||
"subagent_tool_start",
|
||||
"subagent_tool_end",
|
||||
"subagent_tool_error",
|
||||
"subagent_response",
|
||||
"subagent_status",
|
||||
]
|
||||
first_status_payload = emitter.events[0][1]
|
||||
assert first_status_payload["subagent_label"] == "CAD Code Specialist"
|
||||
assert first_status_payload["validation_command"].startswith(
|
||||
"uv run python part/model.py"
|
||||
)
|
||||
assert "ls part/*.stl" in first_status_payload["validation_command"]
|
||||
assert (
|
||||
"(ls part/*.step || ls part/*.stp)"
|
||||
in first_status_payload["validation_command"]
|
||||
)
|
||||
assert "STATUS: SUCCESS" in result
|
||||
assert "Latest code" in result
|
||||
response_payload = emitter.events[4][1]
|
||||
assert response_payload["delta_text"] == "STATUS: SUCCESS\nSUMMARY: ok"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cad_code_generator_retries_when_first_attempt_does_not_produce_code(
|
||||
monkeypatch,
|
||||
):
|
||||
monkeypatch.chdir(WORKSPACE_ROOT)
|
||||
calls = {"count": 0}
|
||||
|
||||
async def fake_specialist(**kwargs):
|
||||
calls["count"] += 1
|
||||
if calls["count"] == 1:
|
||||
assert kwargs["history"] == []
|
||||
assert "Current working directory:" in kwargs["message"]
|
||||
assert (
|
||||
"validation_command: uv run python part/model.py" in kwargs["message"]
|
||||
)
|
||||
assert "references/docs/api/README.md" in kwargs["message"]
|
||||
yield ResponseYield(
|
||||
response=(
|
||||
"STATUS: SUCCESS\n"
|
||||
"SUMMARY: planned but no file write\n"
|
||||
"```python\n"
|
||||
"print('draft only')\n"
|
||||
"```"
|
||||
),
|
||||
messages=[],
|
||||
)
|
||||
return
|
||||
|
||||
assert "must write the required Python script directly" in kwargs["message"]
|
||||
assert "part/model.py" in kwargs["message"]
|
||||
assert kwargs["history"][0]["role"] == "user"
|
||||
assert "target_file: part/model.py" in kwargs["history"][0]["content"]
|
||||
assert kwargs["history"][1]["role"] == "assistant"
|
||||
assert "planned but no file write" in kwargs["history"][1]["content"]
|
||||
yield ResponseYield(
|
||||
response="STATUS: SUCCESS\nCODE_WRITTEN: YES\nSUMMARY: wrote code",
|
||||
messages=[],
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
code_tools_module,
|
||||
"cad_code_generator_specialist",
|
||||
fake_specialist,
|
||||
)
|
||||
|
||||
code_reads = {"count": 0}
|
||||
|
||||
def fake_read_latest_code(path):
|
||||
code_reads["count"] += 1
|
||||
if code_reads["count"] == 1:
|
||||
return None
|
||||
return "print('ok')\n"
|
||||
|
||||
monkeypatch.setattr(code_tools_module, "_read_latest_code", fake_read_latest_code)
|
||||
|
||||
result = await cad_code_generator(
|
||||
task="Create a cube as a new file. This is a create-new-file task.",
|
||||
target_file_path="part/model.py",
|
||||
event_emitter=None,
|
||||
)
|
||||
|
||||
assert calls["count"] == 2
|
||||
assert "CODE_WRITTEN: YES" in result
|
||||
assert "Latest code" in result
|
||||
Reference in New Issue
Block a user