import { AxAIOpenAIModel, ai, ax, optimize } from '@ax-llm/ax';

const apiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!apiKey) {
  throw new Error('Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.');
}

const llm = ai({
  name: 'openai',
  apiKey,
  config: {
    model: AxAIOpenAIModel.GPT54Mini,
    temperature: 0,
  },
});

const program = ax(
  'emailText:string -> priority:class "high, normal, low", rationale:string'
);

const baseline = await program.forward(llm, {
  emailText: 'Production checkout is failing for enterprise customers.',
});

const train = [
  {
    emailText: 'URGENT: checkout is down',
    priority: 'high',
    rationale: 'Production checkout outage blocks customers.',
  },
  {
    emailText: 'Weekly newsletter',
    priority: 'low',
    rationale: 'Informational update with no action needed.',
  },
  {
    emailText: 'Reminder to submit timesheets',
    priority: 'normal',
    rationale: 'Routine request with a clear deadline.',
  },
];

const metric = ({ prediction, example }: { prediction: any; example: any }) =>
  prediction.priority === example.priority ? 1 : 0;

const result = await optimize(program, train, metric, {
  studentAI: llm,
  teacherAI: llm,
  numTrials: 1,
  maxMetricCalls: 4,
});

if (!result.optimizedProgram) {
  throw new Error('Optimizer did not return an optimized program.');
}

program.applyOptimization(result.optimizedProgram);
const after = await program.forward(llm, {
  emailText: 'Production checkout is failing for enterprise customers.',
});

console.log(
  JSON.stringify({ baseline, after, bestScore: result.bestScore }, null, 2)
);

import { AxAIOpenAIModel, ai, ax, optimize } from '@ax-llm/ax';

const apiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!apiKey) {
  throw new Error('Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.');
}

const llm = ai({
  name: 'openai',
  apiKey,
  config: {
    model: AxAIOpenAIModel.GPT54Mini,
    temperature: 0,
  },
});

const program = ax(
  'emailText:string -> priority:class "high, normal, low", rationale:string'
);

const baseline = await program.forward(llm, {
  emailText: 'Production checkout is failing for enterprise customers.',
});

const train = [
  {
    emailText: 'URGENT: checkout is down',
    priority: 'high',
    rationale: 'Production checkout outage blocks customers.',
  },
  {
    emailText: 'Weekly newsletter',
    priority: 'low',
    rationale: 'Informational update with no action needed.',
  },
  {
    emailText: 'Reminder to submit timesheets',
    priority: 'normal',
    rationale: 'Routine request with a clear deadline.',
  },
];

const metric = ({ prediction, example }: { prediction: any; example: any }) =>
  prediction.priority === example.priority ? 1 : 0;

const result = await optimize(program, train, metric, {
  studentAI: llm,
  teacherAI: llm,
  numTrials: 1,
  maxMetricCalls: 4,
});

if (!result.optimizedProgram) {
  throw new Error('Optimizer did not return an optimized program.');
}

program.applyOptimization(result.optimizedProgram);
console.log(JSON.stringify({ baseline, bestScore: result.bestScore }, null, 2));

import { AxAIOpenAIModel, type AxMetricFn, ai, ax, playbook } from '@ax-llm/ax';

const apiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!apiKey) {
  throw new Error('Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.');
}

const studentAI = ai({
  name: 'openai',
  apiKey,
  config: { model: AxAIOpenAIModel.GPT54Mini, temperature: 0.2 },
});

// A generator we want to improve without hand-editing its prompt.
const triage = ax('ticket:string -> urgency:class "p0, p1, p2"');
triage.setDescription('Classify the support ticket urgency.');

// Labeled examples capture the nuance we want the playbook to absorb.
const train = [
  {
    ticket: 'Checkout is down for all customers in the EU region.',
    urgency: 'p0',
  },
  { ticket: 'A single user cannot change their avatar.', urgency: 'p2' },
  {
    ticket: 'Login works but is intermittently slow for many users.',
    urgency: 'p1',
  },
  { ticket: 'Production database returns 500s on every write.', urgency: 'p0' },
  { ticket: 'Typo in the footer copyright year.', urgency: 'p2' },
];

const metric: AxMetricFn = ({ prediction, example }) =>
  (prediction as { urgency?: string }).urgency ===
  (example as { urgency?: string }).urgency
    ? 1
    : 0;

// 1) Grow a playbook offline from the labeled examples (ACE runs under the hood).
const pb = playbook(triage, { studentAI, maxEpochs: 2 });
const { bestScore } = await pb.evolve(train, metric);
pb.applyTo(triage);

console.log(`offline best score: ${bestScore}`);
console.log('\nlearned playbook:\n');
console.log(pb.render());

// 2) Use the improved program.
const live = await triage.forward(studentAI, {
  ticket: 'Password reset emails are delayed ~10 minutes for some users.',
});
console.log('\nlive prediction:', live);

// 3) Keep improving online from feedback — no metric required.
await pb.update({
  example: { ticket: 'The status page itself is unreachable.' },
  prediction: { urgency: 'p2' },
  feedback: 'WRONG: if customers cannot even see status, treat it as p0.',
});
pb.applyTo(triage);

// 4) Persist the playbook and restore it into a fresh program instance.
const snapshot = pb.toJSON();
const restored = playbook(ax('ticket:string -> urgency:class "p0, p1, p2"'), {
  studentAI,
}).load(snapshot);
console.log(
  '\nrestored playbook bullets:',
  restored.getState().playbook.stats.bulletCount
);

import {
  AxAIOpenAIModel,
  ai,
  ax,
  axDeserializeOptimizedProgram,
  axSerializeOptimizedProgram,
  optimize,
} from '@ax-llm/ax';

const apiKey = process.env.OPENAI_API_KEY ?? process.env.OPENAI_APIKEY;
if (!apiKey) {
  throw new Error('Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.');
}

const llm = ai({
  name: 'openai',
  apiKey,
  config: {
    model: AxAIOpenAIModel.GPT54Mini,
    temperature: 0,
  },
});

const program = ax(
  'emailText:string -> priority:class "high, normal, low", rationale:string'
);

const baseline = await program.forward(llm, {
  emailText: 'Production checkout is failing for enterprise customers.',
});

const train = [
  {
    emailText: 'URGENT: checkout is down',
    priority: 'high',
    rationale: 'Production checkout outage blocks customers.',
  },
  {
    emailText: 'Weekly newsletter',
    priority: 'low',
    rationale: 'Informational update with no action needed.',
  },
  {
    emailText: 'Reminder to submit timesheets',
    priority: 'normal',
    rationale: 'Routine request with a clear deadline.',
  },
];

const metric = ({ prediction, example }: { prediction: any; example: any }) =>
  prediction.priority === example.priority ? 1 : 0;

const result = await optimize(program, train, metric, {
  studentAI: llm,
  teacherAI: llm,
  numTrials: 1,
  maxMetricCalls: 4,
});

if (!result.optimizedProgram) {
  throw new Error('Optimizer did not return an optimized program.');
}

const saved = axSerializeOptimizedProgram(result.optimizedProgram);
const restored = axDeserializeOptimizedProgram(saved);

program.applyOptimization(restored);
const after = await program.forward(llm, {
  emailText: 'Production checkout is failing for enterprise customers.',
});

console.log(
  JSON.stringify(
    {
      baseline,
      after,
      bestScore: result.bestScore,
      artifactComponents: Object.keys(saved.componentMap ?? {}),
    },
    null,
    2
  )
);