feat(wren-ui): Reasoning enhancement & SQL pair in asking flow

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

@@ -1086,12 +1086,14 @@ export type ResolveSchemaChangeWhereInput = {
 export type ResultCandidate = {
   __typename?: 'ResultCandidate';
   sql: Scalars['String'];
+  sqlPair?: Maybe<SqlPair>;
   type: ResultCandidateType;
   view?: Maybe<ViewInfo>;
 };
 
 export enum ResultCandidateType {
   LLM = 'LLM',
+  SQL_PAIR = 'SQL_PAIR',
   VIEW = 'VIEW'
 }
 

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

@@ -68,6 +68,12 @@ const COMMON_ASKING_TASK = gql`
         statement
         displayName
       }
+      sqlPair {
+        id
+        question
+        sql
+        projectId
+      }
     }
     error {
       ...CommonError

wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts

@@ -711,6 +711,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
       type: candidate?.type?.toUpperCase() as AskCandidateType,
       sql: candidate.sql,
       viewId: candidate?.viewId ? Number(candidate.viewId) : null,
+      sqlpairId: candidate?.sqlpairId ? Number(candidate.sqlpairId) : null,
     }));
 
     return {

wren-ui/src/apollo/server/models/adaptor.ts

@@ -88,9 +88,11 @@ export enum AskResultType {
 
 // if it's view, viewId will be returned as well. It means the candidate is originally saved in mdl as a view.
 // if it's llm, viewId will not be returned. It means the candidate is generated by AI service.
+// if it's sql_pair, sqlpairId will be returned as well. It means the candidate is generated by sql_pair.
 export enum AskCandidateType {
   VIEW = 'VIEW',
   LLM = 'LLM',
+  SQL_PAIR = 'SQL_PAIR',
 }
 
 export interface AskResponse<R, S> {
@@ -119,6 +121,7 @@ export type AskResult = AskResponse<
     type: AskCandidateType;
     sql: string;
     viewId?: number | null;
+    sqlpairId?: number | null;
   }>,
   AskResultStatus
 > & {

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

@@ -648,10 +648,14 @@ export class AskingResolver {
         const view = response.viewId
           ? await ctx.viewRepository.findOneBy({ id: response.viewId })
           : null;
+        const sqlPair = response.sqlpairId
+          ? await ctx.sqlPairRepository.findOneBy({ id: response.sqlpairId })
+          : null;
         return {
           type: response.type,
           sql: response.sql,
           view,
+          sqlPair,
         };
       }),
     );

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

@@ -571,12 +571,14 @@ export const typeDefs = gql`
   enum ResultCandidateType {
     VIEW # View type candidate is provided basd on a saved view
     LLM # LLM type candidate is created by LLM
+    SQL_PAIR # SQL pair type candidate is created by SQL pair
   }
 
   type ResultCandidate {
     type: ResultCandidateType!
     sql: String!
     view: ViewInfo
+    sqlPair: SqlPair
   }
 
   type AskingTask {

wren-ui/src/components/pages/home/preparation/PreparationStatus.tsx

@@ -56,7 +56,12 @@ export default function PreparationStatus(props: Props) {
     );
   } else if (askingTask.status === AskingTaskStatus.FINISHED) {
     const showView = data.view !== null;
-    return <div className="gray-6">{showView ? '1 step' : '3 steps'}</div>;
+    const showSqlPair = !!askingTask?.candidates[0]?.sqlPair;
+    return (
+      <div className="gray-6">
+        {showView || showSqlPair ? '1 step' : '3 steps'}
+      </div>
+    );
   }
 
   return null;

wren-ui/src/components/pages/home/preparation/PreparationSteps.tsx

@@ -7,6 +7,7 @@ import Organizing from './step/Organizing';
 import Generating from './step/Generating';
 import FixedSQLFinished from './step/FixedSQLFinished';
 import ViewFinished from './step/ViewFinished';
+import SQLPairFinished from './step/SQLPairFinished';
 import { PROCESS_STATE } from '@/utils/enum';
 import {
   ProcessStateMachine,
@@ -47,7 +48,7 @@ const getProcessDot = (processing: boolean) => {
 };
 
 export default function PreparationSteps(props: Props) {
-  const { className, data, askingStreamTask, isAnswerFinished } = props;
+  const { className, data, askingStreamTask, minimized } = props;
   const { askingTask, view, sql } = data;
 
   const processState = useMemo(
@@ -60,12 +61,10 @@ export default function PreparationSteps(props: Props) {
 
   // displays
   const showView = !!view;
-  const showRetrieving =
-    retrievingNextStates.includes(processState) && !showView;
-  const showOrganizing =
-    organizingNextStates.includes(processState) && !showView;
-  const showGenerating =
-    generatingNextStates.includes(processState) && !showView;
+  const showSqlPair = !!askingTask?.candidates[0]?.sqlPair;
+  const showRetrieving = retrievingNextStates.includes(processState);
+  const showOrganizing = organizingNextStates.includes(processState);
+  const showGenerating = generatingNextStates.includes(processState);
 
   // data
   const retrievedTables = askingTask.retrievedTables || [];
@@ -77,10 +76,11 @@ export default function PreparationSteps(props: Props) {
   const organizing = processState === PROCESS_STATE.PLANNING;
   const generating = processState === PROCESS_STATE.GENERATING;
   const correcting = processState === PROCESS_STATE.CORRECTING;
-  const wrapping = !isAnswerFinished;
+  const wrapping = !minimized;
 
   // templates
   if (showView) return <ViewTimelineSteps {...props} />;
+  if (showSqlPair) return <SQLPairTimelineSteps {...props} />;
   if (isFixedSQL) return <FixedSQLTimelineSteps {...props} />;
 
   // default
@@ -139,3 +139,15 @@ function ViewTimelineSteps(props: Props) {
     </Timeline>
   );
 }
+
+function SQLPairTimelineSteps(props: Props) {
+  const { className } = props;
+
+  return (
+    <Timeline className={className}>
+      <Timeline.Item dot={fileDone}>
+        <SQLPairFinished />
+      </Timeline.Item>
+    </Timeline>
+  );
+}

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

@@ -15,19 +15,19 @@ import {
 export type Props = IPromptThreadStore['preparation'] & {
   className?: string;
   data: ThreadResponse;
-  isAnswerFinished?: boolean;
+  minimized?: boolean;
 };
 
 export default function Preparation(props: Props) {
-  const { className, data, isAnswerFinished, onFixSQLStatement } = props;
+  const { className, data, minimized, onFixSQLStatement } = props;
   const { askingTask, id: responseId, sql } = data;
 
   const [isActive, setIsActive] = useState(!sql);
 
   // wrapping up after answer is prepared
   useEffect(() => {
-    setIsActive(!isAnswerFinished);
-  }, [isAnswerFinished]);
+    setIsActive(!minimized);
+  }, [minimized]);
   const error = useMemo(() => {
     return askingTask?.error && !sql
       ? {

wren-ui/src/components/pages/home/preparation/step/SQLPairFinished.tsx

@@ -0,0 +1,16 @@
+import { Typography } from 'antd';
+
+export default function SQLPairFinished() {
+  return (
+    <>
+      <Typography.Text className="gray-8">
+        Using question-SQL pair
+      </Typography.Text>
+      <div className="gray-7 text-sm mt-1">
+        <div>
+          Matching question-SQL pair found. Returning results instantly.
+        </div>
+      </div>
+    </>
+  );
+}

wren-ui/src/components/pages/home/prompt/Result.tsx

@@ -159,7 +159,8 @@ const IntentionFinished = (props: Props) => {
     }
   }, [type]);
 
-  return null;
+  // To keep the UI result keep showing as understanding
+  return <Understanding {...props} />;
 };
 
 const GeneralAnswer = (props: Props) => {

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

@@ -231,17 +231,14 @@ export default function AnswerResult(props: Props) {
     isAnswerPrepared ||
     isBreakdownOnly;
 
-  const isAnswerFinished =
-    isAnswerPrepared && getAnswerIsFinished(answerDetail?.status);
-
   return (
     <div style={resultStyle} data-jsid="answerResult">
       <QuestionTitle className="mb-4" question={question} />
       <Preparation
         className="mb-3"
         {...preparation}
         data={threadResponse}
-        isAnswerFinished={isAnswerFinished}
+        minimized={isAnswerPrepared}
       />
       {showAnswerTabs && (
         <>

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

@@ -13,6 +13,7 @@ import RecommendedQuestionsPrompt from '@/components/pages/home/prompt/Recommend
 import {
   useSuggestedQuestionsQuery,
   useCreateThreadMutation,
+  useThreadLazyQuery,
 } from '@/apollo/client/graphql/home.generated';
 import { useGetSettingsQuery } from '@/apollo/client/graphql/settings.generated';
 import { CreateThreadInput } from '@/apollo/client/graphql/__types__';
@@ -98,6 +99,9 @@ export default function Home() {
   const [createThread, { loading: threadCreating }] = useCreateThreadMutation({
     onCompleted: () => homeSidebar.refetch(),
   });
+  const [preloadThread] = useThreadLazyQuery({
+    fetchPolicy: 'cache-and-network',
+  });
 
   const { data: settingsResult } = useGetSettingsQuery();
   const settings = settingsResult?.settings;
@@ -119,7 +123,9 @@ export default function Home() {
     try {
       askPrompt.onStopPolling();
       const response = await createThread({ variables: { data: payload } });
-      router.push(Path.Home + `/${response.data.createThread.id}`);
+      const threadId = response.data.createThread.id;
+      await preloadThread({ variables: { threadId } });
+      router.push(Path.Home + `/${threadId}`);
     } catch (error) {
       console.error(error);
     }