feat(wren-ui): Fix SQL flow in reasoning

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

@@ -600,6 +600,7 @@ export type Mutation = {
   updateRelation: Scalars['JSON'];
   updateSqlPair: SqlPair;
   updateThread: Thread;
+  updateThreadResponse: ThreadResponse;
   updateViewMetadata: Scalars['Boolean'];
   validateCalculatedField: CalculatedFieldValidationResponse;
   validateView: ViewValidationResponse;
@@ -865,6 +866,12 @@ export type MutationUpdateThreadArgs = {
 };
 
 
+export type MutationUpdateThreadResponseArgs = {
+  data: UpdateThreadResponseInput;
+  where: ThreadResponseUniqueWhereInput;
+};
+
+
 export type MutationUpdateViewMetadataArgs = {
   data: UpdateViewMetadataInput;
   where: ViewWhereUniqueInput;
@@ -1241,6 +1248,10 @@ export type ThreadResponseChartDetail = {
   status: ChartTaskStatus;
 };
 
+export type ThreadResponseUniqueWhereInput = {
+  id: Scalars['Int'];
+};
+
 export type ThreadUniqueWhereInput = {
   id: Scalars['Int'];
 };
@@ -1328,6 +1339,10 @@ export type UpdateThreadInput = {
   summary?: InputMaybe<Scalars['String']>;
 };
 
+export type UpdateThreadResponseInput = {
+  sql?: InputMaybe<Scalars['String']>;
+};
+
 export type UpdateViewColumnMetadataInput = {
   description?: InputMaybe<Scalars['String']>;
   referenceName: Scalars['String'];

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

@@ -92,6 +92,14 @@ export type UpdateThreadMutationVariables = Types.Exact<{
 
 export type UpdateThreadMutation = { __typename?: 'Mutation', updateThread: { __typename?: 'Thread', id: number, summary: string } };
 
+export type UpdateThreadResponseMutationVariables = Types.Exact<{
+  where: Types.ThreadResponseUniqueWhereInput;
+  data: Types.UpdateThreadResponseInput;
+}>;
+
+
+export type UpdateThreadResponseMutation = { __typename?: 'Mutation', updateThreadResponse: { __typename?: 'ThreadResponse', id: number, threadId: number, question: string, sql?: string | null, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, breakdownDetail?: { __typename?: 'ThreadResponseBreakdownDetail', queryId?: string | null, status: Types.AskingTaskStatus, description?: string | null, steps?: Array<{ __typename?: 'DetailStep', summary: string, sql: string, cteName?: string | null }> | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array<string | null> | null } | null } | null, answerDetail?: { __typename?: 'ThreadResponseAnswerDetail', queryId?: string | null, status?: Types.ThreadResponseAnswerStatus | null, content?: string | null, numRowsUsedInLLM?: number | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array<string | null> | null } | null } | null, chartDetail?: { __typename?: 'ThreadResponseChartDetail', queryId?: string | null, status: Types.ChartTaskStatus, description?: string | null, chartType?: Types.ChartType | null, chartSchema?: any | null, adjustment?: boolean | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array<string | null> | null } | null } | null, askingTask?: { __typename?: 'AskingTask', status: Types.AskingTaskStatus, type?: Types.AskingTaskType | null, rephrasedQuestion?: string | null, intentReasoning?: string | null, sqlGenerationReasoning?: string | null, retrievedTables?: Array<string> | null, invalidSql?: string | null, traceId?: string | null, queryId?: string | null, candidates: Array<{ __typename?: 'ResultCandidate', sql: string, type: Types.ResultCandidateType, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null }>, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array<string | null> | null } | null } | null } };
+
 export type DeleteThreadMutationVariables = Types.Exact<{
   where: Types.ThreadUniqueWhereInput;
 }>;
@@ -681,6 +689,40 @@ export function useUpdateThreadMutation(baseOptions?: Apollo.MutationHookOptions
 export type UpdateThreadMutationHookResult = ReturnType<typeof useUpdateThreadMutation>;
 export type UpdateThreadMutationResult = Apollo.MutationResult<UpdateThreadMutation>;
 export type UpdateThreadMutationOptions = Apollo.BaseMutationOptions<UpdateThreadMutation, UpdateThreadMutationVariables>;
+export const UpdateThreadResponseDocument = gql`
+    mutation UpdateThreadResponse($where: ThreadResponseUniqueWhereInput!, $data: UpdateThreadResponseInput!) {
+  updateThreadResponse(where: $where, data: $data) {
+    ...CommonResponse
+  }
+}
+    ${CommonResponseFragmentDoc}`;
+export type UpdateThreadResponseMutationFn = Apollo.MutationFunction<UpdateThreadResponseMutation, UpdateThreadResponseMutationVariables>;
+
+/**
+ * __useUpdateThreadResponseMutation__
+ *
+ * To run a mutation, you first call `useUpdateThreadResponseMutation` within a React component and pass it any options that fit your needs.
+ * When your component renders, `useUpdateThreadResponseMutation` returns a tuple that includes:
+ * - A mutate function that you can call at any time to execute the mutation
+ * - An object with fields that represent the current status of the mutation's execution
+ *
+ * @param baseOptions options that will be passed into the mutation, supported options are listed on: https://www.apollographql.com/docs/react/api/react-hooks/#options-2;
+ *
+ * @example
+ * const [updateThreadResponseMutation, { data, loading, error }] = useUpdateThreadResponseMutation({
+ *   variables: {
+ *      where: // value for 'where'
+ *      data: // value for 'data'
+ *   },
+ * });
+ */
+export function useUpdateThreadResponseMutation(baseOptions?: Apollo.MutationHookOptions<UpdateThreadResponseMutation, UpdateThreadResponseMutationVariables>) {
+        const options = {...defaultOptions, ...baseOptions}
+        return Apollo.useMutation<UpdateThreadResponseMutation, UpdateThreadResponseMutationVariables>(UpdateThreadResponseDocument, options);
+      }
+export type UpdateThreadResponseMutationHookResult = ReturnType<typeof useUpdateThreadResponseMutation>;
+export type UpdateThreadResponseMutationResult = Apollo.MutationResult<UpdateThreadResponseMutation>;
+export type UpdateThreadResponseMutationOptions = Apollo.BaseMutationOptions<UpdateThreadResponseMutation, UpdateThreadResponseMutationVariables>;
 export const DeleteThreadDocument = gql`
     mutation DeleteThread($where: ThreadUniqueWhereInput!) {
   deleteThread(where: $where)

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

@@ -235,6 +235,18 @@ export const UPDATE_THREAD = gql`
   }
 `;
 
+export const UPDATE_THREAD_RESPONSE = gql`
+  mutation UpdateThreadResponse(
+    $where: ThreadResponseUniqueWhereInput!
+    $data: UpdateThreadResponseInput!
+  ) {
+    updateThreadResponse(where: $where, data: $data) {
+      ...CommonResponse
+    }
+  }
+  ${COMMON_RESPONSE}
+`;
+
 export const DELETE_THREAD = gql`
   mutation DeleteThread($where: ThreadUniqueWhereInput!) {
     deleteThread(where: $where)

wren-ui/src/apollo/server/resolvers.ts

@@ -102,6 +102,7 @@ const resolvers = {
     updateThread: askingResolver.updateThread,
     deleteThread: askingResolver.deleteThread,
     createThreadResponse: askingResolver.createThreadResponse,
+    updateThreadResponse: askingResolver.updateThreadResponse,
     previewData: askingResolver.previewData,
     previewBreakdownData: askingResolver.previewBreakdownData,
 

wren-ui/src/apollo/server/resolvers/askingResolver.ts

@@ -84,6 +84,7 @@ export class AskingResolver {
     this.deleteThread = this.deleteThread.bind(this);
     this.listThreads = this.listThreads.bind(this);
     this.createThreadResponse = this.createThreadResponse.bind(this);
+    this.updateThreadResponse = this.updateThreadResponse.bind(this);
     this.getResponse = this.getResponse.bind(this);
     this.previewData = this.previewData.bind(this);
     this.previewBreakdownData = this.previewBreakdownData.bind(this);
@@ -418,6 +419,17 @@ export class AskingResolver {
     }
   }
 
+  public async updateThreadResponse(
+    _root: any,
+    args: { where: { id: number }; data: { sql: string } },
+    ctx: IContext,
+  ): Promise<ThreadResponse> {
+    const { where, data } = args;
+    const askingService = ctx.askingService;
+    const response = await askingService.updateThreadResponse(where.id, data);
+    return response;
+  }
+
   public async rerunAskingTask(
     _root: any,
     args: { responseId: number },

wren-ui/src/apollo/server/schema.ts

@@ -637,6 +637,14 @@ export const typeDefs = gql`
     summary: String
   }
 
+  input ThreadResponseUniqueWhereInput {
+    id: Int!
+  }
+
+  input UpdateThreadResponseInput {
+    sql: String
+  }
+
   input AdjustThreadResponseChartInput {
     chartType: ChartType!
     xAxis: String
@@ -1034,6 +1042,10 @@ export const typeDefs = gql`
       threadId: Int!
       data: CreateThreadResponseInput!
     ): ThreadResponse!
+    updateThreadResponse(
+      where: ThreadResponseUniqueWhereInput!
+      data: UpdateThreadResponseInput!
+    ): ThreadResponse!
     previewData(where: PreviewDataInput!): JSON!
     previewBreakdownData(where: PreviewDataInput!): JSON!
 

wren-ui/src/apollo/server/services/askingService.ts

@@ -132,6 +132,10 @@ export interface IAskingService {
     input: AskingDetailTaskInput,
     threadId: number,
   ): Promise<ThreadResponse>;
+  updateThreadResponse(
+    responseId: number,
+    data: { sql: string },
+  ): Promise<ThreadResponse>;
   getResponsesWithThread(threadId: number): Promise<ThreadResponse[]>;
   getResponse(responseId: number): Promise<ThreadResponse>;
   generateThreadResponseBreakdown(
@@ -711,6 +715,22 @@ export class AskingService implements IAskingService {
     return threadResponse;
   }
 
+  public async updateThreadResponse(
+    responseId: number,
+    data: { sql: string },
+  ): Promise<ThreadResponse> {
+    const threadResponse = await this.threadResponseRepository.findOneBy({
+      id: responseId,
+    });
+    if (!threadResponse) {
+      throw new Error(`Thread response ${responseId} not found`);
+    }
+
+    return await this.threadResponseRepository.updateOne(responseId, {
+      sql: data.sql,
+    });
+  }
+
   public async generateThreadResponseBreakdown(
     threadResponseId: number,
     configurations: { language: string },

wren-ui/src/components/ErrorCollapse.tsx

@@ -1,3 +1,4 @@
+import { useState } from 'react';
 import { Collapse } from 'antd';
 import styled from 'styled-components';
 import CaretRightOutlined from '@ant-design/icons/CaretRightOutlined';
@@ -26,14 +27,21 @@ const StyledCollapse = styled(Collapse)`
 interface Props {
   message: string;
   className?: string;
+  defaultActive?: boolean;
 }
 
 export default function ErrorCollapse(props: Props) {
-  const { message, className } = props;
+  const { message, className, defaultActive } = props;
+  const [activeKey, setActiveKey] = useState<string[]>(
+    defaultActive ? ['1'] : [],
+  );
+
   return (
     <StyledCollapse
       className={className}
       ghost
+      activeKey={activeKey}
+      onChange={(key) => setActiveKey(key as string[])}
       expandIcon={({ isActive }) => (
         <CaretRightOutlined rotate={isActive ? 90 : 0} />
       )}

wren-ui/src/components/modals/FixSQLModal.tsx

@@ -0,0 +1,135 @@
+import { useEffect, useMemo, useState } from 'react';
+import { Button, Form, Modal, Typography, Alert } from 'antd';
+import { ERROR_TEXTS } from '@/utils/error';
+import { ModalAction } from '@/hooks/useModalAction';
+import { attachLoading } from '@/utils/helper';
+import { parseGraphQLError } from '@/utils/errorHandler';
+import SQLEditor from '@/components/editor/SQLEditor';
+import ErrorCollapse from '@/components/ErrorCollapse';
+import PreviewData from '@/components/dataPreview/PreviewData';
+import { usePreviewSqlMutation } from '@/apollo/client/graphql/sql.generated';
+
+type Props = ModalAction<{ sql: string; responseId: number }> & {
+  loading?: boolean;
+};
+
+export function FixSQLModal(props: Props) {
+  const { visible, defaultValue, loading, onSubmit, onClose } = props;
+  const [previewLoading, setPreviewLoading] = useState(false);
+  const [form] = Form.useForm();
+
+  const [previewSqlMutation, previewSqlResult] = usePreviewSqlMutation();
+  const error = useMemo(() => {
+    if (!previewSqlResult.error) return null;
+    const graphQLError = parseGraphQLError(previewSqlResult.error);
+    return { ...graphQLError, shortMessage: 'Invalid SQL syntax' };
+  }, [previewSqlResult.error]);
+
+  useEffect(() => {
+    if (!visible) return;
+    form.setFieldsValue(defaultValue || {});
+  }, [form, defaultValue, visible]);
+
+  const validateSql = async () => {
+    const sql = form.getFieldValue('sql');
+    await previewSqlMutation({
+      variables: { data: { sql, limit: 1, dryRun: true } },
+    });
+  };
+
+  const previewData = async () => {
+    form
+      .validateFields()
+      .then(async (values) => {
+        await attachLoading(
+          previewSqlMutation,
+          setPreviewLoading,
+        )({
+          variables: { data: { sql: values.sql, limit: 50 } },
+        });
+      })
+      .catch(console.error);
+  };
+
+  const reset = () => {
+    form.resetFields();
+    previewSqlResult.reset();
+  };
+
+  const submit = async () => {
+    form
+      .validateFields()
+      .then(async (values) => {
+        await validateSql();
+        await onSubmit(values.sql);
+        onClose();
+      })
+      .catch(console.error);
+  };
+
+  const showPreview = previewSqlResult.data || previewSqlResult.loading;
+
+  return (
+    <Modal
+      title="Fix SQL"
+      width={640}
+      visible={visible}
+      okText="Submit"
+      onOk={submit}
+      onCancel={onClose}
+      confirmLoading={loading}
+      maskClosable={false}
+      destroyOnClose
+      centered
+      afterClose={reset}
+    >
+      <Typography.Text className="d-block gray-7 mb-3">
+        The following SQL statement needs to be fixed:
+      </Typography.Text>
+      <Form form={form} preserve={false} layout="vertical">
+        <Form.Item
+          label="SQL statement"
+          name="sql"
+          required
+          rules={[
+            {
+              required: true,
+              message: ERROR_TEXTS.FIX_SQL.SQL.REQUIRED,
+            },
+          ]}
+        >
+          <SQLEditor autoFocus />
+        </Form.Item>
+      </Form>
+      <div className="my-3">
+        <Typography.Text className="d-block gray-7 mb-2">
+          Data preview (50 rows)
+        </Typography.Text>
+        <Button
+          onClick={previewData}
+          loading={previewLoading}
+          disabled={previewLoading}
+        >
+          Preview data
+        </Button>
+        {showPreview && (
+          <div className="my-3">
+            <PreviewData
+              loading={previewLoading}
+              previewData={previewSqlResult?.data?.previewSql}
+              copyable={false}
+            />
+          </div>
+        )}
+      </div>
+      {!!error && (
+        <Alert
+          showIcon
+          type="error"
+          message={error.shortMessage}
+          description={<ErrorCollapse message={error.message} />}
+        />
+      )}
+    </Modal>
+  );
+}

wren-ui/src/components/modals/QuestionSQLPairModal.tsx

@@ -85,7 +85,7 @@ export default function QuestionSQLPairModal(props: Props) {
 
   const handleError = (error) => {
     const graphQLError = parseGraphQLError(error);
-    setError({ ...graphQLError, shortMessage: 'Invalid SQL Syntax' });
+    setError({ ...graphQLError, shortMessage: 'Invalid SQL syntax' });
     console.error(graphQLError);
   };
 
@@ -274,7 +274,7 @@ export default function QuestionSQLPairModal(props: Props) {
         <Alert
           showIcon
           type="error"
-          message="Invalid SQL Syntax"
+          message={error.shortMessage}
           description={<ErrorCollapse message={error.message} />}
         />
       )}