feat: update prompt content with response from API (a93a693b) · Commits · GitLab.org / gitlab-vscode-extension

src/common/chat/gitlab_chat_api.test.ts

+13 −4

Original line number	Diff line number	Diff line
		@@ -14,7 +14,13 @@ const mockedMutationResponse = {
		const mockedQueryResponse = {
		aiMessages: {
		nodes: [
		{ content: 'test', requestId: '123', errors: ['bar'], timestamp: '2023-01-01 01:01:01' },
		{
		content: 'test',
		requestId: '123',
		role: 'assistant',
		errors: ['bar'],
		timestamp: '2023-01-01 01:01:01',
		},
		],
		},
		};
		@@ -60,13 +66,16 @@ describe('GitLabChatApi', () => {
		};

		describe('getAiMessage', () => {
		it('returns first assistant message with given requestId', async () => {
		it('returns first message with given requestId and role', async () => {
		const manager = createManager(currentProject, mockedQueryResponse);
		const gitlabChatApi = new GitLabChatApi(manager);

		const expectedMessage = mockedQueryResponse.aiMessages.nodes[0];

		const response = await gitlabChatApi.pullAiMessage(expectedMessage.requestId);
		const response = await gitlabChatApi.pullAiMessage(
		expectedMessage.requestId,
		expectedMessage.role,
		);

		assert(response.type === 'message');

		@@ -83,7 +92,7 @@ describe('GitLabChatApi', () => {
		const manager = createManager(currentProject, mockedEmptyQueryResponse);
		const gitlabChatApi = new GitLabChatApi(manager);

		const response = await gitlabChatApi.pullAiMessage('123');
		const response = await gitlabChatApi.pullAiMessage('123', 'assistant');

		expect(response.requestId).toBe('123');
		expect(response.errors).toContainEqual('Reached timeout while fetching response.');

src/common/chat/gitlab_chat_api.ts

+9 −6

Original line number	Diff line number	Diff line
		@@ -19,8 +19,8 @@ export type AiActionResponseType = {
		};

		export const AI_MESSAGES_QUERY = gql`
		query getAiMessages($requestIds: [ID!]) {
		aiMessages(requestIds: $requestIds, roles: [ASSISTANT]) {
		query getAiMessages($requestIds: [ID!], $roles: [AiChatMessageRole!]) {
		aiMessages(requestIds: $requestIds, roles: $roles) {
		nodes {
		requestId
		role
		@@ -98,8 +98,8 @@ export class GitLabChatApi {
		return this.sendAiAction(AI_ACTIONS.chat, { question });
		}

		async pullAiMessage(requestId: string): Promise<AiMessage> {
		const response = await pullHandler(() => this.getAiMessage(requestId));
		async pullAiMessage(requestId: string, role: string): Promise<AiMessage> {
		const response = await pullHandler(() => this.getAiMessage(requestId, role));

		if (!response)
		return new Promise(resolve => {
		@@ -118,11 +118,14 @@ export class GitLabChatApi {
		return platform;
		}

		private async getAiMessage(requestId: string): Promise<AiMessageResponseType \| undefined> {
		private async getAiMessage(
		requestId: string,
		role: string,
		): Promise<AiMessageResponseType \| undefined> {
		const request: GraphQLRequest<AiMessagesResponseType> = {
		type: 'graphql',
		query: AI_MESSAGES_QUERY,
		variables: { requestIds: [requestId] },
		variables: { requestIds: [requestId], roles: [role.toUpperCase()] },
		};
		const platform = await this.currentPlatform();
		const history = await platform.fetchFromApi(request);

src/common/chat/gitlab_chat_controller.test.ts

+19 −7

Original line number	Diff line number	Diff line
		@@ -35,11 +35,12 @@ describe('GitLabChatController', () => {
		},
		});

		apiMock.pullAiMessage = jest.fn().mockReturnValue({
		content: 'api response',
		requestId: 'uniqueId',
		apiMock.pullAiMessage = jest.fn().mockImplementation((requestId: string, role: string) => ({
		content: `api response ${role}`,
		role,
		requestId,
		timestamp: '2023-01-01 01:01:01',
		});
		}));
		});

		describe('resolveWebviewView', () => {
		@@ -94,7 +95,17 @@ describe('GitLabChatController', () => {

		expect(viewMock.updateRecord).toHaveBeenCalledWith(
		expect.objectContaining({
		content: 'api response',
		content: 'api response user',
		state: 'ready',
		role: 'user',
		requestId: 'uniqueId',
		timestamp: Date.parse('2023-01-01 01:01:01'),
		}),
		);

		expect(viewMock.updateRecord).toHaveBeenCalledWith(
		expect.objectContaining({
		content: 'api response assistant',
		state: 'ready',
		role: 'assistant',
		requestId: 'uniqueId',
		@@ -110,14 +121,15 @@ describe('GitLabChatController', () => {

		expect(controller.chatHistory[0]).toEqual(
		expect.objectContaining({
		content: 'hello',
		content: 'api response user',
		state: 'ready',
		role: 'user',
		requestId: 'uniqueId',
		}),
		);
		expect(controller.chatHistory[1]).toEqual(
		expect.objectContaining({
		content: 'api response',
		content: 'api response assistant',
		state: 'ready',
		role: 'assistant',
		requestId: 'uniqueId',

src/common/chat/gitlab_chat_controller.ts

+19 −21

Original line number	Diff line number	Diff line
		@@ -61,26 +61,11 @@ export class GitLabChatController implements vscode.WebviewViewProvider {
		role: 'assistant',
		content: '...',
		state: 'pending',
		requestId: record.requestId,
		});

		await this.addToChat(responseRecord);

		if (!record.requestId) {
		break;
		}

		const response = await this.#api.pullAiMessage(record.requestId);

		if (response.type !== 'error') {
		responseRecord.content = response.content;
		responseRecord.setTimestamp(response.timestamp);
		}

		responseRecord.errors = response.errors;
		responseRecord.requestId = response.requestId;
		responseRecord.state = 'ready';

		await this.#view.updateRecord(responseRecord);
		await Promise.all([this.refreshRecord(record), this.refreshRecord(responseRecord)]);
		}
		}
		}
		@@ -101,9 +86,22 @@ export class GitLabChatController implements vscode.WebviewViewProvider {

		const actionResponse = await this.#api.processNewUserPrompt(record.content);

		// eslint-disable-next-line no-param-reassign
		record.requestId = actionResponse.aiAction.requestId;
		// eslint-disable-next-line no-param-reassign
		record.errors = actionResponse.aiAction.errors;
		record.update(actionResponse.aiAction);
		}

		private async refreshRecord(record: GitLabChatRecord) {
		if (!record.requestId) {
		throw Error('requestId must be present!');
		}

		const apiResponse = await this.#api.pullAiMessage(record.requestId, record.role);

		if (apiResponse.type !== 'error') {
		record.update({ content: apiResponse.content, timestamp: apiResponse.timestamp });
		}

		record.update({ errors: apiResponse.errors, state: 'ready' });

		await this.#view.updateRecord(record);
		}
		}

src/common/chat/gitlab_chat_record.ts

+18 −12

Original line number	Diff line number	Diff line
		@@ -4,6 +4,17 @@ type ChatRecordRole = 'user' \| 'assistant' \| 'system';
		type ChatRecordState = 'pending' \| 'ready';
		type ChatRecordType = 'general' \| 'explainCode' \| 'newConversation';

		type GitLabChatRecordAttributes = {
		type?: ChatRecordType;
		role: ChatRecordRole;
		content?: string;
		requestId?: string;
		state?: ChatRecordState;
		payload?: object;
		errors?: string[];
		timestamp?: string;
		};

		export class GitLabChatRecord {
		role: ChatRecordRole;

		@@ -32,16 +43,7 @@ export class GitLabChatRecord {
		payload,
		errors,
		timestamp,
		}: {
		type?: ChatRecordType;
		role: ChatRecordRole;
		content?: string;
		requestId?: string;
		state?: ChatRecordState;
		payload?: object;
		errors?: string[];
		timestamp?: string;
		}) {
		}: GitLabChatRecordAttributes) {
		this.role = role;
		this.content = content;
		this.type = type ?? this.detectType();
		@@ -53,8 +55,12 @@ export class GitLabChatRecord {
		this.timestamp = timestamp ? Date.parse(timestamp) : Date.now();
		}

		setTimestamp(timestamp?: string) {
		this.timestamp = timestamp ? Date.parse(timestamp) : Date.now();
		update(attributes: Partial<GitLabChatRecordAttributes>) {
		const convertedAttributes = attributes as Partial<GitLabChatRecord>;
		if (attributes.timestamp) {
		convertedAttributes.timestamp = Date.parse(attributes.timestamp);
		}
		Object.assign(this, convertedAttributes);
		}

		private detectType(): ChatRecordType {