Merge pull request #99 from ca-dp/add-model-select

Add commands for selecting OpenAI models

.github/workflows/code-butler.yml

@@ -18,6 +18,7 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           cmd: review
+          model: gpt-4-1106-preview
   chat:
     if: startsWith(github.event.comment.body, '/chat')
     runs-on: ubuntu-latest
@@ -27,4 +28,5 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           cmd: chat
+          model: gpt-4-1106-preview
           comment_body: ${{ github.event.comment.body }}

__tests__/ai.test.ts

@@ -71,4 +71,42 @@ describe('completionRequest', () => {
     )
     consoleLogSpy.mockRestore()
   })
+
+  it('model gpt4 should be used', async () => {
+    const apiKey = 'test-api-key'
+    const systemPrompt = 'System prompt'
+    const userPrompt = 'User prompt'
+    const model = 'gpt-4'
+
+    // Mock the OpenAI API request
+    // Mock the OpenAI API request
+    nock('https://api.openai.com')
+      .post('/v1/chat/completions', {
+        messages: [
+          { role: 'system', content: systemPrompt },
+          { role: 'user', content: userPrompt }
+        ],
+        model: 'gpt-4'
+      })
+      .reply(200, {
+        choices: [
+          {
+            message: {
+              content: 'Expected completed content'
+            }
+          }
+        ]
+      })
+
+    // Call the completionRequest function to test it
+    const result = await completionRequest(
+      apiKey,
+      systemPrompt,
+      userPrompt,
+      model
+    )
+
+    // Assert based on the expected result
+    expect(result).toBe('Expected completed content')
+  })
 })

__tests__/main.test.ts

@@ -37,6 +37,8 @@ describe('run', () => {
           return 'mocked-api-key'
         case 'cmd':
           return 'review'
+        case 'model':
+          return 'mock ai model'
         default:
           return ''
       }
@@ -61,7 +63,8 @@ describe('run', () => {
     expect(completionRequestMock).toHaveBeenCalledWith(
       'mocked-api-key',
       'Mocked system prompt',
-      'Mocked PR diff'
+      'Mocked PR diff',
+      'mock ai model'
     )
     expect(createGitHubCommentMock).toHaveBeenCalledWith('Mocked AI response')
   })
@@ -159,7 +162,8 @@ describe('run', () => {
     expect(completionRequestMock).toHaveBeenCalledWith(
       'mocked-api-key',
       'Mocked system prompt',
-      'Mocked issue comment'
+      'Mocked issue comment',
+      ''
     )
     expect(createGitHubCommentMock).toHaveBeenCalledWith('Mocked AI response')
   })

package.json

@@ -68,7 +68,7 @@
     "@actions/core": "^1.10.1",
     "@actions/github": "^5.1.1",
     "js-tiktoken": "^1.0.7",
-    "openai": "^4.10.0"
+    "openai": "^4.19.1"
   },
   "devDependencies": {
     "@types/jest": "^29.5.8",

src/ai.ts

@@ -3,18 +3,20 @@ import OpenAI from 'openai'
 export async function completionRequest(
   apiKey: string,
   systemPrompt: string,
-  userPrompt: string
+  userPrompt: string,
+  model: string | undefined = undefined
 ): Promise<string> {
   try {
     const openai = new OpenAI({ apiKey })
     const requestTimeout = 300 * 1000 // 5 minutes
+    const targetModel = model || 'gpt-3.5-turbo'
     const response = await openai.chat.completions.create(
       {
         messages: [
           { role: 'system', content: systemPrompt },
           { role: 'user', content: userPrompt }
         ],
-        model: 'gpt-3.5-turbo'
+        model: targetModel
       },
       {
         timeout: requestTimeout

src/main.ts

@@ -22,10 +22,12 @@ export async function run(): Promise<void> {
 
         for (const group of groups) {
           const diffToSend = group.join('')
+          const model = core.getInput('model', { required: false })
           const messagePromise = ai.completionRequest(
             core.getInput('OPENAI_API_KEY', { required: true }),
             sysPrompt,
-            diffToSend
+            diffToSend,
+            model
           )
           const message = await messagePromise
 
@@ -46,10 +48,12 @@ export async function run(): Promise<void> {
         }
 
         const chatSystemPrompt = prompt.getChatSystemPrompt()
+        const model = core.getInput('model', { required: false })
         const responseMessage = ai.completionRequest(
           core.getInput('OPENAI_API_KEY', { required: true }),
           chatSystemPrompt,
-          comment
+          comment,
+          model
         )
         const response = await responseMessage
         if (response === '') {