@prismatic-io/lux 0.0.1

package/src/optimization/reflective.ts

@@ -0,0 +1,319 @@
+import { z } from "zod";
+import {
+  type CandidateEvaluation,
+  type CandidateProposal,
+  defaultHarnessModel,
+  emptyUsage,
+  ModelCliProviderSchema,
+  type Optimizer,
+  ReasoningEffortSchema,
+  runTextModel,
+} from "../core/index.js";
+
+export const ReflectiveOptimizerConfigSchema = z
+  .object({
+    provider: ModelCliProviderSchema.default("claude-code"),
+    model: z.string().optional(),
+    command: z.string().optional(),
+    reasoningEffort: ReasoningEffortSchema.optional(),
+    timeoutMs: z.number().int().positive().default(180_000),
+    maxComponentChars: z.number().int().positive().default(80_000),
+    maxAlternativeChars: z.number().int().positive().default(40_000),
+    maxFeedbackChars: z.number().int().positive().default(50_000),
+    /** Hard cap for the complete serialized optimizer prompt. */
+    maxPromptChars: z.number().int().min(5_000).max(500_000).default(120_000),
+    proposalsPerIteration: z.number().int().min(1).max(4).default(1),
+  })
+  .strict();
+export type ReflectiveOptimizerConfig = z.infer<typeof ReflectiveOptimizerConfigSchema>;
+
+const ProposalResponseSchema = z
+  .object({
+    proposals: z.array(
+      z
+        .object({
+          rationale: z.string().optional(),
+          updates: z.record(z.string(), z.string()),
+        })
+        .strict(),
+    ),
+  })
+  .strict();
+
+const parseResponse = (text: string): z.infer<typeof ProposalResponseSchema> | null => {
+  const direct = (() => {
+    try {
+      return JSON.parse(text);
+    } catch {
+      return undefined;
+    }
+  })();
+  const directResult = ProposalResponseSchema.safeParse(direct);
+  if (directResult.success) return directResult.data;
+  const fenced = /```(?:json)?\s*([\s\S]*?)```/i.exec(text)?.[1];
+  if (!fenced) return null;
+  try {
+    const parsed = ProposalResponseSchema.safeParse(JSON.parse(fenced));
+    return parsed.success ? parsed.data : null;
+  } catch {
+    return null;
+  }
+};
+
+const boundedComponents = (
+  components: Record<string, string>,
+  budget: number,
+): { complete: Record<string, string>; omitted: string[] } => {
+  const out: Record<string, string> = {};
+  const omitted: string[] = [];
+  let used = 0;
+  for (const [path, content] of Object.entries(components).sort(([a], [b]) => a.localeCompare(b))) {
+    const serializedChars = JSON.stringify({ [path]: content }).length;
+    if (serializedChars > budget - used) {
+      omitted.push(path);
+      continue;
+    }
+    out[path] = content;
+    used += serializedChars;
+  }
+  return { complete: out, omitted };
+};
+
+const boundedRecord = (
+  input: Record<string, number>,
+  budget: number,
+): { complete: Record<string, number>; omitted: number } => {
+  const complete: Record<string, number> = {};
+  let used = 0;
+  let omitted = 0;
+  for (const [key, value] of Object.entries(input).sort(([a], [b]) => a.localeCompare(b))) {
+    const size = JSON.stringify({ [key]: value }).length;
+    if (used + size > budget) {
+      omitted++;
+      continue;
+    }
+    complete[key] = value;
+    used += size;
+  }
+  return { complete, omitted };
+};
+
+const boundedResults = (
+  observations: readonly CandidateEvaluation[],
+  budget: number,
+  includeMetrics: boolean,
+): { results: unknown[]; omitted: number } => {
+  const results: unknown[] = [];
+  for (const observation of observations) {
+    const share = Math.max(100, Math.floor(budget / Math.max(1, observations.length)));
+    const casePassRates = boundedRecord(observation.casePassRates, Math.floor(share * 0.55));
+    const metrics = boundedRecord(
+      includeMetrics ? observation.metrics : {},
+      Math.floor(share * 0.3),
+    );
+    const next = {
+      profile: observation.profileId.slice(0, 200),
+      split: observation.split,
+      passRate: observation.passRate,
+      casePassRates: casePassRates.complete,
+      omittedCasePassRates: casePassRates.omitted,
+      ...(includeMetrics ? { metrics: metrics.complete, omittedMetrics: metrics.omitted } : {}),
+    };
+    if (JSON.stringify([...results, next]).length > budget) break;
+    results.push(next);
+  }
+  return { results, omitted: observations.length - results.length };
+};
+
+const boundedFeedback = (
+  observations: readonly CandidateEvaluation[],
+  budget: number,
+): { failures: unknown[]; omitted: number } => {
+  const failures = observations.flatMap((observation) =>
+    observation.failedCases.map((failure) => ({
+      profile: observation.profileId,
+      split: observation.split,
+      caseId: failure.caseId,
+      score: failure.score,
+      feedback: failure.feedback,
+    })),
+  );
+  const included: unknown[] = [];
+  for (const failure of failures) {
+    const next = [...included, failure];
+    if (JSON.stringify(next).length > budget) break;
+    included.push(failure);
+  }
+  return { failures: included, omitted: failures.length - included.length };
+};
+
+const systemPrompt = (kind: "agent" | "grader"): string => `You improve text components of a
+${kind === "grader" ? "grading system from disagreements with human expert labels" : "coding-agent system from evaluation results"}.
+Propose the smallest changes that fix recurring failures while preserving successful behavior
+and reducing instruction/context size where possible.
+
+Return only JSON:
+{"proposals":[{"rationale":"short explanation","updates":{"relative/path":"complete new file contents"}}]}
+
+Rules:
+- Update only paths present in the supplied components.
+- Each update value is the complete replacement file, never a patch.
+- ${
+  kind === "grader"
+    ? "Align rubric criteria, evidence bounds, thresholds, and weights to the human labels. Do not make a grader merely more lenient or leak annotation text, IDs, or expected labels into it."
+    : "Do not weaken tests, graders, safety rules, or required user-owned decisions."
+}
+- Prefer deletion, routing, explicit compact contracts, and mechanical checks over added prose.
+- Do not mention the evaluation cases in the proposed component text.
+${kind === "grader" ? "- Preserve stable assertion IDs so labels and historical grades remain comparable." : ""}`;
+
+const optimizerPrompt = (
+  input: Parameters<Optimizer["propose"]>[0],
+  config: ReflectiveOptimizerConfig,
+): { prompt: string; visibleComponents: Set<string> } => {
+  const cap = config.maxPromptChars;
+  const currentComponents = boundedComponents(
+    input.candidate.components,
+    Math.min(config.maxComponentChars, Math.floor(cap * 0.45)),
+  );
+  const resultBudget = Math.floor(cap * 0.12);
+  const currentResults = boundedResults(input.observations, resultBudget, true);
+  const objectives =
+    JSON.stringify(input.campaign.objectives).length <= Math.floor(cap * 0.08)
+      ? input.campaign.objectives
+      : {
+          omitted: true,
+          constraints: input.campaign.objectives.constraints.length,
+          minimize: input.campaign.objectives.minimize.length,
+        };
+  const alternatives = (input.alternatives ?? []).slice(0, 2).map((candidate) => {
+    const observations = input.alternativeObservations?.[candidate.id] ?? [];
+    return {
+      id: candidate.id.slice(0, 200),
+      sourceBytes: candidate.sourceBytes,
+      trainResults: boundedResults(
+        observations,
+        Math.floor((cap * 0.08) / Math.max(1, input.alternatives?.length ?? 1)),
+        false,
+      ),
+      componentView: boundedComponents(
+        candidate.components,
+        Math.min(
+          Math.floor(config.maxAlternativeChars / Math.max(1, input.alternatives?.length ?? 1)),
+          Math.floor((cap * 0.08) / Math.max(1, input.alternatives?.length ?? 1)),
+        ),
+      ),
+    };
+  });
+  const payload: Record<string, unknown> = {
+    campaign: {
+      id: input.campaign.id.slice(0, 200),
+      objectives,
+      remainingCandidates: input.remainingCandidates,
+    },
+    currentCandidate: {
+      id: input.candidate.id.slice(0, 200),
+      sourceBytes: input.candidate.sourceBytes,
+      trainResults: currentResults,
+      components: currentComponents.complete,
+      omittedComponents: currentComponents.omitted,
+    },
+    frontierAlternatives: alternatives,
+    omittedFrontierAlternatives: Math.max(
+      0,
+      (input.alternatives?.length ?? 0) - alternatives.length,
+    ),
+    failedEvaluations: boundedFeedback(
+      input.observations,
+      Math.min(config.maxFeedbackChars, Math.floor(cap * 0.12)),
+    ),
+    proposalCount: Math.min(config.proposalsPerIteration, input.remainingCandidates),
+  };
+  let prompt = JSON.stringify(payload, null, 2);
+  if (prompt.length > cap) {
+    // Preserve valid JSON and the component visibility contract while
+    // shedding reproducibly ordered context from least essential sections.
+    payload.frontierAlternatives = [];
+    payload.omittedFrontierAlternatives = input.alternatives?.length ?? 0;
+    payload.failedEvaluations = { failures: [], omitted: input.observations.length };
+    (payload.campaign as Record<string, unknown>).objectives = {
+      omitted: true,
+      constraints: input.campaign.objectives.constraints.length,
+      minimize: input.campaign.objectives.minimize.length,
+    };
+    (payload.currentCandidate as Record<string, unknown>).trainResults = {
+      results: [],
+      omitted: input.observations.length,
+    };
+    prompt = JSON.stringify(payload, null, 2);
+  }
+  while (prompt.length > cap) {
+    const paths = Object.keys(currentComponents.complete).sort();
+    const path = paths.at(-1);
+    if (!path) break;
+    delete currentComponents.complete[path];
+    currentComponents.omitted.push(path);
+    prompt = JSON.stringify(payload, null, 2);
+  }
+  if (prompt.length > cap) {
+    throw new Error(`reflective optimizer prompt metadata exceeds maxPromptChars=${cap}`);
+  }
+  return { prompt, visibleComponents: new Set(Object.keys(currentComponents.complete)) };
+};
+
+const reflectiveOptimizer = (config: ReflectiveOptimizerConfig): Optimizer => ({
+  propose: async (input): Promise<CandidateProposal[]> => {
+    const model = config.model ?? defaultHarnessModel(config.provider);
+    const builtPrompt = optimizerPrompt(input, config);
+    let response: Awaited<ReturnType<typeof runTextModel>>;
+    try {
+      await input.recordModelAttemptStart?.("proposal");
+      response = await runTextModel({
+        provider: config.provider,
+        model,
+        command: config.command,
+        reasoningEffort: config.reasoningEffort,
+        timeoutMs: config.timeoutMs,
+        system: systemPrompt(input.campaign.evaluation.kind),
+        signal: input.abortSignal,
+        prompt: builtPrompt.prompt,
+      });
+    } catch (error) {
+      // A failed provider call has unknown usage/price. The campaign ledger
+      // must still consume one explicit unpriced call.
+      await input.recordModelCall?.({ usage: emptyUsage() });
+      await input.completeModelAttempts?.();
+      throw error;
+    }
+    await input.recordModelCall?.({
+      usage: response.usage,
+      ...(response.costUsd !== undefined ? { costUsd: response.costUsd } : {}),
+    });
+    await input.completeModelAttempts?.();
+    const parsed = parseResponse(response.text);
+    if (!parsed) {
+      throw new Error(`reflective optimizer returned invalid JSON: ${response.text.slice(0, 300)}`);
+    }
+    const visibleComponents = builtPrompt.visibleComponents;
+    return parsed.proposals
+      .slice(0, Math.min(config.proposalsPerIteration, input.remainingCandidates))
+      .map((proposal) => {
+        const hiddenUpdates = Object.keys(proposal.updates).filter(
+          (path) => !visibleComponents.has(path),
+        );
+        if (hiddenUpdates.length > 0) {
+          throw new Error(
+            `reflective optimizer attempted to replace components it was not shown in full: ${hiddenUpdates.join(", ")}`,
+          );
+        }
+        return {
+          updates: proposal.updates,
+          ...(proposal.rationale ? { rationale: proposal.rationale } : {}),
+          usage: response.usage,
+        };
+      });
+  },
+});
+
+export const createReflectiveOptimizer = (rawConfig: unknown = {}): Optimizer =>
+  reflectiveOptimizer(ReflectiveOptimizerConfigSchema.parse(rawConfig));

package/src/orchestrator/annotation-loader.ts

@@ -0,0 +1,97 @@
+import { readFile } from "node:fs/promises";
+import { dirname, extname, resolve, sep } from "node:path";
+import { pathToFileURL } from "node:url";
+import { z } from "zod";
+import {
+  type ExperimentSelector,
+  formatZodIssues,
+  type GraderAnnotation,
+  GraderAnnotationDatasetSchema,
+  GraderAnnotationSchema,
+  valueHash,
+} from "../core/index.js";
+
+export const LoadedGraderAnnotationsSchema = z
+  .object({
+    path: z.string(),
+    directory: z.string(),
+    dataset: GraderAnnotationDatasetSchema,
+    annotations: z.array(GraderAnnotationSchema),
+    hash: z.string(),
+  })
+  .strict();
+export type LoadedGraderAnnotations = z.infer<typeof LoadedGraderAnnotationsSchema>;
+
+const loadAuthoredDataset = async (path: string): Promise<unknown> => {
+  if (extname(path).toLowerCase() === ".json") {
+    return JSON.parse(await readFile(path, "utf8"));
+  }
+  const mod = (await import(pathToFileURL(path).href)) as { default?: unknown };
+  if (mod.default === undefined) {
+    throw new Error(
+      `Lux grader annotations at ${path} have no default export. Use \`export default defineGraderAnnotations({ ... })\`.`,
+    );
+  }
+  return mod.default;
+};
+
+export const loadGraderAnnotations = async (filePath: string): Promise<LoadedGraderAnnotations> => {
+  const path = resolve(filePath);
+  const parsed = GraderAnnotationDatasetSchema.safeParse(await loadAuthoredDataset(path));
+  if (!parsed.success) {
+    throw new Error(`Invalid Lux grader annotations at ${path}: ${formatZodIssues(parsed.error)}`);
+  }
+  const directory = dirname(path);
+  return LoadedGraderAnnotationsSchema.parse({
+    path,
+    directory,
+    dataset: parsed.data,
+    annotations: parsed.data.annotations.map((annotation) => ({
+      ...annotation,
+      runDir: resolve(directory, annotation.runDir.split(/[\\/]/).join(sep)),
+    })),
+    hash: valueHash(parsed.data),
+  });
+};
+
+export const selectGraderAnnotations = (
+  annotations: readonly GraderAnnotation[],
+  selector: ExperimentSelector,
+): GraderAnnotation[] => {
+  const patterns = selector.patterns.map((pattern) => pattern.toLowerCase());
+  const tags = selector.tags.map((tag) => tag.toLowerCase());
+  return annotations.filter((annotation) => {
+    const haystack = `${annotation.id}\n${annotation.rubricPath}`.toLowerCase();
+    const patternMatch =
+      patterns.length === 0 || patterns.some((pattern) => haystack.includes(pattern));
+    const annotationTags = new Set(annotation.tags.map((tag) => tag.toLowerCase()));
+    const tagMatch = tags.length === 0 || tags.some((tag) => annotationTags.has(tag));
+    return patternMatch && tagMatch;
+  });
+};
+
+export const assertDisjointAnnotationSplits = (
+  splits: Array<{ name: string; annotations: readonly GraderAnnotation[] }>,
+): void => {
+  const idOwners = new Map<string, string>();
+  const runOwners = new Map<string, string>();
+  for (const split of splits) {
+    for (const annotation of split.annotations) {
+      const idOwner = idOwners.get(annotation.id);
+      if (idOwner) {
+        throw new Error(
+          `grader annotation '${annotation.id}' appears in both ${idOwner} and ${split.name} splits`,
+        );
+      }
+      idOwners.set(annotation.id, split.name);
+      const runIdentity = resolve(annotation.runDir);
+      const runOwner = runOwners.get(runIdentity);
+      if (runOwner) {
+        throw new Error(
+          `grader evidence run '${annotation.runDir}' appears in both ${runOwner} and ${split.name} splits`,
+        );
+      }
+      runOwners.set(runIdentity, split.name);
+    }
+  }
+};

package/src/orchestrator/annotation-store.ts

@@ -0,0 +1,112 @@
+import { randomUUID } from "node:crypto";
+import { existsSync } from "node:fs";
+import { mkdir, readFile, rename, rm, writeFile } from "node:fs/promises";
+import { hostname } from "node:os";
+import { dirname, resolve } from "node:path";
+import { setTimeout as delay } from "node:timers/promises";
+import {
+  type GraderAnnotation,
+  type GraderAnnotationDataset,
+  GraderAnnotationDatasetSchema,
+} from "../core/index.js";
+import { claimLockForRecovery, publishOwnedLock, releaseOwnedLock } from "./owned-lock.js";
+
+const LOCK_WAIT_MS = 5_000;
+const STALE_LOCK_MS = 30_000;
+
+type AnnotationLockOwner = {
+  pid: number;
+  hostname: string;
+  acquiredAt: string;
+};
+
+const processIsAlive = (pid: number): boolean => {
+  try {
+    process.kill(pid, 0);
+    return true;
+  } catch (error) {
+    return (error as NodeJS.ErrnoException).code !== "ESRCH";
+  }
+};
+
+const recoverStaleLock = async (lockPath: string): Promise<boolean> => {
+  const claim = await claimLockForRecovery(lockPath);
+  if (!claim) return false;
+  let owner: AnnotationLockOwner;
+  try {
+    try {
+      owner = JSON.parse(await readFile(claim.path, "utf8")) as AnnotationLockOwner;
+    } catch {
+      return false;
+    }
+    const ageMs = Date.now() - Date.parse(owner.acquiredAt);
+    if (
+      owner.hostname !== hostname() ||
+      !Number.isFinite(ageMs) ||
+      ageMs < STALE_LOCK_MS ||
+      processIsAlive(owner.pid)
+    ) {
+      return false;
+    }
+    return await claim.removeLock();
+  } finally {
+    await claim.release();
+  }
+};
+
+const acquireAnnotationLock = async (path: string): Promise<() => Promise<void>> => {
+  const lockPath = `${path}.lock`;
+  const deadline = Date.now() + LOCK_WAIT_MS;
+  while (true) {
+    try {
+      const owner = await publishOwnedLock(lockPath, {
+        pid: process.pid,
+        hostname: hostname(),
+        acquiredAt: new Date().toISOString(),
+      });
+      return async () => {
+        await releaseOwnedLock(lockPath, owner.lockToken);
+      };
+    } catch (error) {
+      if ((error as NodeJS.ErrnoException).code !== "EEXIST") throw error;
+      if (await recoverStaleLock(lockPath)) continue;
+      if (Date.now() >= deadline) {
+        throw new Error(`timed out waiting to update grader annotations: ${path}`);
+      }
+      await delay(10);
+    }
+  }
+};
+
+export const upsertGraderAnnotation = async (
+  filePath: string,
+  annotation: GraderAnnotation,
+): Promise<GraderAnnotationDataset> => {
+  const path = resolve(filePath);
+  await mkdir(dirname(path), { recursive: true });
+  const release = await acquireAnnotationLock(path);
+  const temporary = `${path}.${process.pid}.${randomUUID()}.tmp`;
+  try {
+    const dataset: GraderAnnotationDataset = existsSync(path)
+      ? GraderAnnotationDatasetSchema.parse(JSON.parse(await readFile(path, "utf8")))
+      : { schemaVersion: 1, annotations: [] };
+    const existing = dataset.annotations.findIndex((entry) => entry.id === annotation.id);
+    if (existing === -1) dataset.annotations.push(annotation);
+    else {
+      const current = dataset.annotations[existing];
+      if (current?.runDir !== annotation.runDir || current.rubricPath !== annotation.rubricPath) {
+        throw new Error(
+          `annotation id '${annotation.id}' already targets different evidence; choose a unique id`,
+        );
+      }
+      dataset.annotations[existing] = annotation;
+    }
+    dataset.annotations.sort((a, b) => a.id.localeCompare(b.id));
+    await writeFile(temporary, `${JSON.stringify(dataset, null, 2)}\n`, { flag: "wx" });
+    await rename(temporary, path);
+    return dataset;
+  } finally {
+    await rm(temporary, { force: true });
+    await release();
+  }
+};