feat(wren-ui): API enhancement

docker/.env.example

@@ -22,7 +22,7 @@ OPENAI_API_KEY=
 # CHANGE THIS TO THE LATEST VERSION
 WREN_PRODUCT_VERSION=0.23.0
 WREN_ENGINE_VERSION=0.16.1
-WREN_AI_SERVICE_VERSION=0.23.2
+WREN_AI_SERVICE_VERSION=custom-instruction-v2
 IBIS_SERVER_VERSION=0.16.1
 WREN_UI_VERSION=0.28.0
 WREN_BOOTSTRAP_VERSION=0.1.5

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

@@ -77,9 +77,22 @@ export type ApiHistoryResponse = {
 };
 
 export enum ApiType {
+  ASK = 'ASK',
+  ASYNC_ASK = 'ASYNC_ASK',
+  ASYNC_GENERATE_SQL = 'ASYNC_GENERATE_SQL',
+  CREATE_INSTRUCTION = 'CREATE_INSTRUCTION',
+  CREATE_SQL_PAIR = 'CREATE_SQL_PAIR',
+  DELETE_INSTRUCTION = 'DELETE_INSTRUCTION',
+  DELETE_SQL_PAIR = 'DELETE_SQL_PAIR',
   GENERATE_SQL = 'GENERATE_SQL',
+  GENERATE_SUMMARY = 'GENERATE_SUMMARY',
   GENERATE_VEGA_CHART = 'GENERATE_VEGA_CHART',
-  RUN_SQL = 'RUN_SQL'
+  GET_INSTRUCTIONS = 'GET_INSTRUCTIONS',
+  GET_MODELS = 'GET_MODELS',
+  GET_SQL_PAIRS = 'GET_SQL_PAIRS',
+  RUN_SQL = 'RUN_SQL',
+  UPDATE_INSTRUCTION = 'UPDATE_INSTRUCTION',
+  UPDATE_SQL_PAIR = 'UPDATE_SQL_PAIR'
 }
 
 export type AskingTask = {
@@ -294,9 +307,9 @@ export type DashboardSchedule = {
   __typename?: 'DashboardSchedule';
   cron?: Maybe<Scalars['String']>;
   day?: Maybe<CacheScheduleDayEnum>;
-  frequency: ScheduleFrequencyEnum;
-  hour: Scalars['Int'];
-  minute: Scalars['Int'];
+  frequency?: Maybe<ScheduleFrequencyEnum>;
+  hour?: Maybe<Scalars['Int']>;
+  minute?: Maybe<Scalars['Int']>;
   timezone?: Maybe<Scalars['String']>;
 };
 
@@ -317,8 +330,8 @@ export enum DataSourceName {
   CLICK_HOUSE = 'CLICK_HOUSE',
   DUCKDB = 'DUCKDB',
   MSSQL = 'MSSQL',
-  ORACLE = 'ORACLE',
   MYSQL = 'MYSQL',
+  ORACLE = 'ORACLE',
   POSTGRES = 'POSTGRES',
   SNOWFLAKE = 'SNOWFLAKE',
   TRINO = 'TRINO'
@@ -384,7 +397,7 @@ export type DetailedDashboard = {
   items: Array<DashboardItem>;
   name: Scalars['String'];
   nextScheduledAt?: Maybe<Scalars['String']>;
-  schedule: DashboardSchedule;
+  schedule?: Maybe<DashboardSchedule>;
 };
 
 export type DetailedModel = {

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

@@ -56,7 +56,7 @@ export type SetDashboardScheduleMutation = { __typename?: 'Mutation', setDashboa
 export type DashboardQueryVariables = Types.Exact<{ [key: string]: never; }>;
 
 
-export type DashboardQuery = { __typename?: 'Query', dashboard: { __typename?: 'DetailedDashboard', id: number, name: string, description?: string | null, cacheEnabled: boolean, nextScheduledAt?: string | null, schedule: { __typename?: 'DashboardSchedule', frequency: Types.ScheduleFrequencyEnum, hour: number, minute: number, day?: Types.CacheScheduleDayEnum | null, timezone?: string | null, cron?: string | null }, items: Array<{ __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> } };
+export type DashboardQuery = { __typename?: 'Query', dashboard: { __typename?: 'DetailedDashboard', id: number, name: string, description?: string | null, cacheEnabled: boolean, nextScheduledAt?: string | null, schedule?: { __typename?: 'DashboardSchedule', frequency?: Types.ScheduleFrequencyEnum | null, hour?: number | null, minute?: number | null, day?: Types.CacheScheduleDayEnum | null, timezone?: string | null, cron?: string | null } | null, items: Array<{ __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> } };
 
 export const CommonDashboardItemFragmentDoc = gql`
     fragment CommonDashboardItem on DashboardItem {

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

@@ -1,11 +1,30 @@
-import { camelCase, isPlainObject, mapKeys, mapValues } from 'lodash';
+import {
+  camelCase,
+  isPlainObject,
+  mapKeys,
+  mapValues,
+  snakeCase,
+} from 'lodash';
 import { BaseRepository, IBasicRepository } from './baseRepository';
 import { Knex } from 'knex';
 
 export enum ApiType {
   GENERATE_SQL = 'GENERATE_SQL',
   RUN_SQL = 'RUN_SQL',
   GENERATE_VEGA_CHART = 'GENERATE_VEGA_CHART',
+  GENERATE_SUMMARY = 'GENERATE_SUMMARY',
+  ASK = 'ASK',
+  GET_INSTRUCTIONS = 'GET_INSTRUCTIONS',
+  CREATE_INSTRUCTION = 'CREATE_INSTRUCTION',
+  UPDATE_INSTRUCTION = 'UPDATE_INSTRUCTION',
+  DELETE_INSTRUCTION = 'DELETE_INSTRUCTION',
+  GET_SQL_PAIRS = 'GET_SQL_PAIRS',
+  CREATE_SQL_PAIR = 'CREATE_SQL_PAIR',
+  UPDATE_SQL_PAIR = 'UPDATE_SQL_PAIR',
+  DELETE_SQL_PAIR = 'DELETE_SQL_PAIR',
+  GET_MODELS = 'GET_MODELS',
+  ASYNC_ASK = 'ASYNC_ASK',
+  ASYNC_GENERATE_SQL = 'ASYNC_GENERATE_SQL',
 }
 
 export interface ApiHistory {
@@ -147,6 +166,20 @@ export class ApiHistoryRepository
     return formattedData;
   };
 
+  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));
+  };
+
   /**
    * Convert camelCase to snake_case for DB column names
    */

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

@@ -8,6 +8,19 @@ export const typeDefs = gql`
     GENERATE_SQL
     RUN_SQL
     GENERATE_VEGA_CHART
+    GENERATE_SUMMARY
+    ASK
+    GET_INSTRUCTIONS
+    CREATE_INSTRUCTION
+    UPDATE_INSTRUCTION
+    DELETE_INSTRUCTION
+    CREATE_SQL_PAIR
+    UPDATE_SQL_PAIR
+    DELETE_SQL_PAIR
+    GET_SQL_PAIRS
+    GET_MODELS
+    ASYNC_ASK
+    ASYNC_GENERATE_SQL
   }
 
   input ApiHistoryFilterInput {

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

@@ -11,6 +11,7 @@ import * as Errors from '@server/utils/error';
 import { GeneralErrorCodes } from '@server/utils/error';
 export interface IInstructionService {
   getInstructions(projectId: number): Promise<Instruction[]>;
+  getInstruction(id: number): Promise<Instruction>;
   createInstruction(instruction: InstructionInput): Promise<Instruction>;
   createInstructions(instructions: InstructionInput[]): Promise<Instruction[]>;
   updateInstruction(instruction: UpdateInstructionInput): Promise<Instruction>;
@@ -30,10 +31,15 @@ export class InstructionService implements IInstructionService {
     this.instructionRepository = instructionRepository;
     this.wrenAIAdaptor = wrenAIAdaptor;
   }
+
   public async getInstructions(projectId: number): Promise<Instruction[]> {
     return this.instructionRepository.findAllBy({ projectId });
   }
 
+  public async getInstruction(id: number): Promise<Instruction> {
+    return this.instructionRepository.findOneBy({ id });
+  }
+
   public async createInstruction(
     input: InstructionInput,
   ): Promise<Instruction> {

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

@@ -1,11 +1,143 @@
 import { NextApiResponse } from 'next';
 import { v4 as uuidv4 } from 'uuid';
-import { ApiType } from '@server/repositories/apiHistoryRepository';
+import { ApiType, ApiHistory } from '@server/repositories/apiHistoryRepository';
 import * as Errors from '@server/utils/error';
 import { components } from '@/common';
+import {
+  AskResult,
+  AskResultStatus,
+  AskResultType,
+  WrenAIError,
+  TextBasedAnswerResult,
+  TextBasedAnswerStatus,
+} from '@/apollo/server/models/adaptor';
 
 const { apiHistoryRepository } = components;
 
+export const MAX_WAIT_TIME = 1000 * 60 * 3; // 3 minutes
+
+export const isAskResultFinished = (result: AskResult) => {
+  return (
+    result.status === AskResultStatus.FINISHED ||
+    result.status === AskResultStatus.FAILED ||
+    result.status === AskResultStatus.STOPPED ||
+    result.error
+  );
+};
+
+/**
+ * Validates the AI result and throws appropriate errors for different failure cases
+ * @param result The AI result to validate
+ * @param taskQueryId The query ID of the task (used for explanation queries)
+ * @throws ApiError if result contains errors or is of an invalid type
+ */
+export const validateAskResult = (
+  result: AskResult,
+  taskQueryId: string,
+): void => {
+  // Check for error in result
+  if (result.error) {
+    const errorMessage =
+      (result.error as WrenAIError).message || 'Unknown error';
+    const additionalData: Record<string, any> = {};
+
+    // Include invalid SQL if available
+    if (result.invalidSql) {
+      additionalData.invalidSql = result.invalidSql;
+    }
+
+    throw new ApiError(errorMessage, 400, result.error.code, additionalData);
+  }
+
+  // Check for misleading query type
+  if (result.type === AskResultType.MISLEADING_QUERY) {
+    throw new ApiError(
+      result.intentReasoning ||
+        Errors.errorMessages[Errors.GeneralErrorCodes.NON_SQL_QUERY],
+      400,
+      Errors.GeneralErrorCodes.NON_SQL_QUERY,
+    );
+  }
+
+  // Check for general type response
+  if (result.type === AskResultType.GENERAL) {
+    throw new ApiError(
+      result.intentReasoning ||
+        Errors.errorMessages[Errors.GeneralErrorCodes.NON_SQL_QUERY],
+      400,
+      Errors.GeneralErrorCodes.NON_SQL_QUERY,
+      { explanationQueryId: taskQueryId },
+    );
+  }
+};
+
+/**
+ * Validates the summary generation result and checks for errors
+ * @param result The summary result to validate
+ * @throws ApiError if the result has errors or is in a failed state
+ */
+export const validateSummaryResult = (result: TextBasedAnswerResult): void => {
+  // Check for errors or failed status
+  if (result.status === TextBasedAnswerStatus.FAILED || result.error) {
+    throw new ApiError(
+      result.error?.message || 'Failed to generate summary',
+      400,
+      Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR,
+    );
+  }
+
+  // Verify that the status is succeeded
+  if (result.status !== TextBasedAnswerStatus.SUCCEEDED) {
+    throw new ApiError('Summary generation is still in progress', 500);
+  }
+};
+
+export const transformHistoryInput = (histories: ApiHistory[]) => {
+  if (!histories) {
+    return [];
+  }
+  return histories
+    .filter(
+      (history) =>
+        history.responsePayload?.sql && history.requestPayload?.question,
+    )
+    .map((history) => ({
+      question: history.requestPayload?.question,
+      sql: history.responsePayload?.sql,
+    }));
+};
+
+/**
+ * Validates SQL syntax and compatibility with the project's manifest.
+ * Throws an ApiError if the SQL is invalid or cannot be previewed.
+ * @param sql The SQL string to validate
+ * @param project The project object (must have id)
+ * @param deployService The deployment service instance
+ * @param queryService The query service instance
+ */
+export const validateSql = async (
+  sql: string,
+  project: any,
+  deployService: any,
+  queryService: any,
+) => {
+  const lastDeployment = await deployService.getLastDeployment(project.id);
+  const manifest = lastDeployment.manifest;
+  try {
+    await queryService.preview(sql, {
+      manifest,
+      project,
+      dryRun: true,
+    });
+  } catch (err: any) {
+    throw new ApiError(
+      err.message || 'Invalid SQL',
+      400,
+      Errors.GeneralErrorCodes.INVALID_SQL_ERROR,
+    );
+  }
+};
+
 /**
  * Common error class for API endpoints
  */
@@ -71,6 +203,45 @@ export const respondWith = async ({
   });
 };
 
+/**
+ * Simple response handler for API endpoints that don't need responseId or threadId
+ * Used for simple CRUD operations like instructions
+ */
+export const respondWithSimple = async ({
+  res,
+  statusCode,
+  responsePayload,
+  projectId,
+  apiType,
+  headers,
+  requestPayload,
+  startTime,
+}: {
+  res: NextApiResponse;
+  statusCode: number;
+  responsePayload: any;
+  projectId: number;
+  apiType: ApiType;
+  startTime: number;
+  requestPayload?: Record<string, any>;
+  headers?: Record<string, string>;
+}) => {
+  const durationMs = startTime ? Date.now() - startTime : undefined;
+  const responseId = uuidv4();
+  await apiHistoryRepository.createOne({
+    id: responseId,
+    projectId,
+    apiType,
+    headers,
+    requestPayload,
+    responsePayload,
+    statusCode,
+    durationMs,
+  });
+
+  return res.status(statusCode).json(responsePayload);
+};
+
 /**
  * Common error handler for API endpoints
  */

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

@@ -57,6 +57,9 @@ export enum GeneralErrorCodes {
   // vega schema error
   FAILED_TO_GENERATE_VEGA_SCHEMA = 'FAILED_TO_GENERATE_VEGA_SCHEMA',
   POLLING_TIMEOUT = 'POLLING_TIMEOUT',
+
+  // sql execution error
+  SQL_EXECUTION_ERROR = 'SQL_EXECUTION_ERROR',
 }
 
 export const errorMessages = {
@@ -122,6 +125,9 @@ export const errorMessages = {
   [GeneralErrorCodes.FAILED_TO_GENERATE_VEGA_SCHEMA]:
     'Failed to generate Vega spec',
   [GeneralErrorCodes.POLLING_TIMEOUT]: 'Polling timeout',
+
+  // sql execution error
+  [GeneralErrorCodes.SQL_EXECUTION_ERROR]: 'SQL execution error',
 };
 
 export const shortMessages = {
@@ -154,6 +160,7 @@ export const shortMessages = {
   [GeneralErrorCodes.FAILED_TO_GENERATE_VEGA_SCHEMA]:
     'Failed to generate Vega spec',
   [GeneralErrorCodes.POLLING_TIMEOUT]: 'Polling timeout',
+  [GeneralErrorCodes.SQL_EXECUTION_ERROR]: 'SQL execution error',
 };
 
 export const create = (

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

@@ -6,3 +6,5 @@ export * from './docker';
 export * from './model';
 export * from './helper';
 export * from './regex';
+export * from './sseTypes';
+export * from './sseUtils';