@prismatic-io/lux 0.0.1

package/src/orchestrator/experiment-corpus.ts

@@ -0,0 +1,416 @@
+import { readFile } from "node:fs/promises";
+import { dirname, join } from "node:path";
+import { RubricConfigSchema } from "../assertions/rubric/index.js";
+import {
+  bindSubjectRoot,
+  type Candidate,
+  contentHash,
+  type ExperimentCampaign,
+  type ExperimentSelector,
+  valueHash,
+} from "../core/index.js";
+import { assertDisjointAnnotationSplits, selectGraderAnnotations } from "./annotation-loader.js";
+import { discoverCases, selectCases } from "./discover.js";
+import type { CampaignExecutionOptions, EvaluationCorpus } from "./experiment-context.js";
+import { profileOrThrow } from "./experiment-context.js";
+import {
+  claudeConfigRoot,
+  claudeUserBehaviorHash,
+  cliIdentity,
+  configuredClaudeSourceHashes,
+  referencedAnswererCli,
+  referencedCli,
+} from "./experiment-runtime-identity.js";
+import { fixtureIdentity } from "./fixture-identity.js";
+import { loadEvalCase } from "./loader.js";
+import { resolveAnswererReference, resolveDriverReference } from "./orchestrator.js";
+import { loadFrozenRegradableRun, releaseFrozenRegradableRun } from "./regrade.js";
+import { loadRubricDefinition } from "./rubric-loader.js";
+
+type ExperimentSplit = "train" | "validation" | "test";
+
+export const selectorsBySplit = (
+  campaign: ExperimentCampaign,
+): Array<{ name: ExperimentSplit; selector: ExperimentSelector }> => [
+  ...(campaign.splits.train ? [{ name: "train" as const, selector: campaign.splits.train }] : []),
+  { name: "validation", selector: campaign.splits.validation },
+  ...(campaign.splits.test ? [{ name: "test" as const, selector: campaign.splits.test }] : []),
+];
+
+const fixtureEvidenceHash = async (
+  execution: CampaignExecutionOptions,
+  casePath: string,
+  fixtures: Awaited<ReturnType<typeof loadEvalCase>>["fixtures"],
+): Promise<string | null> => {
+  if (!fixtures) return null;
+  return (await fixtureIdentity(fixtures, execution.config.fixturesRoot ?? dirname(casePath))).hash;
+};
+
+const agentCaseEvidenceHash = async (
+  execution: CampaignExecutionOptions,
+  casePath: string,
+): Promise<string> => {
+  const authored = await loadEvalCase(casePath);
+  const { id: _caseId, meta: _meta, fixtures, assertions = [], ...behavior } = authored;
+  const normalizedAssertions = assertions.map(
+    ({ id: _assertionId, name: _assertionName, ...assertion }) => assertion,
+  );
+  return valueHash({
+    case: { ...behavior, assertions: normalizedAssertions },
+    fixtures: await fixtureEvidenceHash(execution, casePath, fixtures),
+  });
+};
+
+export const caseRubricCliReferences = (
+  execution: CampaignExecutionOptions,
+  authored: Awaited<ReturnType<typeof loadEvalCase>>,
+): Array<{
+  assertionId: string;
+  profileId: string;
+  provider: "claude-code" | "codex";
+  command: string;
+}> => {
+  const rubrics = (authored.assertions ?? []).flatMap((assertion, index) =>
+    assertion.type === "rubric" ? [{ index, config: RubricConfigSchema.parse(assertion) }] : [],
+  );
+  return execution.loaded.campaign.profiles.flatMap((profileId) => {
+    const profile = profileOrThrow(execution.config, profileId);
+    const harness = profile.harness ?? execution.config.harness;
+    return rubrics.map(({ index, config }) => {
+      const provider = config.provider ?? harness?.provider ?? "claude-code";
+      const inheritsHarness =
+        config.provider === undefined || config.provider === harness?.provider;
+      return {
+        assertionId: config.id ?? `assertion-${index}`,
+        profileId,
+        provider,
+        command:
+          config.command ??
+          (inheritsHarness ? harness?.command : undefined) ??
+          (provider === "codex" ? "codex" : "claude"),
+      };
+    });
+  });
+};
+
+const caseRubricCliIdentities = (
+  execution: CampaignExecutionOptions,
+  authored: Awaited<ReturnType<typeof loadEvalCase>>,
+): Promise<
+  Array<{
+    assertionId: string;
+    profileId: string;
+    provider: "claude-code" | "codex";
+    cliIdentity: Awaited<ReturnType<typeof cliIdentity>>;
+  }>
+> =>
+  Promise.all(
+    caseRubricCliReferences(execution, authored).map(async ({ command, ...reference }) => ({
+      ...reference,
+      cliIdentity: await cliIdentity(command, execution.config.rootDir),
+    })),
+  );
+
+export const agentCorpusHash = async (
+  execution: CampaignExecutionOptions,
+  cases: Awaited<ReturnType<typeof discoverCases>>["cases"],
+  splits: ReturnType<typeof selectorsBySplit>,
+): Promise<string> => {
+  const selected = [
+    ...new Map(
+      splits
+        .flatMap((split) => selectCases(cases, split.selector))
+        .map((evalCase) => [evalCase.path, evalCase]),
+    ).values(),
+  ];
+  const caseBehaviorHashes = new Map<string, Promise<string | null>>();
+  const configuredSourceHashes = new Map<string, Promise<string>>();
+  const definitions = await Promise.all(
+    selected.map(async (evalCase) => {
+      const authored = await loadEvalCase(evalCase.path);
+      const bound = bindSubjectRoot(authored, execution.loaded.subjectRoot);
+      const profileCliIdentities = await Promise.all(
+        execution.loaded.campaign.profiles.map(async (profileId) => {
+          const profile = profileOrThrow(execution.config, profileId);
+          const harness = profile.harness ?? execution.config.harness;
+          const driver = resolveDriverReference(
+            bound,
+            profile.driver ?? execution.config.defaultDriver,
+          );
+          const answerer = resolveAnswererReference(
+            bound,
+            profile.answerer ?? execution.config.defaultAnswerer,
+            harness,
+          );
+          const driverCli = referencedCli(driver);
+          const answererCli = referencedAnswererCli(answerer);
+          const configuredRoot =
+            driverCli?.name === "claude-code" ? claudeConfigRoot(driver) : null;
+          let driverClaudeBehaviorHash: string | null = null;
+          if (configuredRoot) {
+            let hash = caseBehaviorHashes.get(configuredRoot);
+            if (!hash) {
+              hash = claudeUserBehaviorHash(configuredRoot);
+              caseBehaviorHashes.set(configuredRoot, hash);
+            }
+            driverClaudeBehaviorHash = await hash;
+          }
+          return {
+            profileId,
+            driver,
+            driverCliIdentity: driverCli
+              ? await cliIdentity(driverCli.command, execution.config.rootDir)
+              : null,
+            driverClaudeBehaviorHash,
+            configuredDriverSources: await configuredClaudeSourceHashes(
+              driver,
+              execution.config.rootDir,
+              configuredSourceHashes,
+            ),
+            answerer,
+            answererCliIdentity: answererCli
+              ? await cliIdentity(answererCli.command, execution.config.rootDir)
+              : null,
+          };
+        }),
+      );
+      return {
+        path: evalCase.path,
+        sourceHash: contentHash(await readFile(evalCase.path)),
+        case: authored,
+        fixturesHash: await fixtureEvidenceHash(execution, evalCase.path, authored.fixtures),
+        profileCliIdentities,
+        caseRubricCliIdentities: await caseRubricCliIdentities(execution, bound),
+      };
+    }),
+  );
+  return valueHash(definitions);
+};
+
+export const graderCorpusHash = async (
+  execution: CampaignExecutionOptions,
+  splits: ReturnType<typeof selectorsBySplit>,
+): Promise<string> => {
+  const loaded = execution.loaded.graderAnnotations;
+  if (!loaded) throw new Error("grader experiment is missing its annotation dataset");
+  const selected = [
+    ...new Map(
+      splits
+        .flatMap((split) => selectGraderAnnotations(loaded.annotations, split.selector))
+        .map((annotation) => [annotation.id, annotation]),
+    ).values(),
+  ];
+  const evidence = await Promise.all(
+    selected.map(async (annotation) => {
+      const run = await loadFrozenRegradableRun(annotation.runDir);
+      try {
+        const rubric = await loadRubricDefinition(
+          join(execution.loaded.subjectRoot, annotation.rubricPath),
+        );
+        if (rubric.id !== run.case.id) {
+          throw new Error(
+            `rubric '${annotation.rubricPath}' id '${rubric.id}' does not match persisted case '${run.case.id}'`,
+          );
+        }
+        const artifacts = await Promise.all(
+          run.artifacts.map(async (artifact) => {
+            const root = artifact.root ?? join(annotation.runDir, "artifacts");
+            try {
+              return {
+                ...artifact,
+                root: undefined,
+                contentHash: contentHash(await readFile(join(root, artifact.path))),
+              };
+            } catch (error) {
+              if ((error as NodeJS.ErrnoException).code === "ENOENT") {
+                return { ...artifact, root: undefined, contentHash: null };
+              }
+              throw error;
+            }
+          }),
+        );
+        return {
+          annotationId: annotation.id,
+          runDir: annotation.runDir,
+          evidenceHash: valueHash({
+            caseId: run.case.id,
+            prompt: run.case.prompt,
+            events: run.events.flatMap((event) => {
+              if (event.type === "lifecycle" || event.type === "ready") return [];
+              const { ts: _timestamp, ...visible } = event;
+              return [visible];
+            }),
+            artifacts,
+          }),
+          rubric,
+        };
+      } finally {
+        await releaseFrozenRegradableRun(run);
+      }
+    }),
+  );
+  const splitByAnnotation = new Map(
+    splits.flatMap((split) =>
+      selectGraderAnnotations(loaded.annotations, split.selector).map(
+        (annotation) => [annotation.id, split.name] as const,
+      ),
+    ),
+  );
+  const ownerByEvidence = new Map<string, { annotationId: string; split: ExperimentSplit }>();
+  for (const item of evidence) {
+    const split = splitByAnnotation.get(item.annotationId);
+    if (!split) continue;
+    const owner = ownerByEvidence.get(item.evidenceHash);
+    if (owner) {
+      const scope = owner.split === split ? `duplicate evidence within ${split}` : "split leakage";
+      throw new Error(
+        `grader ${scope}: annotations '${owner.annotationId}' (${owner.split}) and '${item.annotationId}' (${split}) contain identical frozen evidence`,
+      );
+    }
+    ownerByEvidence.set(item.evidenceHash, { annotationId: item.annotationId, split });
+  }
+  return valueHash({ annotationsHash: loaded.hash, evidence });
+};
+
+export const prepareCorpus = async (
+  execution: CampaignExecutionOptions,
+): Promise<EvaluationCorpus> => {
+  const campaign = execution.loaded.campaign;
+  const splits = selectorsBySplit(campaign);
+  if (campaign.evaluation.kind === "grader") {
+    const loaded = execution.loaded.graderAnnotations;
+    if (!loaded) throw new Error("grader experiment is missing its annotation dataset");
+    const selectedSplits = splits.map((split) => ({
+      name: split.name,
+      annotations: selectGraderAnnotations(loaded.annotations, split.selector),
+    }));
+    for (const split of selectedSplits) {
+      if (split.annotations.length === 0) {
+        throw new Error(
+          `experiment '${campaign.id}' ${split.name} selector matched no grader annotations`,
+        );
+      }
+    }
+    assertDisjointAnnotationSplits(selectedSplits);
+    return { kind: "grader", hash: await graderCorpusHash(execution, splits) };
+  }
+  const discovery = await discoverCases(
+    execution.config.casesRoot ?? join(execution.config.rootDir, "cases"),
+  );
+  if (discovery.problems.length > 0) {
+    throw new Error(
+      `experiment case discovery failed: ${discovery.problems.map((problem) => `${problem.path}: ${problem.reason}`).join("; ")}`,
+    );
+  }
+  const owners = new Map<string, string>();
+  const caseIds = new Map<string, string>();
+  const evidenceOwners = new Map<
+    string,
+    { caseId: string; path: string; split: ExperimentSplit }
+  >();
+  for (const split of splits) {
+    const selected = selectCases(discovery.cases, split.selector);
+    if (selected.length === 0) {
+      throw new Error(`experiment '${campaign.id}' ${split.name} selector matched no cases`);
+    }
+    for (const evalCase of selected) {
+      const existingPath = caseIds.get(evalCase.id);
+      if (existingPath && existingPath !== evalCase.path) {
+        throw new Error(
+          `experiment case id '${evalCase.id}' is duplicated by ${existingPath} and ${evalCase.path}`,
+        );
+      }
+      caseIds.set(evalCase.id, evalCase.path);
+      const owner = owners.get(evalCase.path);
+      if (owner) {
+        throw new Error(
+          `eval case '${evalCase.id}' appears in both ${owner} and ${split.name} splits`,
+        );
+      }
+      owners.set(evalCase.path, split.name);
+      const evidenceHash = await agentCaseEvidenceHash(execution, evalCase.path);
+      const evidenceOwner = evidenceOwners.get(evidenceHash);
+      if (evidenceOwner) {
+        const scope =
+          evidenceOwner.split === split.name
+            ? `duplicate evidence within ${split.name}`
+            : "split leakage";
+        throw new Error(
+          `agent ${scope}: cases '${evidenceOwner.caseId}' (${evidenceOwner.split}) and '${evalCase.id}' (${split.name}) contain identical behavioral evidence`,
+        );
+      }
+      evidenceOwners.set(evidenceHash, {
+        caseId: evalCase.id,
+        path: evalCase.path,
+        split: split.name,
+      });
+    }
+  }
+  return {
+    kind: "agent",
+    cases: discovery.cases,
+    hash: await agentCorpusHash(execution, discovery.cases, splits),
+  };
+};
+
+export const assertGraderSubject = async (
+  execution: CampaignExecutionOptions,
+  baseline: Candidate,
+): Promise<void> => {
+  if (execution.loaded.campaign.evaluation.kind !== "grader") return;
+  const loadedAnnotations = execution.loaded.graderAnnotations?.annotations ?? [];
+  const annotations = [
+    ...new Map(
+      selectorsBySplit(execution.loaded.campaign)
+        .flatMap((split) => selectGraderAnnotations(loadedAnnotations, split.selector))
+        .map((annotation) => [annotation.id, annotation]),
+    ).values(),
+  ];
+  const rubricPaths = [...new Set(annotations.map((annotation) => annotation.rubricPath))].sort();
+  const componentPaths = Object.keys(baseline.components).sort();
+  if (
+    rubricPaths.length !== componentPaths.length ||
+    rubricPaths.some((path, index) => path !== componentPaths[index])
+  ) {
+    throw new Error(
+      "grader experiment mutable components must exactly match the annotation rubric paths",
+    );
+  }
+  for (const rubricPath of rubricPaths) {
+    const rubric = await loadRubricDefinition(join(execution.loaded.subjectRoot, rubricPath));
+    const baselineIds = new Set(rubric.assertions.map((assertion) => assertion.id));
+    const relevant = annotations.filter((annotation) => annotation.rubricPath === rubricPath);
+    const allowedAdditions = new Set(
+      execution.loaded.campaign.evaluation.allowedAssertionAdditions[rubricPath] ?? [],
+    );
+    const allLabeledIds = new Set(
+      relevant.flatMap((annotation) => Object.keys(annotation.expected.assertions)),
+    );
+    const undeclaredAdditions = [...allLabeledIds].filter(
+      (id) => !baselineIds.has(id) && !allowedAdditions.has(id),
+    );
+    if (undeclaredAdditions.length > 0) {
+      throw new Error(
+        `annotation assertion IDs absent from '${rubricPath}' must be declared in evaluation.allowedAssertionAdditions: ${undeclaredAdditions.join(", ")}`,
+      );
+    }
+    const incompleteAdditions = [...allowedAdditions].filter(
+      (id) =>
+        baselineIds.has(id) ||
+        !relevant.every((annotation) => Object.hasOwn(annotation.expected.assertions, id)),
+    );
+    if (incompleteAdditions.length > 0) {
+      throw new Error(
+        `allowed assertion additions for '${rubricPath}' must be absent from baseline and labeled on every frozen run: ${incompleteAdditions.join(", ")}`,
+      );
+    }
+  }
+  const unknownRubricPaths = Object.keys(
+    execution.loaded.campaign.evaluation.allowedAssertionAdditions,
+  ).filter((path) => !rubricPaths.includes(path));
+  if (unknownRubricPaths.length > 0) {
+    throw new Error(
+      `allowed assertion additions reference rubric paths absent from annotations: ${unknownRubricPaths.join(", ")}`,
+    );
+  }
+};