@swapai/core 0.2.0 → 0.2.2

package/README.md

@@ -95,6 +95,14 @@ init({
 
 The string result is inferred as `"rabbit" | "fish" | "pig"`.
 
+Needle is a classifier, not a regression engine. For bounded numbers, SwapAI
+converts reference scores into a small, closed set of internal labels, then
+converts the selected label back to a number. When the training set contains
+only a few scores, those scores stay exact. Larger score sets are quantized to
+a bounded set based on `acceptableError`. The held-out test still compares the
+decoded number with the original reference score, so quantization consumes the
+same error budget and cannot bypass the accuracy gate.
+
 ## Effect
 
 Effect is optional and lives in a separate import:

package/dist/index.js

@@ -46,6 +46,9 @@ function validateResultConfig(result) {
     if (result.max <= result.min) {
       invalidConfiguration("number result max must be greater than min");
     }
+    if (!Number.isFinite(result.max - result.min)) {
+      invalidConfiguration("number result range must be finite");
+    }
     return;
   }
   if (result.type === "string") {
@@ -122,6 +125,7 @@ function averageError(config, comparisons) {
 // src/runtime.ts
 import { createHash, randomUUID } from "crypto";
 import { spawn } from "child_process";
+import { readFileSync, statSync } from "fs";
 import {
   chmod,
   copyFile,
@@ -137,16 +141,66 @@ import { dirname, join, resolve, sep } from "path";
 import { fileURLToPath } from "url";
 
 // src/needle.ts
-function createNeedleTool(config) {
+var MAX_NEEDLE_NUMBER_LABELS = 64;
+function createNeedleNumberLabels(values, config, acceptableError = 0) {
+  if (typeof acceptableError !== "number" || !Number.isFinite(acceptableError) || acceptableError < 0 || acceptableError > 1) {
+    throw new TypeError("acceptableError must be between 0 and 1.");
+  }
+  const valid = values.map(
+    (value) => validateNeedleResult(value, config)
+  );
+  const unique = [...new Set(valid)].sort((left, right) => left - right);
+  const desiredLabelCount = acceptableError === 0 ? MAX_NEEDLE_NUMBER_LABELS : Math.min(
+    MAX_NEEDLE_NUMBER_LABELS,
+    Math.max(2, Math.floor(1 / (2 * acceptableError)) + 1)
+  );
+  if (unique.length <= desiredLabelCount) return { values: unique };
+  const width = (config.max - config.min) / desiredLabelCount;
+  const occupied = /* @__PURE__ */ new Set();
+  for (const value of unique) {
+    occupied.add(Math.min(
+      desiredLabelCount - 1,
+      Math.floor((value - config.min) / width)
+    ));
+  }
+  return {
+    values: [...occupied].sort((left, right) => left - right).map((index) => config.min + (index + 0.5) * width)
+  };
+}
+function numberLabel(index, count) {
+  if (count === 1) return "only_score";
+  if (count === 2) return index === 0 ? "lower_score" : "higher_score";
+  return `score_bucket_${index}`;
+}
+function encodeNumberResult(value, labels) {
+  if (labels.values.length === 0) {
+    throw new TypeError("Numeric result has no Needle label.");
+  }
+  let closestIndex = 0;
+  let closestDistance = Math.abs(value - labels.values[0]);
+  for (let index = 1; index < labels.values.length; index += 1) {
+    const distance = Math.abs(value - labels.values[index]);
+    if (distance < closestDistance) {
+      closestIndex = index;
+      closestDistance = distance;
+    }
+  }
+  return numberLabel(closestIndex, labels.values.length);
+}
+function createNeedleTool(config, numberLabels) {
   const description = "The classification result.";
   let result;
   switch (config.type) {
     case "number":
+      if (!numberLabels || numberLabels.values.length === 0) {
+        throw new TypeError("Numeric Needle tools require at least one result label.");
+      }
       result = {
-        type: "number",
-        minimum: config.min,
-        maximum: config.max,
-        description
+        type: "string",
+        enum: numberLabels.values.map(
+          (_, index) => numberLabel(index, numberLabels.values.length)
+        ),
+        description: `The numeric classification bucket from ${config.min} to ${config.max}.`
       };
       break;
     case "boolean":
@@ -167,12 +221,13 @@ function createNeedleTool(config) {
     }
   };
 }
-function createTrainingLine(input, value, config) {
+function createTrainingLine(input, value, config, numberLabels) {
   const result = validateNeedleResult(value, config);
+  const needleResult = config.type === "number" ? encodeNumberResult(result, numberLabels ?? { values: [] }) : result;
   return JSON.stringify({
     query: input,
-    tools: [createNeedleTool(config)],
-    answers: [{ name: "classify", arguments: { result } }]
+    tools: [createNeedleTool(config, numberLabels)],
+    answers: [{ name: "classify", arguments: { result: needleResult } }]
   });
 }
 function validateNeedleResult(value, config) {
@@ -196,11 +251,62 @@ function validateNeedleResult(value, config) {
       return value;
   }
 }
+function decodeNeedleResult(result, config, numberLabels) {
+  if (config.type === "number") {
+    if (typeof result !== "string") {
+      throw new TypeError(
+        `Needle returned a number outside ${config.min} to ${config.max}.`
+      );
+    }
+    const index = numberLabels?.values.findIndex(
+      (_, candidateIndex) => result === numberLabel(candidateIndex, numberLabels.values.length)
+    ) ?? -1;
+    const decoded = numberLabels?.values[index];
+    if (index < 0 || decoded === void 0) {
+      throw new TypeError(
+        `Needle returned a number outside ${config.min} to ${config.max}.`
+      );
+    }
+    return validateNeedleResult(decoded, config);
+  }
+  return validateNeedleResult(result, config);
+}
 
 // src/runtime.ts
 var NEEDLE_VERSION = "2.0.14";
+var NEEDLE_NUMBER_MODEL_VERSION = `${NEEDLE_VERSION}/number-buckets-v1`;
 var UV_VERSION = "0.11.4";
 var PYTHON_VERSION = "3.12";
+function needleModelVersion(resultConfig) {
+  return resultConfig.type === "number" ? NEEDLE_NUMBER_MODEL_VERSION : NEEDLE_VERSION;
+}
+var NeedleModelArtifactError = class extends SwapAIError {
+  constructor(message, cause) {
+    super(
+      "service_unavailable",
+      message,
+      cause === void 0 ? void 0 : { cause }
+    );
+    this.name = "NeedleModelArtifactError";
+  }
+};
+function hasNeedleModelArtifacts(options) {
+  try {
+    if (!statSync(options.modelPath).isFile()) return false;
+    if (options.resultConfig.type !== "number") return true;
+    const saved = JSON.parse(
+      readFileSync(`${options.modelPath}.numbers.json`, "utf8")
+    );
+    if (saved.format !== 1 || !Array.isArray(saved.values) || saved.values.length === 0) return false;
+    const savedValues = saved.values;
+    const labels = createNeedleNumberLabels(savedValues, options.resultConfig);
+    return labels.values.length === savedValues.length && labels.values.every(
+      (value, index) => value === savedValues[index]
+    );
+  } catch {
+    return false;
+  }
+}
 var NeedleRuntimeError = class extends SwapAIError {
   constructor(code, message, cause) {
     super(code, message, cause === void 0 ? void 0 : { cause });
@@ -285,8 +391,18 @@ var ManagedNeedleRuntime = class {
     const modelPath = join(candidateDirectory, "model.cact");
     await mkdir(checkpointDirectory, { recursive: true, mode: 448 });
     await mkdir(candidateDirectory, { recursive: true, mode: 448 });
+    const numberLabels = options.resultConfig.type === "number" ? createNeedleNumberLabels(
+      options.examples.map((example) => example.result),
+      options.resultConfig,
+      options.acceptableError
+    ) : void 0;
     const lines = options.examples.map(
-      (example) => createTrainingLine(example.input, example.result, options.resultConfig)
+      (example) => createTrainingLine(
+        example.input,
+        example.result,
+        options.resultConfig,
+        numberLabels
+      )
     );
     const lockPath = join(
       this.#dataDirectory,
@@ -294,33 +410,39 @@ var ManagedNeedleRuntime = class {
       `${createHash("sha256").update(options.classifierName).digest("hex")}.sqlite`
     );
     try {
+      const trainArguments = [
+        this.#workerPath,
+        "train",
+        "--training-data",
+        trainingPath,
+        "--output",
+        modelPath,
+        "--checkpoint-dir",
+        checkpointDirectory,
+        "--epochs",
+        String(options.epochs ?? 10),
+        "--artifact-lock-database",
+        lockPath,
+        "--main-database",
+        join(this.#dataDirectory, "swapai.sqlite"),
+        "--classifier-name",
+        options.classifierName,
+        "--expected-epoch",
+        String(options.expectedEpoch)
+      ];
+      if (numberLabels) {
+        trainArguments.push("--numeric-labels-from-stdin");
+      }
       await this.#runCommand(
         this.#pythonPath,
-        [
-          this.#workerPath,
-          "train",
-          "--training-data",
-          trainingPath,
-          "--output",
-          modelPath,
-          "--checkpoint-dir",
-          checkpointDirectory,
-          "--epochs",
-          String(options.epochs ?? 10),
-          "--artifact-lock-database",
-          lockPath,
-          "--main-database",
-          join(this.#dataDirectory, "swapai.sqlite"),
-          "--classifier-name",
-          options.classifierName,
-          "--expected-epoch",
-          String(options.expectedEpoch)
-        ],
+        trainArguments,
         {
           cwd: candidateDirectory,
           env: this.#environment,
           timeoutMs: 6 * 60 * 60 * 1e3,
-          stdin: `${lines.join("\n")}
+          stdin: numberLabels ? `${JSON.stringify({ format: 1, values: numberLabels.values })}
+${lines.join("\n")}
+` : `${lines.join("\n")}
 `
         }
       );
@@ -328,19 +450,51 @@ var ManagedNeedleRuntime = class {
     } catch (error) {
       throw toRuntimeError("Needle training failed.", error);
     }
-    return { modelPath, needleVersion: NEEDLE_VERSION };
+    return {
+      modelPath,
+      needleVersion: needleModelVersion(options.resultConfig)
+    };
   }
   async loadModel(options) {
     await this.ready();
     try {
       await stat(options.modelPath);
     } catch (error) {
-      throw toRuntimeError(`Needle model does not exist: ${options.modelPath}`, error);
+      throw new NeedleModelArtifactError(
+        `Needle model does not exist: ${options.modelPath}`,
+        error
+      );
+    }
+    let numberLabels;
+    if (options.resultConfig.type === "number") {
+      try {
+        const saved = JSON.parse(
+          await readFile(`${options.modelPath}.numbers.json`, "utf8")
+        );
+        if (saved.format !== 1 || !Array.isArray(saved.values) || saved.values.length === 0) {
+          throw new TypeError("Numeric label map has an invalid format.");
+        }
+        const savedValues = saved.values;
+        numberLabels = createNeedleNumberLabels(
+          savedValues,
+          options.resultConfig
+        );
+        if (numberLabels.values.length !== savedValues.length || numberLabels.values.some(
+          (value, index) => value !== savedValues[index]
+        )) {
+          throw new TypeError("Numeric label map is not sorted and unique.");
+        }
+      } catch (error) {
+        throw new NeedleModelArtifactError(
+          "Could not load Needle numeric labels.",
+          error
+        );
+      }
     }
     const schemaPath = `${options.modelPath}.schema.json`;
     await writeFile(
       schemaPath,
-      JSON.stringify([createNeedleTool(options.resultConfig)]),
+      JSON.stringify([createNeedleTool(options.resultConfig, numberLabels)]),
       "utf8"
     );
     const process2 = this.#spawnProcess(
@@ -359,6 +513,7 @@ var ManagedNeedleRuntime = class {
     const model = new NeedleModelProcess(
       process2,
       options.resultConfig,
+      numberLabels,
       (error) => this.#onBackgroundError?.(error),
       () => this.#models.delete(model)
     );
@@ -599,6 +754,7 @@ var ManagedNeedleRuntime = class {
 var NeedleModelProcess = class {
   #process;
   #resultConfig;
+  #numberLabels;
   #onBackgroundError;
   #onClose;
   #pending = /* @__PURE__ */ new Map();
@@ -612,9 +768,10 @@ var NeedleModelProcess = class {
   #exited = false;
   #closing = false;
   #stderr = "";
-  constructor(process2, resultConfig, onBackgroundError, onClose) {
+  constructor(process2, resultConfig, numberLabels, onBackgroundError, onClose) {
     this.#process = process2;
     this.#resultConfig = resultConfig;
+    this.#numberLabels = numberLabels;
     this.#onBackgroundError = onBackgroundError;
     this.#onClose = onClose;
     this.#readyPromise = new Promise((resolve2, reject) => {
@@ -765,7 +922,13 @@ var NeedleModelProcess = class {
       return;
     }
     try {
-      pending.resolve(validateNeedleResult(message.result, this.#resultConfig));
+      pending.resolve(
+        decodeNeedleResult(
+          message.result,
+          this.#resultConfig,
+          this.#numberLabels
+        )
+      );
     } catch (error) {
       pending.reject(
         new NeedleRuntimeError(
@@ -1359,7 +1522,7 @@ function openStorage(options) {
         WHERE name = ? AND training_lease_owner = ?
       `).run(Date.now(), options.name, owner);
     },
-    archiveAndReset(expectedDataEpoch) {
+    archiveAndReset(expectedDataEpoch, resetOptions) {
       assertOpen(closed);
       let nextGeneration = 0;
       const changedAt = Date.now();
@@ -1385,11 +1548,19 @@ function openStorage(options) {
             classifier_name, generation, status, created_at
           ) VALUES (?, ?, 'active', ?)
         `).run(options.name, nextGeneration, changedAt);
+        let retainedExampleCount = 0;
+        if (resetOptions?.retainExamples) {
+          retainedExampleCount = Number(database.prepare(`
+            UPDATE examples
+            SET generation = ?
+            WHERE classifier_name = ? AND generation = ?
+          `).run(nextGeneration, options.name, currentGeneration).changes);
+        }
         database.prepare(`
           UPDATE classifiers
           SET active_generation = ?,
               data_epoch = data_epoch + 1,
-              new_examples_since_training = 0,
+              new_examples_since_training = ?,
               training_attempts = 0,
               examples_used_for_training = 0,
               training_lease_owner = NULL,
@@ -1399,7 +1570,12 @@ function openStorage(options) {
               consecutive_retest_failures = 0,
               updated_at = ?
           WHERE name = ?
-        `).run(nextGeneration, changedAt, options.name);
+        `).run(
+          nextGeneration,
+          retainedExampleCount,
+          changedAt,
+          options.name
+        );
         reset = true;
       });
       return reset ? generationByNumber(database, options.name, nextGeneration) : null;
@@ -1804,8 +1980,11 @@ function createClassifier(config, storage, runtime) {
   let dataEpoch = initialState.dataEpoch;
   let clearedEpoch = initialState.clearPending ? initialState.dataEpoch - 1 : initialState.dataEpoch;
   const saved = initialState;
-  if (!saved.clearPending && saved.trained && (saved.modelPath === null || saved.needleVersion !== NEEDLE_VERSION)) {
-    if (storage.archiveAndReset(saved.dataEpoch) !== null) {
+  if (!saved.clearPending && saved.trained && (saved.modelPath === null || saved.needleVersion !== needleModelVersion(config.result) || !hasNeedleModelArtifacts({
+    modelPath: saved.modelPath,
+    resultConfig: config.result
+  }))) {
+    if (storage.archiveAndReset(saved.dataEpoch, { retainExamples: true }) !== null) {
       const reset = storage.snapshot();
       dataEpoch = reset.dataEpoch;
       clearedEpoch = reset.dataEpoch;
@@ -1827,13 +2006,22 @@ function createClassifier(config, storage, runtime) {
     }
     try {
       await runtime.ready();
-      const snapshot = storage.snapshot();
-      if (snapshot.clearPending) {
-        trained = false;
-        dataEpoch = snapshot.dataEpoch;
-        return;
-      }
-      if (snapshot.trained && snapshot.modelPath !== null) {
+    } catch (error) {
+      rememberBackgroundFailure(toSwapAIError(
+        error,
+        "service_unavailable",
+        "Needle could not start"
+      ));
+      return;
+    }
+    const snapshot = storage.snapshot();
+    if (snapshot.clearPending) {
+      trained = false;
+      dataEpoch = snapshot.dataEpoch;
+      return;
+    }
+    if (snapshot.trained && snapshot.modelPath !== null) {
+      try {
         const restoredModel = await runtime.loadModel({
           modelPath: snapshot.modelPath,
           resultConfig: config.result
@@ -1846,15 +2034,28 @@ function createClassifier(config, storage, runtime) {
         } else {
           await restoredModel.close();
         }
+      } catch (error) {
+        if (error instanceof NeedleModelArtifactError && storage.archiveAndReset(snapshot.dataEpoch, {
+          retainExamples: true
+        }) !== null) {
+          const reset = refreshStoredState();
+          dataEpoch = reset.dataEpoch;
+          clearedEpoch = reset.dataEpoch;
+          trained = false;
+        }
+        const loadError = toSwapAIError(
+          error,
+          "service_unavailable",
+          error instanceof NeedleModelArtifactError ? `Could not load classifier "${config.name}"; retraining from saved examples` : `Could not load classifier "${config.name}"`
+        );
+        if (error instanceof NeedleModelArtifactError) {
+          reportBackgroundError(config, loadError);
+        } else {
+          rememberBackgroundFailure(loadError);
+        }
       }
-    } catch (error) {
-      const swapAIError = toSwapAIError(
-        error,
-        "service_unavailable",
-        "Needle could not start"
-      );
-      rememberBackgroundFailure(swapAIError);
     }
+    scheduleTraining();
   }
   function rememberBackgroundFailure(error) {
     queuedFailure = error;
@@ -1971,7 +2172,8 @@ function createClassifier(config, storage, runtime) {
         generation: snapshot.activeGeneration,
         expectedEpoch: trainingEpoch,
         examples: trainingExamples.map(({ input, result }) => ({ input, result })),
-        resultConfig: config.result
+        resultConfig: config.result,
+        acceptableError: config.acceptableError
       });
       const candidateModel = await runtime.loadModel({
         modelPath: candidate.modelPath,
@@ -1980,10 +2182,18 @@ function createClassifier(config, storage, runtime) {
       try {
         const comparisons = [];
         for (const example of heldOutExamples) {
-          const candidateResult = validateResult(
-            config.result,
-            await candidateModel.classify(example.input)
-          );
+          let candidateResult;
+          try {
+            candidateResult = validateResult(
+              config.result,
+              await candidateModel.classify(example.input)
+            );
+          } catch (error) {
+            if (error instanceof SwapAIError && error.code === "classification_failed") {
+              return true;
+            }
+            throw error;
+          }
           comparisons.push({
             reference: validateResult(config.result, example.result),
             candidate: candidateResult
@@ -2062,10 +2272,22 @@ function createClassifier(config, storage, runtime) {
         );
       }
       await runtime.ready();
-      const restored = await runtime.loadModel({
-        modelPath: current.modelPath,
-        resultConfig: config.result
-      });
+      let restored;
+      try {
+        restored = await runtime.loadModel({
+          modelPath: current.modelPath,
+          resultConfig: config.result
+        });
+      } catch (error) {
+        if (error instanceof NeedleModelArtifactError && storage.archiveAndReset(expectedEpoch, { retainExamples: true }) !== null) {
+          const reset = refreshStoredState();
+          dataEpoch = reset.dataEpoch;
+          clearedEpoch = reset.dataEpoch;
+          trained = false;
+          scheduleTraining();
+        }
+        throw error;
+      }
       const afterLoad = refreshStoredState();
       if (afterLoad.dataEpoch !== expectedEpoch || afterLoad.clearPending || !afterLoad.trained) {
         await restored.close();