import json
import os

from axllm import OptimizerEngine, ai, ax, novel_f1_score


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-6-luna"),
    model_config={"temperature": 0},
)
program = ax('emailText:string -> priority:class "high, normal, low", rationale:string', {"id": "priority", "instruction": "Classify the email priority."})
baseline = program.forward(client, {"emailText": "Production checkout is failing for enterprise customers."})


class ExampleOptimizer(OptimizerEngine):
    name = "example"
    version = "1"

    def optimize(self, request, evaluator=None):
        return {"componentMap": {"priority::instruction": "Classify operational risk. Use high for production-impacting urgency."}, "metadata": {"source": "axgen"}}


artifact = program.optimize_with(ExampleOptimizer(), [{"emailText": "URGENT: checkout is down", "priority": "high"}], {"apply": False})
program.apply_optimization(json.dumps(artifact))
after = program.forward(client, {"emailText": "Production checkout is failing for enterprise customers."})
score = novel_f1_score("", after["rationale"], "Production checkout is failing for enterprise customers.")
print(json.dumps({"baseline": baseline, "after": after, "rationaleNovelF1": score}, indent=2, sort_keys=True))

import json
import os

from axllm import AxGEPA, OptimizerEvaluator, ai, ax


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-6-luna"),
)
program = ax('emailText:string -> priority:class "high, normal, low", rationale:string', {"id": "priority", "instruction": "Classify the email priority."})
baseline = program.forward(client, {"emailText": "Production checkout is failing for enterprise customers."})


class PriorityEvaluator(OptimizerEvaluator):
    def evaluate(self, candidate_map, options=None):
        result = program.evaluate_optimization(client, options["dataset"], candidate_map, options)
        for row in result["rows"]:
            expected = row["input"]["expectedOutput"]["priority"]
            actual = row["prediction"].get("output", {}).get("priority")
            score = float(actual == expected)
            row.update(scores={"accuracy": score}, scalar=score)
        result["avg"] = sum(row["scalar"] for row in result["rows"]) / result["count"]
        return result


def task(text, priority):
    return {"input": {"emailText": text}, "expectedOutput": {"priority": priority}}


request = {
    "programKind": "axgen",
    "components": program.get_optimizable_components(),
    "dataset": {
        "train": [task("URGENT: checkout is down", "high"), task("Weekly product newsletter", "low")],
        "validation": [task("Production payments are failing", "high"), task("Optional community meetup invite", "low")],
    },
    "options": {"numTrials": 1, "minibatchSize": 2, "maxMetricCalls": 12, "bootstrap": False, "seed": 7},
}
artifact = AxGEPA(client, seed=7).optimize(request, PriorityEvaluator())
original = artifact["baselineScore"]
best = artifact["metadata"]["bestScore"]
print(json.dumps({"baselinePrediction": baseline, "baselineScore": original, "baselineScores": artifact["baselineScores"], "optimizedScore": best, "delta": best - original}, indent=2))

import json
import os

from axllm import OptimizerEngine, ai, ax


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
program = ax('emailText:string -> priority:class "high, normal, low", rationale:string', {"id": "priority", "instruction": "Classify the email priority."})
baseline = program.forward(client, {"emailText": "Production checkout is failing for enterprise customers."})


class ExampleOptimizer(OptimizerEngine):
    name = "example"
    version = "1"

    def optimize(self, request, evaluator=None):
        return {"componentMap": {"priority::instruction": "Classify operational risk. Use high for production-impacting urgency."}, "metadata": {"source": "artifact"}}


artifact = program.optimize_with(ExampleOptimizer(), [{"emailText": "URGENT: checkout is down", "priority": "high"}], {"apply": False})
program.apply_optimization(json.dumps(artifact))
after = program.forward(client, {"emailText": "Production checkout is failing for enterprise customers."})
print(json.dumps({"baseline": baseline, "after": after}, indent=2, sort_keys=True))

import json
import os

from axllm import agent, ai
from axllm.runtime_quickjs import AxQuickJsCodeRuntime

api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-6-luna"),
)

seed = {
    "playbook": {
        "version": 1,
        "sections": {
            "failures_to_avoid": [
                {
                    "id": "failures-to-avoid-00001",
                    "section": "failures_to_avoid",
                    "content": "Check the available evidence before answering.",
                    "helpfulCount": 0,
                    "harmfulCount": 0,
                    "createdAt": "2026-07-15T00:00:00.000Z",
                    "updatedAt": "2026-07-15T00:00:00.000Z",
                }
            ]
        },
        "updatedAt": "2026-07-15T00:00:00.000Z",
    },
    "artifact": {"feedback": [], "history": []},
}

observed_citations = []
playbook_updates = []
assistant = agent(
    "question:string -> answer:string",
    {
        "ai": client,
        "contextFields": [],
        "runtime": {"language": "JavaScript"},
        "playbook": {"playbook": seed, "seed": 7, "onUpdate": playbook_updates.append},
        "citations": {
            "surface": "hidden",
            "onCitations": observed_citations.append,
        },
    },
)
assistant.set_instruction("Answer from evidence and state uncertainty plainly.")
assistant.add_actor_instruction("Before finishing, verify the answer against the collected evidence.")

runtime = AxQuickJsCodeRuntime()
answer = assistant.forward(
    client,
    {"question": "What should a support agent verify before answering?"},
    {"runtime": runtime, "max_actor_steps": 8},
)

# A score below scoreThreshold becomes a deterministic failure cluster. The
# default verification gate re-runs the task and rolls back a proposal that
# does not improve the held-in score. Add validation for production workloads.
evolution = assistant.playbook().evolve(
    {
        "train": [
            {
                "input": {"question": "Give a concise evidence-first answer."},
                "score": 0,
            }
        ]
    },
    {"verify": True, "maxProposals": 1, "runtime": runtime},
)

# Each train task replays on the runtime. A replay that fails scores 0 and
# would be mined as the weakness, so stop instead.
for record in evolution["records"]:
    prediction = record.get("prediction") or {}
    if record.get("error") or prediction.get("completionType") == "error":
        raise SystemExit(f"evolve replay failed: {record.get('error') or prediction.get('error')}")

print(json.dumps(answer, indent=2, sort_keys=True))
print("replays:", [(record.get("prediction") or {}).get("completionType") for record in evolution["records"]])
print("citations:", observed_citations[-1] if observed_citations else [])
print("run-end updates:", len(playbook_updates))
print("outcomes:", evolution["outcomes"])
print(assistant.get_playbook().render())