agent_runner.py

"""
agent_runner.py

Main agent loop. Claude proposes design parameters, calls
check_constraints / estimate_resources / run_simulation / query_history,
and iterates toward the optimization goal within a fixed simulation budget.

Requires:
    pip install anthropic
    export ANTHROPIC_API_KEY=...

Run from inside agent_pipeline/ (or with this directory + your HCres.py /
SQDMetal install on PYTHONPATH):

    python agent_runner.py
"""

import json

import anthropic

import config
from agent_tools import TOOLS, dispatch


# Sonnet is a reasonable default for cost; switch to an Opus model if the
# agent seems to be making weak strategic choices over a long run.
MODEL = "claude-sonnet-4-6"
MAX_SIMULATIONS = 20
MAX_TURNS = 200  # hard safety cap on total API calls, independent of MAX_SIMULATIONS


SYSTEM_PROMPT = f"""You are optimizing the design of a high-capacitance, \
low-geometric-inductance superconducting resonator ("HCres") intended for \
use as a flux-tunable coupler.

You can adjust 11 geometric design parameters, all in SI units (meters, or \
a dimensionless count for fingers_n), preferably adjusting the metal_width, geo_ind_thick, geo_ind_length, res_height, cap_section_height, and finger_length before trying to change other parameters:

{json.dumps(config.PARAM_BOUNDS, indent=2)}

For reference, here is a previously-tested point in this space:
{json.dumps(config.DEMO_PARAMS, indent=2)}

Goals, in priority order:
1. Maximize coupling_max (the peak coupling coefficient over the spatial
   sweep), returned by run_simulation.
2. Bring impedance, which in this code is saved as Z,  as close as possible to 50 ohms.
3. Stay within the fabrication / geometric feasibility bounds enforced by
   check_constraints.
4. Never let a simulation crash the machine -- always call
   estimate_resources before run_simulation and pass its recommended_cores
   as num_cpus.

Workflow for every proposed design:
1. check_constraints(params). If violations are returned, fix the
   parameters and re-check -- do not spend a simulation on an infeasible
   design.
2. estimate_resources(params). Use the returned recommended_cores as
   num_cpus in the next step.
3. run_simulation(params, num_cpus=...). This is expensive (three Palace
   solves) -- you have a budget of {MAX_SIMULATIONS} calls to this tool for
   the whole session.
4. Briefly note what you learned (which parameters seem to move
   coupling_max and impedance, and in which direction) before proposing the
   next design.

Call query_history early (it may be empty on the first run) and again
periodically to ground your reasoning in what's actually been tried,
instead of re-deriving everything from scratch each turn.

When your run_simulation budget is exhausted, stop calling tools and write
a final summary: the best design found (its parameters, coupling_max, and
impedance), how it compares to the reference point above, and what you'd
try next with a larger budget.
"""


def run_agent():
    client = anthropic.Anthropic()
    messages = [{
        "role": "user",
        "content": (
            "Begin the design optimization. Start with query_history (it "
            "may be empty), then propose, check, and evaluate your first "
            "design."
        ),
    }]

    n_simulations = 0
    response = None

    for _turn in range(MAX_TURNS):
        response = client.messages.create(
            model=MODEL,
            max_tokens=4096,
            system=SYSTEM_PROMPT,
            tools=TOOLS,
            messages=messages,
        )
        messages.append({"role": "assistant", "content": response.content})

        for block in response.content:
            if block.type == "text" and block.text.strip():
                print(f"\n--- agent ---\n{block.text}\n")

        if response.stop_reason != "tool_use":
            if n_simulations >= MAX_SIMULATIONS:
                break
            messages.append({
                "role": "user",
                "content": "Continue: propose, check, and evaluate your next design.",
            })
            continue

        tool_results = []
        for block in response.content:
            if block.type != "tool_use":
                continue

            if block.name == "run_simulation":
                n_simulations += 1
                print(f"[run_simulation {n_simulations}/{MAX_SIMULATIONS}] {block.input}")

            result = dispatch(block.name, block.input)

            if block.name == "run_simulation":
                print(f"  -> status={result.get('status')} "
                      f"coupling_max={result.get('coupling_max')} "
                      f"impedance={result.get('Z')}")

            tool_results.append({
                "type": "tool_result",
                "tool_use_id": block.id,
                "content": json.dumps(result, default=str),
            })

        messages.append({"role": "user", "content": tool_results})

        if n_simulations >= MAX_SIMULATIONS:
            messages.append({
                "role": "user",
                "content": (
                    f"Your simulation budget of {MAX_SIMULATIONS} is now "
                    "exhausted. Do not call run_simulation again. Write your "
                    "final summary."
                ),
            })
            response = client.messages.create(
                model=MODEL, max_tokens=2048, system=SYSTEM_PROMPT,
                tools=TOOLS, messages=messages,
            )
            for block in response.content:
                if block.type == "text":
                    print(f"\n--- final summary ---\n{block.text}\n")
            break

    return messages


if __name__ == "__main__":
    run_agent()