feat(wren-ui): Adjust steps integration

wren-ui/src/apollo/client/graphql/__types__.ts

@@ -23,6 +23,21 @@ export type AdjustThreadResponseChartInput = {
   yAxis?: InputMaybe<Scalars['String']>;
 };
 
+export type AdjustThreadResponseInput = {
+  sql?: InputMaybe<Scalars['String']>;
+  sqlGenerationReasoning?: InputMaybe<Scalars['String']>;
+  tables?: InputMaybe<Array<Scalars['String']>>;
+};
+
+export type AdjustmentTask = {
+  __typename?: 'AdjustmentTask';
+  error?: Maybe<Error>;
+  queryId?: Maybe<Scalars['String']>;
+  sql?: Maybe<Scalars['String']>;
+  status?: Maybe<AskingTaskStatus>;
+  traceId?: Maybe<Scalars['String']>;
+};
+
 export type AskingTask = {
   __typename?: 'AskingTask';
   candidates: Array<ResultCandidate>;
@@ -547,7 +562,9 @@ export type ModelWhereInput = {
 
 export type Mutation = {
   __typename?: 'Mutation';
+  adjustThreadResponse: ThreadResponse;
   adjustThreadResponseChart: ThreadResponse;
+  cancelAdjustmentTask: Scalars['Boolean'];
   cancelAskingTask: Scalars['Boolean'];
   createAskingTask: Task;
   createCalculatedField: Scalars['JSON'];
@@ -581,6 +598,7 @@ export type Mutation = {
   previewModelData: Scalars['JSON'];
   previewSql: Scalars['JSON'];
   previewViewData: Scalars['JSON'];
+  rerunAdjustmentTask: Scalars['Boolean'];
   rerunAskingTask: Task;
   resetCurrentProject: Scalars['Boolean'];
   resolveSchemaChange: Scalars['Boolean'];
@@ -607,12 +625,23 @@ export type Mutation = {
 };
 
 
+export type MutationAdjustThreadResponseArgs = {
+  data: AdjustThreadResponseInput;
+  responseId: Scalars['Int'];
+};
+
+
 export type MutationAdjustThreadResponseChartArgs = {
   data: AdjustThreadResponseChartInput;
   responseId: Scalars['Int'];
 };
 
 
+export type MutationCancelAdjustmentTaskArgs = {
+  taskId: Scalars['String'];
+};
+
+
 export type MutationCancelAskingTaskArgs = {
   taskId: Scalars['String'];
 };
@@ -774,6 +803,11 @@ export type MutationPreviewViewDataArgs = {
 };
 
 
+export type MutationRerunAdjustmentTaskArgs = {
+  responseId: Scalars['Int'];
+};
+
+
 export type MutationRerunAskingTaskArgs = {
   responseId: Scalars['Int'];
 };
@@ -933,7 +967,7 @@ export type PreviewItemSqlInput = {
 export type PreviewSqlDataInput = {
   dryRun?: InputMaybe<Scalars['Boolean']>;
   limit?: InputMaybe<Scalars['Int']>;
-  projectId?: InputMaybe<Scalars['Int']>;
+  projectId?: InputMaybe<Scalars['String']>;
   sql: Scalars['String'];
 };
 
@@ -957,6 +991,7 @@ export enum ProjectLanguage {
 
 export type Query = {
   __typename?: 'Query';
+  adjustmentTask?: Maybe<AdjustmentTask>;
   askingTask?: Maybe<AskingTask>;
   autoGenerateRelation: Array<RecommendRelations>;
   dashboardItems: Array<DashboardItem>;
@@ -985,6 +1020,11 @@ export type Query = {
 };
 
 
+export type QueryAdjustmentTaskArgs = {
+  taskId: Scalars['String'];
+};
+
+
 export type QueryAskingTaskArgs = {
   taskId: Scalars['String'];
 };
@@ -1200,6 +1240,8 @@ export type Thread = {
 
 export type ThreadResponse = {
   __typename?: 'ThreadResponse';
+  adjustment?: Maybe<ThreadResponseAdjustment>;
+  adjustmentTask?: Maybe<AdjustmentTask>;
   answerDetail?: Maybe<ThreadResponseAnswerDetail>;
   askingTask?: Maybe<AskingTask>;
   breakdownDetail?: Maybe<ThreadResponseBreakdownDetail>;
@@ -1211,6 +1253,17 @@ export type ThreadResponse = {
   view?: Maybe<ViewInfo>;
 };
 
+export type ThreadResponseAdjustment = {
+  __typename?: 'ThreadResponseAdjustment';
+  payload?: Maybe<Scalars['JSON']>;
+  type: ThreadResponseAdjustmentType;
+};
+
+export enum ThreadResponseAdjustmentType {
+  APPLY_SQL = 'APPLY_SQL',
+  REASONING = 'REASONING'
+}
+
 export type ThreadResponseAnswerDetail = {
   __typename?: 'ThreadResponseAnswerDetail';
   content?: Maybe<Scalars['String']>;

wren-ui/src/apollo/client/graphql/home.ts

@@ -113,12 +113,26 @@ const COMMON_RESPONSE = gql`
     askingTask {
       ...CommonAskingTask
     }
+    adjustment {
+      type
+      payload
+    }
+    adjustmentTask {
+      queryId
+      status
+      error {
+        ...CommonError
+      }
+      sql
+      traceId
+    }
   }
 
   ${COMMON_BREAKDOWN_DETAIL}
   ${COMMON_ANSWER_DETAIL}
   ${COMMON_CHART_DETAIL}
   ${COMMON_ASKING_TASK}
+  ${COMMON_ERROR}
 `;
 
 const COMMON_RECOMMENDED_QUESTIONS_TASK = gql`
@@ -253,6 +267,19 @@ export const UPDATE_THREAD_RESPONSE = gql`
   ${COMMON_RESPONSE}
 `;
 
+// For adjust reasoning steps or SQL
+export const ADJUST_THREAD_RESPONSE = gql`
+  mutation AdjustThreadResponse(
+    $responseId: Int!
+    $data: AdjustThreadResponseInput!
+  ) {
+    adjustThreadResponse(responseId: $responseId, data: $data) {
+      ...CommonResponse
+    }
+  }
+  ${COMMON_RESPONSE}
+`;
+
 export const DELETE_THREAD = gql`
   mutation DeleteThread($where: ThreadUniqueWhereInput!) {
     deleteThread(where: $where)
@@ -369,3 +396,32 @@ export const ADJUST_THREAD_RESPONSE_CHART = gql`
   }
   ${COMMON_RESPONSE}
 `;
+
+export const ADJUSTMENT_TASK = gql`
+  query AdjustmentTask($taskId: String!) {
+    adjustmentTask(taskId: $taskId) {
+      queryId
+      status
+      error {
+        code
+        shortMessage
+        message
+        stacktrace
+      }
+      sql
+      traceId
+    }
+  }
+`;
+
+export const CANCEL_ADJUSTMENT_TASK = gql`
+  mutation CancelAdjustmentTask($taskId: String!) {
+    cancelAdjustmentTask(taskId: $taskId)
+  }
+`;
+
+export const RERUN_ADJUSTMENT_TASK = gql`
+  mutation RerunAdjustmentTask($responseId: Int!) {
+    rerunAdjustmentTask(responseId: $responseId)
+  }
+`;

wren-ui/src/components/diagram/CustomDropdown.tsx

@@ -1,4 +1,5 @@
 import React from 'react';
+import styled from 'styled-components';
 import { Dropdown, Menu } from 'antd';
 import { ItemType } from 'antd/lib/menu/hooks/useItems';
 import { MORE_ACTION, NODE_TYPE } from '@/utils/enum';
@@ -18,9 +19,15 @@ import {
   DeleteInstructionModal,
 } from '@/components/modals/DeleteModal';
 
+const StyledMenu = styled(Menu)`
+  .ant-dropdown-menu-item:not(.ant-dropdown-menu-item-disabled) {
+    color: var(--gray-8);
+  }
+`;
+
 interface Props {
   [key: string]: any;
-  onMoreClick: (type: MORE_ACTION) => void;
+  onMoreClick: (type: MORE_ACTION | { type: MORE_ACTION; data: any }) => void;
   onMenuEnter?: (event: React.MouseEvent) => void;
   children: React.ReactNode;
   onDropdownVisibleChange?: (visible: boolean) => void;
@@ -37,7 +44,7 @@ const makeDropdown =
         trigger={['click']}
         overlayStyle={{ minWidth: 100, userSelect: 'none' }}
         overlay={
-          <Menu
+          <StyledMenu
             onClick={(e) => e.domEvent.stopPropagation()}
             items={items}
             onMouseEnter={onMenuEnter}
@@ -57,7 +64,7 @@ export const ModelDropdown = makeDropdown((props: Props) => {
     {
       label: (
         <>
-          <EditOutlined className="gray-8 mr-2" />
+          <EditOutlined className="mr-2" />
           Update Columns
         </>
       ),
@@ -106,7 +113,7 @@ export const ColumnDropdown = makeDropdown((props: Props) => {
     {
       label: (
         <>
-          <EditOutlined className="gray-8 mr-2" />
+          <EditOutlined className="mr-2" />
           Edit
         </>
       ),
@@ -132,12 +139,12 @@ export const DashboardItemDropdown = makeDropdown((props: Props) => {
     {
       label: isHideLegend ? (
         <>
-          <EyeOutlined className="gray-8 mr-2" />
+          <EyeOutlined className="mr-2" />
           Show categories
         </>
       ) : (
         <>
-          {<EyeInvisibleOutlined className="gray-8 mr-2" />}
+          {<EyeInvisibleOutlined className="mr-2" />}
           Hide categories
         </>
       ),
@@ -147,7 +154,7 @@ export const DashboardItemDropdown = makeDropdown((props: Props) => {
     {
       label: (
         <>
-          <ReloadOutlined className="gray-8 mr-2" />
+          <ReloadOutlined className="mr-2" />
           Refresh
         </>
       ),
@@ -179,7 +186,7 @@ export const SQLPairDropdown = makeDropdown(
       {
         label: (
           <>
-            <EyeOutlined className="gray-8 mr-2" />
+            <EyeOutlined className="mr-2" />
             View
           </>
         ),
@@ -193,7 +200,7 @@ export const SQLPairDropdown = makeDropdown(
       {
         label: (
           <>
-            <EditOutlined className="gray-8 mr-2" />
+            <EditOutlined className="mr-2" />
             Edit
           </>
         ),
@@ -238,7 +245,7 @@ export const InstructionDropdown = makeDropdown(
       {
         label: (
           <>
-            <EyeOutlined className="gray-8 mr-2" />
+            <EyeOutlined className="mr-2" />
             View
           </>
         ),
@@ -252,7 +259,7 @@ export const InstructionDropdown = makeDropdown(
       {
         label: (
           <>
-            <EditOutlined className="gray-8 mr-2" />
+            <EditOutlined className="mr-2" />
             Edit
           </>
         ),
@@ -296,8 +303,8 @@ export const AdjustAnswerDropdown = makeDropdown(
     const items: ItemType[] = [
       {
         label: 'Adjust steps',
-        className: 'gray-8',
-        icon: <EditSVG className="gray-8" />,
+        icon: <EditSVG />,
+        disabled: !data.sqlGenerationReasoning,
         key: 'adjust-steps',
         onClick: () =>
           onMoreClick({
@@ -307,8 +314,7 @@ export const AdjustAnswerDropdown = makeDropdown(
       },
       {
         label: 'Adjust SQL',
-        className: 'gray-8',
-        icon: <CodeFilled className="gray-8 text-base" />,
+        icon: <CodeFilled className="text-base" />,
         key: 'adjust-sql',
         onClick: () =>
           onMoreClick({