feat(wren-ui): Implement asking intention routing (#882)

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

@@ -19,6 +19,7 @@ export type AskingTask = {
   candidates: Array<ResultCandidate>;
   error?: Maybe<Error>;
   status: AskingTaskStatus;
+  type?: Maybe<AskingTaskType>;
 };
 
 export type AskingTaskInput = {
@@ -35,6 +36,12 @@ export enum AskingTaskStatus {
   UNDERSTANDING = 'UNDERSTANDING'
 }
 
+export enum AskingTaskType {
+  GENERAL = 'GENERAL',
+  MISLEADING_QUERY = 'MISLEADING_QUERY',
+  TEXT_TO_SQL = 'TEXT_TO_SQL'
+}
+
 export type CalculatedFieldInput = {
   diagram?: InputMaybe<Scalars['JSON']>;
   expression: Scalars['String'];
@@ -382,6 +389,10 @@ export type GetMdlResult = {
   mdl?: Maybe<Scalars['String']>;
 };
 
+export type InstantRecommendedQuestionsInput = {
+  previousQuestions?: InputMaybe<Array<Scalars['String']>>;
+};
+
 export type LearningRecord = {
   __typename?: 'LearningRecord';
   paths: Array<Scalars['String']>;
@@ -422,6 +433,7 @@ export type Mutation = {
   cancelAskingTask: Scalars['Boolean'];
   createAskingTask: Task;
   createCalculatedField: Scalars['JSON'];
+  createInstantRecommendedQuestions: Task;
   createModel: Scalars['JSON'];
   createRelation: Scalars['JSON'];
   createThread: Thread;
@@ -433,6 +445,8 @@ export type Mutation = {
   deleteThread: Scalars['Boolean'];
   deleteView: Scalars['Boolean'];
   deploy: Scalars['JSON'];
+  generateProjectRecommendationQuestions: Scalars['Boolean'];
+  generateThreadRecommendationQuestions: Scalars['Boolean'];
   previewData: Scalars['JSON'];
   previewModelData: Scalars['JSON'];
   previewSql: Scalars['JSON'];
@@ -473,6 +487,11 @@ export type MutationCreateCalculatedFieldArgs = {
 };
 
 
+export type MutationCreateInstantRecommendedQuestionsArgs = {
+  data: InstantRecommendedQuestionsInput;
+};
+
+
 export type MutationCreateModelArgs = {
   data: CreateModelInput;
 };
@@ -529,6 +548,11 @@ export type MutationDeployArgs = {
 };
 
 
+export type MutationGenerateThreadRecommendationQuestionsArgs = {
+  threadId: Scalars['Int'];
+};
+
+
 export type MutationPreviewDataArgs = {
   where: PreviewDataInput;
 };
@@ -703,6 +727,9 @@ export type Query = {
   autoGenerateRelation: Array<RecommendRelations>;
   diagram: Diagram;
   getMDL: GetMdlResult;
+  getProjectRecommendationQuestions: RecommendedQuestionsTask;
+  getThreadRecommendationQuestions: RecommendedQuestionsTask;
+  instantRecommendedQuestions: RecommendedQuestionsTask;
   learningRecord: LearningRecord;
   listDataSourceTables: Array<CompactTable>;
   listModels: Array<ModelInfo>;
@@ -731,6 +758,16 @@ export type QueryGetMdlArgs = {
 };
 
 
+export type QueryGetThreadRecommendationQuestionsArgs = {
+  threadId: Scalars['Int'];
+};
+
+
+export type QueryInstantRecommendedQuestionsArgs = {
+  taskId: Scalars['String'];
+};
+
+
 export type QueryModelArgs = {
   where: ModelWhereInput;
 };
@@ -763,6 +800,20 @@ export type RecommendRelations = {
   relations: Array<Maybe<Relation>>;
 };
 
+export type RecommendedQuestionsTask = {
+  __typename?: 'RecommendedQuestionsTask';
+  error?: Maybe<Error>;
+  questions: Array<ResultQuestion>;
+  status: RecommendedQuestionsTaskStatus;
+};
+
+export enum RecommendedQuestionsTaskStatus {
+  FAILED = 'FAILED',
+  FINISHED = 'FINISHED',
+  GENERATING = 'GENERATING',
+  NOT_STARTED = 'NOT_STARTED'
+}
+
 export type Relation = {
   __typename?: 'Relation';
   fromColumnId: Scalars['Int'];
@@ -808,6 +859,13 @@ export enum ResultCandidateType {
   VIEW = 'VIEW'
 }
 
+export type ResultQuestion = {
+  __typename?: 'ResultQuestion';
+  category: Scalars['String'];
+  explanation: Scalars['String'];
+  question: Scalars['String'];
+};
+
 export type SampleDatasetInput = {
   name: SampleDatasetName;
 };

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

@@ -17,7 +17,7 @@ export type AskingTaskQueryVariables = Types.Exact<{
 }>;
 
 
-export type AskingTaskQuery = { __typename?: 'Query', askingTask: { __typename?: 'AskingTask', status: Types.AskingTaskStatus, candidates: Array<{ __typename?: 'ResultCandidate', sql: string, summary: 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 } };
+export type AskingTaskQuery = { __typename?: 'Query', askingTask: { __typename?: 'AskingTask', status: Types.AskingTaskStatus, type?: Types.AskingTaskType | null, candidates: Array<{ __typename?: 'ResultCandidate', sql: string, summary: 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 } };
 
 export type ThreadsQueryVariables = Types.Exact<{ [key: string]: never; }>;
 
@@ -96,6 +96,20 @@ export type GetNativeSqlQueryVariables = Types.Exact<{
 
 export type GetNativeSqlQuery = { __typename?: 'Query', nativeSql: string };
 
+export type CreateInstantRecommendedQuestionsMutationVariables = Types.Exact<{
+  data: Types.InstantRecommendedQuestionsInput;
+}>;
+
+
+export type CreateInstantRecommendedQuestionsMutation = { __typename?: 'Mutation', createInstantRecommendedQuestions: { __typename?: 'Task', id: string } };
+
+export type InstantRecommendedQuestionsQueryVariables = Types.Exact<{
+  taskId: Types.Scalars['String'];
+}>;
+
+
+export type InstantRecommendedQuestionsQuery = { __typename?: 'Query', instantRecommendedQuestions: { __typename?: 'RecommendedQuestionsTask', status: Types.RecommendedQuestionsTaskStatus, questions: Array<{ __typename?: 'ResultQuestion', question: string, category: string, explanation: string }>, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array<string | null> | null } | null } };
+
 export const CommonErrorFragmentDoc = gql`
     fragment CommonError on Error {
   code
@@ -168,6 +182,7 @@ export const AskingTaskDocument = gql`
     query AskingTask($taskId: String!) {
   askingTask(taskId: $taskId) {
     status
+    type
     candidates {
       sql
       summary
@@ -600,4 +615,80 @@ export function useGetNativeSqlLazyQuery(baseOptions?: Apollo.LazyQueryHookOptio
         }
 export type GetNativeSqlQueryHookResult = ReturnType<typeof useGetNativeSqlQuery>;
 export type GetNativeSqlLazyQueryHookResult = ReturnType<typeof useGetNativeSqlLazyQuery>;
-export type GetNativeSqlQueryResult = Apollo.QueryResult<GetNativeSqlQuery, GetNativeSqlQueryVariables>;
+export type GetNativeSqlQueryResult = Apollo.QueryResult<GetNativeSqlQuery, GetNativeSqlQueryVariables>;
+export const CreateInstantRecommendedQuestionsDocument = gql`
+    mutation CreateInstantRecommendedQuestions($data: InstantRecommendedQuestionsInput!) {
+  createInstantRecommendedQuestions(data: $data) {
+    id
+  }
+}
+    `;
+export type CreateInstantRecommendedQuestionsMutationFn = Apollo.MutationFunction<CreateInstantRecommendedQuestionsMutation, CreateInstantRecommendedQuestionsMutationVariables>;
+
+/**
+ * __useCreateInstantRecommendedQuestionsMutation__
+ *
+ * To run a mutation, you first call `useCreateInstantRecommendedQuestionsMutation` within a React component and pass it any options that fit your needs.
+ * When your component renders, `useCreateInstantRecommendedQuestionsMutation` 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 [createInstantRecommendedQuestionsMutation, { data, loading, error }] = useCreateInstantRecommendedQuestionsMutation({
+ *   variables: {
+ *      data: // value for 'data'
+ *   },
+ * });
+ */
+export function useCreateInstantRecommendedQuestionsMutation(baseOptions?: Apollo.MutationHookOptions<CreateInstantRecommendedQuestionsMutation, CreateInstantRecommendedQuestionsMutationVariables>) {
+        const options = {...defaultOptions, ...baseOptions}
+        return Apollo.useMutation<CreateInstantRecommendedQuestionsMutation, CreateInstantRecommendedQuestionsMutationVariables>(CreateInstantRecommendedQuestionsDocument, options);
+      }
+export type CreateInstantRecommendedQuestionsMutationHookResult = ReturnType<typeof useCreateInstantRecommendedQuestionsMutation>;
+export type CreateInstantRecommendedQuestionsMutationResult = Apollo.MutationResult<CreateInstantRecommendedQuestionsMutation>;
+export type CreateInstantRecommendedQuestionsMutationOptions = Apollo.BaseMutationOptions<CreateInstantRecommendedQuestionsMutation, CreateInstantRecommendedQuestionsMutationVariables>;
+export const InstantRecommendedQuestionsDocument = gql`
+    query InstantRecommendedQuestions($taskId: String!) {
+  instantRecommendedQuestions(taskId: $taskId) {
+    status
+    questions {
+      question
+      category
+      explanation
+    }
+    error {
+      ...CommonError
+    }
+  }
+}
+    ${CommonErrorFragmentDoc}`;
+
+/**
+ * __useInstantRecommendedQuestionsQuery__
+ *
+ * To run a query within a React component, call `useInstantRecommendedQuestionsQuery` and pass it any options that fit your needs.
+ * When your component renders, `useInstantRecommendedQuestionsQuery` returns an object from Apollo Client that contains loading, error, and data properties
+ * you can use to render your UI.
+ *
+ * @param baseOptions options that will be passed into the query, supported options are listed on: https://www.apollographql.com/docs/react/api/react-hooks/#options;
+ *
+ * @example
+ * const { data, loading, error } = useInstantRecommendedQuestionsQuery({
+ *   variables: {
+ *      taskId: // value for 'taskId'
+ *   },
+ * });
+ */
+export function useInstantRecommendedQuestionsQuery(baseOptions: Apollo.QueryHookOptions<InstantRecommendedQuestionsQuery, InstantRecommendedQuestionsQueryVariables>) {
+        const options = {...defaultOptions, ...baseOptions}
+        return Apollo.useQuery<InstantRecommendedQuestionsQuery, InstantRecommendedQuestionsQueryVariables>(InstantRecommendedQuestionsDocument, options);
+      }
+export function useInstantRecommendedQuestionsLazyQuery(baseOptions?: Apollo.LazyQueryHookOptions<InstantRecommendedQuestionsQuery, InstantRecommendedQuestionsQueryVariables>) {
+          const options = {...defaultOptions, ...baseOptions}
+          return Apollo.useLazyQuery<InstantRecommendedQuestionsQuery, InstantRecommendedQuestionsQueryVariables>(InstantRecommendedQuestionsDocument, options);
+        }
+export type InstantRecommendedQuestionsQueryHookResult = ReturnType<typeof useInstantRecommendedQuestionsQuery>;
+export type InstantRecommendedQuestionsLazyQueryHookResult = ReturnType<typeof useInstantRecommendedQuestionsLazyQuery>;
+export type InstantRecommendedQuestionsQueryResult = Apollo.QueryResult<InstantRecommendedQuestionsQuery, InstantRecommendedQuestionsQueryVariables>;

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

@@ -48,6 +48,7 @@ export const ASKING_TASK = gql`
   query AskingTask($taskId: String!) {
     askingTask(taskId: $taskId) {
       status
+      type
       candidates {
         sql
         summary
@@ -180,3 +181,30 @@ export const GET_NATIVE_SQL = gql`
     nativeSql(responseId: $responseId)
   }
 `;
+
+export const CREATE_INSTANT_RECOMMENDED_QUESTIONS = gql`
+  mutation CreateInstantRecommendedQuestions(
+    $data: InstantRecommendedQuestionsInput!
+  ) {
+    createInstantRecommendedQuestions(data: $data) {
+      id
+    }
+  }
+`;
+
+export const INSTANT_RECOMMENDED_QUESTIONS = gql`
+  query InstantRecommendedQuestions($taskId: String!) {
+    instantRecommendedQuestions(taskId: $taskId) {
+      status
+      questions {
+        question
+        category
+        explanation
+      }
+      error {
+        ...CommonError
+      }
+    }
+  }
+  ${COMMON_ERROR}
+`;

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

@@ -1,4 +1,5 @@
 import axios from 'axios';
+import { Readable } from 'stream';
 import { Manifest } from '@server/mdl/type';
 import { getLogger } from '@server/utils';
 import * as Errors from '@server/utils/error';
@@ -83,6 +84,12 @@ export enum AskResultStatus {
   STOPPED = 'STOPPED',
 }
 
+export enum AskResultType {
+  GENERAL = 'GENERAL',
+  TEXT_TO_SQL = 'TEXT_TO_SQL',
+  MISLEADING_QUERY = 'MISLEADING_QUERY',
+}
+
 // 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.
 export enum AskCandidateType {
@@ -91,6 +98,7 @@ export enum AskCandidateType {
 }
 
 export interface AskResponse<R, S> {
+  type: AskResultType | null;
   status: S;
   response: R | null;
   error: WrenAIError | null;
@@ -179,6 +187,7 @@ export interface IWrenAIAdaptor {
   ask(input: AskInput): Promise<AsyncQueryResponse>;
   cancelAsk(queryId: string): Promise<void>;
   getAskResult(queryId: string): Promise<AskResult>;
+  getAskStreamingResult(queryId: string): Promise<Readable>;
 
   /**
    * After you choose a candidate, you can request AI service to generate the detail.
@@ -256,6 +265,25 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
     }
   }
 
+  public async getAskStreamingResult(queryId: string): Promise<Readable> {
+    // make GET request /v1/asks/:query_id/streaming-result to get the streaming result
+    try {
+      const res = await axios.get(
+        `${this.wrenAIBaseEndpoint}/v1/asks/${queryId}/streaming-result`,
+        { responseType: 'stream' },
+      );
+      return res.data;
+    } catch (err: any) {
+      logger.debug(
+        `Got error when getting ask streaming result: ${getAISerciceError(err)}`,
+      );
+      // throw err;
+      throw Errors.create(Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, {
+        originalError: err,
+      });
+    }
+  }
+
   /**
    * After you choose a candidate, you can request AI service to generate the detail.
    */
@@ -419,6 +447,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
   }
 
   private transformAskResult(body: any): AskResult {
+    const { type } = body;
     const { status, error } = this.transformStatusAndError(body);
     const candidates = (body?.response || []).map((candidate: any) => ({
       type: candidate?.type?.toUpperCase() as AskCandidateType,
@@ -428,6 +457,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
     }));
 
     return {
+      type,
       status,
       error,
       response: candidates,
@@ -446,6 +476,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
   }
 
   private transformAskDetailResult(body: any): AskDetailResult {
+    const { type } = body;
     const { status, error } = this.transformStatusAndError(body);
 
     // snake_case to camelCase
@@ -456,6 +487,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
     }));
 
     return {
+      type,
       status,
       error,
       response: {

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

@@ -27,6 +27,7 @@ const resolvers = {
     // Ask
     askingTask: askingResolver.getAskingTask,
     suggestedQuestions: askingResolver.getSuggestedQuestions,
+    instantRecommendedQuestions: askingResolver.getInstantRecommendedQuestions,
 
     // Thread
     thread: askingResolver.getThread,
@@ -79,6 +80,8 @@ const resolvers = {
     // Ask
     createAskingTask: askingResolver.createAskingTask,
     cancelAskingTask: askingResolver.cancelAskingTask,
+    createInstantRecommendedQuestions:
+      askingResolver.createInstantRecommendedQuestions,
 
     // Thread
     createThread: askingResolver.createThread,

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

@@ -2,6 +2,8 @@ import {
   WrenAIError,
   WrenAILanguage,
   AskResultStatus,
+  AskResultType,
+  RecommendationQuestionStatus,
 } from '../adaptors/wrenAIAdaptor';
 import { Thread } from '../repositories/threadRepository';
 import {
@@ -35,6 +37,7 @@ export interface Task {
 }
 
 export interface AskingTask {
+  type: AskResultType | null;
   status: AskResultStatus;
   candidates: Array<{
     sql: string;
@@ -51,6 +54,16 @@ export interface DetailedThread {
   responses: ThreadResponse[];
 }
 
+export interface RecommendedQuestionsTask {
+  questions: {
+    question: string;
+    category: string;
+    explanation: string;
+  }[];
+  status: RecommendationQuestionStatus;
+  error: WrenAIError | null;
+}
+
 export class AskingResolver {
   constructor() {
     this.createAskingTask = this.createAskingTask.bind(this);
@@ -63,7 +76,12 @@ export class AskingResolver {
     this.listThreads = this.listThreads.bind(this);
     this.createThreadResponse = this.createThreadResponse.bind(this);
     this.getResponse = this.getResponse.bind(this);
+    this.previewData = this.previewData.bind(this);
     this.getSuggestedQuestions = this.getSuggestedQuestions.bind(this);
+    this.createInstantRecommendedQuestions =
+      this.createInstantRecommendedQuestions.bind(this);
+    this.getInstantRecommendedQuestions =
+      this.getInstantRecommendedQuestions.bind(this);
     this.generateThreadRecommendationQuestions =
       this.generateThreadRecommendationQuestions.bind(this);
     this.generateProjectRecommendationQuestions =
@@ -196,6 +214,7 @@ export class AskingResolver {
     );
 
     return {
+      type: askResult.type,
       status: askResult.status,
       error: askResult.error,
       candidates,
@@ -396,6 +415,31 @@ export class AskingResolver {
     return data;
   }
 
+  public async createInstantRecommendedQuestions(
+    _root: any,
+    args: { data: { previousQuestions?: string[] } },
+    ctx: IContext,
+  ): Promise<Task> {
+    const { data } = args;
+    const askingService = ctx.askingService;
+    return askingService.createInstantRecommendedQuestions(data);
+  }
+
+  public async getInstantRecommendedQuestions(
+    _root: any,
+    args: { taskId: string },
+    ctx: IContext,
+  ): Promise<RecommendedQuestionsTask> {
+    const { taskId } = args;
+    const askingService = ctx.askingService;
+    const result = await askingService.getInstantRecommendedQuestions(taskId);
+    return {
+      questions: result.response?.questions || [],
+      status: result.status,
+      error: result.error,
+    };
+  }
+
   /**
    * Nested resolvers
    */

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

@@ -542,6 +542,12 @@ export const typeDefs = gql`
     STOPPED
   }
 
+  enum AskingTaskType {
+    GENERAL
+    TEXT_TO_SQL
+    MISLEADING_QUERY
+  }
+
   enum ResultCandidateType {
     VIEW # View type candidate is provided basd on a saved view
     LLM # LLM type candidate is created by LLM
@@ -556,10 +562,34 @@ export const typeDefs = gql`
 
   type AskingTask {
     status: AskingTaskStatus!
+    type: AskingTaskType
     error: Error
     candidates: [ResultCandidate!]!
   }
 
+  input InstantRecommendedQuestionsInput {
+    previousQuestions: [String!]
+  }
+
+  enum RecommendedQuestionsTaskStatus {
+    NOT_STARTED
+    GENERATING
+    FINISHED
+    FAILED
+  }
+
+  type ResultQuestion {
+    question: String!
+    category: String!
+    explanation: String!
+  }
+
+  type RecommendedQuestionsTask {
+    status: RecommendedQuestionsTaskStatus!
+    questions: [ResultQuestion!]!
+    error: Error
+  }
+
   # Thread
   input CreateThreadInput {
     question: String
@@ -717,31 +747,6 @@ export const typeDefs = gql`
     path: String!
   }
 
-  enum RecommendQuestionResultStatus {
-    NOT_STARTED
-    GENERATING
-    FINISHED
-    FAILED
-  }
-
-  type RecommendationQuestion {
-    question: String!
-    category: String!
-    sql: String!
-  }
-
-  type ThreadRecommendQuestionResult {
-    status: RecommendQuestionResultStatus!
-    questions: [RecommendationQuestion!]!
-    error: Error
-  }
-
-  type ProjectRecommendationQuestionsResult {
-    status: RecommendQuestionResultStatus!
-    questions: [RecommendationQuestion!]!
-    error: Error
-  }
-
   # Query and Mutation
   type Query {
     # On Boarding Steps
@@ -777,10 +782,10 @@ export const typeDefs = gql`
     # Learning
     learningRecord: LearningRecord!
 
-    getThreadRecommendationQuestions(
-      threadId: Int!
-    ): ThreadRecommendQuestionResult!
-    getProjectRecommendationQuestions: ProjectRecommendationQuestionsResult!
+    # Recommendation questions
+    getThreadRecommendationQuestions(threadId: Int!): RecommendedQuestionsTask!
+    getProjectRecommendationQuestions: RecommendedQuestionsTask!
+    instantRecommendedQuestions(taskId: String!): RecommendedQuestionsTask!
   }
 
   type Mutation {
@@ -864,5 +869,8 @@ export const typeDefs = gql`
     # Recommendation questions
     generateThreadRecommendationQuestions(threadId: Int!): Boolean!
     generateProjectRecommendationQuestions: Boolean!
+    createInstantRecommendedQuestions(
+      data: InstantRecommendedQuestionsInput!
+    ): Task!
   }
 `;

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

@@ -3,6 +3,7 @@ import {
   IWrenAIAdaptor,
   AskResultStatus,
   AskHistory,
+  RecommendationQuestionsResult,
   RecommendationQuestionsInput,
   RecommendationQuestion,
   WrenAIError,
@@ -66,6 +67,10 @@ export interface ThreadRecommendQuestionResult {
   error?: WrenAIError;
 }
 
+export interface InstantRecommendedQuestionsInput {
+  previousQuestions?: string[];
+}
+
 export interface IAskingService {
   /**
    * Asking task.
@@ -103,13 +108,22 @@ export interface IAskingService {
     stepIndex?: number,
     limit?: number,
   ): Promise<PreviewDataResponse>;
-  deleteAllByProjectId(projectId: number): Promise<void>;
 
-  // recommendation questions
+  /**
+   * Recommendation questions
+   */
+  createInstantRecommendedQuestions(
+    input: InstantRecommendedQuestionsInput,
+  ): Promise<Task>;
+  getInstantRecommendedQuestions(
+    queryId: string,
+  ): Promise<RecommendationQuestionsResult>;
   generateThreadRecommendationQuestions(threadId: number): Promise<void>;
   getThreadRecommendationQuestions(
     threadId: number,
   ): Promise<ThreadRecommendQuestionResult>;
+
+  deleteAllByProjectId(projectId: number): Promise<void>;
 }
 
 /**
@@ -651,6 +665,31 @@ export class AskingService implements IAskingService {
     }
   }
 
+  public async createInstantRecommendedQuestions(
+    input: InstantRecommendedQuestionsInput,
+  ): Promise<Task> {
+    const project = await this.projectService.getCurrentProject();
+    const { manifest } = await this.deployService.getLastDeployment(project.id);
+
+    const response = await this.wrenAIAdaptor.generateRecommendationQuestions({
+      manifest,
+      projectId: project.id.toString(),
+      previousQuestions: input.previousQuestions,
+      maxCategories: 3,
+      maxQuestions: 3,
+      configuration: { language: project.language },
+    });
+    return { id: response.queryId };
+  }
+
+  public async getInstantRecommendedQuestions(
+    queryId: string,
+  ): Promise<RecommendationQuestionsResult> {
+    const response =
+      await this.wrenAIAdaptor.getRecommendationQuestionsResult(queryId);
+    return response;
+  }
+
   public async deleteAllByProjectId(projectId: number): Promise<void> {
     // delete all threads
     await this.threadRepository.deleteAllBy({ projectId });

wren-ui/src/apollo/server/utils/error.ts

@@ -4,7 +4,6 @@ import { WrenService } from '../telemetry/telemetry';
 export enum GeneralErrorCodes {
   INTERNAL_SERVER_ERROR = 'INTERNAL_SERVER_ERROR',
   // AI service errors
-  MISLEADING_QUERY = 'MISLEADING_QUERY',
   NO_RELEVANT_DATA = 'NO_RELEVANT_DATA',
   NO_RELEVANT_SQL = 'NO_RELEVANT_SQL',
   RESOURCE_NOT_FOUND = 'RESOURCE_NOT_FOUND',
@@ -41,12 +40,10 @@ export const errorMessages = {
   [GeneralErrorCodes.INTERNAL_SERVER_ERROR]: 'Internal server error',
 
   // AI service errors
-  [GeneralErrorCodes.MISLEADING_QUERY]:
-    'The query provided is misleading and may not yield accurate results. Please refine your query.',
   [GeneralErrorCodes.NO_RELEVANT_DATA]:
-    'No relevant data found for the given query. Please try a different query.',
+    'I can’t find the exact data you’re looking for, but feel free to ask about other available topics.',
   [GeneralErrorCodes.NO_RELEVANT_SQL]:
-    'No relevant SQL found for the given query. Please check your query and try again.',
+    "Could you please provide more details or specify the information you're seeking?",
 
   // Connector errors
   [GeneralErrorCodes.CONNECTION_ERROR]: 'Can not connect to data source',
@@ -80,9 +77,8 @@ export const errorMessages = {
 
 export const shortMessages = {
   [GeneralErrorCodes.INTERNAL_SERVER_ERROR]: 'Internal server error',
-  [GeneralErrorCodes.MISLEADING_QUERY]: 'Misleading query',
-  [GeneralErrorCodes.NO_RELEVANT_DATA]: 'No relevant data',
-  [GeneralErrorCodes.NO_RELEVANT_SQL]: 'No relevant SQL',
+  [GeneralErrorCodes.NO_RELEVANT_DATA]: 'Try a different query',
+  [GeneralErrorCodes.NO_RELEVANT_SQL]: 'Clarification needed',
   [GeneralErrorCodes.CONNECTION_ERROR]: 'Failed to connect',
   [GeneralErrorCodes.IBIS_SERVER_ERROR]: 'Data connection error',
   [GeneralErrorCodes.INIT_SQL_ERROR]: 'Invalid initializing SQL',

wren-ui/src/components/editor/MarkdownBlock.tsx

@@ -0,0 +1,62 @@
+import styled from 'styled-components';
+import ReactMarkdown from 'react-markdown';
+import remarkGfm from 'remark-gfm';
+
+const ReactMarkdownBlock = styled(ReactMarkdown)`
+  h1,
+  h2,
+  h3,
+  h4,
+  h5,
+  h6 {
+    color: var(--gray-10);
+    margin-bottom: 8px;
+  }
+  h1 {
+    font-size: 20px;
+  }
+  h2 {
+    font-size: 18px;
+  }
+  h3 {
+    font-size: 16px;
+  }
+  h4 {
+    font-size: 14px;
+  }
+  hr {
+    border-top: 1px solid var(--gray-5);
+    border-bottom: none;
+    border-left: none;
+    border-right: none;
+    margin: 18px 0;
+  }
+  pre {
+    background-color: var(--gray-2);
+    border: 1px var(--gray-4) solid;
+    padding: 16px;
+    border-radius: 4px;
+  }
+  table td,
+  table th {
+    border: 1px solid var(--gray-4);
+    padding: 4px 8px;
+  }
+  table th {
+    background-color: var(--gray-2);
+    font-weight: 600;
+  }
+  table {
+    border: 1px solid var(--gray-4);
+    border-collapse: collapse;
+    margin-bottom: 16px;
+  }
+`;
+
+export default function MarkdownBlock(props: { content: string }) {
+  return (
+    <ReactMarkdownBlock remarkPlugins={[remarkGfm]}>
+      {props.content}
+    </ReactMarkdownBlock>
+  );
+}

wren-ui/src/components/pages/home/RecommendedQuestions.tsx

@@ -8,6 +8,12 @@ import { makeIterable } from '@/utils/iteration';
 interface Props {
   items: string[];
   loading?: boolean;
+  error?: {
+    shortMessage?: string;
+    code?: string;
+    message?: string;
+    stacktrace?: string[];
+  };
   className?: string;
   onSelect?: (question: string) => void;
 }

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

@@ -1,4 +1,4 @@
-import { ReactNode, useEffect } from 'react';
+import { ReactNode, useEffect, useRef } from 'react';
 import { Button } from 'antd';
 import styled from 'styled-components';
 import { PROCESS_STATE } from '@/utils/enum';
@@ -7,9 +7,18 @@ import StopOutlined from '@ant-design/icons/StopFilled';
 import LoadingOutlined from '@ant-design/icons/LoadingOutlined';
 import CloseCircleFilled from '@ant-design/icons/CloseCircleFilled';
 import WarningOutlined from '@ant-design/icons/WarningOutlined';
+import MessageOutlined from '@ant-design/icons/MessageOutlined';
 import ErrorCollapse from '@/components/ErrorCollapse';
+import InfoCircleOutlined from '@ant-design/icons/InfoCircleOutlined';
 import useAskProcessState from '@/hooks/useAskProcessState';
-import { AskingTask } from '@/apollo/client/graphql/__types__';
+import RecommendedQuestions from '@/components/pages/home/RecommendedQuestions';
+import MarkdownBlock from '@/components/editor/MarkdownBlock';
+import {
+  AskingTask,
+  AskingTaskType,
+  RecommendedQuestionsTask,
+  RecommendedQuestionsTaskStatus,
+} from '@/apollo/client/graphql/__types__';
 
 const StyledResult = styled.div`
   position: absolute;
@@ -24,9 +33,16 @@ const StyledResult = styled.div`
 
 interface Props {
   processState: ReturnType<typeof useAskProcessState>;
-  data: AskingTask['candidates'];
+  data: {
+    type: AskingTaskType;
+    originalQuestion: string;
+    candidates: AskingTask['candidates'];
+    askingStreamTask: string;
+    recommendedQuestions: RecommendedQuestionsTask;
+  };
   error?: any;
-  onSelect: (payload: { sql: string; summary: string }) => void;
+  onSelectResult: (payload: { sql: string; summary: string }) => void;
+  onSelectQuestion: (question: string) => void;
   onClose: () => void;
   onStop: () => void;
   loading?: boolean;
@@ -40,6 +56,20 @@ const Wrapper = ({ children }) => {
   );
 };
 
+const getRecommendedQuestionProps = (data: RecommendedQuestionsTask) => {
+  if (!data) return { show: false };
+  const questions = (data?.questions || []).map((item) => item.question);
+  const loading = data?.status === RecommendedQuestionsTaskStatus.GENERATING;
+  return {
+    show: loading || questions.length > 0,
+    state: {
+      items: questions,
+      loading,
+      error: data?.error,
+    },
+  };
+};
+
 const makeProcessing = (text: string) => (props: Props) => {
   const { onStop } = props;
   return (
@@ -66,9 +96,14 @@ const makeProcessing = (text: string) => (props: Props) => {
 const makeProcessingError =
   (config: { icon: ReactNode; title?: string; description?: string }) =>
   (props: Props) => {
-    const { onClose, error } = props;
+    const { onClose, onSelectQuestion, data, error } = props;
     const { message, shortMessage, stacktrace } = error || {};
     const hasStacktrace = !!stacktrace;
+
+    const recommendedQuestionProps = getRecommendedQuestionProps(
+      data?.recommendedQuestions,
+    );
+
     return (
       <Wrapper>
         <div className="d-flex justify-space-between text-medium mb-2">
@@ -90,6 +125,14 @@ const makeProcessingError =
         {hasStacktrace && (
           <ErrorCollapse className="mt-2" message={stacktrace.join('\n')} />
         )}
+
+        {recommendedQuestionProps.show && (
+          <RecommendedQuestions
+            className="mt-2"
+            {...recommendedQuestionProps.state}
+            onSelect={onSelectQuestion}
+          />
+        )}
       </Wrapper>
     );
   };
@@ -110,25 +153,121 @@ const Understanding = makeProcessing('Understanding question');
 const Searching = makeProcessing('Searching data');
 const Generating = makeProcessing('Generating answer');
 const Finished = (props: Props) => {
-  const { data, onSelect } = props;
+  const { data, onSelectResult } = props;
+  // only one candidate
+  const { candidates } = data;
 
   useEffect(() => {
-    if (data.length) {
-      const [result] = data;
-      onSelect && onSelect({ sql: result.sql, summary: result.summary });
+    if (candidates.length) {
+      const [result] = candidates;
+      onSelectResult &&
+        onSelectResult({ sql: result.sql, summary: result.summary });
     }
   }, [data]);
 
-  if (data.length === 0)
+  if (candidates.length === 0)
     return (
       <Wrapper>
         <NoResult {...props} />
       </Wrapper>
     );
+
   return null;
 };
 
-const getProcessStateComponent = (state: PROCESS_STATE) => {
+const GeneralAnswer = (props: Props) => {
+  const { onClose, onSelectQuestion, data, loading } = props;
+  const $wrapper = useRef<HTMLDivElement>(null);
+
+  const { originalQuestion, askingStreamTask, recommendedQuestions } = data;
+  const isDone = askingStreamTask && !loading;
+
+  const scrollBottom = () => {
+    if ($wrapper.current) {
+      $wrapper.current.scrollTo({
+        top: $wrapper.current.scrollHeight,
+      });
+    }
+  };
+
+  useEffect(() => {
+    scrollBottom();
+  }, [askingStreamTask]);
+
+  useEffect(() => {
+    if (isDone) scrollBottom();
+  }, [isDone]);
+
+  const recommendedQuestionProps =
+    getRecommendedQuestionProps(recommendedQuestions);
+
+  return (
+    <div>
+      <div className="d-flex justify-space-between">
+        <div className="d-flex align-start">
+          <MessageOutlined className="mr-2 mt-1 geekblue-6" />
+          <b className="text-semi-bold">{originalQuestion}</b>
+        </div>
+        <Button
+          className="adm-btn-no-style gray-7 bg-gray-3 text-sm px-2"
+          type="text"
+          size="small"
+          onClick={onClose}
+        >
+          <CloseOutlined className="-mr-1" />
+          Close
+        </Button>
+      </div>
+      <div className="py-3">
+        <div
+          ref={$wrapper}
+          className="py-2 px-3"
+          style={{ maxHeight: 'calc(100vh - 420px)', overflowY: 'auto' }}
+        >
+          <MarkdownBlock content={askingStreamTask} />
+          {isDone && (
+            <div className="gray-6">
+              <InfoCircleOutlined className="mr-2" />
+              For the most accurate semantics, please visit the modeling page.
+            </div>
+          )}
+        </div>
+      </div>
+
+      {recommendedQuestionProps.show && (
+        <RecommendedQuestions
+          {...recommendedQuestionProps.state}
+          onSelect={onSelectQuestion}
+        />
+      )}
+    </div>
+  );
+};
+
+const MisleadingQuery = makeProcessingError({
+  icon: <WarningOutlined className="mr-2 text-lg gold-6" />,
+  title: 'Clarification needed',
+  description:
+    "Could you please provide more details or specify the information you're seeking?",
+});
+
+const getGeneralAnswerStateComponent = (state: PROCESS_STATE) => {
+  return (
+    {
+      [PROCESS_STATE.FINISHED]: GeneralAnswer,
+    }[state] || null
+  );
+};
+
+const getMisleadingQueryStateComponent = (state: PROCESS_STATE) => {
+  return (
+    {
+      [PROCESS_STATE.FINISHED]: MisleadingQuery,
+    }[state] || null
+  );
+};
+
+const getDefaultStateComponent = (state: PROCESS_STATE) => {
   return (
     {
       [PROCESS_STATE.UNDERSTANDING]: Understanding,
@@ -140,9 +279,19 @@ const getProcessStateComponent = (state: PROCESS_STATE) => {
   );
 };
 
-export default function PromptResult(props: Props) {
-  const { processState } = props;
+const makeProcessStateStrategy = (type: AskingTaskType) => {
+  // note that the asking task type only has value when the asking status was finished
+  // by default, we use the default state component (also the text to sql state component)
+  if (type === AskingTaskType.GENERAL) return getGeneralAnswerStateComponent;
+  if (type === AskingTaskType.MISLEADING_QUERY)
+    return getMisleadingQueryStateComponent;
+  return getDefaultStateComponent;
+};
 
+export default function PromptResult(props: Props) {
+  const { processState, data } = props;
+
+  const getProcessStateComponent = makeProcessStateStrategy(data?.type);
   const StateComponent = getProcessStateComponent(processState.currentState);
 
   if (StateComponent === null) return null;

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

@@ -13,9 +13,11 @@ import PromptResult from '@/components/pages/home/prompt/Result';
 import useAskProcessState, {
   getIsProcessing,
 } from '@/hooks/useAskProcessState';
+import { AskPromptData } from '@/hooks/useAskPrompt';
 import {
   AskingTask,
   AskingTaskStatus,
+  AskingTaskType,
 } from '@/apollo/client/graphql/__types__';
 
 interface Props {
@@ -27,7 +29,11 @@ interface Props {
   }) => Promise<void>;
   onStop: () => void;
   onSubmit: (value: string) => Promise<void>;
-  data?: AskingTask;
+  onStopPolling: () => void;
+  onStopStreaming: () => void;
+  onStopRecommend: () => void;
+  data: AskPromptData;
+  loading: boolean;
 }
 
 interface Attributes {
@@ -60,7 +66,11 @@ const convertAskingTaskToProcessState = (data: AskingTask) => {
     [AskingTaskStatus.FINISHED]: PROCESS_STATE.FINISHED,
   }[data.status];
 
-  if (processState === PROCESS_STATE.FINISHED && data.candidates.length === 0) {
+  if (
+    data?.type === AskingTaskType.TEXT_TO_SQL &&
+    processState === PROCESS_STATE.FINISHED &&
+    data.candidates.length === 0
+  ) {
     return PROCESS_STATE.NO_RESULT;
   }
   return processState;
@@ -68,12 +78,36 @@ const convertAskingTaskToProcessState = (data: AskingTask) => {
 
 export default forwardRef<Attributes, Props>(function Prompt(props, ref) {
   const $promptInput = useRef<HTMLTextAreaElement>(null);
-  const { data, onSubmit, onStop, onSelect } = props;
+  const {
+    data,
+    loading,
+    onSubmit,
+    onStop,
+    onSelect,
+    onStopStreaming,
+    onStopRecommend,
+  } = props;
   const [inputValue, setInputValue] = useState('');
   const askProcessState = useAskProcessState();
 
-  const candidates = useMemo(() => data?.candidates || [], [data?.candidates]);
-  const error = useMemo(() => data?.error || null, [data?.error]);
+  const {
+    originalQuestion,
+    askingTask,
+    askingStreamTask,
+    recommendedQuestions,
+  } = data;
+
+  const result = useMemo(
+    () => ({
+      type: askingTask?.type, // question's type
+      originalQuestion, // original question
+      candidates: askingTask?.candidates || [], // for text to sql answer, only one candidate
+      askingStreamTask, // for general answer
+      recommendedQuestions, // guiding user to ask
+    }),
+    [data],
+  );
+  const error = useMemo(() => askingTask?.error || null, [askingTask?.error]);
   const question = useMemo(() => inputValue.trim(), [inputValue]);
   const isProcessing = useMemo(
     () => getIsProcessing(askProcessState.currentState),
@@ -85,11 +119,11 @@ export default forwardRef<Attributes, Props>(function Prompt(props, ref) {
   }, [isProcessing]);
 
   useEffect(() => {
-    if (data) {
-      const processState = convertAskingTaskToProcessState(data);
+    if (askingTask) {
+      const processState = convertAskingTaskToProcessState(askingTask);
       askProcessState.setState(processState);
     }
-  }, [data]);
+  }, [askingTask]);
 
   useEffect(() => {
     if (error) {
@@ -97,6 +131,13 @@ export default forwardRef<Attributes, Props>(function Prompt(props, ref) {
     }
   }, [error]);
 
+  const selectQuestion = async (value: string) => {
+    setInputValue(value);
+    onStopStreaming && onStopStreaming();
+    askProcessState.resetState();
+    onSubmit && (await onSubmit(value));
+  };
+
   const selectResult = async (payload) => {
     const isSavedViewCandidate = !!payload.viewId;
 
@@ -121,6 +162,8 @@ export default forwardRef<Attributes, Props>(function Prompt(props, ref) {
   const closeResult = () => {
     askProcessState.resetState();
     setInputValue('');
+    onStopStreaming && onStopStreaming();
+    onStopRecommend && onStopRecommend();
   };
 
   const stopProcess = () => {
@@ -178,10 +221,12 @@ export default forwardRef<Attributes, Props>(function Prompt(props, ref) {
       </PromptButton>
 
       <PromptResult
-        data={candidates}
+        data={result}
         error={error}
+        loading={loading}
         processState={askProcessState}
-        onSelect={selectResult}
+        onSelectQuestion={selectQuestion}
+        onSelectResult={selectResult}
         onClose={closeResult}
         onStop={stopProcess}
       />

wren-ui/src/hooks/useAskPrompt.tsx

@@ -1,10 +1,27 @@
-import { useEffect, useMemo } from 'react';
-import { AskingTaskStatus } from '@/apollo/client/graphql/__types__';
+import { useCallback, useEffect, useMemo, useState } from 'react';
+import { uniq } from 'lodash';
+import {
+  AskingTask,
+  AskingTaskStatus,
+  AskingTaskType,
+  RecommendedQuestionsTask,
+  RecommendedQuestionsTaskStatus,
+} from '@/apollo/client/graphql/__types__';
 import {
   useAskingTaskLazyQuery,
   useCancelAskingTaskMutation,
   useCreateAskingTaskMutation,
+  useCreateInstantRecommendedQuestionsMutation,
+  useInstantRecommendedQuestionsLazyQuery,
 } from '@/apollo/client/graphql/home.generated';
+import useAskingStreamTask from './useAskingStreamTask';
+
+export interface AskPromptData {
+  originalQuestion: string;
+  askingTask?: AskingTask;
+  askingStreamTask?: string;
+  recommendedQuestions?: RecommendedQuestionsTask;
+}
 
 export const getIsFinished = (status: AskingTaskStatus) =>
   [
@@ -13,22 +30,96 @@ export const getIsFinished = (status: AskingTaskStatus) =>
     AskingTaskStatus.STOPPED,
   ].includes(status);
 
+const isRecommendedFinished = (status: RecommendedQuestionsTaskStatus) =>
+  [
+    RecommendedQuestionsTaskStatus.FINISHED,
+    RecommendedQuestionsTaskStatus.FAILED,
+  ].includes(status);
+
+const isNeedRecommendedQuestions = (askingTask: AskingTask) => {
+  return (
+    [AskingTaskType.GENERAL, AskingTaskType.MISLEADING_QUERY].includes(
+      askingTask?.type,
+    ) || askingTask?.status === AskingTaskStatus.FAILED
+  );
+};
+
 export default function useAskPrompt(threadId?: number) {
+  const [originalQuestion, setOriginalQuestion] = useState<string>('');
+  const [threadQuestions, setThreadQuestions] = useState<string[]>([]);
   const [createAskingTask, createAskingTaskResult] =
     useCreateAskingTaskMutation();
   const [cancelAskingTask] = useCancelAskingTaskMutation();
   const [fetchAskingTask, askingTaskResult] = useAskingTaskLazyQuery({
     pollInterval: 1000,
   });
-  const data = useMemo(
+  const [fetchAskingStreamTask, askingStreamTaskResult] = useAskingStreamTask();
+  const [createInstantRecommendedQuestions] =
+    useCreateInstantRecommendedQuestionsMutation();
+  const [fetchInstantRecommendedQuestions, instantRecommendedQuestionsResult] =
+    useInstantRecommendedQuestionsLazyQuery({
+      pollInterval: 1000,
+    });
+
+  const askingTask = useMemo(
     () => askingTaskResult.data?.askingTask || null,
     [askingTaskResult.data],
   );
-  const isFinished = useMemo(() => getIsFinished(data?.status), [data]);
+  const askingTaskType = useMemo(() => askingTask?.type, [askingTask?.type]);
+  const askingStreamTask = askingStreamTaskResult.data;
+  const recommendedQuestions = useMemo(
+    () =>
+      instantRecommendedQuestionsResult.data?.instantRecommendedQuestions ||
+      null,
+    [instantRecommendedQuestionsResult.data],
+  );
+
+  const loading = askingStreamTaskResult.loading;
+
+  const data = useMemo(
+    () => ({
+      originalQuestion,
+      askingTask,
+      askingStreamTask,
+      recommendedQuestions,
+    }),
+    [originalQuestion, askingTask, askingStreamTask, recommendedQuestions],
+  );
+
+  const startRecommendedQuestions = useCallback(async () => {
+    const previousQuestions = [
+      // slice the last 5 questions in threadQuestions
+      ...uniq(threadQuestions).slice(-5),
+      originalQuestion,
+    ];
+    const response = await createInstantRecommendedQuestions({
+      variables: { data: { previousQuestions } },
+    });
+    fetchInstantRecommendedQuestions({
+      variables: { taskId: response.data.createInstantRecommendedQuestions.id },
+    });
+  }, [originalQuestion]);
 
   useEffect(() => {
+    const isFinished = getIsFinished(askingTask?.status);
     if (isFinished) askingTaskResult.stopPolling();
-  }, [isFinished]);
+
+    if (isNeedRecommendedQuestions(askingTask)) {
+      startRecommendedQuestions();
+    }
+  }, [askingTask]);
+
+  useEffect(() => {
+    if (isRecommendedFinished(recommendedQuestions?.status))
+      instantRecommendedQuestionsResult.stopPolling();
+  }, [recommendedQuestions]);
+
+  useEffect(() => {
+    const taskId = createAskingTaskResult.data?.createAskingTask.id;
+    if (taskId && askingTaskType === AskingTaskType.GENERAL) {
+      fetchAskingStreamTask(taskId);
+    }
+  }, [askingTaskType, createAskingTaskResult.data]);
 
   const onStop = () => {
     const taskId = createAskingTaskResult.data?.createAskingTask.id;
@@ -40,6 +131,7 @@ export default function useAskPrompt(threadId?: number) {
   };
 
   const onSubmit = async (value) => {
+    setOriginalQuestion(value);
     try {
       const response = await createAskingTask({
         variables: { data: { question: value, threadId } },
@@ -52,12 +144,22 @@ export default function useAskPrompt(threadId?: number) {
     }
   };
 
-  const stopPolling = () => askingTaskResult.stopPolling();
+  const onStopPolling = () => askingTaskResult.stopPolling();
+
+  const onStopStreaming = () => askingStreamTaskResult.reset();
+
+  const onStopRecommend = () => instantRecommendedQuestionsResult.stopPolling();
+
+  const onStoreThreadQuestions = (questions: string[]) => setThreadQuestions(questions);
 
   return {
     data,
+    loading,
     onStop,
     onSubmit,
-    stopPolling,
+    onStopPolling,
+    onStopStreaming,
+    onStopRecommend,
+    onStoreThreadQuestions,
   };
 }

wren-ui/src/hooks/useAskingStreamTask.tsx

@@ -0,0 +1,52 @@
+import { useRef, useState } from 'react';
+
+type useAskingStreamTaskReturn = [
+  (queryId: string) => void,
+  { data: string; loading: boolean; reset: () => void },
+];
+
+export default function useAskingStreamTask() {
+  const eventSourceRef = useRef<EventSource | null>(null);
+  const [loading, setLoading] = useState<boolean>(false);
+  const [data, setData] = useState<string>('');
+
+  const reset = () => {
+    if (eventSourceRef.current) {
+      eventSourceRef.current?.close();
+      eventSourceRef.current = null;
+    }
+    setData('');
+  };
+
+  const fetchAskingStreamingTask = (queryId: string) => {
+    setLoading(true);
+    reset();
+
+    const eventSource = new EventSource(
+      `/api/ask_task/streaming?queryId=${queryId}`,
+    );
+
+    eventSource.onmessage = (event) => {
+      const eventData = JSON.parse(event.data);
+      if (eventData.done) {
+        eventSource.close();
+        setLoading(false);
+      } else {
+        setData((state) => state + (eventData?.message || ''));
+      }
+    };
+
+    eventSource.onerror = (error) => {
+      console.error(error);
+      eventSource.close();
+      setLoading(false);
+    };
+
+    eventSourceRef.current = eventSource;
+  };
+
+  return [
+    fetchAskingStreamingTask,
+    { data, loading, reset },
+  ] as useAskingStreamTaskReturn;
+}

wren-ui/src/pages/api/ask_task/streaming.ts

@@ -0,0 +1,37 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { components } from '@/common';
+
+const { wrenAIAdaptor } = components;
+
+export default async function handler(
+  req: NextApiRequest,
+  res: NextApiResponse,
+) {
+  res.setHeader('Content-Type', 'text/event-stream');
+  res.setHeader('Cache-Control', 'no-cache, no-transform');
+  res.setHeader('Connection', 'keep-alive');
+  res.flushHeaders();
+
+  const { queryId } = req.query;
+  try {
+    const stream = await wrenAIAdaptor.getAskStreamingResult(queryId as string);
+
+    stream.on('data', (chunk) => {
+      // pass the chunk directly to the client
+      res.write(chunk);
+    });
+
+    stream.on('end', () => {
+      res.write(`data: ${JSON.stringify({ done: true })}\n\n`);
+      res.end();
+    });
+
+    // destroy the stream if the client closes the connection
+    req.on('close', () => {
+      stream.destroy();
+    });
+  } catch (error) {
+    console.error(error);
+    res.status(500).end();
+  }
+}

wren-ui/src/pages/home/[id].tsx

@@ -79,7 +79,7 @@ export default function HomeThread() {
 
   // stop all requests when change thread
   useEffect(() => {
-    askPrompt.stopPolling();
+    askPrompt.onStopPolling();
     threadResponseResult.stopPolling();
     $prompt.current?.close();
   }, [threadId]);
@@ -92,6 +92,10 @@ export default function HomeThread() {
     if (unfinishedRespose) {
       fetchThreadResponse({ variables: { responseId: unfinishedRespose.id } });
     }
+
+    // store thread questions for instant recommended questions
+    const questions = thread?.responses.flatMap((res) => res.question || [])
+    if(questions) askPrompt.onStoreThreadQuestions(questions);
   }, [thread]);
 
   useEffect(() => {
@@ -100,7 +104,7 @@ export default function HomeThread() {
 
   const onSelect = async (payload) => {
     try {
-      askPrompt.stopPolling();
+      askPrompt.onStopPolling();
       const response = await createThreadResponse({
         variables: { threadId: thread.id, data: payload },
       });
@@ -119,13 +123,7 @@ export default function HomeThread() {
         onOpenSaveAsViewModal={saveAsViewModal.openModal}
       />
       <div className="py-12" />
-      <Prompt
-        ref={$prompt}
-        data={askPrompt.data}
-        onSubmit={askPrompt.onSubmit}
-        onStop={askPrompt.onStop}
-        onSelect={onSelect}
-      />
+      <Prompt ref={$prompt} {...askPrompt} onSelect={onSelect} />
       <SaveAsViewModal
         {...saveAsViewModal.state}
         loading={creating}

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

@@ -41,7 +41,7 @@ export default function Home() {
 
   const onSelect = async (payload) => {
     try {
-      askPrompt.stopPolling();
+      askPrompt.onStopPolling();
       const response = await createThread({ variables: { data: payload } });
       router.push(Path.Home + `/${response.data.createThread.id}`);
     } catch (error) {
@@ -64,13 +64,7 @@ export default function Home() {
           <DemoPrompt demo={sampleQuestions} onSelect={onDemoSelect} />
         )}
       </div>
-      <Prompt
-        ref={$prompt}
-        data={askPrompt.data}
-        onSubmit={askPrompt.onSubmit}
-        onStop={askPrompt.onStop}
-        onSelect={onSelect}
-      />
+      <Prompt ref={$prompt} {...askPrompt} onSelect={onSelect} />
     </SiderLayout>
   );
 }