Chore: add migration files and update adaptors for feedback loop feature (#506)

* chore: Add migration for creating thread_response_explain table

* chore: Add analysisSql method to WrenEngineAdaptor

* adding explain api to ai service adapter:

* fix migration errer

* chore(wren-ui): Add regenerations API to wrenAIAdaptor

---------

Co-authored-by: andreashimin <[email protected]>

Author: onlyjackfrost

wren-ui/migrations/20240711021133_create_thread_response_explain_table.js

@@ -0,0 +1,32 @@
+/**
+ * @param { import("knex").Knex } knex
+ * @returns { Promise<void> }
+ */
+exports.up = function (knex) {
+  return knex.schema.createTable('thread_response_explain', (table) => {
+    table.increments('id').comment('ID');
+    table
+      .integer('thread_response_id')
+      .comment('Reference to thread_response.id');
+    table
+      .foreign('thread_response_id')
+      .references('thread_response.id')
+      .onDelete('CASCADE');
+
+    table.string('query_id').notNullable();
+    table.string('status').notNullable();
+    table.jsonb('detail').notNullable();
+    table.jsonb('error').notNullable();
+
+    // timestamps
+    table.timestamps(true, true);
+  });
+};
+
+/**
+ * @param { import("knex").Knex } knex
+ * @returns { Promise<void> }
+ */
+exports.down = function (knex) {
+  return knex.schema.dropTable('thread_response_explain');
+};

wren-ui/migrations/20240711082655_update_thread_response_table.js

@@ -0,0 +1,22 @@
+/**
+ * @param { import("knex").Knex } knex
+ * @returns { Promise<void> }
+ */
+exports.up = function (knex) {
+  return knex.schema.alterTable('thread_response', (table) => {
+    table
+      .jsonb('corrections')
+      .nullable()
+      .comment('the corrections of the previous thread response'); // [{type, id, correct}, ...]
+  });
+};
+
+/**
+ * @param { import("knex").Knex } knex
+ * @returns { Promise<void> }
+ */
+exports.down = function (knex) {
+  return knex.schema.alterTable('thread_response', (table) => {
+    table.dropColumn('corrections');
+  });
+};

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

@@ -55,6 +55,15 @@ export interface AsyncQueryResponse {
   queryId: string;
 }
 
+export type ExplainResult = AIServiceResponse<any, ExplainPipelineStatus>;
+
+export enum ExplainPipelineStatus {
+  UNDERSTANDING = 'UNDERSTANDING',
+  GENERATING = 'GENERATING',
+  FINISHED = 'FINISHED',
+  FAILED = 'FAILED',
+}
+
 export enum AskResultStatus {
   UNDERSTANDING = 'UNDERSTANDING',
   SEARCHING = 'SEARCHING',
@@ -71,7 +80,21 @@ export enum AskCandidateType {
   LLM = 'LLM',
 }
 
-export interface AskResponse<R, S> {
+export enum ExplainType {
+  FILTER = 'filter',
+  SELECT_ITEMS = 'selectItems',
+  RELATION = 'relation',
+  GROUP_BY_KEYS = 'groupByKeys',
+  SORTINGS = 'sortings',
+}
+
+// UI currently only support nl_expression
+export enum ExpressionType {
+  SQL_EXPRESSION = 'sql_expression',
+  NL_EXPRESSION = 'nl_expression',
+}
+
+export interface AIServiceResponse<R, S> {
   status: S;
   response: R | null;
   error: WrenAIError | null;
@@ -83,15 +106,15 @@ export interface AskDetailInput {
   summary: string;
 }
 
-export type AskDetailResult = AskResponse<
+export type AskDetailResult = AIServiceResponse<
   {
     description: string;
     steps: AskStep[];
   },
   AskResultStatus
 >;
 
-export type AskResult = AskResponse<
+export type AskResult = AIServiceResponse<
   Array<{
     type: AskCandidateType;
     sql: string;
@@ -101,7 +124,26 @@ export type AskResult = AskResponse<
   AskResultStatus
 >;
 
-const getAISerciceError = (error: any) => {
+export interface CorrectionObject<T> {
+  type: T;
+  value: string;
+}
+
+export interface AskCorrection {
+  before: CorrectionObject<ExplainType>;
+  after: CorrectionObject<ExpressionType>;
+}
+
+export interface AskStepWithCorrections extends AskStep {
+  corrections: AskCorrection[];
+}
+
+export interface RegenerateAskDetailInput {
+  description: string;
+  steps: AskStepWithCorrections[];
+}
+
+const getAIServiceError = (error: any) => {
   const { data } = error.response || {};
   return data?.detail
     ? `${error.message}, detail: ${data.detail}`
@@ -129,6 +171,12 @@ export interface IWrenAIAdaptor {
    */
   generateAskDetail(input: AskDetailInput): Promise<AsyncQueryResponse>;
   getAskDetailResult(queryId: string): Promise<AskDetailResult>;
+  explain(question: string, analysisResults: any): Promise<AsyncQueryResponse>;
+  getExplainResult(queryId: string): Promise<ExplainResult>;
+  regenerateAskDetail(
+    input: RegenerateAskDetailInput,
+  ): Promise<AsyncQueryResponse>;
+  getRegeneratedAskDetailResult(queryId: string): Promise<AskDetailResult>;
 }
 
 export class WrenAIAdaptor implements IWrenAIAdaptor {
@@ -148,11 +196,11 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
       const res = await axios.post(`${this.wrenAIBaseEndpoint}/v1/asks`, {
         query: input.query,
         id: input.deployId,
-        history: this.transfromHistoryInput(input.history),
+        history: this.transformHistoryInput(input.history),
       });
       return { queryId: res.data.query_id };
     } catch (err: any) {
-      logger.debug(`Got error when asking wren AI: ${getAISerciceError(err)}`);
+      logger.debug(`Got error when asking wren AI: ${getAIServiceError(err)}`);
       throw err;
     }
   }
@@ -164,7 +212,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
         status: 'stopped',
       });
     } catch (err: any) {
-      logger.debug(`Got error when canceling ask: ${getAISerciceError(err)}`);
+      logger.debug(`Got error when canceling ask: ${getAIServiceError(err)}`);
       throw err;
     }
   }
@@ -178,7 +226,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
       return this.transformAskResult(res.data);
     } catch (err: any) {
       logger.debug(
-        `Got error when getting ask result: ${getAISerciceError(err)}`,
+        `Got error when getting ask result: ${getAIServiceError(err)}`,
       );
       // throw err;
       throw Errors.create(Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, {
@@ -187,6 +235,44 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
     }
   }
 
+  public async explain(
+    question: string,
+    analysisResults: any,
+  ): Promise<AsyncQueryResponse> {
+    try {
+      const res = await axios.post(
+        `${this.wrenAIBaseEndpoint}/v1/sql-explanations`,
+        {
+          question,
+          steps_with_analysis_results: analysisResults,
+        },
+      );
+      return { queryId: res.data.query_id };
+    } catch (err: any) {
+      logger.debug(`Got error when explaining: ${getAIServiceError(err)}`);
+      throw err;
+    }
+  }
+
+  public async getExplainResult(queryId: string): Promise<ExplainResult> {
+    // make GET request /v1/sql-explanations/:query_id/result to get the result
+    try {
+      const res = await axios.get(
+        `${this.wrenAIBaseEndpoint}/v1/sql-explanations/${queryId}/result`,
+      );
+      return {
+        status: res.data.status as ExplainPipelineStatus,
+        response: res.data.response,
+        error: this.transformStatusAndError(res.data).error,
+      };
+    } catch (err: any) {
+      logger.debug(
+        `Got error when getting explain result: ${getAIServiceError(err)}`,
+      );
+      throw err;
+    }
+  }
+
   /**
    * After you choose a candidate, you can request AI service to generate the detail.
    */
@@ -202,7 +288,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
       return { queryId: res.data.query_id };
     } catch (err: any) {
       logger.debug(
-        `Got error when generating ask detail: ${getAISerciceError(err)}`,
+        `Got error when generating ask detail: ${getAIServiceError(err)}`,
       );
       throw err;
     }
@@ -217,7 +303,37 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
       return this.transformAskDetailResult(res.data);
     } catch (err: any) {
       logger.debug(
-        `Got error when getting ask detail result: ${getAISerciceError(err)}`,
+        `Got error when getting ask detail result: ${getAIServiceError(err)}`,
+      );
+      throw err;
+    }
+  }
+
+  public async regenerateAskDetail(input: RegenerateAskDetailInput) {
+    try {
+      const res = await axios.post(
+        `${this.wrenAIBaseEndpoint}/v1/sql-regenerations`,
+        input,
+      );
+      return { queryId: res.data.query_id };
+    } catch (err: any) {
+      logger.debug(
+        `Got error when regenerating ask detail: ${getAIServiceError(err)}`,
+      );
+      throw err;
+    }
+  }
+
+  public async getRegeneratedAskDetailResult(queryId: string) {
+    // make GET request /v1/sql-regenerations/:query_id/result to get the result
+    try {
+      const res = await axios.get(
+        `${this.wrenAIBaseEndpoint}/v1/sql-regenerations/${queryId}/result`,
+      );
+      return this.transformAskDetailResult(res.data);
+    } catch (err: any) {
+      logger.debug(
+        `Got error when getting regenerated ask detail result: ${getAIServiceError(err)}`,
       );
       throw err;
     }
@@ -309,7 +425,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
     }));
 
     return {
-      status,
+      status: status as AskResultStatus,
       error,
       response: candidates,
     };
@@ -326,7 +442,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
     }));
 
     return {
-      status,
+      status: status as AskResultStatus,
       error,
       response: {
         description: body?.response?.description,
@@ -336,17 +452,19 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
   }
 
   private transformStatusAndError(body: any): {
-    status: AskResultStatus;
+    status: AskResultStatus | ExplainPipelineStatus;
     error?: {
       code: Errors.GeneralErrorCodes;
       message: string;
       shortMessage: string;
     } | null;
   } {
     // transform status to enum
-    const status = AskResultStatus[
-      body?.status?.toUpperCase()
-    ] as AskResultStatus;
+    const status =
+      (AskResultStatus[body?.status?.toUpperCase()] as AskResultStatus) ||
+      (ExplainPipelineStatus[
+        body.status
+      ]?.toUpperCase() as ExplainPipelineStatus);
 
     if (!status) {
       throw new Error(`Unknown ask status: ${body?.status}`);
@@ -380,7 +498,7 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
     };
   }
 
-  private transfromHistoryInput(history: AskHistory) {
+  private transformHistoryInput(history: AskHistory) {
     if (!history) {
       return null;
     }

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

@@ -94,6 +94,9 @@ export interface IWrenEngineAdaptor {
     sql: string,
     options: WrenEngineDryRunOption,
   ): Promise<DryRunResponse[]>;
+
+  // analysis
+  analysisSql(sql: string, mdl: Manifest): Promise<any>;
 }
 
 export class WrenEngineAdaptor implements IWrenEngineAdaptor {
@@ -105,6 +108,7 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor {
   private dryPlanUrlPath = '/v1/mdl/dry-plan';
   private dryRunUrlPath = '/v1/mdl/dry-run';
   private validateUrlPath = '/v1/mdl/validate';
+  private analysisUrlPath = '/v1/analysis/sql';
 
   constructor({ wrenEngineEndpoint }: { wrenEngineEndpoint: string }) {
     this.wrenEngineBaseEndpoint = wrenEngineEndpoint;
@@ -315,16 +319,20 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor {
     }
   }
 
-  private async getDeployStatus(): Promise<WrenEngineDeployStatusResponse> {
+  public async analysisSql(sql: string, mdl: Manifest) {
     try {
-      const res = await axios.get(
-        `${this.wrenEngineBaseEndpoint}/v1/mdl/status`,
+      const url = new URL(this.analysisUrlPath, this.wrenEngineBaseEndpoint);
+      const headers = {
+        'Content-Type': 'application/json',
+      };
+      const res = await axios.post(
+        url.href,
+        { sql, manifest: mdl },
+        { headers },
       );
-      return res.data as WrenEngineDeployStatusResponse;
+      return res.data;
     } catch (err: any) {
-      logger.debug(
-        `WrenEngine: Got error when getting deploy status: ${err.message}`,
-      );
+      logger.debug(`Got error when analyzing sql: ${err.message}`);
       throw err;
     }
   }