Add support for maxTokens option in CLI and run functions (#288)

microsoft · Mar 21, 2024 · 9ecc60f · 9ecc60f
1 parent c8e69c8
commit 9ecc60f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/docs/src/content/docs/reference/cli/commands.md b/docs/src/content/docs/reference/cli/commands.md
@@ -42,6 +42,7 @@ Options:
   -t, --temperature <number>         temperature for the run
   -tp, --top-p <number>              top-p for the run
   -m, --model <string>               model for the run
+  -mt, --max-tokens <number>         maximum tokens for the run
   -se, --seed <number>               seed for the run
   --no-cache                         disable LLM result cache
   --cs, --csv-separator <string>     csv separator (default: "\t")
@@ -73,6 +74,7 @@ Options:
   -t, --temperature <number>         temperature for the run
   -tp, --top-p <number>              top-p for the run
   -m, --model <string>               model for the run
+  -mt, --max-tokens <number>         maximum tokens for the run
   -se, --seed <number>               seed for the run
   --no-cache                         disable LLM result cache
   -ae, --apply-edits                 apply file edits

diff --git a/packages/cli/src/main.ts b/packages/cli/src/main.ts
@@ -153,6 +153,7 @@ async function batch(
         temperature: string
         topP: string
         seed: string
+        maxTokens: string
         model: string
         cache: boolean
         applyEdits: boolean
@@ -185,6 +186,7 @@ async function batch(
     const temperature = normalizeFloat(options.temperature)
     const topP = normalizeFloat(options.topP)
     const seed = normalizeFloat(options.seed)
+    const maxTokens = normalizeInt(options.maxTokens)
 
     const toolFiles: string[] = []
     if (scriptRx.test(tool)) toolFiles.push(tool)
@@ -264,6 +266,7 @@ async function batch(
                     temperature,
                     topP,
                     seed,
+                    maxTokens,
                     model,
                     retry,
                     retryDelay,
@@ -394,6 +397,7 @@ async function run(
         temperature: string
         topP: string
         seed: string
+        maxTokens: string
         model: string
         csvSeparator: string
         cache: boolean
@@ -419,6 +423,7 @@ async function run(
     const temperature = normalizeFloat(options.temperature)
     const topP = normalizeFloat(options.topP)
     const seed = normalizeFloat(options.seed)
+    const maxTokens = normalizeInt(options.maxTokens)
     const cache = !!options.cache
     const applyEdits = !!options.applyEdits
     const model = options.model
@@ -522,6 +527,7 @@ ${Array.from(files)
         temperature,
         topP,
         seed,
+        maxTokens,
         model,
         retry,
         retryDelay,
@@ -895,6 +901,7 @@ async function main() {
         .option("-t, --temperature <number>", "temperature for the run")
         .option("-tp, --top-p <number>", "top-p for the run")
         .option("-m, --model <string>", "model for the run")
+        .option("-mt, --max-tokens <number>", "maximum tokens for the run")
         .option("-se, --seed <number>", "seed for the run")
         .option("--no-cache", "disable LLM result cache")
         .option("--cs, --csv-separator <string>", "csv separator", "\t")
@@ -931,6 +938,7 @@ async function main() {
         .option("-t, --temperature <number>", "temperature for the run")
         .option("-tp, --top-p <number>", "top-p for the run")
         .option("-m, --model <string>", "model for the run")
+        .option("-mt, --max-tokens <number>", "maximum tokens for the run")
         .option("-se, --seed <number>", "seed for the run")
         .option("--no-cache", "disable LLM result cache")
         .option("-ae, --apply-edits", "apply file edits")