Merge pull request #173 from GitPaulo/main

Support passing max_tokens and max_completion_tokens

README.md

@@ -123,6 +123,33 @@ supplied via the `input` parameter in YAML format. Additionally, you can
 provide file-based variables via `file_input`, where each key maps to a file
 path.
 
+### Prompt.yml with model parameters
+
+You can specify model parameters directly in your `.prompt.yml` files using the
+`modelParameters` key:
+
+```yaml
+messages:
+  - role: system
+    content: Be as concise as possible
+  - role: user
+    content: 'Compare {{a}} and {{b}}, please'
+model: openai/gpt-4o
+modelParameters:
+  maxCompletionTokens: 500
+  temperature: 0.7
+```
+
+| Key                   | Type   | Description                                           |
+| --------------------- | ------ | ----------------------------------------------------- |
+| `maxCompletionTokens` | number | The maximum number of tokens to generate              |
+| `maxTokens`           | number | The maximum number of tokens to generate (deprecated) |
+| `temperature`         | number | The sampling temperature to use (0-1)                 |
+| `topP`                | number | The nucleus sampling parameter to use (0-1)           |
+
+> ![Note]
+> Parameters set in `modelParameters` take precedence over the corresponding action inputs.
+
 ### Using a system prompt file
 
 In addition to the regular prompt, you can provide a system prompt file instead
@@ -276,21 +303,24 @@ perform actions like searching issues and PRs.
 Various inputs are defined in [`action.yml`](action.yml) to let you configure
 the action:
 
-| Name                 | Description                                                                                                                                                                                                        | Default                              |
-| -------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------ |
-| `token`              | Token to use for inference. Typically the GITHUB_TOKEN secret                                                                                                                                                      | `github.token`                       |
-| `prompt`             | The prompt to send to the model                                                                                                                                                                                    | N/A                                  |
-| `prompt-file`        | Path to a file containing the prompt (supports .txt and .prompt.yml formats). If both `prompt` and `prompt-file` are provided, `prompt-file` takes precedence                                                      | `""`                                 |
-| `input`              | Template variables in YAML format for .prompt.yml files (e.g., `var1: value1` on separate lines)                                                                                                                   | `""`                                 |
-| `file_input`         | Template variables in YAML where values are file paths. The file contents are read and used for templating                                                                                                         | `""`                                 |
-| `system-prompt`      | The system prompt to send to the model                                                                                                                                                                             | `"You are a helpful assistant"`      |
-| `system-prompt-file` | Path to a file containing the system prompt. If both `system-prompt` and `system-prompt-file` are provided, `system-prompt-file` takes precedence                                                                  | `""`                                 |
-| `model`              | The model to use for inference. Must be available in the [GitHub Models](https://github.com/marketplace?type=models) catalog                                                                                       | `openai/gpt-4o`                      |
-| `endpoint`           | The endpoint to use for inference. If you're running this as part of an org, you should probably use the org-specific Models endpoint                                                                              | `https://models.github.ai/inference` |
-| `max-tokens`         | The max number of tokens to generate                                                                                                                                                                               | 200                                  |
-| `enable-github-mcp`  | Enable Model Context Protocol integration with GitHub tools                                                                                                                                                        | `false`                              |
-| `github-mcp-token`   | Token to use for GitHub MCP server (defaults to the main token if not specified).                                                                                                                                  | `""`                                 |
-| `custom-headers`     | Custom HTTP headers to include in API requests. Supports both YAML format (`header1: value1`) and JSON format (`{"header1": "value1"}`). Useful for API Management platforms, rate limiting, and request tracking. | `""`                                 |
+| Name                    | Description                                                                                                                                                                                                        | Default                              |
+| ----------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------ |
+| `token`                 | Token to use for inference. Typically the GITHUB_TOKEN secret                                                                                                                                                      | `github.token`                       |
+| `prompt`                | The prompt to send to the model                                                                                                                                                                                    | N/A                                  |
+| `prompt-file`           | Path to a file containing the prompt (supports .txt and .prompt.yml formats). If both `prompt` and `prompt-file` are provided, `prompt-file` takes precedence                                                      | `""`                                 |
+| `input`                 | Template variables in YAML format for .prompt.yml files (e.g., `var1: value1` on separate lines)                                                                                                                   | `""`                                 |
+| `file_input`            | Template variables in YAML where values are file paths. The file contents are read and used for templating                                                                                                         | `""`                                 |
+| `system-prompt`         | The system prompt to send to the model                                                                                                                                                                             | `"You are a helpful assistant"`      |
+| `system-prompt-file`    | Path to a file containing the system prompt. If both `system-prompt` and `system-prompt-file` are provided, `system-prompt-file` takes precedence                                                                  | `""`                                 |
+| `model`                 | The model to use for inference. Must be available in the [GitHub Models](https://github.com/marketplace?type=models) catalog                                                                                       | `openai/gpt-4o`                      |
+| `endpoint`              | The endpoint to use for inference. If you're running this as part of an org, you should probably use the org-specific Models endpoint                                                                              | `https://models.github.ai/inference` |
+| `max-tokens`            | The maximum number of tokens to generate (deprecated, use `max-completion-tokens` instead)                                                                                                                         | 200                                  |
+| `max-completion-tokens` | The maximum number of tokens to generate                                                                                                                                                                           | `""`                                 |
+| `temperature`           | The sampling temperature to use (0-1)                                                                                                                                                                              | `""`                                 |
+| `top-p`                 | The nucleus sampling parameter to use (0-1)                                                                                                                                                                        | `""`                                 |
+| `enable-github-mcp`     | Enable Model Context Protocol integration with GitHub tools                                                                                                                                                        | `false`                              |
+| `github-mcp-token`      | Token to use for GitHub MCP server (defaults to the main token if not specified).                                                                                                                                  | `""`                                 |
+| `custom-headers`        | Custom HTTP headers to include in API requests. Supports both YAML format (`header1: value1`) and JSON format (`{"header1": "value1"}`). Useful for API Management platforms, rate limiting, and request tracking. | `""`                                 |
 
 ## Outputs
 

__tests__/helpers-inference.test.ts

@@ -109,6 +109,7 @@ describe('helpers.ts - inference request building', () => {
         undefined,
         undefined,
         100,
+        undefined,
         'https://api.test.com',
         'test-token',
       )
@@ -122,6 +123,7 @@ describe('helpers.ts - inference request building', () => {
         temperature: undefined,
         topP: undefined,
         maxTokens: 100,
+        maxCompletionTokens: undefined,
         endpoint: 'https://api.test.com',
         token: 'test-token',
         responseFormat: {
@@ -143,6 +145,7 @@ describe('helpers.ts - inference request building', () => {
         undefined,
         undefined,
         100,
+        undefined,
         'https://api.test.com',
         'test-token',
       )
@@ -156,6 +159,7 @@ describe('helpers.ts - inference request building', () => {
         temperature: undefined,
         topP: undefined,
         maxTokens: 100,
+        maxCompletionTokens: undefined,
         endpoint: 'https://api.test.com',
         token: 'test-token',
         responseFormat: undefined,

__tests__/inference.test.ts

@@ -31,7 +31,7 @@ describe('inference.ts', () => {
       {role: 'user' as const, content: 'Hello, AI!'},
     ],
     modelName: 'gpt-4',
-    maxTokens: 100,
+    maxCompletionTokens: 100,
     endpoint: 'https://api.test.com',
     token: 'test-token',
   }
@@ -72,7 +72,7 @@ describe('inference.ts', () => {
             content: 'Hello, AI!',
           },
         ],
-        max_tokens: 100,
+        max_completion_tokens: 100,
         model: 'gpt-4',
       })
 
@@ -176,7 +176,7 @@ describe('inference.ts', () => {
             content: 'Hello, AI!',
           },
         ],
-        max_tokens: 100,
+        max_completion_tokens: 100,
         model: 'gpt-4',
         response_format: requestWithResponseFormat.responseFormat,
       })
@@ -228,7 +228,7 @@ describe('inference.ts', () => {
       expect(callArgs.tools).toEqual(mockMcpClient.tools)
       expect(callArgs.response_format).toBeUndefined()
       expect(callArgs.model).toBe('gpt-4')
-      expect(callArgs.max_tokens).toBe(100)
+      expect(callArgs.max_completion_tokens).toBe(100)
 
       // Verify OpenAI client was initialized with empty custom headers
       expect(mockOpenAIClient).toHaveBeenCalledWith({
@@ -633,4 +633,64 @@ describe('inference.ts', () => {
       expect(result).toBe('{"immediate": "result"}')
     })
   })
+
+  describe('token param routing', () => {
+    it('sends max_tokens when only maxTokens is set', async () => {
+      const requestWithMaxTokens = {
+        ...mockRequest,
+        maxCompletionTokens: undefined,
+        maxTokens: 100,
+      }
+
+      const mockResponse = {
+        choices: [
+          {
+            message: {
+              content: 'Direct max_tokens response',
+            },
+          },
+        ],
+      }
+
+      mockCreate.mockResolvedValueOnce(mockResponse)
+
+      const result = await simpleInference(requestWithMaxTokens)
+
+      expect(result).toBe('Direct max_tokens response')
+      expect(mockCreate).toHaveBeenCalledTimes(1)
+
+      // Should have sent max_tokens directly
+      expect(mockCreate.mock.calls[0][0]).toHaveProperty('max_tokens', 100)
+      expect(mockCreate.mock.calls[0][0]).not.toHaveProperty('max_completion_tokens')
+    })
+
+    it('sends neither token param when both are undefined', async () => {
+      const requestWithNoTokens = {
+        ...mockRequest,
+        maxCompletionTokens: undefined,
+        maxTokens: undefined,
+      }
+
+      const mockResponse = {
+        choices: [
+          {
+            message: {
+              content: 'No token limit response',
+            },
+          },
+        ],
+      }
+
+      mockCreate.mockResolvedValueOnce(mockResponse)
+
+      const result = await simpleInference(requestWithNoTokens)
+
+      expect(result).toBe('No token limit response')
+      expect(mockCreate).toHaveBeenCalledTimes(1)
+
+      const params = mockCreate.mock.calls[0][0]
+      expect(params).not.toHaveProperty('max_tokens')
+      expect(params).not.toHaveProperty('max_completion_tokens')
+    })
+  })
 })

__tests__/main.test.ts

@@ -168,6 +168,7 @@ describe('main.ts', () => {
       ],
       modelName: 'gpt-4',
       maxTokens: 100,
+      maxCompletionTokens: undefined,
       endpoint: 'https://api.test.com',
       token: 'fake-token',
       responseFormat: undefined,
@@ -259,6 +260,7 @@ describe('main.ts', () => {
       ],
       modelName: 'gpt-4',
       maxTokens: 100,
+      maxCompletionTokens: undefined,
       endpoint: 'https://api.test.com',
       token: 'fake-token',
       responseFormat: undefined,

action.yml

@@ -43,9 +43,21 @@ inputs:
     required: false
     default: ''
   max-tokens:
-    description: The maximum number of tokens to generate
+    description: The maximum number of tokens to generate (deprecated)
     required: false
     default: '200'
+  max-completion-tokens:
+    description: The maximum number of tokens to generate
+    required: false
+    default: ''
+  temperature:
+    description: The sampling temperature to use (0-1)
+    required: false
+    default: ''
+  top-p:
+    description: The nucleus sampling parameter to use (0-1)
+    required: false
+    default: ''
   token:
     description: The token to use
     required: false

dist/index.js

@@ -58300,6 +58300,19 @@ OpenAI.Responses = Responses;
 OpenAI.Evals = Evals;
 OpenAI.Containers = Containers;
 
+/**
+ * Build according to what input was passed, default to max_tokens.
+ * Only one of max_tokens or max_completion_tokens will be set.
+ */
+function buildMaxTokensParam(request) {
+    if (request.maxCompletionTokens != null) {
+        return { max_completion_tokens: request.maxCompletionTokens };
+    }
+    if (request.maxTokens != null) {
+        return { max_tokens: request.maxTokens };
+    }
+    return {};
+}
 /**
  * Simple one-shot inference without tools
  */
@@ -58312,10 +58325,10 @@ async function simpleInference(request) {
     });
     const chatCompletionRequest = {
         messages: request.messages,
-        max_tokens: request.maxTokens,
         model: request.modelName,
         temperature: request.temperature,
         top_p: request.topP,
+        ...buildMaxTokensParam(request), // Note: solution around models using different underlying max tokens properties
     };
     // Add response format if specified
     if (request.responseFormat) {
@@ -58349,10 +58362,10 @@ async function mcpInference(request, githubMcpClient) {
         coreExports.info(`MCP inference iteration ${iterationCount}`);
         const chatCompletionRequest = {
             messages: messages,
-            max_tokens: request.maxTokens,
             model: request.modelName,
             temperature: request.temperature,
             top_p: request.topP,
+            ...buildMaxTokensParam(request), // Note: solution around models using different underlying max tokens properties
         };
         // Add response format if specified (only on final iteration to avoid conflicts with tool calls)
         if (finalMessage && request.responseFormat) {
@@ -61382,7 +61395,8 @@ function validateAndMaskHeaders(headers) {
 /**
  * Build complete InferenceRequest from prompt config and inputs
  */
-function buildInferenceRequest(promptConfig, systemPrompt, prompt, modelName, temperature, topP, maxTokens, endpoint, token, customHeaders) {
+function buildInferenceRequest(promptConfig, systemPrompt, prompt, modelName, temperature, topP, maxTokens, // Deprecated
+maxCompletionTokens, endpoint, token, customHeaders) {
     const messages = buildMessages(promptConfig, systemPrompt, prompt);
     const responseFormat = buildResponseFormat(promptConfig);
     return {
@@ -61390,7 +61404,8 @@ function buildInferenceRequest(promptConfig, systemPrompt, prompt, modelName, te
         modelName,
         temperature,
         topP,
-        maxTokens,
+        maxTokens, // Deprecated
+        maxCompletionTokens,
         endpoint,
         token,
         responseFormat,
@@ -61536,10 +61551,11 @@ async function run() {
         }
         // Get common parameters
         const modelName = promptConfig?.model || coreExports.getInput('model');
-        let maxTokens = promptConfig?.modelParameters?.maxTokens ?? coreExports.getInput('max-tokens');
-        if (typeof maxTokens === 'string') {
-            maxTokens = parseInt(maxTokens, 10);
-        }
+        // Parse token limit inputs
+        const maxCompletionTokensInput = promptConfig?.modelParameters?.maxCompletionTokens ?? coreExports.getInput('max-completion-tokens');
+        const maxCompletionTokens = maxCompletionTokensInput ? Number(maxCompletionTokensInput) : undefined;
+        const maxTokensInput = promptConfig?.modelParameters?.maxTokens ?? coreExports.getInput('max-tokens');
+        const maxTokens = maxCompletionTokens != null ? undefined : maxTokensInput ? Number(maxTokensInput) : undefined;
         const token = process.env['GITHUB_TOKEN'] || coreExports.getInput('token');
         if (token === undefined) {
             throw new Error('GITHUB_TOKEN is not set');
@@ -61548,11 +61564,16 @@ async function run() {
         const githubMcpToken = coreExports.getInput('github-mcp-token') || token;
         const githubMcpToolsets = coreExports.getInput('github-mcp-toolsets');
         const endpoint = coreExports.getInput('endpoint');
+        // Get temperature and topP (prompt YAML modelParameters takes precedence over action inputs)
+        const temperatureInput = coreExports.getInput('temperature');
+        const topPInput = coreExports.getInput('top-p');
+        const temperature = promptConfig?.modelParameters?.temperature ?? (temperatureInput ? parseFloat(temperatureInput) : undefined);
+        const topP = promptConfig?.modelParameters?.topP ?? (topPInput ? parseFloat(topPInput) : undefined);
         // Parse custom headers
         const customHeadersInput = coreExports.getInput('custom-headers');
         const customHeaders = parseCustomHeaders(customHeadersInput);
         // Build the inference request with pre-processed messages and response format
-        const inferenceRequest = buildInferenceRequest(promptConfig, systemPrompt, prompt, modelName, promptConfig?.modelParameters?.temperature, promptConfig?.modelParameters?.topP, maxTokens, endpoint, token, customHeaders);
+        const inferenceRequest = buildInferenceRequest(promptConfig, systemPrompt, prompt, modelName, temperature, topP, maxTokens, maxCompletionTokens, endpoint, token, customHeaders);
         const enableMcp = coreExports.getBooleanInput('enable-github-mcp') || false;
         let modelResponse = null;
         if (enableMcp) {

package-lock.json

@@ -11,6 +11,7 @@
       "dependencies": {
         "@actions/core": "^1.11.1",
         "@modelcontextprotocol/sdk": "^1.25.2",
+        "@rollup/rollup-linux-x64-gnu": "4.57.1",
         "@types/tmp": "^0.2.6",
         "js-yaml": "^4.1.1",
         "openai": "^5.11.0",
@@ -2389,9 +2390,9 @@
       ]
     },
     "node_modules/@rollup/rollup-linux-x64-gnu": {
-      "version": "4.52.5",
-      "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-x64-gnu/-/rollup-linux-x64-gnu-4.52.5.tgz",
-      "integrity": "sha512-hXGLYpdhiNElzN770+H2nlx+jRog8TyynpTVzdlc6bndktjKWyZyiCsuDAlpd+j+W+WNqfcyAWz9HxxIGfZm1Q==",
+      "version": "4.57.1",
+      "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-x64-gnu/-/rollup-linux-x64-gnu-4.57.1.tgz",
+      "integrity": "sha512-ABca4ceT4N+Tv/GtotnWAeXZUZuM/9AQyCyKYyKnpk4yoA7QIAuBt6Hkgpw8kActYlew2mvckXkvx0FfoInnLg==",
       "cpu": [
         "x64"
       ],
@@ -6653,9 +6654,9 @@
       }
     },
     "node_modules/lodash": {
-      "version": "4.17.21",
-      "resolved": "https://registry.npmjs.org/lodash/-/lodash-4.17.21.tgz",
-      "integrity": "sha512-v2kDEe57lecTulaDIuNTPy3Ry4gLGJ6Z1O3vE1krgXZNrsQ+LFTGHVxVjcXPs17LhbZVGedAJv8XZ1tvj5FvSg==",
+      "version": "4.17.23",
+      "resolved": "https://registry.npmjs.org/lodash/-/lodash-4.17.23.tgz",
+      "integrity": "sha512-LgVTMpQtIopCi79SJeDiP0TfWi5CNEc/L/aRdTh3yIvmZXTnheWpKjSZhnvMl8iXbC1tFg9gdHHDMLoV7CnG+w==",
       "dev": true,
       "license": "MIT"
     },

src/helpers.ts

@@ -162,7 +162,8 @@ export function buildInferenceRequest(
   modelName: string,
   temperature: number | undefined,
   topP: number | undefined,
-  maxTokens: number,
+  maxTokens: number | undefined, // Deprecated
+  maxCompletionTokens: number | undefined,
   endpoint: string,
   token: string,
   customHeaders?: Record<string, string>,
@@ -175,7 +176,8 @@ export function buildInferenceRequest(
     modelName,
     temperature,
     topP,
-    maxTokens,
+    maxTokens, // Deprecated
+    maxCompletionTokens,
     endpoint,
     token,
     responseFormat,

src/inference.ts

@@ -12,7 +12,8 @@ interface ChatMessage {
 export interface InferenceRequest {
   messages: Array<{role: 'system' | 'user' | 'assistant' | 'tool'; content: string}>
   modelName: string
-  maxTokens: number
+  maxTokens?: number // Deprecated
+  maxCompletionTokens?: number
   endpoint: string
   token: string
   temperature?: number
@@ -33,6 +34,20 @@ export interface InferenceResponse {
   }>
 }
 
+/**
+ * Build according to what input was passed, default to max_tokens.
+ * Only one of max_tokens or max_completion_tokens will be set.
+ */
+function buildMaxTokensParam(request: InferenceRequest): {max_tokens?: number; max_completion_tokens?: number} {
+  if (request.maxCompletionTokens != null) {
+    return {max_completion_tokens: request.maxCompletionTokens}
+  }
+  if (request.maxTokens != null) {
+    return {max_tokens: request.maxTokens}
+  }
+  return {}
+}
+
 /**
  * Simple one-shot inference without tools
  */
@@ -47,10 +62,10 @@ export async function simpleInference(request: InferenceRequest): Promise<string
 
   const chatCompletionRequest: OpenAI.Chat.Completions.ChatCompletionCreateParams = {
     messages: request.messages as OpenAI.Chat.Completions.ChatCompletionMessageParam[],
-    max_tokens: request.maxTokens,
     model: request.modelName,
     temperature: request.temperature,
     top_p: request.topP,
+    ...buildMaxTokensParam(request), // Note: solution around models using different underlying max tokens properties
   }
 
   // Add response format if specified
@@ -95,10 +110,10 @@ export async function mcpInference(
 
     const chatCompletionRequest: OpenAI.Chat.Completions.ChatCompletionCreateParams = {
       messages: messages as OpenAI.Chat.Completions.ChatCompletionMessageParam[],
-      max_tokens: request.maxTokens,
       model: request.modelName,
       temperature: request.temperature,
       top_p: request.topP,
+      ...buildMaxTokensParam(request), // Note: solution around models using different underlying max tokens properties
     }
 
     // Add response format if specified (only on final iteration to avoid conflicts with tool calls)

src/main.ts

@@ -48,11 +48,14 @@ export async function run(): Promise<void> {
 
     // Get common parameters
     const modelName = promptConfig?.model || core.getInput('model')
-    let maxTokens = promptConfig?.modelParameters?.maxTokens ?? core.getInput('max-tokens')
 
-    if (typeof maxTokens === 'string') {
-      maxTokens = parseInt(maxTokens, 10)
-    }
+    // Parse token limit inputs
+    const maxCompletionTokensInput =
+      promptConfig?.modelParameters?.maxCompletionTokens ?? core.getInput('max-completion-tokens')
+    const maxCompletionTokens = maxCompletionTokensInput ? Number(maxCompletionTokensInput) : undefined
+
+    const maxTokensInput = promptConfig?.modelParameters?.maxTokens ?? core.getInput('max-tokens')
+    const maxTokens = maxCompletionTokens != null ? undefined : maxTokensInput ? Number(maxTokensInput) : undefined
 
     const token = process.env['GITHUB_TOKEN'] || core.getInput('token')
     if (token === undefined) {
@@ -65,6 +68,13 @@ export async function run(): Promise<void> {
 
     const endpoint = core.getInput('endpoint')
 
+    // Get temperature and topP (prompt YAML modelParameters takes precedence over action inputs)
+    const temperatureInput = core.getInput('temperature')
+    const topPInput = core.getInput('top-p')
+    const temperature =
+      promptConfig?.modelParameters?.temperature ?? (temperatureInput ? parseFloat(temperatureInput) : undefined)
+    const topP = promptConfig?.modelParameters?.topP ?? (topPInput ? parseFloat(topPInput) : undefined)
+
     // Parse custom headers
     const customHeadersInput = core.getInput('custom-headers')
     const customHeaders = parseCustomHeaders(customHeadersInput)
@@ -75,9 +85,10 @@ export async function run(): Promise<void> {
       systemPrompt,
       prompt,
       modelName,
-      promptConfig?.modelParameters?.temperature,
-      promptConfig?.modelParameters?.topP,
+      temperature,
+      topP,
       maxTokens,
+      maxCompletionTokens,
       endpoint,
       token,
       customHeaders,

src/prompt.ts

@@ -8,7 +8,8 @@ export interface PromptMessage {
 }
 
 export interface ModelParameters {
-  maxTokens?: number
+  maxTokens?: number // Deprecated
+  maxCompletionTokens?: number
   temperature?: number
   topP?: number
 }