WIP: dataset unique across different runs

Author: Michal Warda

README.md

@@ -221,6 +221,29 @@ const schema = () => [
 
 `oneOf` accepts plain schema entries or `{ value, weight }` objects. Provide any subset of weights (summing to ≤ 1) and the remaining probability is spread evenly across unweighted entries.
 
+#### Unique draws across a dataset
+
+Pass a `uniqueBy` configuration when you need each option to be used at most once across every row/schema during generation:
+
+```ts
+const toolOptions = [
+  weatherTool.toolFunction(),
+  calendarTool.toolFunction(),
+  flightTool.toolFunction(),
+] as const;
+
+const schema = () => [
+  oneOf(toolOptions, {
+    uniqueBy: {
+      collection: "tools",
+      itemId: "name",
+    },
+  }),
+];
+```
+
+The `collection` name identifies the shared pool (so multiple `oneOf` calls can coordinate), and `itemId` can be either a property key or a function that returns a stable identifier. Omit `itemId` to default to the common `id` field. Torque throws if the pool is exhausted, making it easy to guarantee perfect round-robin coverage.
+
 > 💡 See weighted example: [`examples/weighted-one-of.ts`](examples/weighted-one-of.ts)  
 > 💡 Full utilities demo: [`examples/composition-utilities.ts`](examples/composition-utilities.ts) | [▶️ Try in Browser](https://stackblitz.com/github/qforge-dev/torque/tree/main/stackblitz-templates/composition-utilities)
 

bun.lock

@@ -230,6 +230,29 @@ const schema = () => [
 
 `oneOf` accepts plain schema entries or `{ value, weight }` objects. Provide any subset of weights (summing to ≤ 1) and the remaining probability is spread evenly across unweighted entries.
 
+#### Unique draws across a dataset
+
+Pass a `uniqueBy` configuration when you need each option to be used at most once across every row/schema during generation:
+
+```ts
+const toolOptions = [
+  weatherTool.toolFunction(),
+  calendarTool.toolFunction(),
+  flightTool.toolFunction(),
+] as const;
+
+const schema = () => [
+  oneOf(toolOptions, {
+    uniqueBy: {
+      collection: "tools",
+      itemId: "name",
+    },
+  }),
+];
+```
+
+The `collection` name identifies the shared pool (so multiple `oneOf` calls can coordinate), and `itemId` can be either a property key or a function that returns a stable identifier. Omit `itemId` to default to the common `id` field. Torque throws if the pool is exhausted, making it easy to guarantee perfect round-robin coverage.
+
 > 💡 See weighted example: [`examples/weighted-one-of.ts`](examples/weighted-one-of.ts)  
 > 💡 Full utilities demo: [`examples/composition-utilities.ts`](examples/composition-utilities.ts) | [▶️ Try in Browser](https://stackblitz.com/github/qforge-dev/torque/tree/main/stackblitz-templates/composition-utilities)
 

packages/torque/README.md

@@ -32,6 +32,8 @@ import { createWriter } from "./writer";
 import { createFormatter } from "./formatter";
 import { TokenCounterPool } from "./token-counting/tokenCounterPool";
 import { hoistSystemMessages } from "./ai-message-order";
+import { runWithMessageContext } from "./message-context";
+import { runWithUniqueSelectionScope } from "./unique-selection";
 
 const DEFAULT_TOKEN_COUNTER_WORKERS = 3;
 
@@ -417,7 +419,9 @@ async function checkMessageSchemaStructure(
     generationContext,
   };
 
-  const message = await messageFactory(checkContext);
+  const message = await runWithMessageContext(checkContext, () =>
+    messageFactory(checkContext)
+  );
   if (message === null) return structure;
 
   if (Array.isArray(message)) {
@@ -458,7 +462,9 @@ async function checkMessageSchemaStructure(
         generationId: string;
       }[] = [];
       for (const tc of message.toolCalls) {
-        const toolCall = await tc(checkContext);
+        const toolCall = await runWithMessageContext(checkContext, () =>
+          tc(checkContext)
+        );
         toolCallStructures.push({
           toolCallId: toolCall.toolCallId,
           toolName: toolCall.toolName,
@@ -609,7 +615,9 @@ async function convertMessageSchemaToDatasetMessage(
     generationContext,
   };
 
-  const message = await messageFactory(context);
+  const message = await runWithMessageContext(context, () =>
+    messageFactory(context)
+  );
 
   if (message === null) return acc;
   if (Array.isArray(message)) {
@@ -733,7 +741,9 @@ async function convertMessageSchemaToDatasetMessage(
       if (message.toolCalls && message.toolCalls.length > 0) {
         const toolCallParts: IToolCallSchema<any>[] = [];
         for (const tc of message.toolCalls) {
-          const toolCall = await tc(context);
+          const toolCall = await runWithMessageContext(context, () =>
+            tc(context)
+          );
           toolCallParts.push(toolCall);
         }
 

packages/torque/src/dataset.ts

@@ -0,0 +1,22 @@
+import { AsyncLocalStorage } from "async_hooks";
+import type { IMessageSchemaContext } from "./types";
+
+type MaybePromise<T> = T | Promise<T>;
+
+const messageContextStorage = new AsyncLocalStorage<IMessageSchemaContext>();
+
+export function runWithMessageContext<T>(
+  context: IMessageSchemaContext,
+  fn: () => MaybePromise<T>
+): MaybePromise<T> {
+  const currentStore = messageContextStorage.getStore();
+  if (currentStore === context) {
+    return fn();
+  }
+
+  return messageContextStorage.run(context, fn);
+}
+
+export function getCurrentMessageContext(): IMessageSchemaContext | undefined {
+  return messageContextStorage.getStore();
+}

packages/torque/src/message-context.ts

@@ -1,7 +1,35 @@
 import type { WeightedOneOfOption } from "./schema";
 import { isWeightedOption } from "./schema";
+import { getCurrentMessageContext } from "./message-context";
+import { hasValueBeenUsed, markValueAsUsed } from "./unique-selection";
 import { random } from "./utils";
 
+type PrimitiveKey = string | number | boolean;
+type UniqueIdentifierResolver<T> =
+  | keyof any
+  | ((value: T) => PrimitiveKey | null | undefined);
+
+export interface OneOfUniqueBy<T> {
+  collection: string;
+  itemId?: UniqueIdentifierResolver<T>;
+}
+
+export interface OneOfOptions<T> {
+  unique?: boolean;
+  uniqueBy?: OneOfUniqueBy<T>;
+}
+
+type NormalizedWeightedOption<T> = {
+  value: T;
+  weight?: number;
+  uniqueKey?: string;
+};
+
+type ResolvedUniqueBy<T> = {
+  collection: string;
+  itemId: UniqueIdentifierResolver<T>;
+};
+
 export function optional<T>(message: T): T | (() => []) {
   return random() < 0.5 ? message : () => [];
 }
@@ -31,13 +59,27 @@ export function randomSample<T>(n: number, array: T[]): T[] {
 
   return result;
 }
-export function oneOf<T>(options: Array<WeightedOneOfOption<T>>): T {
+export function oneOf<T>(
+  options: Array<WeightedOneOfOption<T>>,
+  config?: OneOfOptions<T>
+): T {
   if (options.length === 0) {
     throw new Error("oneOf requires at least one option");
   }
 
+  const rawUniqueBy = config?.uniqueBy;
+  const enforceUnique = Boolean(config?.unique ?? rawUniqueBy);
+
+  if (enforceUnique && !rawUniqueBy) {
+    throw new Error(
+      "oneOf unique mode requires a uniqueBy option with a collection name"
+    );
+  }
+
+  const uniqueBy = resolveUniqueBy(rawUniqueBy);
+
   const normalized = options.map((option) =>
-    isWeightedOption(option) ? option : { value: option }
+    isWeightedOption(option) ? { ...option } : { value: option }
   );
 
   let providedWeightTotal = 0;
@@ -81,34 +123,145 @@ export function oneOf<T>(options: Array<WeightedOneOfOption<T>>): T {
     }
   }
 
-  const totalWeight = normalized.reduce(
+  const candidateBase = normalized.map<NormalizedWeightedOption<T>>((option) => ({
+    value: option.value,
+    weight: option.weight,
+  }));
+
+  let candidateOptions = candidateBase;
+
+  if (enforceUnique && uniqueBy) {
+    candidateOptions = candidateBase
+      .map((option) => ({
+        ...option,
+        uniqueKey: buildUniqueKey(option.value, uniqueBy),
+      }))
+      .filter(
+        (option) => !hasValueBeenUsed(uniqueBy.collection, option.uniqueKey!)
+      );
+
+    if (candidateOptions.length === 0) {
+      throw new Error(
+        `oneOf uniqueBy collection "${uniqueBy.collection}" is exhausted`
+      );
+    }
+  }
+
+  const totalWeight = candidateOptions.reduce(
     (sum, option) => sum + (option.weight ?? 0),
     0
   );
 
   if (totalWeight <= 0) {
-    const fallback = normalized[Math.floor(random() * normalized.length)];
+    const fallback =
+      candidateOptions[Math.floor(random() * candidateOptions.length)];
     if (!fallback) {
       throw new Error("oneOf failed to select a fallback option");
     }
+    recordUniqueSelection(fallback, uniqueBy, enforceUnique);
     return fallback.value;
   }
 
   const needle = random() * totalWeight;
   let cumulative = 0;
 
-  for (const option of normalized) {
+  for (const option of candidateOptions) {
     cumulative += option.weight ?? 0;
 
     if (needle <= cumulative) {
+      recordUniqueSelection(option, uniqueBy, enforceUnique);
       return option.value;
     }
   }
 
-  const lastOption = normalized[normalized.length - 1];
+  const lastOption = candidateOptions[candidateOptions.length - 1];
   if (!lastOption) {
     throw new Error("oneOf failed to resolve a selection");
   }
 
+  recordUniqueSelection(lastOption, uniqueBy, enforceUnique);
   return lastOption.value;
 }
+
+function recordUniqueSelection<T>(
+  option: NormalizedWeightedOption<T>,
+  uniqueBy: ResolvedUniqueBy<T> | undefined,
+  enforceUnique: boolean
+): void {
+  if (!enforceUnique || !uniqueBy || !option.uniqueKey) {
+    return;
+  }
+
+  const phase = getCurrentMessageContext()?.phase ?? "generate";
+  if (phase !== "generate") {
+    return;
+  }
+
+  markValueAsUsed(uniqueBy.collection, option.uniqueKey);
+}
+
+function buildUniqueKey<T>(
+  value: T,
+  uniqueBy: ResolvedUniqueBy<T>
+): string {
+  const rawValue =
+    typeof uniqueBy.itemId === "function"
+      ? uniqueBy.itemId(value)
+      : readProperty(value, uniqueBy.itemId);
+
+  if (
+    rawValue === undefined ||
+    rawValue === null ||
+    (typeof rawValue !== "string" &&
+      typeof rawValue !== "number" &&
+      typeof rawValue !== "boolean")
+  ) {
+    throw new Error(
+      `oneOf uniqueBy.itemId "${String(
+        uniqueBy.itemId
+      )}" must resolve to a string, number, or boolean`
+    );
+  }
+
+  const prefix = typeof rawValue;
+  return `${prefix}:${String(rawValue)}`;
+}
+
+function readProperty<T>(value: T, key: keyof any): unknown {
+  if (
+    value === null ||
+    (typeof value !== "object" && typeof value !== "function")
+  ) {
+    throw new Error(
+      `oneOf uniqueBy.itemId "${String(
+        key
+      )}" requires the option value to be an object or function`
+    );
+  }
+
+  return (value as any)[key];
+}
+
+function resolveUniqueBy<T>(
+  uniqueBy?: OneOfUniqueBy<T>
+): ResolvedUniqueBy<T> | undefined {
+  if (!uniqueBy) {
+    return undefined;
+  }
+
+  if (typeof uniqueBy.collection !== "string") {
+    throw new Error("oneOf uniqueBy.collection must be a non-empty string");
+  }
+
+  const collection = uniqueBy.collection.trim();
+  if (!collection) {
+    throw new Error("oneOf uniqueBy.collection must be a non-empty string");
+  }
+
+  const itemId = uniqueBy.itemId ?? ("id" as UniqueIdentifierResolver<T>);
+
+  return {
+    collection,
+    itemId,
+  };
+}