v0

agent_loop.py

@@ -1,9 +1,15 @@
+import json
 import os
 import subprocess
+import sys
+import textwrap
 from pathlib import Path
 
-from anthropic import Anthropic
-from dotenv import load_dotenv
+try:
+    from dotenv import load_dotenv
+except ImportError:
+    def load_dotenv(*args, **kwargs):
+        return False
 
 load_dotenv(override=True)
 
@@ -11,10 +17,26 @@ if os.getenv("ANTHROPIC_BASE_URL"):
     os.environ.pop("ANTHROPIC_AUTH_TOKEN", None)
 
 WORKDIR = Path.cwd()
-client = Anthropic(base_url=os.getenv("ANTHROPIC_BASE_URL"))
-MODEL = os.environ["MODEL_ID"]
 
-SYSTEM = f"You are a coding agent at {WORKDIR}. Use tools to solve tasks. Act, don't explain."
+SYSTEM = f"""
+You are a CAD generation agent at {WORKDIR}.
+
+Your only tool is execute. For CAD tasks:
+- Write CadQuery Python code.
+- The code must assign the final model to a variable named result.
+- You may use cadquery as cq; it is pre-imported by the tool.
+- Call execute with the code and optional output_path.
+- If execute returns ok=false, inspect the structured error, fix the code, and call execute again.
+- When execute returns ok=true, report the output_path to the user.
+
+Do not ask the user to run commands or create files manually.
+""".strip()
+
+
+def get_client():
+    from anthropic import Anthropic
+
+    return Anthropic(base_url=os.getenv("ANTHROPIC_BASE_URL"))
 
 
 def safe_path(p: str) -> Path:
@@ -24,76 +46,269 @@ def safe_path(p: str) -> Path:
     return path
 
 
-def run_bash(command: str) -> str:
-    dangerous = ["rm -rf /", "sudo", "shutdown", "reboot", "> /dev/"]
-    if any(d in command for d in dangerous):
-        return "Error: Dangerous command blocked"
-    try:
-        r = subprocess.run(command, shell=True, cwd=WORKDIR,
-                           capture_output=True, text=True, timeout=120)
-        out = (r.stdout + r.stderr).strip()
-        return out[:50000] if out else "(no output)"
-    except subprocess.TimeoutExpired:
-        return "Error: Timeout (120s)"
+CADQUERY_RUNNER = r"""
+import contextlib
+import io
+import json
+import sys
+import traceback
+from pathlib import Path
+
+stdout_buffer = io.StringIO()
+stderr_buffer = io.StringIO()
+
+
+def tail(text, limit=4000):
+    if not text:
+        return ""
+    return text[-limit:]
+
+
+def fail(stage, exc=None, error_type=None, error=None):
+    payload = {
+        "ok": False,
+        "stage": stage,
+        "error_type": error_type or (type(exc).__name__ if exc else "Error"),
+        "error": error or (str(exc) if exc else ""),
+        "traceback_tail": tail(traceback.format_exc() if exc else ""),
+        "stdout": tail(stdout_buffer.getvalue()),
+        "stderr": tail(stderr_buffer.getvalue()),
+    }
+    print(json.dumps(payload, ensure_ascii=False))
+
+
+try:
+    request = json.loads(sys.stdin.read())
+    code = request["code"]
+    output_path = Path(request["output_path"])
+except Exception as exc:
+    fail("input", exc)
+    sys.exit(0)
+
+try:
+    import cadquery as cq
+except Exception as exc:
+    fail("import", exc)
+    sys.exit(0)
+
+try:
+    compiled = compile(code, "<cadquery_code>", "exec")
+except Exception as exc:
+    fail("syntax", exc)
+    sys.exit(0)
+
+namespace = {
+    "__name__": "__cadquery_user_code__",
+    "cq": cq,
+}
 
+try:
+    with contextlib.redirect_stdout(stdout_buffer), contextlib.redirect_stderr(stderr_buffer):
+        exec(compiled, namespace)
+except Exception as exc:
+    fail("execution", exc)
+    sys.exit(0)
 
-def run_read(path: str, limit: int = None) -> str:
+if "result" not in namespace:
+    fail("result", error_type="MissingResult", error="CadQuery code must assign the final model to variable 'result'.")
+    sys.exit(0)
+
+result = namespace["result"]
+if result is None:
+    fail("result", error_type="InvalidResult", error="'result' is None.")
+    sys.exit(0)
+
+exportable_types = tuple(
+    value
+    for name in ("Workplane", "Shape", "Assembly", "Sketch")
+    for value in [getattr(cq, name, None)]
+    if isinstance(value, type)
+)
+
+is_exportable = isinstance(result, exportable_types)
+if isinstance(result, (list, tuple)) and result:
+    is_exportable = all(isinstance(item, exportable_types) for item in result)
+
+if not is_exportable:
+    fail(
+        "result",
+        error_type="InvalidResultType",
+        error=f"'result' must be a CadQuery object or a non-empty list/tuple of CadQuery objects, got {type(result).__name__}.",
+    )
+    sys.exit(0)
+
+try:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    cq.exporters.export(result, str(output_path))
+except Exception as exc:
+    fail("export", exc)
+    sys.exit(0)
+
+payload = {
+    "ok": True,
+    "stage": "done",
+    "output_path": str(output_path),
+    "stdout": tail(stdout_buffer.getvalue()),
+    "stderr": tail(stderr_buffer.getvalue()),
+}
+print(json.dumps(payload, ensure_ascii=False))
+"""
+
+
+def _json_response(payload: dict) -> str:
+    return json.dumps(payload, ensure_ascii=False, indent=2)
+
+
+def run_execute(code: str, output_path: str = "outputs/model.step") -> str:
     try:
-        text = safe_path(path).read_text()
-        lines = text.splitlines()
-        if limit and limit < len(lines):
-            lines = lines[:limit] + [f"... ({len(lines) - limit} more lines)"]
-        return "\n".join(lines)[:50000]
-    except Exception as e:
-        return f"Error: {e}"
+        if not code or not code.strip():
+            return _json_response({
+                "ok": False,
+                "stage": "input",
+                "error_type": "EmptyCode",
+                "error": "execute requires non-empty CadQuery code.",
+                "traceback_tail": "",
+                "stdout": "",
+                "stderr": "",
+            })
 
+        output = safe_path(output_path)
+        if output.suffix.lower() not in (".step", ".stp"):
+            return _json_response({
+                "ok": False,
+                "stage": "input",
+                "error_type": "UnsupportedOutputFormat",
+                "error": "Only STEP output is supported in this version. Use a .step or .stp path.",
+                "traceback_tail": "",
+                "stdout": "",
+                "stderr": "",
+            })
+    except Exception as exc:
+        return _json_response({
+            "ok": False,
+            "stage": "input",
+            "error_type": type(exc).__name__,
+            "error": str(exc),
+            "traceback_tail": "",
+            "stdout": "",
+            "stderr": "",
+        })
+
+    request = {
+        "code": code,
+        "output_path": str(output),
+    }
 
-def run_write(path: str, content: str) -> str:
     try:
-        fp = safe_path(path)
-        fp.parent.mkdir(parents=True, exist_ok=True)
-        fp.write_text(content)
-        return f"Wrote {len(content)} bytes to {path}"
-    except Exception as e:
-        return f"Error: {e}"
+        process = subprocess.run(
+            [sys.executable, "-c", CADQUERY_RUNNER],
+            input=json.dumps(request, ensure_ascii=False),
+            cwd=WORKDIR,
+            capture_output=True,
+            text=True,
+            timeout=120,
+        )
+    except subprocess.TimeoutExpired as exc:
+        return _json_response({
+            "ok": False,
+            "stage": "execution",
+            "error_type": "TimeoutExpired",
+            "error": "CadQuery execution timed out after 120 seconds.",
+            "traceback_tail": "",
+            "stdout": exc.stdout[-4000:] if exc.stdout else "",
+            "stderr": exc.stderr[-4000:] if exc.stderr else "",
+        })
+    except Exception as exc:
+        return _json_response({
+            "ok": False,
+            "stage": "subprocess",
+            "error_type": type(exc).__name__,
+            "error": str(exc),
+            "traceback_tail": "",
+            "stdout": "",
+            "stderr": "",
+        })
 
+    raw_output = process.stdout.strip()
+    if not raw_output:
+        return _json_response({
+            "ok": False,
+            "stage": "subprocess",
+            "error_type": "NoRunnerOutput",
+            "error": "CadQuery runner returned no JSON output.",
+            "traceback_tail": "",
+            "stdout": "",
+            "stderr": process.stderr[-4000:],
+        })
 
-def run_edit(path: str, old_text: str, new_text: str) -> str:
     try:
-        fp = safe_path(path)
-        content = fp.read_text()
-        if old_text not in content:
-            return f"Error: Text not found in {path}"
-        fp.write_text(content.replace(old_text, new_text, 1))
-        return f"Edited {path}"
-    except Exception as e:
-        return f"Error: {e}"
+        payload = json.loads(raw_output.splitlines()[-1])
+    except Exception as exc:
+        return _json_response({
+            "ok": False,
+            "stage": "subprocess",
+            "error_type": type(exc).__name__,
+            "error": f"Failed to parse CadQuery runner output: {exc}",
+            "traceback_tail": "",
+            "stdout": raw_output[-4000:],
+            "stderr": process.stderr[-4000:],
+        })
+
+    if process.returncode != 0 and payload.get("ok") is not False:
+        payload = {
+            "ok": False,
+            "stage": "subprocess",
+            "error_type": "RunnerFailed",
+            "error": f"CadQuery runner exited with code {process.returncode}.",
+            "traceback_tail": "",
+            "stdout": raw_output[-4000:],
+            "stderr": process.stderr[-4000:],
+        }
+
+    if payload.get("ok") and payload.get("output_path"):
+        payload["output_path"] = str(Path(payload["output_path"]).relative_to(WORKDIR))
+
+    return _json_response(payload)
 
 
 # -- The dispatch map: {tool_name: handler} --
 TOOL_HANDLERS = {
-    "bash":       lambda **kw: run_bash(kw["command"]),
-    "read_file":  lambda **kw: run_read(kw["path"], kw.get("limit")),
-    "write_file": lambda **kw: run_write(kw["path"], kw["content"]),
-    "edit_file":  lambda **kw: run_edit(kw["path"], kw["old_text"], kw["new_text"]),
+    "execute": lambda **kw: run_execute(kw["code"], kw.get("output_path", "outputs/model.step")),
 }
 
 TOOLS = [
-    {"name": "bash", "description": "Run a shell command.",
-     "input_schema": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}},
-    {"name": "read_file", "description": "Read file contents.",
-     "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "limit": {"type": "integer"}}, "required": ["path"]}},
-    {"name": "write_file", "description": "Write content to file.",
-     "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, "required": ["path", "content"]}},
-    {"name": "edit_file", "description": "Replace exact text in file.",
-     "input_schema": {"type": "object", "properties": {"path": {"type": "string"}, "old_text": {"type": "string"}, "new_text": {"type": "string"}}, "required": ["path", "old_text", "new_text"]}},
+    {
+        "name": "execute",
+        "description": textwrap.dedent("""
+            Execute CadQuery Python code and export the result as a STEP file.
+            The code must assign the final CadQuery model to a variable named result.
+            cadquery is pre-imported as cq. output_path is optional and defaults to outputs/model.step.
+        """).strip(),
+        "input_schema": {
+            "type": "object",
+            "properties": {
+                "code": {
+                    "type": "string",
+                    "description": "CadQuery Python code. It must assign the final model to variable 'result'.",
+                },
+                "output_path": {
+                    "type": "string",
+                    "description": "Optional workspace-relative STEP path. Defaults to outputs/model.step.",
+                },
+            },
+            "required": ["code"],
+        },
+    },
 ]
 
 
 def agent_loop(messages: list):
+    client = get_client()
+    model = os.environ["MODEL_ID"]
+
     while True:
         response = client.messages.create(
-            model=MODEL, system=SYSTEM, messages=messages,
+            model=model, system=SYSTEM, messages=messages,
             tools=TOOLS, max_tokens=8000,
         )
         messages.append({"role": "assistant", "content": response.content})
@@ -113,7 +328,7 @@ if __name__ == "__main__":
     history = []
     while True:
         try:
-            query = input("\033[36ms02 >> \033[0m")
+            query = input("\033[36mcad agent >> \033[0m")
         except (EOFError, KeyboardInterrupt):
             break
         if query.strip().lower() in ("q", "exit", ""):

environment.yml

@@ -0,0 +1,10 @@
+name: aigc
+channels:
+  - conda-forge
+dependencies:
+  - python
+  - cadquery
+  - pip
+  - pip:
+      - anthropic
+      - python-dotenv