Skip to content

Model selection for code review tool #2196

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ashishrp-aws merged 6 commits into from
Sep 3, 2025
Merged

Model selection for code review tool #2196

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
6a244bf
feat: enabling model selection for code review tool
mkovelam Sep 2, 2025
67473ae
feat: added unit test for model selection for code review tool
mkovelam Sep 2, 2025
bb96f18
Merge branch 'main' into modelSelectionForCodeReviewTool
laileni-aws Sep 2, 2025
91eb7ca
fix: making modelId required for codeReviewTool in input schema valid…
mkovelam Sep 2, 2025
9a3c147
fix: fixing unit tests
mkovelam Sep 3, 2025
9b68466
Merge branch 'main' into modelSelectionForCodeReviewTool
laileni-aws Sep 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions server/aws-lsp-codewhisperer/src/client/token/bearer-token-service.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4822,6 +4822,12 @@
},
"profileArn": {
"shape": "ProfileArn"
},
"languageModelId": {
"shape": "ModelId"
},
"clientType": {
Copy link
Contributor

@laileni-aws laileni-aws Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For my understanding, Why do we need clientType?
Do we pass IDE || CLI etc in this ? to separate CLI requests from IDE?

Copy link
Contributor Author

@mkovelam mkovelam Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah we want to send the clientType to distinguish between IDE and CLI. To use this as a tenant after getting the capacity for ReviewBird.

"shape": "Origin"
}
}
},
Expand Down
4 changes: 3 additions & 1 deletion server/aws-lsp-codewhisperer/src/client/token/codewhispererbearertokenclient.d.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1539,6 +1539,8 @@ declare namespace CodeWhispererBearerTokenClient {
codeScanName?: CodeScanName;
codeDiffMetadata?: CodeDiffMetadata;
profileArn?: ProfileArn;
languageModelId?: ModelId;
clientType?: Origin;
Comment on lines +1542 to +1543
Copy link
Contributor

@BlakeLazarine BlakeLazarine Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you cant modify the d.ts I am pretty sure. This is a built file I believe

Copy link
Contributor Author

@mkovelam mkovelam Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So, I didn't modify it directly. I used another script to generate this. Should this not be included in the PR/commit ?

Copy link
Contributor

@BlakeLazarine BlakeLazarine Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you need to modify the RTS model file, and re-build this (Im pretty sure that is the process)

Copy link
Contributor

@BlakeLazarine BlakeLazarine Sep 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is fine then

Copy link
Contributor

@ashishrp-aws ashishrp-aws Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you need to modify it. This is fine

}
export type StartCodeAnalysisRequestClientTokenString = string;
export interface StartCodeAnalysisResponse {
Expand Down Expand Up @@ -2186,4 +2188,4 @@ declare namespace CodeWhispererBearerTokenClient {
}
export = CodeWhispererBearerTokenClient;


1 change: 1 addition & 0 deletions server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@
* Will be deleted or merged.
*/

import * as crypto from 'crypto'

Check warning on line 6 in server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

View workflow job for this annotation

GitHub Actions / Test 

Do not import Node.js builtin module "crypto"

Check warning on line 6 in server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

View workflow job for this annotation

GitHub Actions / Test (Windows) 

Do not import Node.js builtin module "crypto"
import * as path from 'path'

Check warning on line 7 in server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

View workflow job for this annotation

GitHub Actions / Test 

Do not import Node.js builtin module "path"

Check warning on line 7 in server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

View workflow job for this annotation

GitHub Actions / Test (Windows) 

Do not import Node.js builtin module "path"
import * as os from 'os'

Check warning on line 8 in server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

View workflow job for this annotation

GitHub Actions / Test 

Do not import Node.js builtin module "os"

Check warning on line 8 in server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

View workflow job for this annotation

GitHub Actions / Test (Windows) 

Do not import Node.js builtin module "os"
import {
ChatTriggerType,
Origin,
Expand Down Expand Up @@ -1989,6 +1989,7 @@
.filter(c => c.type === 'rule')
.map(c => ({ path: c.path }))
}
initialInput['modelId'] = session.modelId
toolUse.input = initialInput
} catch (e) {
this.#features.logging.warn(`could not parse CodeReview tool input: ${e}`)
Expand Down
61 changes: 61 additions & 0 deletions ...-lsp-codewhisperer/src/language-server/agenticChat/tools/qCodeAnalysis/codeReview.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import * as path from 'path'
import { expect } from 'chai'
import { CancellationError } from '@aws/lsp-core'
import * as JSZip from 'jszip'
import { Origin } from '@amzn/codewhisperer-streaming'

describe('CodeReview', () => {
let sandbox: sinon.SinonSandbox
Expand Down Expand Up @@ -103,6 +104,7 @@ describe('CodeReview', () => {
folderLevelArtifacts: [],
ruleArtifacts: [],
scopeOfReview: FULL_REVIEW,
modelId: 'claude-4-sonnet',
}
})

Expand Down Expand Up @@ -143,6 +145,61 @@ describe('CodeReview', () => {
expect(result.output.kind).to.equal('json')
})

it('should execute successfully and pass languageModelId and clientType to startCodeAnalysis', async () => {
const inputWithModelId = {
...validInput,
modelId: 'test-model-789',
}

// Setup mocks for successful execution
mockCodeWhispererClient.createUploadUrl.resolves({
uploadUrl: 'https://upload.com',
uploadId: 'upload-123',
requestHeaders: {},
})

mockCodeWhispererClient.startCodeAnalysis.resolves({
jobId: 'job-123',
status: 'Pending',
})

mockCodeWhispererClient.getCodeAnalysis.resolves({
status: 'Completed',
})

mockCodeWhispererClient.listCodeAnalysisFindings.resolves({
codeAnalysisFindings: '[]',
nextToken: undefined,
})

sandbox.stub(CodeReviewUtils, 'uploadFileToPresignedUrl').resolves()
sandbox.stub(codeReview as any, 'prepareFilesAndFoldersForUpload').resolves({
zipBuffer: Buffer.from('test'),
md5Hash: 'hash123',
isCodeDiffPresent: false,
programmingLanguages: new Set(['javascript']),
})
sandbox.stub(codeReview as any, 'parseFindings').returns([])

const result = await codeReview.execute(inputWithModelId, context)

expect(result.output.success).to.be.true
expect(result.output.kind).to.equal('json')

// Verify that startCodeAnalysis was called with the correct parameters
expect(mockCodeWhispererClient.startCodeAnalysis.calledOnce).to.be.true
const startAnalysisCall = mockCodeWhispererClient.startCodeAnalysis.getCall(0)
const callArgs = startAnalysisCall.args[0]

expect(callArgs).to.have.property('languageModelId', 'test-model-789')
expect(callArgs).to.have.property('clientType', Origin.IDE)
expect(callArgs).to.have.property('artifacts')
expect(callArgs).to.have.property('programmingLanguage')
expect(callArgs).to.have.property('clientToken')
expect(callArgs).to.have.property('codeScanName')
expect(callArgs).to.have.property('scope', 'AGENTIC')
})

it('should handle missing client error', async () => {
context.codeWhispererClient = undefined

Expand All @@ -160,6 +217,7 @@ describe('CodeReview', () => {
folderLevelArtifacts: [],
ruleArtifacts: [],
scopeOfReview: FULL_REVIEW,
modelId: 'claude-4-sonnet',
}

try {
Expand Down Expand Up @@ -279,6 +337,7 @@ describe('CodeReview', () => {
folderLevelArtifacts: [],
ruleArtifacts: [],
scopeOfReview: FULL_REVIEW,
modelId: 'claude-4-sonnet',
}

const context = {
Expand All @@ -303,6 +362,7 @@ describe('CodeReview', () => {
folderLevelArtifacts: [{ path: '/test/folder' }],
ruleArtifacts: [],
scopeOfReview: CODE_DIFF_REVIEW,
modelId: 'claude-4-sonnet',
}

const context = {
Expand Down Expand Up @@ -614,6 +674,7 @@ describe('CodeReview', () => {
folderLevelArtifacts: [],
ruleArtifacts: [],
scopeOfReview: FULL_REVIEW,
modelId: 'claude-4-sonnet',
}

// Make prepareFilesAndFoldersForUpload throw an error
Expand Down
11 changes: 10 additions & 1 deletion ...r/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/qCodeAnalysis/codeReview.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ import {
SuccessMetricName,
} from './codeReviewTypes'
import { CancellationError } from '@aws/lsp-core'
import { Origin } from '@amzn/codewhisperer-streaming'

export class CodeReview {
private static readonly CUSTOMER_CODE_BASE_PATH = 'customerCodeBaseFolder'
Expand Down Expand Up @@ -158,6 +159,7 @@ export class CodeReview {
const fileArtifacts = validatedInput.fileLevelArtifacts || []
const folderArtifacts = validatedInput.folderLevelArtifacts || []
const ruleArtifacts = validatedInput.ruleArtifacts || []
const modelId = validatedInput.modelId

if (fileArtifacts.length === 0 && folderArtifacts.length === 0) {
CodeReviewUtils.emitMetric(
Expand All @@ -182,7 +184,7 @@ export class CodeReview {
const programmingLanguage = 'java'
const scanName = 'Standard-' + randomUUID()

this.logging.info(`Agentic scan name: ${scanName}`)
this.logging.info(`Agentic scan name: ${scanName} selectedModel: ${modelId}`)

return {
fileArtifacts,
Expand All @@ -192,6 +194,7 @@ export class CodeReview {
programmingLanguage,
scanName,
ruleArtifacts,
modelId,
}
}

Expand Down Expand Up @@ -274,6 +277,8 @@ export class CodeReview {
codeScanName: setup.scanName,
scope: CodeReview.SCAN_SCOPE,
codeDiffMetadata: uploadResult.isCodeDiffPresent ? { codeDiffPath: '/code_artifact/codeDiff/' } : undefined,
languageModelId: setup.modelId,
clientType: Origin.IDE,
})

if (!createResponse.jobId) {
Expand All @@ -293,6 +298,7 @@ export class CodeReview {
customRules: setup.ruleArtifacts.length,
programmingLanguages: Array.from(uploadResult.programmingLanguages),
scope: setup.isFullReviewRequest ? FULL_REVIEW : CODE_DIFF_REVIEW,
modelId: setup.modelId,
},
},
this.logging,
Expand Down Expand Up @@ -349,6 +355,7 @@ export class CodeReview {
programmingLanguages: Array.from(uploadResult.programmingLanguages),
scope: setup.isFullReviewRequest ? FULL_REVIEW : CODE_DIFF_REVIEW,
status: status,
modelId: setup.modelId,
},
},
this.logging,
Expand Down Expand Up @@ -382,6 +389,7 @@ export class CodeReview {
programmingLanguages: Array.from(uploadResult.programmingLanguages),
scope: setup.isFullReviewRequest ? FULL_REVIEW : CODE_DIFF_REVIEW,
status: status,
modelId: setup.modelId,
},
},
this.logging,
Expand Down Expand Up @@ -426,6 +434,7 @@ export class CodeReview {
programmingLanguages: Array.from(uploadResult.programmingLanguages),
scope: setup.isFullReviewRequest ? FULL_REVIEW : CODE_DIFF_REVIEW,
latency: Date.now() - this.toolStartTime,
modelId: setup.modelId,
},
},
this.logging,
Expand Down
1 change: 1 addition & 0 deletions ...sp-codewhisperer/src/language-server/agenticChat/tools/qCodeAnalysis/codeReviewSchemas.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,7 @@ export const Z_CODE_REVIEW_INPUT_SCHEMA = z.object({
})
)
.optional(),
modelId: z.string(),
})

/**
Expand Down
1 change: 1 addition & 0 deletions ...-lsp-codewhisperer/src/language-server/agenticChat/tools/qCodeAnalysis/codeReviewTypes.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ export type ValidateInputAndSetupResult = {
programmingLanguage: string
scanName: string
ruleArtifacts: RuleArtifacts
modelId?: string
}

export type PrepareAndUploadArtifactsResult = {
Expand Down
Loading