Skip to content

Commit e09e659

Browse files
stephaniegiangstephaniegiang
authored
Merge pull request #173 from GitPaulo/main
Support passing max_tokens and max_completion_tokens
2 parents a380166 + e608d2b commit e09e659

File tree

11 files changed

+175
-39
lines changed

11 files changed

+175
-39
lines changed

‎README.md‎

Lines changed: 45 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,33 @@ supplied via the `input` parameter in YAML format. Additionally, you can
123123
provide file-based variables via `file_input`, where each key maps to a file
124124
path.
125125

126+
### Prompt.yml with model parameters
127+
128+
You can specify model parameters directly in your `.prompt.yml` files using the
129+
`modelParameters` key:
130+
131+
```yaml
132+
messages:
133+
- role: system
134+
content: Be as concise as possible
135+
- role: user
136+
content: 'Compare {{a}} and {{b}}, please'
137+
model: openai/gpt-4o
138+
modelParameters:
139+
maxCompletionTokens: 500
140+
temperature: 0.7
141+
```
142+
143+
| Key | Type | Description |
144+
| --------------------- | ------ | ----------------------------------------------------- |
145+
| `maxCompletionTokens` | number | The maximum number of tokens to generate |
146+
| `maxTokens` | number | The maximum number of tokens to generate (deprecated) |
147+
| `temperature` | number | The sampling temperature to use (0-1) |
148+
| `topP` | number | The nucleus sampling parameter to use (0-1) |
149+
150+
> ![Note]
151+
> Parameters set in `modelParameters` take precedence over the corresponding action inputs.
152+
126153
### Using a system prompt file
127154

128155
In addition to the regular prompt, you can provide a system prompt file instead
@@ -276,23 +303,24 @@ perform actions like searching issues and PRs.
276303
Various inputs are defined in [`action.yml`](action.yml) to let you configure
277304
the action:
278305

279-
| Name | Description | Default |
280-
| -------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------ |
281-
| `token` | Token to use for inference. Typically the GITHUB_TOKEN secret | `github.token` |
282-
| `prompt` | The prompt to send to the model | N/A |
283-
| `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 | `""` |
284-
| `input` | Template variables in YAML format for .prompt.yml files (e.g., `var1: value1` on separate lines) | `""` |
285-
| `file_input` | Template variables in YAML where values are file paths. The file contents are read and used for templating | `""` |
286-
| `system-prompt` | The system prompt to send to the model | `"You are a helpful assistant"` |
287-
| `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 | `""` |
288-
| `model` | The model to use for inference. Must be available in the [GitHub Models](http://github.com/marketplace?type=models) catalog | `openai/gpt-4o` |
289-
| `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` |
290-
| `max-tokens` | The max number of tokens to generate | 200 |
291-
| `temperature` | The sampling temperature to use (0-1) | `""` |
292-
| `top-p` | The nucleus sampling parameter to use (0-1) | `""` |
293-
| `enable-github-mcp` | Enable Model Context Protocol integration with GitHub tools | `false` |
294-
| `github-mcp-token` | Token to use for GitHub MCP server (defaults to the main token if not specified). | `""` |
295-
| `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. | `""` |
306+
| Name | Description | Default |
307+
| ----------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------ |
308+
| `token` | Token to use for inference. Typically the GITHUB_TOKEN secret | `github.token` |
309+
| `prompt` | The prompt to send to the model | N/A |
310+
| `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 | `""` |
311+
| `input` | Template variables in YAML format for .prompt.yml files (e.g., `var1: value1` on separate lines) | `""` |
312+
| `file_input` | Template variables in YAML where values are file paths. The file contents are read and used for templating | `""` |
313+
| `system-prompt` | The system prompt to send to the model | `"You are a helpful assistant"` |
314+
| `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 | `""` |
315+
| `model` | The model to use for inference. Must be available in the [GitHub Models](http://github.com/marketplace?type=models) catalog | `openai/gpt-4o` |
316+
| `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` |
317+
| `max-tokens` | The maximum number of tokens to generate (deprecated, use `max-completion-tokens` instead) | 200 |
318+
| `max-completion-tokens` | The maximum number of tokens to generate | `""` |
319+
| `temperature` | The sampling temperature to use (0-1) | `""` |
320+
| `top-p` | The nucleus sampling parameter to use (0-1) | `""` |
321+
| `enable-github-mcp` | Enable Model Context Protocol integration with GitHub tools | `false` |
322+
| `github-mcp-token` | Token to use for GitHub MCP server (defaults to the main token if not specified). | `""` |
323+
| `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. | `""` |
296324

297325
## Outputs
298326

‎__tests__/helpers-inference.test.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ describe('helpers.ts - inference request building', () => {
109109
undefined,
110110
undefined,
111111
100,
112+
undefined,
112113
'https://api.test.com',
113114
'test-token',
114115
)
@@ -122,6 +123,7 @@ describe('helpers.ts - inference request building', () => {
122123
temperature: undefined,
123124
topP: undefined,
124125
maxTokens: 100,
126+
maxCompletionTokens: undefined,
125127
endpoint: 'https://api.test.com',
126128
token: 'test-token',
127129
responseFormat: {
@@ -143,6 +145,7 @@ describe('helpers.ts - inference request building', () => {
143145
undefined,
144146
undefined,
145147
100,
148+
undefined,
146149
'https://api.test.com',
147150
'test-token',
148151
)
@@ -156,6 +159,7 @@ describe('helpers.ts - inference request building', () => {
156159
temperature: undefined,
157160
topP: undefined,
158161
maxTokens: 100,
162+
maxCompletionTokens: undefined,
159163
endpoint: 'https://api.test.com',
160164
token: 'test-token',
161165
responseFormat: undefined,

‎__tests__/inference.test.ts‎

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ describe('inference.ts', () => {
3131
{role: 'user' as const, content: 'Hello, AI!'},
3232
],
3333
modelName: 'gpt-4',
34-
maxTokens: 100,
34+
maxCompletionTokens: 100,
3535
endpoint: 'https://api.test.com',
3636
token: 'test-token',
3737
}
@@ -633,4 +633,64 @@ describe('inference.ts', () => {
633633
expect(result).toBe('{"immediate": "result"}')
634634
})
635635
})
636+
637+
describe('token param routing', () => {
638+
it('sends max_tokens when only maxTokens is set', async () => {
639+
const requestWithMaxTokens = {
640+
...mockRequest,
641+
maxCompletionTokens: undefined,
642+
maxTokens: 100,
643+
}
644+
645+
const mockResponse = {
646+
choices: [
647+
{
648+
message: {
649+
content: 'Direct max_tokens response',
650+
},
651+
},
652+
],
653+
}
654+
655+
mockCreate.mockResolvedValueOnce(mockResponse)
656+
657+
const result = await simpleInference(requestWithMaxTokens)
658+
659+
expect(result).toBe('Direct max_tokens response')
660+
expect(mockCreate).toHaveBeenCalledTimes(1)
661+
662+
// Should have sent max_tokens directly
663+
expect(mockCreate.mock.calls[0][0]).toHaveProperty('max_tokens', 100)
664+
expect(mockCreate.mock.calls[0][0]).not.toHaveProperty('max_completion_tokens')
665+
})
666+
667+
it('sends neither token param when both are undefined', async () => {
668+
const requestWithNoTokens = {
669+
...mockRequest,
670+
maxCompletionTokens: undefined,
671+
maxTokens: undefined,
672+
}
673+
674+
const mockResponse = {
675+
choices: [
676+
{
677+
message: {
678+
content: 'No token limit response',
679+
},
680+
},
681+
],
682+
}
683+
684+
mockCreate.mockResolvedValueOnce(mockResponse)
685+
686+
const result = await simpleInference(requestWithNoTokens)
687+
688+
expect(result).toBe('No token limit response')
689+
expect(mockCreate).toHaveBeenCalledTimes(1)
690+
691+
const params = mockCreate.mock.calls[0][0]
692+
expect(params).not.toHaveProperty('max_tokens')
693+
expect(params).not.toHaveProperty('max_completion_tokens')
694+
})
695+
})
636696
})

‎__tests__/main.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,7 @@ describe('main.ts', () => {
168168
],
169169
modelName: 'gpt-4',
170170
maxTokens: 100,
171+
maxCompletionTokens: undefined,
171172
endpoint: 'https://api.test.com',
172173
token: 'fake-token',
173174
responseFormat: undefined,
@@ -259,6 +260,7 @@ describe('main.ts', () => {
259260
],
260261
modelName: 'gpt-4',
261262
maxTokens: 100,
263+
maxCompletionTokens: undefined,
262264
endpoint: 'https://api.test.com',
263265
token: 'fake-token',
264266
responseFormat: undefined,

‎action.yml‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,13 @@ inputs:
4343
required: false
4444
default: ''
4545
max-tokens:
46-
description: The maximum number of tokens to generate
46+
description: The maximum number of tokens to generate (deprecated)
4747
required: false
4848
default: '200'
49+
max-completion-tokens:
50+
description: The maximum number of tokens to generate
51+
required: false
52+
default: ''
4953
temperature:
5054
description: The sampling temperature to use (0-1)
5155
required: false

‎dist/index.js‎

Lines changed: 25 additions & 9 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎dist/index.js.map‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)