import os import re import sys import base64 import shutil import mimetypes import subprocess import tempfile from pathlib import Path import time # ========================= # Paths / environment # ========================= BASE_DIR = Path.cwd().resolve() GRADIO_TEMP_DIR = BASE_DIR / "gradio_tmp" ARTIFACT_DIR = BASE_DIR / "outputs" GRADIO_TEMP_DIR.mkdir(parents=True, exist_ok=True) ARTIFACT_DIR.mkdir(parents=True, exist_ok=True) os.environ["GRADIO_TEMP_DIR"] = str(GRADIO_TEMP_DIR) import gradio as gr from anthropic import Anthropic from openai import OpenAI # ========================= # Clients # ========================= ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY") OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") if not ANTHROPIC_API_KEY: raise RuntimeError("ANTHROPIC_API_KEY is not set") if not OPENAI_API_KEY: raise RuntimeError("OPENAI_API_KEY is not set") anthropic_client = Anthropic(api_key=ANTHROPIC_API_KEY) openai_client = OpenAI(api_key=OPENAI_API_KEY) # ========================= # Model options # ========================= CLAUDE_MODELS = [ "claude-opus-4-7", "claude-opus-4-6", "claude-sonnet-4-6", "claude-sonnet-4-5", "claude-haiku-4-5-20251001" ] OPENAI_MODELS = [ "gpt-5.4-pro-2026-03-05", "gpt-5.2-2025-12-11", "gpt-5-nano-2025-08-07", ] DEFAULT_CLAUDE_MODEL = "claude-opus-4-7" DEFAULT_OPENAI_MODEL = "gpt-5-nano-2025-08-07" # ========================= # Helpers # ========================= def strip_code_fences(text: str) -> str: text = text.strip() match = re.search(r"```(?:python)?\s*(.*?)```", text, re.DOTALL | re.IGNORECASE) if match: return match.group(1).strip() return text def detect_media_type(image_path: str) -> str: media_type, _ = mimetypes.guess_type(image_path) return media_type or "application/octet-stream" def copy_to_artifact(src_path: Path, suffix: str) -> str: dst = ARTIFACT_DIR / f"{next(tempfile._get_candidate_names())}{suffix}" shutil.copy2(src_path, dst) return str(dst.resolve()) def truncate_text(text: str, max_chars: int = 4000) -> str: if not text: return "" return text[-max_chars:] def append_log(log: str, message: str) -> str: return (log + message.rstrip() + "\n").strip() + "\n" # ========================= # LLM steps # ========================= def generate_cad_code(text_prompt, image_path, step_path, claude_model): content = [] spec_text = (text_prompt or "").strip() if spec_text: content.append( { "type": "text", "text": f"""Generate executable CadQuery Python code to build a 3D CAD model. Requirements: - Use CadQuery only - Save files in the current working directory - Export BOTH final files: - ./output.step - ./output.stl - ALSO export progressive STL snapshots during construction: - ./snap_001.stl - ./snap_002.stl - ./snap_003.stl - ... - The script must be directly executable with `python model.py` - Do not use markdown fences - Output ONLY valid Python code - At the end, ensure the exports are actually executed Important implementation rules: - Build the model in multiple clear steps - Keep the current shape in a variable named `result` - After every major modeling step, export `result` to a new snapshot STL - Use zero-padded numbering like snap_001.stl - At the end, export the final `result` to output.step and output.stl Specification: {spec_text} """, } ) if image_path: with open(image_path, "rb") as f: img_b64 = base64.b64encode(f.read()).decode("utf-8") content.append( { "type": "image", "source": { "type": "base64", "media_type": detect_media_type(image_path), "data": img_b64, }, } ) if step_path: uploaded_file = anthropic_client.beta.files.upload( file=( Path(step_path).name, open(step_path, "rb"), "text/plain", ) ) content.append( { "type": "document", "source": { "type": "file", "file_id": uploaded_file.id, }, "title": "input.step", "context": "This is a STEP CAD file in plain text format.", } ) if not content: raise gr.Error("Text CAD specification or image is required.") response = anthropic_client.beta.messages.create( model=claude_model, max_tokens=16000, betas=["files-api-2025-04-14"], messages=[ { "role": "user", "content": content } ], ) code = "".join( block.text for block in response.content if getattr(block, "type", None) == "text" ).strip() code = strip_code_fences(code) if not code: raise gr.Error("Claude returned empty code.") return code def translate_to_japanese(text, openai_model): if not text or not text.strip(): return "" response = openai_client.responses.create( model=openai_model, instructions=( "You are a professional translator. " "Translate the user's English text into natural, concise Japanese. " "Preserve structure, numbering, and technical meaning. " "Return only the Japanese translation." ), input=text, ) return response.output_text.strip() def parse_plan_and_code(response_text: str): plan_text = "" code_text = response_text.strip() if "[PLAN]" in response_text and "[CODE]" in response_text: before_code, after_code = response_text.split("[CODE]", 1) _, plan_part = before_code.split("[PLAN]", 1) plan_text = plan_part.strip() code_text = after_code.strip() code_text = strip_code_fences(code_text) return plan_text, code_text def generate_cad_artifacts(text_prompt, image_path, step_path, claude_model): content = [] spec_text = (text_prompt or "").strip() if spec_text: content.append( { "type": "text", "text": f"""Generate a 3D CAD model in CadQuery. Return output in exactly this format: [PLAN] Write 4 to 8 concrete build steps for the model. Each step should explain: - what geometry or feature is being created - why it is needed - any important design intent, constraint, or tradeoff The PLAN should be specific and descriptive, suitable for showing as a "Build Progress" log. [CODE] Return executable Python code using CadQuery only. Requirements for the code: - Use CadQuery only - Save files in the current working directory - Export BOTH final files: - ./output.step - ./output.stl - ALSO export progressive STL snapshots during construction: - ./snap_001.stl - ./snap_002.stl - ./snap_003.stl - ... - The script must be directly executable with `python model.py` - Do not use markdown fences - Output ONLY the PLAN block and the CODE block - At the end, ensure the exports are actually executed Important implementation rules: - Build the model in multiple clear steps - Keep the current shape in a variable named `result` - After every major modeling step, export `result` to a new snapshot STL - Use zero-padded numbering like snap_001.stl - At the end, export the final `result` to output.step and output.stl - The number and order of snapshot exports should roughly align with the PLAN steps Specification: {spec_text} """, } ) if image_path: with open(image_path, "rb") as f: img_b64 = base64.b64encode(f.read()).decode("utf-8") content.append( { "type": "image", "source": { "type": "base64", "media_type": detect_media_type(image_path), "data": img_b64, }, } ) if step_path: uploaded_file = anthropic_client.beta.files.upload( file=( Path(step_path).name, open(step_path, "rb"), "text/plain", ) ) content.append( { "type": "document", "source": { "type": "file", "file_id": uploaded_file.id, }, "title": "input.step", "context": "This is a STEP CAD file in plain text format.", } ) if not content: raise gr.Error("Text CAD specification or image is required.") response = anthropic_client.beta.messages.create( model=claude_model, max_tokens=16000, betas=["files-api-2025-04-14"], messages=[ { "role": "user", "content": content } ], ) raw_text = "".join( block.text for block in response.content if getattr(block, "type", None) == "text" ).strip() if not raw_text: raise gr.Error("Claude returned empty output.") plan_text, code = parse_plan_and_code(raw_text) if not code: raise gr.Error("Claude returned empty code.") return plan_text, code def run_cadquery(code): with tempfile.TemporaryDirectory(dir=str(GRADIO_TEMP_DIR)) as tmpdir: tmpdir = Path(tmpdir) script_path = tmpdir / "model.py" with open(script_path, "w", encoding="utf-8") as f: f.write(code) result = subprocess.run( [sys.executable, str(script_path)], cwd=str(tmpdir), capture_output=True, text=True, timeout=180, ) if result.returncode != 0: raise gr.Error( "CadQuery execution failed.\n\n" f"STDOUT:\n{truncate_text(result.stdout)}\n\n" f"STDERR:\n{truncate_text(result.stderr)}" ) step_path = tmpdir / "output.step" stl_path = tmpdir / "output.stl" files = [p.name for p in tmpdir.iterdir()] if not step_path.exists() and not stl_path.exists(): raise gr.Error( "CAD script finished but did not generate output.step or output.stl.\n\n" f"Files found in temp dir: {files}\n\n" "Make sure the generated CadQuery code exports exactly " "./output.step and ./output.stl" ) if not step_path.exists(): raise gr.Error( "output.step was not created.\n\n" f"Files found in temp dir: {files}" ) if not stl_path.exists(): raise gr.Error( "output.stl was not created.\n\n" f"Files found in temp dir: {files}" ) final_step = copy_to_artifact(step_path, ".step") final_stl = copy_to_artifact(stl_path, ".stl") return final_step, final_stl def run_cadquery_progressive(code): with tempfile.TemporaryDirectory(dir=str(GRADIO_TEMP_DIR)) as tmpdir: tmpdir = Path(tmpdir) script_path = tmpdir / "model.py" with open(script_path, "w", encoding="utf-8") as f: f.write(code) proc = subprocess.Popen( [sys.executable, str(script_path)], cwd=str(tmpdir), stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) seen = set() last_snapshot_path = None while True: snap_files = sorted(tmpdir.glob("snap_*.stl")) for snap in snap_files: if snap.name not in seen: seen.add(snap.name) copied = copy_to_artifact(snap, ".stl") last_snapshot_path = copied yield { "viewer": copied, "status": f"Rendering step {snap.stem}...", "step_file": None, "review": None, "done": False, } ret = proc.poll() if ret is not None: break time.sleep(0.4) stdout, stderr = proc.communicate() if proc.returncode != 0: raise gr.Error( "CadQuery execution failed.\n\n" f"STDOUT:\n{truncate_text(stdout)}\n\n" f"STDERR:\n{truncate_text(stderr)}" ) step_path = tmpdir / "output.step" stl_path = tmpdir / "output.stl" files = [p.name for p in tmpdir.iterdir()] if not step_path.exists() or not stl_path.exists(): raise gr.Error( "CAD script finished but final output.step/output.stl was missing.\n\n" f"Files found: {files}" ) final_step = copy_to_artifact(step_path, ".step") final_stl = copy_to_artifact(stl_path, ".stl") yield { "viewer": final_stl, "status": "Final model completed.", "step_file": final_step, "review": None, "done": True, } import base64 import tempfile from pathlib import Path def _stl_to_png(stl_path): try: import trimesh mesh = trimesh.load_mesh(stl_path) png = mesh.scene().save_image(resolution=(512, 512)) # 👇 STEP/STLと同じ場所に保存 out_path = Path(stl_path).with_suffix(".png") with open(out_path, "wb") as f: f.write(png) return str(out_path) except Exception: return None def _image_to_data_url(path): with open(path, "rb") as f: b64 = base64.b64encode(f.read()).decode("utf-8") return f"data:image/png;base64,{b64}" def gpt_check(step_path, openai_model): if not step_path or not os.path.exists(step_path): return "STEP file not found, so review was skipped." size = os.path.getsize(step_path) # 🔹 STEP 読み込み（軽く truncate） try: with open(step_path, "r", encoding="utf-8", errors="ignore") as f: step_text = f.read() except Exception: step_text = "" # 🔹 STL を推定して PNG 化 stl_path = Path(step_path).with_suffix(".stl") preview_png = _stl_to_png(str(stl_path)) if stl_path.exists() else None # 🔹 GPT に渡す内容 content = [ { "type": "input_text", "text": f"""A CAD STEP file was generated. File size: {size} bytes Below is part of the STEP file: {step_text} Tasks: - Check if the STEP structure looks syntactically valid - Point out possible corruption or missing sections - Infer possible geometry issues from structure - If an image is provided, also check visual problems (holes, broken shapes, weird proportions) Keep it concise. 出力は日本語で。 """, } ] # 🔹 画像も渡す（あれば） if preview_png: content.append( { "type": "input_image", "image_url": _image_to_data_url(preview_png), } ) response = openai_client.responses.create( model=openai_model, input=[{"role": "user", "content": content}], ) return response.output_text.strip() # ========================= # Pipeline # ========================= def pipeline(text_prompt, image, step_file, claude_model, openai_model): log = "" if text_prompt: log += f"🧠 Request: {text_prompt[:120]}\n\n" else: log += "🧠 Request: image-based CAD generation\n\n" log += "🧠 Analyzing request...\n" log += "🧠 Thinking...\n Please wait for 10 seconds ... \n" yield None, None, log, None, None yield None, None, log, None, None plan_text, code = generate_cad_artifacts(text_prompt, image, step_file, claude_model) plan_text = translate_to_japanese(plan_text, openai_model) def extract_plan_steps(plan_text: str): steps = [] for line in plan_text.splitlines(): line = line.strip() if not line: continue if line[0].isdigit() or line.startswith("-"): steps.append(line) return steps plan_steps = extract_plan_steps(plan_text) log += "🏗 Build Progress\n" if not plan_steps: log += "- No explicit design plan was returned.\n" log += "\n✅ CAD code generated\n" log += "⚙️ Starting CAD build...\n" yield None, code, log, None, None final_step_path = None final_stl_path = None step_idx = 0 for update in run_cadquery_progressive(code): viewer_value = update["viewer"] step_file_value = update["step_file"] if update["done"]: log += "▶ Final model completed.\n" else: step_idx += 1 if step_idx <= len(plan_steps): log += f"▶ Step {step_idx}: {plan_steps[step_idx - 1]}\n" else: log += f"▶ Executing build snapshot {step_idx}\n" yield viewer_value, code, log, step_file_value, None if update["done"]: final_step_path = step_file_value final_stl_path = viewer_value time.sleep(5) review = "無し" yield final_stl_path, code, log, final_step_path, review # ========================= # UI # ========================= with gr.Blocks(delete_cache=(86400, 86400)) as demo: gr.Markdown("## 3D CAD AIエージェント🛠") with gr.Row(): with gr.Column(scale=1): claude_model = gr.Radio( choices=CLAUDE_MODELS, value=DEFAULT_CLAUDE_MODEL, label="Claude model for CAD code generation", visible=False ) openai_model = gr.Radio( choices=OPENAI_MODELS, value=DEFAULT_OPENAI_MODEL, label="OpenAI model for STEP review", visible=False ) templates = { "自由入力": "", "ボルト": "A bolt with a hexagonal socket head and a cylindrical shaft", "筐体": "A 60mm x 40mm x 10mm enclosure with 4 corner holes", "修正": "修正箇所：" } text_prompt = gr.Textbox( label="Text CAD specification\n仕様を記載してください", lines=3, placeholder="自由入力、またはテンプレートを選択" ) template = gr.Dropdown( choices=list(templates.keys()), value="自由入力", label="テンプレート" ) template.change( lambda x: templates[x], inputs=template, outputs=text_prompt, ) with gr.Accordion("", open=False): image_input = gr.Image( label="2D drawing (optional)\n2D図面をアップロードすることも可能です", type="filepath" ) step_input = gr.File( label="STEPファイルをアップロード", file_types=[".step", ".stp"], type="filepath" ) run_btn = gr.Button("Generate CAD") with gr.Column(scale=4): viewer = gr.Model3D(label="CAD Viewer (STL)", height=560, camera_position=(20, 30, 100)) with gr.Column(scale=1): status_box = gr.Textbox(label="🤖", lines=6, autoscroll=True) review_box = gr.Textbox(label="Model Check", lines=8) with gr.Accordion("コード", open=False): code_box = gr.Code(label="Generated CAD Code", language="python") step_file = gr.File(label="Download STEP") edit_button = gr.Button("編集する") edit_message = gr.Markdown("") def use_output_as_input(step_output_path): if step_output_path is None: return None src = Path(step_output_path) output_dir = Path(tempfile.mkdtemp()) copied_path = output_dir / src.name shutil.copy(src, copied_path) return str(copied_path), "✅ 編集可能になりました" edit_button.click( fn=use_output_as_input, inputs=[step_file], outputs=[step_input, edit_message] ) run_btn.click( fn=pipeline, inputs=[text_prompt, image_input, step_input, claude_model, openai_model], outputs=[viewer, code_box, status_box, step_file, review_box] ) demo.launch( allowed_paths=[str(ARTIFACT_DIR), str(GRADIO_TEMP_DIR)] )