agent_core/agent.py

from agent_core.config import DEFAULT_CADQUERY_OUTPUT_PATH, get_client, get_model_id
from agent_core.prompts import SYSTEM
from agent_core.tools.cadquery_tool import TOOL_SCHEMA as CADQUERY_TOOL_SCHEMA
from agent_core.tools.cadquery_tool import run_execute_cadquery
from agent_core.tools.lux3d_tool import TOOL_SCHEMA as LUX3D_TOOL_SCHEMA
from agent_core.tools.lux3d_tool import run_generate_3d_model


TOOL_HANDLERS = {
    "execute_cadquery": lambda **kw: run_execute_cadquery(
        kw["code"],
        kw.get("output_path", DEFAULT_CADQUERY_OUTPUT_PATH),
    ),
    "generate_3d_model": lambda **kw: run_generate_3d_model(
        kw["image_path"],
        kw.get("output_path"),
    ),
}

TOOLS = [
    CADQUERY_TOOL_SCHEMA,
    LUX3D_TOOL_SCHEMA,
]


def latest_user_prompt(messages: list) -> str | None:
    for message in reversed(messages):
        if message.get("role") == "user" and isinstance(message.get("content"), str):
            return message["content"]
    return None


def agent_loop(messages: list):
    client = get_client()
    model = get_model_id()

    while True:
        response = client.messages.create(
            model=model,
            system=SYSTEM,
            messages=messages,
            tools=TOOLS,
            max_tokens=8000,
        )
        messages.append({"role": "assistant", "content": response.content})
        if response.stop_reason != "tool_use":
            return

        results = []
        prompt = latest_user_prompt(messages)
        for block in response.content:
            if block.type == "tool_use":
                if block.name == "execute_cadquery":
                    output = run_execute_cadquery(
                        code=block.input["code"],
                        output_path=block.input.get("output_path", DEFAULT_CADQUERY_OUTPUT_PATH),
                        prompt=prompt,
                    )
                elif block.name == "generate_3d_model":
                    output = run_generate_3d_model(
                        image_path=block.input["image_path"],
                        output_path=block.input.get("output_path"),
                        prompt=prompt,
                    )
                else:
                    handler = TOOL_HANDLERS.get(block.name)
                    output = handler(**block.input) if handler else f"Unknown tool: {block.name}"

                print(f"> {block.name}: {output[:200]}")
                results.append({"type": "tool_result", "tool_use_id": block.id, "content": output})
        messages.append({"role": "user", "content": results})