Merge pull request #62 from FormulaMonks/add/schema-constrained-tokens

feat: add `forceSchemaConstrainedTokens` to `KurtSamplingOptions`
FormulaMonks · Dec 6, 2024 · 9ba8bdb · 9ba8bdb
2 parents 65a209d + f0f384c
commit 9ba8bdb
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 0 deletions.
diff --git a/...he/test-retain/stub-b8a5a99fa499ef332c4a599d5b1eff433fc0ee7cd6995e86fb7dbfd8b9ffe999.yaml b/...he/test-retain/stub-b8a5a99fa499ef332c4a599d5b1eff433fc0ee7cd6995e86fb7dbfd8b9ffe999.yaml
@@ -5,6 +5,7 @@ sampling:
   maxOutputTokens: 4096
   temperature: 0.5
   topP: 0.95
+  forceSchemaConstrainedTokens: false
 tools:
   structured_data:
     name: structured_data

diff --git a/packages/kurt/spec/Kurt.spec.ts b/packages/kurt/spec/Kurt.spec.ts
@@ -30,6 +30,7 @@ describe("Kurt", () => {
         maxOutputTokens: 1024,
         temperature: 0.1,
         topP: 0.5,
+        forceSchemaConstrainedTokens: true,
       }
       const { adapter, kurt } = createV1({ sampling })
 
@@ -45,6 +46,7 @@ describe("Kurt", () => {
         maxOutputTokens: 1024,
         temperature: 0.1,
         topP: 0.5,
+        forceSchemaConstrainedTokens: true,
       }
       kurt.generateNaturalLanguage({ prompt: "Hello!", sampling })
 

diff --git a/packages/kurt/src/Kurt.ts b/packages/kurt/src/Kurt.ts
@@ -353,6 +353,25 @@ export const KurtSamplingOptionsDefault = {
    * without any "top tokens" filtering being applied before sampling.
    */
   topP: 0.95,
+
+  /**
+   * When true, structured data generation and/or tool calls will always use
+   * schema-aware constrained token sampling, rather than conditionally
+   * doing so, or falling back to relying only on the LLM's training for
+   * responding with structured data schemas.
+   *
+   * This gives a hard guarantee that the output will match the schema,
+   * but for some underlying LLM providers this may introduce more constraints
+   * on the kinds of JSON Schemas that are supported, which explains why
+   * it might not be the default behavior for a particular adapter.
+   *
+   * If this feature is unavailable for a particular adapter, the adapter
+   * should throw a `KurtCapabilityError` when this option is set to true.
+   *
+   * If this feature is available without downsides, the adapter should silently
+   * enable this feature regardless of the option being set to true or false.
+   */
+  forceSchemaConstrainedTokens: false,
 }
 
 export interface KurtGenerateNaturalLanguageOptions {

diff --git a/packages/kurt/src/KurtError.ts b/packages/kurt/src/KurtError.ts
@@ -14,6 +14,26 @@ export abstract class KurtError extends Error {
   abstract readonly adapter: KurtAdapter
 }
 
+/**
+ * Base class for all Kurt errors thrown for rejected requests.
+ */
+export abstract class KurtRequestError extends KurtError {}
+
+/**
+ * Thrown by a Kurt adapter when the request tried to use a capability or
+ * feature that isn't supported by this particular adapter (but might be
+ * supported by other adapters, or by the same adapter class if instantiated
+ * with different configuration, such as selecting a different LLM snapshot).
+ */
+export class KurtCapabilityError extends KurtError {
+  constructor(
+    readonly adapter: KurtAdapter,
+    readonly missingCapability: string
+  ) {
+    super(`This adapter doesn't support: ${missingCapability}`)
+  }
+}
+
 /**
  * Base class for all Kurt errors thrown when handling a result stream.
  */