implement passing two action input properties to cover all model scenarios

__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',
   }
@@ -633,4 +633,35 @@ 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')
+    })
+  })
 })

__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,13 @@ inputs:
     required: false
     default: ''
   max-tokens:
-    description: The maximum number of tokens to generate
+    description: Maximum tokens to generate (deprecated)
     required: false
     default: '200'
+  max-completion-tokens:
+    description: Maximum tokens to generate
+    required: false
+    default: ''
   temperature:
     description: The sampling temperature to use (0-1)
     required: false

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 the token limit params for a chat completion request.
+ * 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_completion_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_completion_tokens: request.maxTokens,
       model: request.modelName,
       temperature: request.temperature,
       top_p: request.topP,
+      ...buildMaxTokensParam(request),
     }
 
     // Add response format if specified (only on final iteration to avoid conflicts with tool calls)
@@ -162,9 +177,7 @@ export async function mcpInference(
 }
 
 /**
- * Wrapper around OpenAI chat.completions.create with defensive handling for cases where
- * the SDK returns a raw string (e.g., unexpected content-type or streaming body) instead of
- * a parsed object. Ensures an object with a 'choices' array is returned or throws a descriptive error.
+ * Wrapper around OpenAI chat.completions.create with response validation.
  */
 async function chatCompletion(
   client: OpenAI,

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) {
@@ -85,6 +88,7 @@ export async function run(): Promise<void> {
       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
 }