cad-Integration/CADDesigner-Code-main/tools/requirements_tools.py

"""Requirement refinement tool implemented as a specialist subagent."""

from __future__ import annotations

import base64
from pathlib import Path
from typing import Any, Optional, Union

from SimpleLLMFunc import llm_chat, tool
from SimpleLLMFunc.type import HistoryList

from context.conversation_manager import get_current_sketch_pad
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 .reference_image import resolve_reference_image_path
from .sketch_tools import sketch_pad_operations
from .subagent_utils import run_subagent_with_events


def create_requirement_refinement_subagent_tools(
    workspace: Optional[str | Path] = None,
) -> list[Any]:
    """Create the low-level tools owned by the requirement specialist.
    Specialist only produces formatted text; SketchPad storage is done by the caller.
    """
    return [
        execute_command,
        sketch_pad_operations,
        *create_builtin_file_tools(workspace),
    ]


_REQUIRED_SECTIONS = [
    "## API Reference",
    "## Refined User Requirements",
    "## Parameter Table",
    "## Modeling Process",
    "## Notes",
]


def _normalize_requirement_output(text: str) -> str:
    """Trim preamble and ensure all required sections exist."""
    t = text.strip()
    for h in _REQUIRED_SECTIONS:
        idx = t.find(h)
        if idx >= 0:
            t = t[idx:]
            break
    for h in _REQUIRED_SECTIONS:
        if h not in t:
            t += f"\n\n{h}\n"
    return t.strip()


def _build_requirement_request(
    *,
    query: str,
    query_image_path: Optional[str],
) -> str:
    parts = [
        "Generate a detailed modeling specification. Output must include: ## API Reference, ## Refined User Requirements, ## Parameter Table, ## Modeling Process, ## Notes.",
        "Use the workspace facts below.",
        "Before you write the final answer, you MUST use file tools to read the preferred skill root's `SKILL.md`, then `references/docs/api/README.md`, then the exact API Markdown pages you cite.",
        "Do not answer from memory. If you have not read those files yet, continue using tools.",
        "In `## API Reference`, cite the concrete file paths you read and only recommend APIs whose exact Markdown pages you actually opened.",
        "If the task mentions SketchPad keys, use `sketch_pad_operations` to retrieve them before refining the requirement.",
        "",
        build_simplecad_workspace_fact_block(),
        "",
        "[User Query]",
        query.strip(),
    ]
    if query_image_path and query_image_path.strip():
        parts.append("\n[Reference image attached below]")
    return "\n".join(parts)


def _image_path_to_base64_data_url(image_path: str) -> Optional[str]:
    """Read image file and return data URL for OpenAI API."""
    p = Path(image_path)
    if not p.exists() or not p.is_file():
        return None
    ext = p.suffix.lower()
    mime_map = {
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".png": "image/png",
        ".gif": "image/gif",
        ".webp": "image/webp",
    }
    mime = mime_map.get(ext, "image/jpeg")
    try:
        b64 = base64.b64encode(p.read_bytes()).decode("utf-8")
        return f"data:{mime};base64,{b64}"
    except Exception:
        return None


def _build_message_with_image(
    text: str,
    query_image_path: Optional[str],
) -> Union[str, list[dict[str, Any]]]:
    """Build message: text only, or text + image as OpenAI content array."""
    if not query_image_path or not query_image_path.strip():
        return text
    data_url = _image_path_to_base64_data_url(query_image_path.strip())
    if not data_url:
        raise RuntimeError(
            f"Failed to load reference image: {query_image_path.strip()}"
        )
    return [
        {"type": "text", "text": text},
        {"type": "image_url", "image_url": {"url": data_url}},
    ]


@tool(
    name="make_user_query_more_detailed",
    description=(
        "Refine and expand the user's modeling requirement through a specialist subagent. "
        "The specialist can inspect local skill docs, inspect APIs, read local files, and consult SketchPad "
        "before producing a structured modeling specification."
    ),
    best_practices=[
        "Pass the complete user request in `query`, not only a short delta fragment.",
        "If the user provided a reference image, pass its workspace-local path in `query_image_path`.",
        "If `query_image_path` is omitted, the tool will automatically reuse the latest uploaded image from the active conversation when available.",
        "Use this tool when the modeling request is vague, underspecified, or needs a step-by-step plan before coding.",
        "The specialist will read local skill docs directly with its file tools to ground the refinement.",
        "The final result should include a structured modeling process, not only rewritten prose.",
    ],
)
async def make_user_query_more_detailed(
    query: str,
    query_image_path: Optional[str] = None,
    event_emitter: Any = None,
) -> str:
    """Refine the user's modeling request via a requirement specialist subagent.

    The refined requirement is always stored in SketchPad for downstream tools to reference.

    Args:
        query: The user's original request. This may also mention SketchPad ids that the
            specialist should inspect.
        query_image_path: Optional workspace-local reference image path, typically something
            like `./uploads/<conversation_id>/query_image_001.png`.
        event_emitter: Optional tool event emitter used to forward nested specialist activity.

    Returns:
        str: Refined requirement text with SketchPad key for reference.
    """

    print_tool_output(
        title="Requirement Refinement Started",
        content=f"Request: {query}",
    )

    requested_query_image_path = (
        query_image_path.strip()
        if isinstance(query_image_path, str) and query_image_path.strip()
        else None
    )
    resolved_query_image_path = resolve_reference_image_path(query_image_path)

    if requested_query_image_path is not None and resolved_query_image_path is None:
        raise RuntimeError(f"Reference image not found: {requested_query_image_path}")

    if resolved_query_image_path is not None:
        print_tool_output(
            title="Reference Image Attached",
            content=f"Using reference image: {resolved_query_image_path}",
        )

    text_content = _build_requirement_request(
        query=query,
        query_image_path=resolved_query_image_path,
    )
    message = _build_message_with_image(text_content, resolved_query_image_path)

    result_text = await run_subagent_with_events(
        specialist_callable=requirement_refinement_specialist,
        specialist_kwargs={
            "message": message,
            "history": [],
        },
        subagent_label="Requirement Refinement Specialist",
        event_emitter=event_emitter,
        status_payload={
            "query": query,
            "query_image_path": resolved_query_image_path,
        },
        response_transform=_normalize_requirement_output,
    )

    final_text = result_text.strip()
    print_tool_output(title="Refined User Requirements", content=final_text)

    sketch_pad = get_current_sketch_pad()
    if sketch_pad is None:
        raise RuntimeError(
            "The refined requirement must be written to SketchPad, but there is no active conversation context."
        )

    import uuid

    sketch_key = f"req_{uuid.uuid4().hex[:8]}"
    try:
        await sketch_pad.set_item(
            key=sketch_key,
            value=final_text,
            ttl=None,
            summary=None,
            tags={"detailed_query", "requirements", "expanded"},
        )

        print_tool_output(
            title="💾 Stored In SketchPad",
            content=f"Key: {sketch_key}\nThe refined requirement has been saved for downstream tools.",
        )

        return (
            "Detailed requirements generated and stored in SketchPad:\n\n"
            f"🔑 SketchPad Key: {sketch_key}\n"
            "# Tags: detailed_query, requirements, expanded\n"
            f'💡 Tip: You can now reference key "{sketch_key}" in later tool calls, for example:\n'
            "- include it in the natural-language task for `cad_code_generator`\n"
            "- store it alongside other constraints or debugging notes in SketchPad\n"
            "- create the target folder first, then use `echo_into` to write a file if needed\n"
        )
    except Exception as exc:
        print_tool_output(
            "❌ SketchPad Store Failed", f"Failed to store in SketchPad: {exc}"
        )
        raise RuntimeError(
            f"The refined requirement must be written to SketchPad, but storage failed: {exc}"
        ) from exc


@llm_chat(
    llm_interface=get_config().MULTIMODALITY_INTERFACE,
    toolkit=create_requirement_refinement_subagent_tools(),
    max_tool_calls=SUBAGENT_MAX_TOOL_CALLS,
    stream=True,
    enable_event=True,
    timeout=600,
    temperature=1.0,
)
async def requirement_refinement_specialist(
    message: Union[str, list[dict[str, Any]]],
    history: HistoryList | None = None,
) -> None:  # type: ignore[misc]
    """Generate a detailed modeling specification. Output: ## API Reference, ## Refined User Requirements, ## Parameter Table, ## Modeling Process, ## Notes.

    Use the workspace facts included in the user message. Read `SKILL.md`, then the API index, then the exact API Markdown pages you cite.
    REQUIRED: The detailed query MUST use exactly correct API names and code snippets.
    Tools: execute_command, sketch_pad_operations, read_file, grep, sed, echo_into.
    You MUST read SKILL.md and the API index before choosing APIs. Retrieve SketchPad artifacts when task mentions keys.

    <EXAMPLE>
    User: "Create a 7.62mm rifle cartridge model"

    ## Refined User Requirements
    1. **Target Object**: A standard 7.62mm caliber rifle cartridge (Full Metal Jacket type).
    2. **Components**: The model consists of four parts: the bullet tip (projectile), the cartridge case (neck, shoulder, body), the rim/extractor groove, and a primer base.
    3. **Dimensions**:
        - **Projectile**: Diameter 7.62mm, ogive shape with a rounded tip.
        - **Case Body**: Maximum diameter approx 11.3mm, total case length 51mm (based on 7.62x51mm NATO standard).
        - **Shoulder/Neck**: Tapered transition from body to 7.62mm neck.
    4. **Output**: A single combined solid representing the exterior geometry of the cartridge.

    ## Parameter Table
    | Parameter | Type | Default Value | Calculation Logic |
    |---|---|---|---|
    | bullet_dia | float | 7.62 | Nominal caliber |
    | case_body_dia | float | 11.3 | Max diameter of the case body |
    | total_length | float | 71.0 | Full cartridge length including projectile |
    | body_length | float | 38.0 | Length from base to shoulder |
    | shoulder_length | float | 3.5 | Length of the tapered shoulder |
    | neck_length | float | 8.0 | Length of the neck holding the bullet |
    | case_length | float | 51.0 | body_length + shoulder_length + neck_length |
    | rim_dia | float | 11.5 | Diameter of the base rim |

    ## Modeling Process
    1. **Create Case Main Body**
        - **Purpose**: Create the main cylindrical propellant chamber.
        - **API**: `make_cylinder_rsolid`
        - **Spatial Reasoning**: Cylinder radius `case_body_dia/2`, height `body_length`, base at (0,0,0).

    2. **Create Shoulder and Neck**
        - **Purpose**: Model the tapered transition and casing neck.
        - **API**: `make_cone_rsolid`, `make_cylinder_rsolid`, `translate_shape`
        - **Spatial Reasoning**: Shoulder cone bottom radius `case_body_dia/2`, top `bullet_dia/2`, height `shoulder_length`, translate to Z=body_length. Neck cylinder radius `bullet_dia/2`, height `neck_length`, translate to Z=body_length+shoulder_length.

    3. **Create Projectile**
        - **Purpose**: Form the aerodynamic tip.
        - **API**: `make_cone_rsolid`, `union_rsolidlist`
        - **Spatial Reasoning**: Cone base radius `bullet_dia/2`, height `total_length-case_length`, translate to Z=case_length.

    4. **Add Extractor Groove and Rim**
        - **Purpose**: Model the base where extractor grips.
        - **API**: `make_cylinder_rsolid`, `cut_rsolidlist`
        - **Spatial Reasoning**: Rim cylinder radius `rim_dia/2`, height 1.5. Cut groove with smaller cylinder.

    5. **Final Assembly**
        - **Purpose**: Combine into single manifold solid.
        - **API**: `union_rsolidlist`
        - **Spatial Reasoning**: Boolean union on body, shoulder, neck, projectile, rim.

    ## Notes
    Based on 7.62x51mm NATO standard. APIs must be verified against SKILL.md.
    </EXAMPLE>
    """

    pass


__all__ = [
    "make_user_query_more_detailed",
    "requirement_refinement_specialist",
    "create_requirement_refinement_subagent_tools",
]