From 08b3f1a836293ba1c6b049b0caeae5b7c61103c5 Mon Sep 17 00:00:00 2001
From: Marc Saleiko <msaleiko@gitlab.com>
Date: Thu, 22 Feb 2024 16:34:18 +0000
Subject: [PATCH] Adds Ai agents debug chat interface
---
.../ai/graphql/chat.mutation.graphql | 8 +-
.../assets/javascripts/ml/ai_agents/index.js | 3 +-
.../ml/ai_agents/views/show_agent.vue | 92 ++++++++++++-
.../assets/stylesheets/components/_index.scss | 1 +
.../stylesheets/components/ai_agent_chat.scss | 19 +++
.../views/projects/ml/agents/index.html.haml | 2 +-
.../ml/ai_agents/views/show_agent_spec.js | 130 +++++++++++++++---
locale/gitlab.pot | 12 ++
8 files changed, 243 insertions(+), 24 deletions(-)
create mode 100644 ee/app/assets/stylesheets/components/ai_agent_chat.scss
diff --git a/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql b/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql
index 5b88d5d6b6cf1..64eaeba06d23e 100644
--- a/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql
+++ b/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql
@@ -1,12 +1,18 @@
mutation chat(
$question: String!
$resourceId: AiModelID!
+ $agentVersionId: AiAgentVersionID
$clientSubscriptionId: String
$currentFileContext: AiCurrentFileInput
) {
aiAction(
input: {
- chat: { resourceId: $resourceId, content: $question, currentFile: $currentFileContext }
+ chat: {
+ resourceId: $resourceId
+ content: $question
+ agentVersionId: $agentVersionId
+ currentFile: $currentFileContext
+ }
clientSubscriptionId: $clientSubscriptionId
}
) {
diff --git a/ee/app/assets/javascripts/ml/ai_agents/index.js b/ee/app/assets/javascripts/ml/ai_agents/index.js
index 1d62f12aa75ee..1ad2d3b2c41cb 100644
--- a/ee/app/assets/javascripts/ml/ai_agents/index.js
+++ b/ee/app/assets/javascripts/ml/ai_agents/index.js
@@ -13,7 +13,7 @@ export default () => {
return false;
}
- const { basePath, projectPath } = el.dataset;
+ const { basePath, projectPath, userId } = el.dataset;
const apolloProvider = new VueApollo({
defaultClient: createDefaultClient(),
@@ -28,6 +28,7 @@ export default () => {
router,
provide: {
projectPath,
+ userId,
},
render(h) {
return h(BaseApp);
diff --git a/ee/app/assets/javascripts/ml/ai_agents/views/show_agent.vue b/ee/app/assets/javascripts/ml/ai_agents/views/show_agent.vue
index ea1712171e19f..2ad6579a7952a 100644
--- a/ee/app/assets/javascripts/ml/ai_agents/views/show_agent.vue
+++ b/ee/app/assets/javascripts/ml/ai_agents/views/show_agent.vue
@@ -1,25 +1,97 @@
<script>
-import { GlExperimentBadge } from '@gitlab/ui';
+import { GlExperimentBadge, GlDuoChat } from '@gitlab/ui';
+import aiResponseSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response.subscription.graphql';
import TitleArea from '~/vue_shared/components/registry/title_area.vue';
+import { renderMarkdown } from '~/notes/utils';
+import { renderGFM } from '~/behaviors/markdown/render_gfm';
import { sprintf, s__ } from '~/locale';
+import { GENIE_CHAT_MODEL_ROLES } from 'ee/ai/constants';
+import chatMutation from 'ee/ai/graphql/chat.mutation.graphql';
export default {
name: 'ShowAiAgent',
+ // Needed to override the default predefined prompts
+ predefinedPrompts: [],
components: {
TitleArea,
GlExperimentBadge,
+ GlDuoChat,
},
provide() {
return {
projectPath: this.projectPath,
+ userId: this.userId,
+ renderMarkdown,
+ renderGFM,
+ };
+ },
+ inject: ['projectPath', 'userId'],
+ apollo: {
+ // https://apollo.vuejs.org/guide-option/subscriptions.html#simple-subscription
+ $subscribe: {
+ aiCompletionResponse: {
+ query: aiResponseSubscription,
+ variables() {
+ return {
+ userId: this.userId,
+ agentVersionId: `gid://gitlab/Ai::AgentVersion/${this.$route.params.agentId}`,
+ aiAction: 'CHAT',
+ };
+ },
+ result({ data }) {
+ const response = data?.aiCompletionResponse;
+
+ if (!response) {
+ return;
+ }
+
+ this.messages.push(response);
+
+ if (response.role.toLowerCase() === GENIE_CHAT_MODEL_ROLES.assistant) {
+ this.isLoading = false;
+ }
+ },
+ error(err) {
+ this.error = err.toString();
+ },
+ },
+ },
+ },
+ data() {
+ return {
+ error: null,
+ messages: [],
+ isLoading: false,
};
},
- inject: ['projectPath'],
computed: {
title() {
return sprintf(s__('AIAgent|AI Agent: %{agentId}'), { agentId: this.$route.params.agentId });
},
},
+ methods: {
+ onSendChatPrompt(question = '') {
+ this.isLoading = true;
+
+ this.$apollo
+ .mutate({
+ mutation: chatMutation,
+ variables: {
+ question,
+ resourceId: this.userId,
+ agentVersionId: `gid://gitlab/Ai::AgentVersion/${this.$route.params.agentId}`,
+ },
+ })
+ .then(() => {
+ // we add the user message in the aiCompletionResponse subscription
+ this.isLoading = true;
+ })
+ .catch((err) => {
+ this.error = err.toString();
+ this.isLoading = false;
+ });
+ },
+ },
};
</script>
@@ -33,5 +105,21 @@ export default {
</div>
</template>
</title-area>
+
+ <gl-duo-chat
+ :messages="messages"
+ :error="error"
+ :is-loading="isLoading"
+ :predefined-prompts="$options.predefinedPrompts"
+ :tool-name="s__('AIAgent|Agent')"
+ class="ai-agent-chat gl-w-full! gl-static gl-border-r gl-border-transparent"
+ :empty-state-title="s__('AIAgent|Try out your agent')"
+ :empty-state-description="
+ s__('AIAgent|Your agent\'s system prompt will be applied to the chat input.')
+ "
+ :chat-prompt-placeholder="s__('AIAgent|Ask your agent')"
+ :show-header="false"
+ @send-chat-prompt="onSendChatPrompt"
+ />
</div>
</template>
diff --git a/ee/app/assets/stylesheets/components/_index.scss b/ee/app/assets/stylesheets/components/_index.scss
index f7f6ff607620c..9b5389e971eba 100644
--- a/ee/app/assets/stylesheets/components/_index.scss
+++ b/ee/app/assets/stylesheets/components/_index.scss
@@ -1,3 +1,4 @@
+@import './ai_agent_chat';
@import './audit_logs/logs_table';
@import './banner';
@import './batch_comments/draft_note';
diff --git a/ee/app/assets/stylesheets/components/ai_agent_chat.scss b/ee/app/assets/stylesheets/components/ai_agent_chat.scss
new file mode 100644
index 0000000000000..3f6ac7ee9f27a
--- /dev/null
+++ b/ee/app/assets/stylesheets/components/ai_agent_chat.scss
@@ -0,0 +1,19 @@
+.ai-agent-chat {
+ // prevents the chat input from jumping all over the place
+ min-height: $gl-spacing-scale-48;
+
+ section {
+ @include gl-bg-transparent;
+ }
+
+ footer {
+ @include gl-border-transparent;
+ @include gl-bg-white;
+ }
+ // Disable feedback
+ // Improve selector once the feedback has its own class name:
+ // https://gitlab.com/gitlab-org/gitlab/-/issues/442715
+ .duo-chat-message > div.gl-display-flex:last-of-type {
+ @include gl-display-none;
+ }
+}
\ No newline at end of file
diff --git a/ee/app/views/projects/ml/agents/index.html.haml b/ee/app/views/projects/ml/agents/index.html.haml
index 615276826cd0b..a5ba2b8c2450c 100644
--- a/ee/app/views/projects/ml/agents/index.html.haml
+++ b/ee/app/views/projects/ml/agents/index.html.haml
@@ -1,4 +1,4 @@
- breadcrumb_title s_('AIAgents|AI Agents')
- page_title s_('AIAgents|AI Agents')
-#js-mount-index-ml-agents{ data: { project_path: @project.full_path, base_path: namespace_project_ml_agents_path(@project.namespace, @project) } }
+#js-mount-index-ml-agents{ data: { project_path: @project.full_path, base_path: namespace_project_ml_agents_path(@project.namespace, @project), user_id: current_user.to_global_id } }
diff --git a/ee/spec/frontend/ml/ai_agents/views/show_agent_spec.js b/ee/spec/frontend/ml/ai_agents/views/show_agent_spec.js
index fb73d3043363f..f5fc1fed5d75c 100644
--- a/ee/spec/frontend/ml/ai_agents/views/show_agent_spec.js
+++ b/ee/spec/frontend/ml/ai_agents/views/show_agent_spec.js
@@ -1,34 +1,126 @@
-import { GlExperimentBadge } from '@gitlab/ui';
+import { GlDuoChat, GlExperimentBadge } from '@gitlab/ui';
+import Vue, { nextTick } from 'vue';
+import VueApollo from 'vue-apollo';
import { shallowMountExtended } from 'helpers/vue_test_utils_helper';
+import createMockApollo from 'helpers/mock_apollo_helper';
+import { getMarkdown } from '~/rest_api';
import ShowAgent from 'ee/ml/ai_agents/views/show_agent.vue';
import TitleArea from '~/vue_shared/components/registry/title_area.vue';
+import aiResponseSubscription from 'ee/graphql_shared/subscriptions/ai_completion_response.subscription.graphql';
+import waitForPromises from 'helpers/wait_for_promises';
+import chatMutation from 'ee/ai/graphql/chat.mutation.graphql';
+import {
+ MOCK_USER_MESSAGE,
+ MOCK_USER_ID,
+ MOCK_TANUKI_SUCCESS_RES,
+ MOCK_TANUKI_BOT_MUTATATION_RES,
+} from '../../../ai/tanuki_bot/mock_data';
-let wrapper;
+Vue.use(VueApollo);
-const createWrapper = () => {
- wrapper = shallowMountExtended(ShowAgent, {
- provide: { projectPath: 'path/to/project' },
- mocks: {
- $route: {
- params: {
- agentId: 2,
+jest.mock('~/rest_api');
+
+describe('ee/ml/ai_agents/views/create_agent', () => {
+ let wrapper;
+
+ const subscriptionHandlerMock = jest.fn().mockResolvedValue(MOCK_TANUKI_SUCCESS_RES);
+ const chatMutationHandlerMock = jest.fn().mockResolvedValue(MOCK_TANUKI_BOT_MUTATATION_RES);
+ const agentId = 2;
+
+ const findTitleArea = () => wrapper.findComponent(TitleArea);
+ const findBadge = () => wrapper.findComponent(GlExperimentBadge);
+ const findGlDuoChat = () => wrapper.findComponent(GlDuoChat);
+
+ const createWrapper = () => {
+ const apolloProvider = createMockApollo([
+ [aiResponseSubscription, subscriptionHandlerMock],
+ [chatMutation, chatMutationHandlerMock],
+ ]);
+
+ wrapper = shallowMountExtended(ShowAgent, {
+ apolloProvider,
+ provide: { projectPath: 'path/to/project', userId: MOCK_USER_ID },
+ mocks: {
+ $route: {
+ params: {
+ agentId,
+ },
},
},
- },
+ });
+ };
+
+ beforeEach(() => {
+ getMarkdown.mockImplementation(({ text }) => Promise.resolve({ data: { html: text } }));
});
-};
-const findTitleArea = () => wrapper.findComponent(TitleArea);
-const findBadge = () => wrapper.findComponent(GlExperimentBadge);
+ describe('rendering', () => {
+ beforeEach(() => {
+ createWrapper();
+ });
-describe('ee/ml/ai_agents/views/create_agent', () => {
- beforeEach(() => createWrapper());
+ it('shows the title', () => {
+ expect(findTitleArea().text()).toContain('AI Agent: 2');
+ });
+
+ it('displays the experiment badge', () => {
+ expect(findBadge().exists()).toBe(true);
+ });
- it('shows the title', () => {
- expect(findTitleArea().text()).toContain('AI Agent: 2');
+ it('renders the DuoChat component', () => {
+ expect(findGlDuoChat().exists()).toBe(true);
+ });
});
- it('displays the experiment badge', () => {
- expect(findBadge().exists()).toBe(true);
+ describe('@send-chat-prompt', () => {
+ beforeEach(() => {
+ createWrapper();
+ });
+
+ it('does set loading to `true` for a user message', async () => {
+ findGlDuoChat().vm.$emit('send-chat-prompt', MOCK_USER_MESSAGE.content);
+ await nextTick();
+ expect(findGlDuoChat().props('isLoading')).toBe(true);
+ });
+
+ it('calls correct GraphQL mutation', async () => {
+ findGlDuoChat().vm.$emit('send-chat-prompt', MOCK_USER_MESSAGE.content);
+ await nextTick();
+ expect(chatMutationHandlerMock).toHaveBeenCalledWith({
+ resourceId: MOCK_USER_ID,
+ agentVersionId: `gid://gitlab/Ai::AgentVersion/${agentId}`,
+ question: MOCK_USER_MESSAGE.content,
+ });
+ });
+ });
+
+ describe('Error conditions', () => {
+ const errorText = 'Fancy foo';
+
+ describe('when subscription fails', () => {
+ beforeEach(async () => {
+ subscriptionHandlerMock.mockRejectedValue(new Error(errorText));
+ createWrapper();
+ await waitForPromises();
+ });
+
+ it('throws error and displays error message', () => {
+ expect(findGlDuoChat().props('error')).toBe(`Error: ${errorText}`);
+ });
+ });
+
+ describe('when mutation fails', () => {
+ beforeEach(async () => {
+ chatMutationHandlerMock.mockRejectedValue(new Error(errorText));
+ createWrapper();
+ await waitForPromises();
+ findGlDuoChat().vm.$emit('send-chat-prompt', MOCK_USER_MESSAGE.content);
+ await waitForPromises();
+ });
+
+ it('throws error and displays error message', () => {
+ expect(findGlDuoChat().props('error')).toBe(`Error: ${errorText}`);
+ });
+ });
});
});
diff --git a/locale/gitlab.pot b/locale/gitlab.pot
index cbc69f94367bf..12cbdafa0afa8 100644
--- a/locale/gitlab.pot
+++ b/locale/gitlab.pot
@@ -1933,6 +1933,18 @@ msgstr ""
msgid "AIAgent|AI Agent: %{agentId}"
msgstr ""
+msgid "AIAgent|Agent"
+msgstr ""
+
+msgid "AIAgent|Ask your agent"
+msgstr ""
+
+msgid "AIAgent|Try out your agent"
+msgstr ""
+
+msgid "AIAgent|Your agent's system prompt will be applied to the chat input."
+msgstr ""
+
msgid "AIPoweredSM|AI-powered features"
msgstr ""
--
GitLab