"""
pycsamt.agents.code_gen
========================
:class:`CodeGenerationAgent` — Emit a reproducible Python script from a
completed workflow configuration and agent results.
The generated script imports pycsamt, replicates every processing step that
was executed, and is self-contained (no agent infrastructure needed to run
it). It is the "save your workflow as code" button.
"""
from __future__ import annotations
import json
import os
import re
import time
from datetime import datetime
from pathlib import Path
from typing import Any
from ._base import AgentResult, BaseAgent
from ._generation import (
GenerationInput,
clarification_for,
exact_artifact_edit,
excludes_correction,
forward_parameters,
)
#: Folder used when no ``output_dir`` is given. Matches the Agent Master
#: default so generated scripts never land in the current working
#: directory (e.g. a repository root).
DEFAULT_OUTPUT_DIR = "pycsamt_agent_output"
_SYSTEM_PROMPT = """\
You are an expert Python developer specialising in geophysics scripting.
Generate or edit a clean Python script for the user's complete request.
Workflow configurations and templates are references, not fixed plans.
Preserve requested selections, quantities, units, filenames, formats, and
exclusions. For an edit, preserve the prior script except for requested changes.
Use pycsamt v2 public API only. Add a one-line comment above each major
block. Do not use the agents/ subpackage — call emtools, forward, and
site functions directly. Never invent missing APIs or scientific parameters.
Output only valid Python code, or CLARIFY: and one focused question when
required information is missing. Do not claim to have executed the script.
"""
# ── static script template fragments ─────────────────────────────────────────
_HEADER = """\
#!/usr/bin/env python
# -*- coding: utf-8 -*-
\"\"\"
{title}
Generated by pycsamt CodeGenerationAgent — {date}
This script reproduces the processing workflow automatically.
\"\"\"
from __future__ import annotations
import numpy as np
import matplotlib.pyplot as plt
"""
_LOAD_BLOCK = """\
# ── load MT data ──────────────────────────────────────────────
from pycsamt.emtools._core import ensure_sites
sites = ensure_sites({path!r}, verbose=0)
print(f"Loaded {{len(list(sites))}} stations")
"""
_QC_BLOCK = """\
# ── quality control ───────────────────────────────────────────
from pycsamt.emtools.qc import (
build_qc_table,
qc_flags,
plot_frequency_confidence_psection,
)
qc_table = build_qc_table(sites)
flags = qc_flags(sites, min_frac_ok=0.6, min_snr_med=2.0)
fig_qc = (
plot_frequency_confidence_psection(sites).get_figure()
)
fig_qc.savefig(
{out!r} + "/qc_confidence.png",
dpi=150, bbox_inches="tight",
)
"""
_SS_ESTIMATE_BLOCK = """\
# Estimate and export factors; do not apply correction.
from pycsamt.emtools.ss import estimate_ss_ama
ss_table = estimate_ss_ama(sites)
if ss_table.empty:
print("No applicable static-shift factors were estimated.")
else:
ss_table.to_csv({out!r} + "/static_shift_factors.csv", index=False)
sites_corr = sites
"""
_SS_BLOCK = """\
# ── static-shift correction ───────────────────────────────────
from pycsamt.emtools.ss import (
estimate_ss_ama,
correct_ss_ama,
plot_ss_summary,
plot_ss_1d_curves,
)
from pycsamt.emtools._core import (
_get_z_block, _name, _iter_items,
)
ss_table = estimate_ss_ama(
sites,
half_window={hw},
sort_by="lon",
weights="tri",
max_skew=6.0,
)
print(
ss_table[
["station", "delta_log10_rho", "fac_z"]
]
)
sites_corr = correct_ss_ama(
sites,
half_window={hw},
sort_by="lon",
)
def _collect_logRho(S):
rows, freqs = [], None
for i, ed in enumerate(_iter_items(S)):
Z, z, fr = _get_z_block(ed)
if Z is None:
continue
rxy = (
0.2 * np.abs(z[:, 0, 1]) ** 2
/ (fr + 1e-24)
)
ryx = (
0.2 * np.abs(z[:, 1, 0]) ** 2
/ (fr + 1e-24)
)
rows.append(
np.log10(np.sqrt(rxy * ryx) + 1e-24)
)
freqs = fr
return np.array(rows), freqs
logRho_b, ss_freqs = _collect_logRho(sites)
logRho_a, _ = _collect_logRho(sites_corr)
ss_labels = [
_name(ed, i)
for i, ed in enumerate(_iter_items(sites))
]
fig_sum = plot_ss_summary(
logRho_b, logRho_a,
freqs=ss_freqs,
station_labels=ss_labels,
)
fig_sum.savefig(
{out!r} + "/ss_summary.png",
dpi=150, bbox_inches="tight",
)
fig_1d = plot_ss_1d_curves(
logRho_b, logRho_a,
freqs=ss_freqs,
station_labels=ss_labels,
)
fig_1d.savefig(
{out!r} + "/ss_1d_curves.png",
dpi=150, bbox_inches="tight",
)
"""
_PT_BLOCK = """\
# ── phase tensor analysis ─────────────────────────────────────
from pycsamt.emtools.tensor import (
build_phase_tensor_table,
plot_phase_tensor_psection,
)
from pycsamt.emtools.strike import (
estimate_strike_consensus,
plot_strike_analysis,
)
from pycsamt.emtools.dimensionality import classify_dimensionality
pt_table = build_phase_tensor_table(sites_corr)
dim_table = classify_dimensionality(
sites_corr, skew_th={skew_th}, ellipt_th={ellipt_th},
)
st_result = estimate_strike_consensus(sites_corr)
strike_angle = st_result.get("angle_deg")
try:
strike_text = f"{{float(strike_angle):.1f}} deg"
except (TypeError, ValueError):
strike_text = "n/a"
print(f"Consensus strike: {{strike_text}}")
fig_pt = plot_phase_tensor_psection(sites_corr).get_figure()
fig_pt.savefig(
{out!r} + "/pt_psection.png",
dpi=150, bbox_inches="tight",
)
fig_strike = plot_strike_analysis(sites_corr)
fig_strike.savefig(
{out!r} + "/strike_analysis.png",
dpi=150, bbox_inches="tight",
)
"""
_FWD_BLOCK = """\
# ── 1-D MT forward model ──────────────────────────────────────
from pycsamt.forward import (
MT1DForward,
LayeredModel,
plot_response_and_model_1d,
)
layered = LayeredModel(
resistivity={rhos},
thickness={ths},
)
freqs = {freqs}
fwd = MT1DForward(freqs=freqs)
resp = fwd.run(layered)
fig_fwd = plot_response_and_model_1d(resp, layered)
fig_fwd.savefig(
{out!r} + "/forward_response.png",
dpi=150, bbox_inches="tight",
)
"""
_PRE_INV_BLOCK = """\
# Occam2D pre-inversion export
from pycsamt.models.occam2d import InputBuilder, OccamConfig
occam_workdir = os.path.join({out!r}, "occam2d_inputs")
os.makedirs(occam_workdir, exist_ok=True)
occam_config = OccamConfig()
occam_config.to_template(os.path.join(occam_workdir, "occam2d.yml"))
occam_builder = InputBuilder(
sites_corr,
workdir=occam_workdir,
config=occam_config,
)
# Expected Occam2D artifacts: OccamDataFile.dat, Occam2DMesh,
# Occam2DModel, and OccamStartup. The builder records the public API
# used for pre-inversion preparation.
print(f"Occam2D pre-inversion configuration written to {{occam_workdir}}")
"""
_TIPPER_BLOCK = """\
# Tipper and induction-arrow analysis
from pycsamt.emtools.tf import (
plot_induction_arrows,
plot_tipper_hodograms,
plot_tipper_polar,
)
try:
ax_tip = plot_induction_arrows(
sites_corr,
periods=({period},),
convention={convention!r},
)
ax_tip.get_figure().savefig(
{out!r} + "/induction_arrows.png",
dpi=150, bbox_inches="tight",
)
except Exception as exc:
print(f"No plottable induction arrows: {{exc}}")
try:
ax_hodo = plot_tipper_hodograms(sites_corr)
ax_hodo.get_figure().savefig(
{out!r} + "/tipper_hodograms.png",
dpi=150, bbox_inches="tight",
)
except Exception as exc:
print(f"Tipper hodogram skipped: {{exc}}")
try:
ax_polar = plot_tipper_polar(sites_corr, component={component!r})
ax_polar.get_figure().savefig(
{out!r} + "/tipper_polar.png",
dpi=150, bbox_inches="tight",
)
except Exception as exc:
print(f"Tipper polar plot skipped: {{exc}}")
"""
_SENSITIVITY_BLOCK = """\
# Depth-of-investigation and sensitivity analysis
from pycsamt.emtools.csumt import (
depth_coverage_table,
plot_depth_section,
)
doi_table = depth_coverage_table(sites_corr)
doi_table.to_csv(
{out!r} + "/depth_of_investigation.csv",
index=False,
)
ax_doi = plot_depth_section(sites_corr)
ax_doi.get_figure().savefig(
{out!r} + "/depth_of_investigation.png",
dpi=150, bbox_inches="tight",
)
print("Depth-of-investigation and sensitivity outputs written.")
"""
_AI_INV_BLOCK = """\
# ── AI 1-D inversion ──────────────────────────────────────────
from pycsamt.ai.inversion import EMInverter1D
from pycsamt.forward.batch import generate_dataset
freqs_ai = {freqs!r}
n_layers = {n_layers}
n_samples = {n_samples}
epochs = {epochs}
ds = generate_dataset(
solver="mt1d",
n_samples=n_samples,
freqs=freqs_ai,
n_layers=n_layers,
noise_level=0.03,
seed=42,
n_jobs=1,
verbose=False,
)
inv1d = EMInverter1D(
arch={arch!r},
n_layers=n_layers,
solver="mt1d",
)
inv1d.fit(ds.X, ds.y, epochs=epochs, batch_size=32, verbose=False)
import numpy as np
obs_features = {{
f"synthetic_{{i:02d}}": feat
for i, feat in enumerate(ds.X[: min(3, len(ds.X))])
}}
predictions = {{}}
for site_name, feat in obs_features.items():
pred = inv1d.predict(feat.reshape(1, -1))
predictions[site_name] = pred[0]
print(f"AI 1-D inversion done: {{len(predictions)}} stations")
"""
_ENSEMBLE_BLOCK = """\
# ── ensemble 1-D inversion with uncertainty ───────────────────
from pycsamt.ai.inversion import EnsembleInverter
ens = EnsembleInverter(
n_members={n_members},
arch={arch!r},
n_layers={n_layers},
solver="mt1d",
)
ens.fit(ds.X, ds.y, epochs={epochs}, verbose=False)
ensemble_preds = {{}}
ensemble_unc = {{}}
for site_name, feat in obs_features.items():
mean, std = ens.predict_with_uncertainty(
feat.reshape(1, -1)
)
ensemble_preds[site_name] = mean[0]
ensemble_unc[site_name] = std[0]
print(
f"Ensemble inversion done: "
f"{{len(ensemble_preds)}} stations"
)
"""
_INV2D_BLOCK = """\
# ── U-Net 2-D profile inversion ───────────────────────────────
from pycsamt.ai.inversion import EMInverter2D
from scipy.ndimage import zoom
n_depth = {n_depth}
n_freqs_2d = {n_freqs}
n_components = {n_components}
n_sta = len(obs_features)
X_obs_list = list(obs_features.values())
X_obs = np.stack(X_obs_list, axis=2) # (n_freqs, n_comp, n_sta)
X_obs = X_obs.transpose(1, 0, 2) # (n_comp, n_freqs, n_sta)
X_obs_4d = X_obs[None, ...]
inv2d = EMInverter2D(
n_components=n_components,
n_depth=n_depth,
n_stations=n_sta,
n_freqs=n_freqs_2d,
arch={arch!r},
)
inv2d.fit(
X_train_2d, y_train_2d,
epochs={epochs},
batch_size=8,
verbose=False,
)
pred_section = inv2d.predict(X_obs_4d)[0] # (n_depth, n_sta)
print(
f"2-D inversion section shape: "
f"{{pred_section.shape}}"
)
"""
_FULL_AI_BLOCK = """\
# ── full AI workflow summary ───────────────────────────────────
import json, hashlib, platform, sys
from importlib.metadata import version as _v
provenance = {{
"workflow": "full_ai_workflow",
"data_path": {path!r},
"output_dir": {out!r},
"python": sys.version,
"platform": platform.platform(),
"packages": {{
"pycsamt": _v("pycsamt"),
"numpy": _v("numpy"),
}},
}}
prov_path = {out!r} + "/provenance.json"
with open(prov_path, "w") as fh:
json.dump(provenance, fh, indent=2)
print(f"Provenance written to {{prov_path}}")
"""
_FOOTER = """\
plt.show()
print("Workflow complete.")
"""
[docs]
class CodeGenerationAgent(BaseAgent):
"""Generate a reproducible Python script from a completed workflow.
Parameters
----------
api_key, model, llm_provider : str
When an API key is provided the LLM refines and annotates the
generated code. Otherwise the agent uses static templates.
script_title : str
Input keys
----------
``workflow_config`` : dict
The config dict produced by :class:`ContextInputAgent`.
``results`` : dict
The agent results dict from :class:`AgentCoordinator`.
``output_dir`` : str, optional
Folder for the script and its ``.validation.json`` report;
defaults to :data:`DEFAULT_OUTPUT_DIR` (``pycsamt_agent_output``).
Output data keys
----------------
``code`` str — Python source code
``script_path`` str or None — path to saved .py file
Examples
--------
>>> agent = CodeGenerationAgent()
>>> result = agent.execute(
... {
... "workflow_config": cfg,
... "results": coord_results,
... "output_dir": "/out",
... }
... )
>>> print(result["script_path"])
/out/workflow_script.py
"""
SYSTEM_PROMPT = _SYSTEM_PROMPT
def __init__(
self,
*,
api_key: str | None = None,
model: str | None = None,
llm_provider: str = "claude",
script_title: str = "pycsamt MT Processing Workflow",
) -> None:
super().__init__(
"CodeGenerationAgent",
api_key=api_key,
model=model,
llm_provider=llm_provider,
)
self.script_title = script_title
[docs]
def execute(self, input_data: dict[str, Any]) -> AgentResult:
self._last_cost = 0.0
t0 = time.time()
validation_started = time.monotonic()
warnings: list[str] = []
generation_data = input_data.get("generation_input")
generation = (
generation_data
if isinstance(generation_data, GenerationInput)
else GenerationInput.from_mapping(generation_data)
if generation_data is not None
else None
)
cfg = dict(
generation.workflow_config
if generation
else input_data.get("workflow_config") or {}
)
results = input_data.get("results") or {}
output_dir = input_data.get("output_dir") or DEFAULT_OUTPUT_DIR
if generation:
cfg.setdefault("output_dir", output_dir)
if cfg.get("workflow") == "forward":
cfg.update(forward_parameters(generation.task_text))
generation.workflow_config = cfg
title = input_data.get("title", self.script_title)
# Optional RAG context (real symbols / recipe) to ground the LLM
# refinement pass; ignored offline.
rag_context = (
generation.retrieved_evidence
if generation
else input_data.get("rag_context") or ""
)
exact_edit = exact_artifact_edit(generation) if generation else None
if generation:
question = clarification_for(generation)
if question:
return AgentResult(
"needs_review",
question,
data={
"clarification": question,
"generation_input": generation.to_dict(),
},
)
if (
not cfg.get("data_path")
and cfg.get("workflow") not in (None, "custom", "forward")
and not generation.previous_code
):
generation.assumptions.append(
"/path/to/EDIs is a placeholder; replace it with your data directory."
)
if generation.omitted_history_turns:
warnings.append(
"Older conversation turns were omitted from the bounded generation context."
)
data_path = cfg.get("data_path") or "/path/to/EDIs"
workflow = cfg.get("workflow") or ("custom" if generation else "qc")
out_dir = cfg.get("output_dir") or output_dir
supported = {
"qc",
"full",
"static_shift",
"phase_analysis",
"forward",
"pre_inversion",
"occam2d",
"tipper",
"tipper_plot",
"sensitivity",
"ai_inversion",
"inv1d",
"ensemble_inversion",
"inv2d",
"full_ai_workflow",
}
if (
generation
and not self.llm_available
and not exact_edit
and (workflow not in supported or generation.previous_code)
):
message = "This request needs a configured local or cloud model to compose or edit the script. Offline templates cannot apply it reliably."
return AgentResult(
"needs_review",
message,
data={
"clarification": message,
"generation_input": generation.to_dict(),
},
)
if generation and self.llm_available and not exact_edit:
from pycsamt.assistant.tools.api_evidence import (
generation_api_evidence,
)
generation.api_evidence = generation_api_evidence(
generation.task_text, workflow, generation.api_symbols
)
if (
workflow not in supported
and not generation.previous_code
and not rag_context
and len(generation.api_evidence) <= 1
):
question = "Which pyCSAMT operation or API should this script use? I could not find enough verified evidence to compose it."
return AgentResult(
"needs_review",
question,
data={
"clarification": question,
"generation_input": generation.to_dict(),
},
)
# ── build script from templates ───────────────────────────────────────
code = _HEADER.format(
title=title,
date=datetime.now().strftime("%Y-%m-%d"),
)
code += f"import os\nos.makedirs({out_dir!r}, exist_ok=True)\n\n"
needs_sites = workflow != "forward" or any(
k in results for k in ("qc", "static_shift", "phase_analysis")
)
if needs_sites:
code += _LOAD_BLOCK.format(path=data_path)
# add blocks based on workflow + available results
if workflow in ("qc", "full") or "qc" in results:
code += _QC_BLOCK.format(out=out_dir)
if generation:
csv = re.search(
r"\b[\w.-]+\.csv\b", generation.task_text, re.I
)
if (
csv
and re.search(
r"\b(save|export)\b", generation.task_text, re.I
)
and not re.search(
r"(?:do not|don't|never)\s+(?:save|export|write)",
generation.task_text,
re.I,
)
):
code += f"qc_table.to_csv({out_dir!r} + '/{csv[0]}', index=False)\n\n"
if workflow in ("static_shift", "full") or "static_shift" in results:
hw = 3
results.get("static_shift")
code += (
_SS_ESTIMATE_BLOCK.format(out=out_dir)
if generation and excludes_correction(generation.task_text)
else _SS_BLOCK.format(hw=hw, out=out_dir)
)
elif needs_sites:
code += "sites_corr = sites # no static-shift correction\n\n"
if (
workflow in ("phase_analysis", "full")
or "phase_analysis" in results
):
results.get("phase_analysis")
skew_th = 5.0
ellipt_th = 0.1
code += _PT_BLOCK.format(
skew_th=skew_th,
ellipt_th=ellipt_th,
out=out_dir,
)
if "forward" in results or workflow == "forward":
fwd_r = results.get("forward")
rhos = cfg.get("resistivity", [100, 10, 1000, 100])
ths = cfg.get("thickness", [500, 1000, 2000])
if fwd_r is not None:
lm = fwd_r.get("layered_model")
if lm is not None:
rhos = list(
getattr(
lm,
"resistivity",
getattr(lm, "resistivities", rhos),
)
)
ths = list(
getattr(
lm,
"thickness",
getattr(lm, "thicknesses", ths),
)
)
grid = cfg.get("frequency_grid")
freq_expr = (
f"np.geomspace({grid[0]!r}, {grid[1]!r}, {grid[2]!r})"
if grid
else "np.logspace(-4, 3, 60)"
)
code += _FWD_BLOCK.format(
rhos=rhos, ths=ths, out=out_dir, freqs=freq_expr
)
if (
workflow in ("pre_inversion", "occam2d")
or "pre_inversion" in results
):
code += _PRE_INV_BLOCK.format(out=out_dir)
if workflow in ("tipper", "tipper_plot") or "tipper" in results:
code += _TIPPER_BLOCK.format(
out=out_dir,
period=float(cfg.get("period", 1.0)),
component=str(cfg.get("tipper_component", "real")),
convention=str(cfg.get("tipper_convention", "park")),
)
if workflow == "sensitivity" or "sensitivity" in results:
code += _SENSITIVITY_BLOCK.format(out=out_dir)
# ── AI inversion blocks ───────────────────────────────────
ai_wf = {
"ai_inversion",
"inv1d",
"ensemble_inversion",
"inv2d",
"full_ai_workflow",
"full",
}
if workflow in ai_wf or "ai_inv" in results:
inv_r = results.get("ai_inv") or {}
n_layers = int(inv_r.get("n_layers", cfg.get("n_layers", 5)))
n_samples = int(cfg.get("n_train_samples", 2000))
epochs = int(cfg.get("epochs", 50))
arch = str(cfg.get("arch", "cnn1d"))
import numpy as np
freqs_ai = list(np.logspace(-4, 3, 32).round(6).tolist())
code += _AI_INV_BLOCK.format(
freqs=freqs_ai,
n_layers=n_layers,
n_samples=n_samples,
epochs=epochs,
arch=arch,
)
if workflow in ("ensemble_inversion",) or "ensemble" in results:
results.get("ensemble") or {}
n_layers = int(cfg.get("n_layers", 5))
arch = str(cfg.get("arch", "cnn1d"))
code += _ENSEMBLE_BLOCK.format(
n_members=int(cfg.get("n_members", 5)),
arch=arch,
n_layers=n_layers,
epochs=int(cfg.get("epochs", 50)),
)
if workflow in ("inv2d", "full_ai_workflow"):
code += _INV2D_BLOCK.format(
n_depth=int(cfg.get("n_depth", 40)),
n_freqs=int(cfg.get("n_freqs", 32)),
n_components=4,
arch=str(cfg.get("arch", "unet")),
epochs=int(cfg.get("epochs", 30)),
)
if workflow in ("full_ai_workflow",):
code += _FULL_AI_BLOCK.format(
path=data_path,
out=out_dir,
)
code += _FOOTER
# ── optional LLM refinement ───────────────────────────────────────────
if exact_edit:
code, updated_output = exact_edit
out_dir = (
updated_output or generation.previous_output_dir or out_dir
)
elif self.llm_available:
grounding = ""
if rag_context:
grounding = (
"Use ONLY the real pyCSAMT symbols and patterns from "
"this retrieved context (do not invent functions):\n"
f"{rag_context}\n\n"
)
prompt = (
f"Refine and annotate the following pycsamt Python script. "
f"Fix any issues, add helpful inline comments, ensure "
f"all imports are correct. Return only valid Python code.\n\n"
f"{grounding}"
f"```python\n{code}\n```"
)
if generation:
template_allowed = workflow in supported and (
workflow != "forward"
or all(
k in cfg
for k in ("resistivity", "thickness", "frequency_grid")
)
)
prompt = generation.prompt(code if template_allowed else "")
if self.llm_provider == "ollama":
from ._local import LocalSettings, current_request
request = current_request()
settings = request.settings if request else LocalSettings()
needed = (
len((self.SYSTEM_PROMPT + prompt).encode("utf-8"))
+ min(2000, settings.output_tokens)
+ 128
)
if needed > settings.context_tokens:
# Drop the optional example, never the request or prior script.
prompt = generation.prompt()
needed = (
len((self.SYSTEM_PROMPT + prompt).encode("utf-8"))
+ min(2000, settings.output_tokens)
+ 128
)
if needed > settings.context_tokens:
generation.compact_evidence()
prompt = generation.prompt()
llm_code = self.query_llm(prompt, max_tokens=2000)
if generation and not llm_code:
return AgentResult(
"failed",
"No generated script received.",
error="The configured model returned no code.",
)
if llm_code:
if generation and llm_code.strip().startswith("CLARIFY:"):
question = (
llm_code.strip()
.removeprefix("CLARIFY:")
.split("```", 1)[0]
.strip()[:800]
)
return AgentResult(
"needs_review",
question,
data={
"clarification": question,
"generation_input": generation.to_dict(),
},
)
# strip markdown fences if present
if "```python" in llm_code:
llm_code = llm_code.split("```python", 1)[1]
if "```" in llm_code:
llm_code = llm_code.rsplit("```", 1)[0]
code = llm_code.strip()
from ._code_validation import validate_and_repair
if generation and generation.assumptions:
code = (
"".join(
"# Assumption: " + a.replace("\n", " ") + "\n"
for a in generation.assumptions
)
+ code
)
code, validation = validate_and_repair(
self,
code,
generation,
started=validation_started,
allow_repair=not exact_edit,
max_repairs=input_data.get("max_repairs", 2),
execute_fixture=bool(input_data.get("execute_fixture", False)),
fixture_image=input_data.get("fixture_image"),
fixture_outputs=input_data.get("fixture_outputs"),
)
if not validation["ok"]:
message = (
"Generated draft failed validation; no script was saved. "
+ " ".join(validation["errors"][:4])
)
return AgentResult(
"needs_review" if validation["syntax_ok"] else "failed",
message,
error=None if validation["syntax_ok"] else message,
data={
"code": code,
"review_reason": message,
"script_path": None,
"validation": validation,
"generation_input": generation.to_dict()
if generation
else None,
"generation": {
"mode": "exact_edit"
if exact_edit
else "model"
if self.llm_available
else "template",
"requirements_verified": False,
"executed": False,
},
},
warnings=validation["warnings"],
)
# ── write file ────────────────────────────────────────────────────────
script_path: str | None = None
from ._request import checkpoint
checkpoint()
validation_path = None
try:
os.makedirs(output_dir, exist_ok=True)
filename = "workflow_script.py"
if generation:
requested = re.search(
r"\b(?:save(?:\s+(?:it|the script|code))?\s+as|name(?:\s+it)?|called|write)\s+[`\"']?([\w.-]+\.py)\b",
generation.task_text,
re.I,
)
filename = requested[1] if requested else filename
# Version every artifact, including the legacy template interface.
destination = Path(output_dir) / filename
suffix = 2
while (
destination.exists()
or destination.with_suffix(".validation.json").exists()
):
destination = (
Path(output_dir) / f"{Path(filename).stem}_{suffix}.py"
)
suffix += 1
with destination.open("x", encoding="utf-8", newline="\n") as stream:
stream.write(code)
script_path = str(destination)
report_path = destination.with_suffix(".validation.json")
with report_path.open("x", encoding="utf-8") as stream:
stream.write(json.dumps(validation, indent=2))
validation_path = str(report_path)
except Exception as exc:
warnings.append(f"Could not write script: {exc}")
script_path = None
validation_path = None
elapsed = time.time() - t0
n_lines = code.count("\n")
tailored = bool(exact_edit or self.llm_available or not generation)
if generation and not tailored:
warnings.append(
"Offline template only: request-specific constraints have not been applied. Configure a language model for a tailored script."
)
return AgentResult(
status="success" if tailored else "needs_review",
summary=(
f"Generated {n_lines}-line Python script for "
f"workflow={workflow!r}. "
+ (f"Saved to {script_path}" if script_path else "")
),
data={
"code": code,
"script_path": script_path,
"validation": validation,
"validation_path": validation_path,
"generation_input": generation.to_dict()
if generation
else None,
"generation": {
"output_dir": out_dir,
"workflow": workflow,
"mode": "exact_edit"
if exact_edit
else "model"
if self.llm_available
else "template",
"requirements_verified": bool(exact_edit),
"executed": False,
},
},
warnings=warnings,
elapsed_seconds=elapsed,
cost_estimate_usd=self._last_cost,
)
__all__ = ["CodeGenerationAgent"]
