feat(wren-ui): Enhance save as view for follow-up questions

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

@@ -243,6 +243,7 @@ export type CreateThreadResponseInput = {
 
 export type CreateViewInput = {
   name: Scalars['String'];
+  rephrasedQuestion: Scalars['String'];
   responseId: Scalars['Int'];
 };
 

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

@@ -802,7 +802,7 @@ export class ModelResolver {
 
   // create view from sql of a response
   public async createView(_root: any, args: any, ctx: IContext) {
-    const { name: displayName, responseId } = args.data;
+    const { name: displayName, responseId, rephrasedQuestion } = args.data;
 
     // validate view name
     const validateResult = await this.validateViewName(displayName, ctx);
@@ -842,7 +842,7 @@ export class ModelResolver {
 
       // properties from the thread response
       responseId, // helpful for mapping back to the thread response
-      question: response.question,
+      question: rephrasedQuestion,
     };
 
     const eventName = TelemetryEvent.HOME_CREATE_VIEW;

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

@@ -436,6 +436,7 @@ export const typeDefs = gql`
   input CreateViewInput {
     name: String!
     responseId: Int!
+    rephrasedQuestion: String!
   }
 
   input ValidateViewInput {

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

@@ -10,10 +10,11 @@ const { Text } = Typography;
 type Props = ModalAction<{ sql: string }> & {
   loading?: boolean;
   defaultValue: { sql: string; responseId: number };
+  payload: { rephrasedQuestion: string };
 };
 
 export default function SaveAsViewModal(props: Props) {
-  const { visible, loading, onSubmit, onClose, defaultValue } = props;
+  const { visible, loading, onSubmit, onClose, defaultValue, payload } = props;
   const [form] = Form.useForm();
   const [validateViewMutation] = useValidateViewMutation({
     fetchPolicy: 'no-cache',
@@ -23,7 +24,11 @@ export default function SaveAsViewModal(props: Props) {
     form
       .validateFields()
       .then(async (values) => {
-        await onSubmit({ responseId: defaultValue.responseId, ...values });
+        await onSubmit({
+          responseId: defaultValue.responseId,
+          ...payload,
+          ...values,
+        });
         onClose();
       })
       .catch(console.error);

wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx

@@ -122,6 +122,7 @@ export interface Props {
   motion: boolean;
   threadResponse: ThreadResponse;
   isLastThreadResponse: boolean;
+  isOpeningQuestion: boolean;
   onInitPreviewDone: () => void;
 }
 
@@ -186,7 +187,7 @@ const isNeedGenerateAnswer = (answerDetail: ThreadResponseAnswerDetail) => {
 };
 
 export default function AnswerResult(props: Props) {
-  const { threadResponse, isLastThreadResponse } = props;
+  const { threadResponse, isLastThreadResponse, isOpeningQuestion } = props;
 
   const {
     onOpenSaveAsViewModal,
@@ -270,6 +271,17 @@ export default function AnswerResult(props: Props) {
     isAnswerPrepared ||
     isBreakdownOnly;
 
+  const rephrasedQuestion =
+    threadResponse?.askingTask?.rephrasedQuestion || question;
+
+  const questionForSaveAsView = useMemo(() => {
+    // use rephrased question for follow-up questions, otherwise use the original question
+
+    if (isOpeningQuestion) return question;
+
+    return rephrasedQuestion;
+  }, [rephrasedQuestion, question, isOpeningQuestion]);
+
   return (
     <div style={resultStyle} data-jsid="answerResult">
       {isAdjustment && <AdjustmentInformation adjustment={adjustment} />}
@@ -334,9 +346,7 @@ export default function AnswerResult(props: Props) {
                 onClick={() =>
                   onOpenSaveToKnowledgeModal(
                     {
-                      question:
-                        threadResponse?.askingTask?.rephrasedQuestion ||
-                        question,
+                      question: rephrasedQuestion,
                       sql,
                     },
                     { isCreateMode: true },
@@ -352,7 +362,14 @@ export default function AnswerResult(props: Props) {
             </Tooltip>
             <ViewBlock
               view={view}
-              onClick={() => onOpenSaveAsViewModal({ sql, responseId: id })}
+              onClick={() =>
+                onOpenSaveAsViewModal(
+                  { sql, responseId: id },
+                  {
+                    rephrasedQuestion: questionForSaveAsView,
+                  },
+                )
+              }
             />
           </div>
           {renderRecommendedQuestions(

wren-ui/src/components/pages/home/promptThread/index.tsx

@@ -57,6 +57,7 @@ const AnswerResultTemplate: React.FC<
       {index > 0 && <Divider />}
       <AnswerResult
         motion={motion}
+        isOpeningQuestion={index === 0}
         isLastThreadResponse={isLastThreadResponse}
         onInitPreviewDone={onInitPreviewDone}
         threadResponse={threadResponse}

wren-ui/src/components/pages/home/promptThread/store.tsx

@@ -21,7 +21,10 @@ export type IPromptThreadStore = {
     onFixSQLStatement?: (responseId: number, sql: string) => Promise<void>;
     fixStatementLoading?: boolean;
   };
-  onOpenSaveAsViewModal: (data: { sql: string; responseId: number }) => void;
+  onOpenSaveAsViewModal: (
+    data: { sql: string; responseId: number },
+    payload: { rephrasedQuestion: string },
+  ) => void;
   onSelectRecommendedQuestion: ({
     question,
     sql,