"""The four specialized agents. Each = system prompt + strict-schema call. No tool calling anywhere (it can't combine with image inputs on Gemma 4) — every agent exchanges validated JSON via structured outputs. """ from __future__ import annotations import json from . import llm from .schemas import ( BestSelection, CritiqueVerdict, FabPlan, GeneratedArtifact, VisionSpec, response_format, ) VISION_SYS = ( "You are a precision vision-analysis agent for digital fabrication. Examine the " "image and produce a rigorous physical specification of the PRIMARY object. " "Estimate real-world dimensions in millimetres using visible scale cues. Note " "materials, salient features, and any defects. Be concrete and quantitative." ) PLANNER_SYS = ( "You are a fabrication planner. Given an object spec, produce a constructive plan that " "reproduces the object. Choose the BEST construction strategy and say which in `fab_method`/" "`notes`:\n" "• Rotationally-symmetric objects (mug, bottle, vase, bowl, cup, lamp, wheel) → a SURFACE OF " "REVOLUTION: describe the radius/height profile.\n" "• Prismatic objects with a constant cross-section (gear, bracket, sign, key) → an EXTRUDED 2D " "profile: describe the cross-section outline.\n" "• Otherwise → a composition of PARAMETRIC PRIMITIVES (box, cylinder, sphere, torus) with " "boolean add/subtract.\n" "Keep it simple and physically buildable. Dimensions in millimetres. box dims=[x,y,z]; " "cylinder=[radius,height]; sphere=[radius]; torus=[major_R,minor_r]. The `primitives` list is " "for the primitive strategy; for revolve/extrude, describe the profile in `steps`/`notes`." ) GENERATOR_SYS = ( "You are a 3D code generator. Output Python that builds the planned object with " "`trimesh` and `np` (numpy), assigning the final mesh to a variable named `result`. " "Use trimesh.creation.box(extents=[x,y,z]) / cylinder(radius,height,sections=64) / " "icosphere(radius=r) / torus(major_radius=R, minor_radius=r), .apply_translation([x,y,z]), " "and trimesh.boolean.union/difference([...]). Units are millimetres. " "To ROTATE a mesh use mesh.apply_transform(trimesh.transformations.rotation_matrix(angle_rad, " "[x,y,z])) — there is NO apply_rotation method.\n" "RICHER BUILDERS for non-primitive shapes (prefer these when they fit):\n" "• Surface of revolution (mug/bottle/vase/bowl/cup/lamp): build a profile of [radius, height] " "points and revolve it — profile=np.array([[r0,h0],[r1,h1],...]); result=trimesh.creation." "revolve(profile, sections=64). The argument is a list of [radius,height] points (radius first); " "omit `angle` for a full 360° solid; close the profile (radius 0 at top/bottom) for watertight.\n" "• Extruded 2D cross-section (gear/bracket/sign/key): from shapely.geometry import Polygon; " "result=trimesh.creation.extrude_polygon(Polygon([(x,y),...]), height=H).\n" "• Also available: trimesh.creation.cone(radius,height), capsule(height,radius), " "annulus(r_min,r_max,height), sweep_polygon(Polygon([...]), path_points).\n" "Allowed imports: trimesh, numpy as np, math, shapely. No file I/O, no printing — only build " "`result`. Keep it watertight.\n" "CODE STYLE: write minimal code. Use FEW comments and keep every comment on ONE line starting " "with '#'. NEVER wrap a comment across two lines (a continuation line without '#' is a syntax error)." ) CRITIC_SYS = ( "You are a fabrication critic. Compare the produced mesh statistics against the " "object spec and plan. Judge geometric fidelity and printability. Approve only if " "the mesh is watertight and reasonably matches the object; otherwise give concrete " "fix instructions for the generator." ) VISUAL_CRITIC_SYS = ( "You are a visual fabrication critic with eyes. You are shown the ORIGINAL object " "photo(s) and a multi-view RENDER of the candidate 3D model the system generated. " "Compare them directly. Do NOT give a vibe check — produce a concrete, checkable diff: " "for each discrepancy in overall shape, proportions, COUNT of features (holes, handles, " "legs, ribs), presence/absence of parts, and relative sizes, state what is wrong and " "which generator change fixes it (e.g. 'handle too thick — reduce torus minor_radius', " "'missing the spout', 'body should be ~30% taller'). Approve only when the render's " "silhouette and feature set clearly match the photo." ) def _parse(text: str, model): return model.model_validate_json(text) async def vision_agent(image_data_uris, goal: str, view_labels: list[str] | None = None): """Analyze 1..5 images of the SAME object. Multiple views (front/side/top) sharply improve depth/proportion estimates — use the side view for depth, top for footprint.""" if isinstance(image_data_uris, str): image_data_uris = [image_data_uris] multi = len(image_data_uris) > 1 prompt = f"Analyze this object for the goal: {goal}." if multi: prompt += (" You are given multiple views of the SAME object. Cross-reference them: " "use the side view to judge depth/thickness and the top view for the footprint. " "Reconcile the views into one consistent specification.") msgs = [ {"role": "system", "content": VISION_SYS}, {"role": "user", "content": llm.multi_image_content(prompt, image_data_uris, view_labels)}, ] text, meta = await llm.acall(msgs, schema=response_format(VisionSpec), max_tokens=1500) return _parse(text, VisionSpec), meta async def visual_critic_agent(spec, original_uris, render_uris, mesh_stats: dict): """Multimodal critic: SEE the original photo(s) + renders of the candidate mesh (shaded + normal-map) and produce a concrete diff. The main fidelity lever (render→VLM→fix loop). The normal-map view exposes curvature/flatness errors flat shading hides.""" if isinstance(original_uris, str): original_uris = [original_uris] if isinstance(render_uris, str): render_uris = [render_uris] render_labels = ["YOUR MODEL render — shaded (4 views)", "YOUR MODEL render — normal map (4 views)"] uris = list(original_uris) + list(render_uris) labels = ([f"ORIGINAL photo {i + 1}" for i in range(len(original_uris))] + render_labels[:len(render_uris)]) prompt = ( f"Target object: {spec.object}. Mesh stats: {json.dumps(mesh_stats)}.\n" "Compare the ORIGINAL photo(s) to YOUR model renders. For each discrepancy, name the " "generator change that fixes it in the form 'problem → primitive/param + direction' " "(e.g. 'handle too thick → reduce torus minor_radius'). Cover shape, proportions, " "feature counts, missing/extra parts, and relative sizes." ) msgs = [ {"role": "system", "content": VISUAL_CRITIC_SYS}, {"role": "user", "content": llm.multi_image_content(prompt, uris, labels)}, ] # NOTE: structured output + images is fine; tool-calling + images is NOT (Gemma 4). text, meta = await llm.acall(msgs, schema=response_format(CritiqueVerdict), max_tokens=1500) return _parse(text, CritiqueVerdict), meta async def planner_agent(spec: VisionSpec, goal: str): msgs = [ {"role": "system", "content": PLANNER_SYS}, {"role": "user", "content": f"Goal: {goal}\n\nObject spec:\n{spec.model_dump_json(indent=2)}"}, ] # NOTE: reasoning_effort is intentionally OFF here. For Gemma 4 the levels are # equivalent, and enabling it destabilizes structured (json_schema) output — # reasoning tokens can crowd out the JSON and return empty content. text, meta = await llm.acall( msgs, schema=response_format(FabPlan), max_tokens=3000) return _parse(text, FabPlan), meta async def generator_agent(plan: FabPlan, spec: VisionSpec, feedback: str | None = None): user = f"Object: {spec.object}\n\nPlan:\n{plan.model_dump_json(indent=2)}" if feedback: user += f"\n\nThe previous attempt failed. Fix it:\n{feedback}" msgs = [ {"role": "system", "content": GENERATOR_SYS}, {"role": "user", "content": user}, ] text, meta = await llm.acall( msgs, schema=response_format(GeneratedArtifact), max_tokens=3000) return _parse(text, GeneratedArtifact), meta SELECTOR_SYS = ( "You select the best 3D-model candidate. You see the ORIGINAL object photo and several " "candidate renders. Pick the 0-based index whose shape, proportions, and feature set best " "match the photo." ) async def select_best_agent(spec, original_uris, candidate_render_uris, hints=None): """One multimodal call: pick the best candidate render index vs the photo (best-of-N). `hints` (optional) adds a numeric note per candidate, e.g. silhouette-IoU scores.""" if isinstance(original_uris, str): original_uris = [original_uris] uris = list(original_uris) + list(candidate_render_uris) labels = (["ORIGINAL photo"] * len(original_uris) + [f"Candidate {i}" + (f" — {hints[i]}" if hints and i < len(hints) else "") for i in range(len(candidate_render_uris))]) prompt = (f"Target object: {spec.object}. Choose the 0-based index of the candidate that best " "matches the photo in shape, proportions, and features.") msgs = [ {"role": "system", "content": SELECTOR_SYS}, {"role": "user", "content": llm.multi_image_content(prompt, uris, labels)}, ] text, meta = await llm.acall(msgs, schema=response_format(BestSelection), max_tokens=500) return _parse(text, BestSelection), meta async def critic_agent(spec: VisionSpec, plan: FabPlan, mesh_stats: dict): msgs = [ {"role": "system", "content": CRITIC_SYS}, {"role": "user", "content": f"Object spec:\n{spec.model_dump_json(indent=2)}\n\n" f"Plan notes: {plan.notes}\n\n" f"Produced mesh stats:\n{json.dumps(mesh_stats, indent=2)}"}, ] text, meta = await llm.acall(msgs, schema=response_format(CritiqueVerdict), max_tokens=1200) return _parse(text, CritiqueVerdict), meta