feat(wren-ui): add reasoning api

wren-ui/migrations/20250509000000_create_asking_task.js

@@ -0,0 +1,34 @@
+/**
+ * @param { import("knex").Knex } knex
+ * @returns { Promise<void> }
+ */
+exports.up = function (knex) {
+  return knex.schema.createTable('asking_task', (table) => {
+    table.increments('id').primary();
+    table.string('query_id').notNullable().unique();
+    table.text('question');
+    table.jsonb('detail').defaultTo('{}');
+
+    table
+      .integer('thread_id')
+      .references('id')
+      .inTable('thread')
+      .onDelete('CASCADE');
+
+    table
+      .integer('thread_response_id')
+      .references('id')
+      .inTable('thread_response')
+      .onDelete('CASCADE');
+
+    table.timestamps(true, true);
+  });
+};
+
+/**
+ * @param { import("knex").Knex } knex
+ * @returns { Promise<void> }
+ */
+exports.down = function (knex) {
+  return knex.schema.dropTable('asking_task');
+};

wren-ui/migrations/20250509000001_add_task_id_to_thread.js

@@ -0,0 +1,25 @@
+/**
+ * @param { import("knex").Knex } knex
+ * @returns { Promise<void> }
+ */
+exports.up = function (knex) {
+  return knex.schema.alterTable('thread_response', (table) => {
+    table
+      .integer('asking_task_id')
+      .nullable()
+      .references('id')
+      .inTable('asking_task')
+      .onDelete('SET NULL');
+  });
+};
+
+/**
+ * @param { import("knex").Knex } knex
+ * @returns { Promise<void> }
+ */
+exports.down = function (knex) {
+  return knex.schema.alterTable('thread_response', (table) => {
+    table.dropForeign('asking_task_id');
+    table.dropColumn('asking_task_id');
+  });
+};

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

@@ -632,7 +632,6 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
   }
 
   private transformAskResult(body: any): AskResult {
-    const { type, intent_reasoning } = body;
     const { status, error } = this.transformStatusAndError(body);
     const candidates = (body?.response || []).map((candidate: any) => ({
       type: candidate?.type?.toUpperCase() as AskCandidateType,
@@ -641,11 +640,16 @@ export class WrenAIAdaptor implements IWrenAIAdaptor {
     }));
 
     return {
-      type,
+      type: body?.type,
       status: status as AskResultStatus,
       error,
       response: candidates,
-      intentReasoning: intent_reasoning,
+      rephrasedQuestion: body?.rephrased_question,
+      intentReasoning: body?.intent_reasoning,
+      sqlGenerationReasoning: body?.sql_generation_reasoning,
+      retrievedTables: body?.retrieved_tables,
+      invalidSql: body?.invalid_sql,
+      traceId: body?.trace_id,
     };
   }
 

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

@@ -122,7 +122,12 @@ export type AskResult = AskResponse<
   }>,
   AskResultStatus
 > & {
+  rephrasedQuestion?: string;
   intentReasoning?: string;
+  sqlGenerationReasoning?: string;
+  retrievedTables?: string[];
+  invalidSql?: string;
+  traceId?: string;
 };
 
 export enum RecommendationQuestionStatus {

wren-ui/src/apollo/server/repositories/askingTaskRepository.ts

@@ -0,0 +1,71 @@
+import { Knex } from 'knex';
+import { BaseRepository, IBasicRepository } from './baseRepository';
+import {
+  camelCase,
+  isPlainObject,
+  mapKeys,
+  mapValues,
+  snakeCase,
+} from 'lodash';
+import { AskResult } from '../models/adaptor';
+
+export interface AskingTask {
+  id: number;
+  queryId: string;
+  question?: string;
+  detail?: AskResult;
+  threadId?: number;
+  threadResponseId?: number;
+  createdAt: Date;
+  updatedAt: Date;
+}
+
+export interface IAskingTaskRepository extends IBasicRepository<AskingTask> {
+  findByQueryId(queryId: string): Promise<AskingTask | null>;
+}
+
+export class AskingTaskRepository
+  extends BaseRepository<AskingTask>
+  implements IAskingTaskRepository
+{
+  private readonly jsonbColumns = ['detail'];
+
+  constructor(knexPg: Knex) {
+    super({ knexPg, tableName: 'asking_task' });
+  }
+
+  public async findByQueryId(queryId: string): Promise<AskingTask | null> {
+    return this.findOneBy({ queryId });
+  }
+
+  protected override transformFromDBData = (data: any) => {
+    if (!isPlainObject(data)) {
+      throw new Error('Unexpected dbdata');
+    }
+    const camelCaseData = mapKeys(data, (_value, key) => camelCase(key));
+    const transformData = mapValues(camelCaseData, (value, key) => {
+      if (this.jsonbColumns.includes(key)) {
+        if (typeof value === 'string') {
+          return value ? JSON.parse(value) : value;
+        }
+        return value;
+      }
+      return value;
+    });
+    return transformData as AskingTask;
+  };
+
+  protected override transformToDBData = (data: any) => {
+    if (!isPlainObject(data)) {
+      throw new Error('Unexpected dbdata');
+    }
+    const transformedData = mapValues(data, (value, key) => {
+      if (this.jsonbColumns.includes(key)) {
+        return JSON.stringify(value);
+      } else {
+        return value;
+      }
+    });
+    return mapKeys(transformedData, (_value, key) => snakeCase(key));
+  };
+}

wren-ui/src/apollo/server/repositories/index.ts

@@ -15,3 +15,4 @@ export * from './schemaChangeRepository';
 export * from './dashboardRepository';
 export * from './dashboardItemRepository';
 export * from './sqlPairRepository';
+export * from './askingTaskRepository';

wren-ui/src/apollo/server/repositories/threadResponseRepository.ts

@@ -41,10 +41,11 @@ export interface ThreadResponseChartDetail {
 
 export interface ThreadResponse {
   id: number; // ID
+  askingTaskId?: number; // Reference to asking_task.id
   viewId?: number; // View ID, if the response is from a view
   threadId: number; // Reference to thread.id
   question: string; // Thread response question
-  sql: string; // SQL query generated by AI service
+  sql?: string; // SQL query generated by AI service
   answerDetail?: ThreadResponseAnswerDetail; // AI generated text-based answer detail
   breakdownDetail?: ThreadResponseBreakdownDetail; // Thread response breakdown detail
   chartDetail?: ThreadResponseChartDetail; // Thread response chart detail
@@ -114,6 +115,8 @@ export class ThreadResponseRepository
     id: string | number,
     data: Partial<{
       status: AskResultStatus;
+      sql: string;
+      viewId: number;
       answerDetail: ThreadResponseAnswerDetail;
       breakdownDetail: ThreadResponseBreakdownDetail;
       chartDetail: ThreadResponseChartDetail;
@@ -122,6 +125,8 @@ export class ThreadResponseRepository
   ) {
     const transformedData = {
       status: data.status ? data.status : undefined,
+      sql: data.sql ? data.sql : undefined,
+      viewId: data.viewId ? data.viewId : undefined,
       answerDetail: data.answerDetail
         ? JSON.stringify(data.answerDetail)
         : undefined,

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

@@ -91,6 +91,7 @@ const resolvers = {
     cancelAskingTask: askingResolver.cancelAskingTask,
     createInstantRecommendedQuestions:
       askingResolver.createInstantRecommendedQuestions,
+    rerunAskingTask: askingResolver.rerunAskingTask,
 
     // Thread
     createThread: askingResolver.createThread,