Jerry
275 lines
10 KiB
Python
275 lines
10 KiB
Python
"""CAD code generation tool implemented as a specialist subagent."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
import re
|
|
import shlex
|
|
from typing import Any, Optional
|
|
|
|
from SimpleLLMFunc import llm_chat, tool
|
|
from SimpleLLMFunc.type import HistoryList
|
|
|
|
from .builtin_file_toolkit import create_builtin_file_tools
|
|
from .command_tools import execute_command
|
|
from .common import (
|
|
SUBAGENT_MAX_TOOL_CALLS,
|
|
build_simplecad_workspace_fact_block,
|
|
get_config,
|
|
print_tool_output,
|
|
)
|
|
from .sketch_tools import sketch_pad_operations
|
|
from .subagent_utils import run_subagent_with_events
|
|
|
|
|
|
def create_codegen_subagent_tools(
|
|
workspace: Optional[str | Path] = None,
|
|
) -> list[Any]:
|
|
"""Create the low-level tools owned by the CAD codegen specialist."""
|
|
|
|
return [
|
|
execute_command,
|
|
sketch_pad_operations,
|
|
*create_builtin_file_tools(workspace),
|
|
]
|
|
|
|
|
|
def _build_codegen_request(
|
|
*,
|
|
task: str,
|
|
target_file_path: str,
|
|
validation_command: Optional[str],
|
|
requirement_sketch_key: Optional[str] = None,
|
|
) -> str:
|
|
parts = [f"target_file: {target_file_path}"]
|
|
if validation_command and validation_command.strip():
|
|
parts.append(f"validation_command: {validation_command.strip()}")
|
|
if requirement_sketch_key and requirement_sketch_key.strip():
|
|
parts.append(f"requirement_key: {requirement_sketch_key.strip()}")
|
|
parts.append("")
|
|
parts.append(build_simplecad_workspace_fact_block())
|
|
parts.append("")
|
|
parts.append(task.strip())
|
|
return "\n".join(parts)
|
|
|
|
|
|
def _extract_python_code_block(text: str) -> Optional[str]:
|
|
match = re.search(r"```python\s*(.*?)```", text, flags=re.DOTALL)
|
|
if match:
|
|
return match.group(1).strip()
|
|
return None
|
|
|
|
|
|
def _build_missing_code_retry_request(
|
|
*,
|
|
target_file_path: str,
|
|
) -> str:
|
|
return "\n".join(
|
|
[
|
|
"<RETRY_AFTER_NO_CODE>",
|
|
f"You ended without writing any script to {target_file_path}.",
|
|
f"You must write the required Python script directly to {target_file_path} before you finish.",
|
|
"Do not stop after planning, describing the approach, or pasting a code block in chat.",
|
|
"Use the file-writing tool now, save the script to disk, then validate/debug until export succeeds.",
|
|
"</RETRY_AFTER_NO_CODE>",
|
|
]
|
|
)
|
|
|
|
|
|
def _build_missing_code_retry_history(
|
|
*,
|
|
original_request: str,
|
|
prior_report: str,
|
|
) -> HistoryList:
|
|
history: HistoryList = [{"role": "user", "content": original_request}]
|
|
if prior_report.strip():
|
|
history.append({"role": "assistant", "content": prior_report.strip()})
|
|
return history
|
|
|
|
|
|
def _default_validation_command(
|
|
target_file_path: str, validation_command: Optional[str]
|
|
) -> Optional[str]:
|
|
if validation_command and validation_command.strip():
|
|
return validation_command.strip()
|
|
|
|
if target_file_path.endswith(".py"):
|
|
script_path = shlex.quote(target_file_path)
|
|
output_dir = shlex.quote(str(Path(target_file_path).parent or Path(".")))
|
|
stl_glob = f"{output_dir}/*.stl"
|
|
step_glob = f"{output_dir}/*.step"
|
|
stp_glob = f"{output_dir}/*.stp"
|
|
return (
|
|
f"uv run python {script_path} && "
|
|
f"ls {stl_glob} && "
|
|
f"(ls {step_glob} || ls {stp_glob})"
|
|
)
|
|
|
|
return None
|
|
|
|
|
|
def _read_latest_code(target_file_path: str) -> Optional[str]:
|
|
candidate = Path(target_file_path)
|
|
if not candidate.is_absolute():
|
|
candidate = Path.cwd() / candidate
|
|
|
|
if not candidate.exists() or not candidate.is_file():
|
|
return None
|
|
|
|
try:
|
|
return candidate.read_text(encoding="utf-8")
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
@tool(
|
|
name="cad_code_generator",
|
|
description=(
|
|
"Delegate CAD script creation or repair to a specialist subagent that owns "
|
|
"builtin file tools and can iteratively debug the target file."
|
|
),
|
|
best_practices=[
|
|
"MUST pass requirement_sketch_key (the req_xxxx from make_user_query_more_detailed) so the specialist retrieves the detailed spec.",
|
|
"Use one complete natural-language task block instead of splitting context across many parameters.",
|
|
"In the task, always say whether this is create-new-file or modify-existing-code.",
|
|
"In the task, include the full user intent, failure context, relevant SketchPad ids, and success criteria.",
|
|
"In the task, explicitly tell the specialist to validate and keep debugging until STL and STEP/STP export succeeds.",
|
|
"Always provide the exact target_file_path for model.py or the file to repair.",
|
|
],
|
|
)
|
|
async def cad_code_generator(
|
|
task: str,
|
|
target_file_path: str,
|
|
requirement_sketch_key: Optional[str] = None,
|
|
event_emitter: Any = None,
|
|
) -> str:
|
|
"""Run the CAD coding specialist as a single-call subagent.
|
|
|
|
Args:
|
|
task: A complete natural-language mission for the coding specialist.
|
|
This should explicitly include:
|
|
- the full user intent,
|
|
- whether the job is create-new-file or modify-existing-code,
|
|
- the concrete modification target or creation goal,
|
|
- any traceback / visual feedback / failure context,
|
|
- any relevant SketchPad ids that the specialist should inspect,
|
|
- any reference-code SketchPad ids if they matter,
|
|
- the expected success criteria,
|
|
- and an explicit instruction to validate and keep debugging until model export succeeds.
|
|
Prefer one complete instruction block instead of splitting context across
|
|
many parameters.
|
|
target_file_path: The exact path of the script file that the specialist owns.
|
|
In the normal workflow this should be the final `model.py` path.
|
|
requirement_sketch_key: REQUIRED. The SketchPad key (e.g. req_xxxx) from make_user_query_more_detailed.
|
|
The specialist will retrieve and follow this detailed requirement.
|
|
event_emitter: Optional tool event emitter used to forward nested specialist
|
|
progress events back to the outer agent event stream.
|
|
|
|
Returns:
|
|
A concise report from the specialist plus the latest code snapshot.
|
|
"""
|
|
|
|
actual_validation_command = _default_validation_command(
|
|
target_file_path,
|
|
None,
|
|
)
|
|
|
|
request_payload = _build_codegen_request(
|
|
task=task.strip(),
|
|
target_file_path=target_file_path,
|
|
validation_command=actual_validation_command,
|
|
requirement_sketch_key=requirement_sketch_key,
|
|
)
|
|
|
|
print_tool_output(
|
|
"🧠 CAD Code Specialist",
|
|
"\n".join(
|
|
[
|
|
f"Target file: {target_file_path}",
|
|
f"Validation command: {actual_validation_command or '(not provided)'}",
|
|
f"Task summary: {task.strip()[:160]}",
|
|
]
|
|
),
|
|
)
|
|
|
|
report = await run_subagent_with_events(
|
|
specialist_callable=cad_code_generator_specialist,
|
|
specialist_kwargs={
|
|
"message": request_payload,
|
|
"history": [],
|
|
},
|
|
subagent_label="CAD Code Specialist",
|
|
event_emitter=event_emitter,
|
|
status_payload={
|
|
"target_file_path": target_file_path,
|
|
"validation_command": actual_validation_command,
|
|
},
|
|
)
|
|
|
|
written_code = _read_latest_code(target_file_path)
|
|
latest_code = written_code or _extract_python_code_block(report)
|
|
if written_code is None:
|
|
print_tool_output(
|
|
"⚠️ CAD Code Specialist",
|
|
"First attempt did not write the target file. Appending a stricter follow-up instruction.",
|
|
)
|
|
retry_report = await run_subagent_with_events(
|
|
specialist_callable=cad_code_generator_specialist,
|
|
specialist_kwargs={
|
|
"message": _build_missing_code_retry_request(
|
|
target_file_path=target_file_path,
|
|
),
|
|
"history": _build_missing_code_retry_history(
|
|
original_request=request_payload,
|
|
prior_report=report,
|
|
),
|
|
},
|
|
subagent_label="CAD Code Specialist",
|
|
event_emitter=event_emitter,
|
|
status_payload={
|
|
"target_file_path": target_file_path,
|
|
"validation_command": actual_validation_command,
|
|
"retry_reason": "no_code_written",
|
|
},
|
|
)
|
|
report = retry_report.strip() or report
|
|
written_code = _read_latest_code(target_file_path)
|
|
latest_code = written_code or _extract_python_code_block(report)
|
|
|
|
if latest_code is None:
|
|
return report.strip()
|
|
|
|
return (
|
|
f"{report.strip()}\n\n"
|
|
f"📁 Target file: {target_file_path}\n"
|
|
f"📄 Latest code:\n```python\n{latest_code.strip()}\n```"
|
|
)
|
|
|
|
|
|
@llm_chat(
|
|
llm_interface=get_config().REASONING_INTERFACE,
|
|
toolkit=create_codegen_subagent_tools(),
|
|
max_tool_calls=SUBAGENT_MAX_TOOL_CALLS,
|
|
stream=True,
|
|
enable_event=True,
|
|
timeout=900,
|
|
temperature=0.8,
|
|
)
|
|
async def cad_code_generator_specialist(
|
|
message: str,
|
|
history: HistoryList | None = None,
|
|
) -> None: # type: ignore[misc]
|
|
"""You are a CAD coding agent. Write Python code directly to the target file with echo_into.
|
|
|
|
Always create/write the target file first. Use the workspace facts included in the user message. Read the chosen skill root's `SKILL.md`, then `references/docs/api/README.md`, then the exact API Markdown pages you use. Use the provided `validation_command` exactly; when you run Python in this repo/workspace, prefer `uv run python ...`. Run validation directly with `execute_command`; that tool already allows up to 600 seconds for a command, so use it as the standard execution path. After a successful script run, verify exported files with `ls` instead of rerunning the same script just to check whether STL/STEP outputs exist. Do not print whole solids, assemblies, or full model objects for inspection; use QL queries and print only the small queried facts you need for grounding/debugging. Keep debugging until the script is executed successfully and exports both STL and STEP/STP (for example, ./model.stl and ./model.step).
|
|
"""
|
|
|
|
pass
|
|
|
|
|
|
__all__ = [
|
|
"cad_code_generator",
|
|
"cad_code_generator_specialist",
|
|
"create_codegen_subagent_tools",
|
|
]
|