feat: enable streaming for the chat (5e5432cf) · Commits · GitLab.org / gitlab-vscode-extension

src/common/api/graphql/ai_completion_response_channel.test.ts

0 → 100644

+53 −0

Original line number	Diff line number	Diff line
		import { AiCompletionResponseChannel } from './ai_completion_response_channel';

		describe('AiCompletionResponseChannel', () => {
		let channel: AiCompletionResponseChannel;
		let emittedEvents: Array<[string, object]>;

		const messageMock = {
		role: 'USER',
		content: 'abc',
		requestId: 'foo',
		timestamp: 'bar',
		errors: [],
		};

		beforeEach(() => {
		channel = new AiCompletionResponseChannel({ userId: 'foo' });
		emittedEvents = [];

		channel.on('newChunk', data => {
		emittedEvents.push(['newChunk', data]);
		});
		channel.on('fullMessage', data => {
		emittedEvents.push(['fullMessage', data]);
		});
		channel.on('systemMessage', data => {
		emittedEvents.push(['systemMessage', data]);
		});
		});

		it('emits "systemMessage" when message has system role', async () => {
		const expectedData = { ...messageMock, role: 'SYSTEM' };

		channel.receive({ result: { data: { aiCompletionResponse: expectedData } }, more: true });

		expect(emittedEvents[0]).toStrictEqual(['systemMessage', expectedData]);
		});

		it('emits "newChunk" when message has chunkId', async () => {
		const expectedData = { ...messageMock, chunkId: 1 };

		channel.receive({ result: { data: { aiCompletionResponse: expectedData } }, more: true });

		expect(emittedEvents[0]).toStrictEqual(['newChunk', expectedData]);
		});

		it('emits "fullMessage" when message has no chunkId', async () => {
		const expectedData = { ...messageMock };

		channel.receive({ result: { data: { aiCompletionResponse: expectedData } }, more: true });

		expect(emittedEvents[0]).toStrictEqual(['fullMessage', expectedData]);
		});
		});

src/common/api/graphql/ai_completion_response_channel.ts

0 → 100644

+105 −0

Original line number	Diff line number	Diff line
		import { Channel, ChannelEvents } from '@anycable/core';
		import { gql } from 'graphql-request';

		type AiCompletionResponseInput = {
		htmlResponse?: boolean;
		userId: string;
		aiAction?: string;
		clientSubscriptionId?: string;
		};

		type AiCompletionResponseParams = {
		channel: 'GraphqlChannel';
		query: string;
		variables: string;
		operationName: 'aiCompletionResponse';
		};

		const AI_MESSAGE_SUBSCRIPTION_QUERY = gql`
		subscription aiCompletionResponse(
		$userId: UserID
		$clientSubscriptionId: String
		$aiAction: AiAction
		$htmlResponse: Boolean = true
		) {
		aiCompletionResponse(
		userId: $userId
		aiAction: $aiAction
		clientSubscriptionId: $clientSubscriptionId
		) {
		id
		requestId
		content
		contentHtml @include(if: $htmlResponse)
		errors
		role
		timestamp
		type
		chunkId
		extras {
		sources
		}
		}
		}
		`;

		export type AiCompletionResponseMessageType = {
		requestId: string;
		role: string;
		content: string;
		contentHtml?: string;
		timestamp: string;
		errors: string[];
		extras?: {
		sources: object[];
		};
		chunkId?: number;
		type?: string;
		};

		type AiCompletionResponseResponseType = {
		result: {
		data: {
		aiCompletionResponse: AiCompletionResponseMessageType;
		};
		};
		more: boolean;
		};

		interface AiCompletionResponseChannelEvents
		extends ChannelEvents<AiCompletionResponseResponseType> {
		systemMessage: (msg: AiCompletionResponseMessageType) => void;
		newChunk: (msg: AiCompletionResponseMessageType) => void;
		fullMessage: (msg: AiCompletionResponseMessageType) => void;
		}

		export class AiCompletionResponseChannel extends Channel<
		AiCompletionResponseParams,
		AiCompletionResponseResponseType,
		AiCompletionResponseChannelEvents
		> {
		static identifier = 'GraphqlChannel';

		constructor(params: AiCompletionResponseInput) {
		super({
		channel: 'GraphqlChannel',
		operationName: 'aiCompletionResponse',
		query: AI_MESSAGE_SUBSCRIPTION_QUERY,
		variables: JSON.stringify(params),
		});
		}

		receive(message: AiCompletionResponseResponseType) {
		if (!message.result.data.aiCompletionResponse) return;

		const data = message.result.data.aiCompletionResponse;

		if (data.role.toLowerCase() === 'system') {
		this.emit('systemMessage', data);
		} else if (data.chunkId) {
		this.emit('newChunk', data);
		} else {
		this.emit('fullMessage', data);
		}
		}
		}

src/common/chat/gitlab_chat_api.test.ts

+1 −1

Original line number	Diff line number	Diff line
		@@ -123,7 +123,7 @@ describe('GitLabChatApi', () => {
		contentBelowCursor: 'after_text',
		};

		const response = await gitlabChatApi.processNewUserPrompt(mockPrompt, fileContext);
		const response = await gitlabChatApi.processNewUserPrompt(mockPrompt, undefined, fileContext);

		expect(response.aiAction).toBe(mockedMutationResponse.aiAction);

src/common/chat/gitlab_chat_api.ts

+47 −7

Original line number	Diff line number	Diff line
		@@ -3,6 +3,11 @@ import { GraphQLRequest } from '../platform/web_ide';
		import { GitLabPlatformManager } from '../platform/gitlab_platform';
		import { pullHandler } from './api/pulling';
		import { GitLabChatFileContext } from './gitlab_chat_file_context';
		import {
		AiCompletionResponseChannel,
		AiCompletionResponseMessageType,
		} from '../api/graphql/ai_completion_response_channel';
		import { extractUserId } from '../platform/gitlab_account';

		export const AI_ACTIONS = {
		chat: gql`
		@@ -10,10 +15,12 @@ export const AI_ACTIONS = {
		$question: String!
		$resourceId: AiModelID
		$currentFileContext: AiCurrentFileInput
		$clientSubscriptionId: String
		) {
		aiAction(
		input: {
		chat: { resourceId: $resourceId, content: $question, currentFile: $currentFileContext }
		clientSubscriptionId: $clientSubscriptionId
		}
		) {
		requestId
		@@ -45,9 +52,7 @@ export const AI_MESSAGES_QUERY = gql`
		}
		`;

		type AiMessagesResponseType = {
		aiMessages: {
		nodes: {
		type AiMessageResponseType = {
		requestId: string;
		role: string;
		content: string;
		@@ -57,10 +62,13 @@ type AiMessagesResponseType = {
		extras?: {
		sources: object[];
		};
		}[];
		};

		type AiMessagesResponseType = {
		aiMessages: {
		nodes: AiMessageResponseType[];
		};
		};
		type AiMessageResponseType = AiMessagesResponseType['aiMessages']['nodes'][0];

		interface ErrorMessage {
		type: 'error';
		@@ -105,9 +113,14 @@ export class GitLabChatApi {

		async processNewUserPrompt(
		question: string,
		currentFileContext: GitLabChatFileContext \| undefined = undefined,
		subscriptionId?: string,
		currentFileContext?: GitLabChatFileContext,
		): Promise<AiActionResponseType> {
		return this.sendAiAction(AI_ACTIONS.chat, { question, currentFileContext });
		return this.sendAiAction(AI_ACTIONS.chat, {
		question,
		currentFileContext,
		clientSubscriptionId: subscriptionId,
		});
		}

		async pullAiMessage(requestId: string, role: string): Promise<AiMessage> {
		@@ -145,6 +158,33 @@ export class GitLabChatApi {
		return history.aiMessages.nodes[0];
		}

		async subscribeToUpdates(
		messageCallback: (message: AiCompletionResponseMessageType) => Promise<void>,
		subscriptionId?: string,
		) {
		const platform = await this.currentPlatform();

		const channel = new AiCompletionResponseChannel({
		htmlResponse: true,
		userId: `gid://gitlab/User/${extractUserId(platform.account.id)}`,
		aiAction: 'CHAT',
		clientSubscriptionId: subscriptionId,
		});

		const cable = await platform.connectToCable();

		channel.on('newChunk', messageCallback);
		channel.on('fullMessage', async message => {
		await messageCallback(message);

		if (subscriptionId) {
		cable.disconnect();
		}
		});

		cable.subscribe(channel);
		}

		private async sendAiAction(
		actionQuery: string,
		variables: object,

src/common/chat/gitlab_chat_controller.test.ts

+10 −2

Original line number	Diff line number	Diff line
		@@ -4,7 +4,11 @@ import { GitLabPlatformManager } from '../platform/gitlab_platform';
		import { GitLabChatRecord } from './gitlab_chat_record';
		import { submitFeedback } from './utils/submit_feedback';

		const apiMock = { processNewUserPrompt: jest.fn(), pullAiMessage: jest.fn() };
		const apiMock = {
		processNewUserPrompt: jest.fn(),
		pullAiMessage: jest.fn(),
		subscribeToUpdates: jest.fn(),
		};

		jest.mock('./gitlab_chat_api', () => ({
		GitLabChatApi: jest.fn().mockImplementation(() => apiMock),
		@@ -196,7 +200,11 @@ describe('GitLabChatController', () => {

		await controller.processNewUserRecord(record);

		expect(apiMock.processNewUserPrompt).toHaveBeenCalledWith('hello', currentFileContext);
		expect(apiMock.processNewUserPrompt).toHaveBeenCalledWith(
		'hello',
		expect.any(String),
		currentFileContext,
		);
		});

		describe('with newChatConversation command', () => {