From 969d0ede3a62b8161692b5d096063510a2034837 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Mon, 23 Jun 2025 16:10:42 +0800 Subject: [PATCH 01/16] add generate_summary api --- .../repositories/apiHistoryRepository.ts | 1 + wren-ui/src/apollo/server/schema.ts | 1 + wren-ui/src/pages/api/v1/generate_summary.ts | 220 ++++++++++++++++++ 3 files changed, 222 insertions(+) create mode 100644 wren-ui/src/pages/api/v1/generate_summary.ts diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index 9e2164bc9e..6f90908592 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -6,6 +6,7 @@ export enum ApiType { GENERATE_SQL = 'GENERATE_SQL', RUN_SQL = 'RUN_SQL', GENERATE_VEGA_CHART = 'GENERATE_VEGA_CHART', + GENERATE_SUMMARY = 'GENERATE_SUMMARY', } export interface ApiHistory { diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 5590055dbc..86b3089956 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -8,6 +8,7 @@ export const typeDefs = gql` GENERATE_SQL RUN_SQL GENERATE_VEGA_CHART + GENERATE_SUMMARY } input ApiHistoryFilterInput { diff --git a/wren-ui/src/pages/api/v1/generate_summary.ts b/wren-ui/src/pages/api/v1/generate_summary.ts new file mode 100644 index 0000000000..2b0626fbcf --- /dev/null +++ b/wren-ui/src/pages/api/v1/generate_summary.ts @@ -0,0 +1,220 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import * as Errors from '@/apollo/server/utils/error'; +import { v4 as uuidv4 } from 'uuid'; +import { + ApiError, + respondWith, + handleApiError, +} from '@/apollo/server/utils/apiUtils'; +import { + TextBasedAnswerInput, + TextBasedAnswerResult, + TextBasedAnswerStatus, + WrenAILanguage, +} from '@/apollo/server/models/adaptor'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('API_GENERATE_SUMMARY'); +logger.level = 'debug'; + +const { projectService, wrenAIAdaptor, deployService, queryService } = + components; + +const MAX_WAIT_TIME = 1000 * 60 * 3; // 3 minutes + +/** + * 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 + */ +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); + } +}; + +interface GenerateSummaryRequest { + question: string; + sql: string; + sampleSize?: number; + language?: string; + threadId?: string; +} + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const { question, sql, sampleSize, language, threadId } = + req.body as GenerateSummaryRequest; + const startTime = Date.now(); + let project; + + try { + project = await projectService.getCurrentProject(); + + // Only allow POST method + if (req.method !== 'POST') { + throw new ApiError('Method not allowed', 405); + } + + // Input validation + if (!question) { + throw new ApiError('Question is required', 400); + } + + if (!sql) { + throw new ApiError('SQL is required', 400); + } + + // Get current project's last deployment + const lastDeploy = await deployService.getLastDeployment(project.id); + if (!lastDeploy) { + throw new ApiError( + 'No deployment found, please deploy your project first', + 400, + Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, + ); + } + + // Create a new thread if it's a new question + const newThreadId = threadId || uuidv4(); + + // Get the data from the SQL + let sqlData; + try { + const queryResult = await queryService.preview(sql, { + project, + limit: sampleSize || 500, + manifest: lastDeploy.manifest, + modelingOnly: false, + }); + sqlData = queryResult; + } catch (queryError) { + throw new ApiError( + queryError.message || 'Error executing SQL query', + 400, + Errors.GeneralErrorCodes.INVALID_SQL_ERROR, + ); + } + + // Create text-based answer input for summary generation + const textBasedAnswerInput: TextBasedAnswerInput = { + query: question, + sql, + sqlData, + threadId: newThreadId, + configurations: { + language: + language || WrenAILanguage[project.language] || WrenAILanguage.EN, + }, + }; + + // Start the summary generation task + const task = + await wrenAIAdaptor.createTextBasedAnswer(textBasedAnswerInput); + + if (!task || !task.queryId) { + throw new ApiError('Failed to start summary generation task', 500); + } + + // Poll for the result + const deadline = Date.now() + MAX_WAIT_TIME; + let result: TextBasedAnswerResult; + while (true) { + result = await wrenAIAdaptor.getTextBasedAnswerResult(task.queryId); + if ( + result.status === TextBasedAnswerStatus.SUCCEEDED || + result.status === TextBasedAnswerStatus.FAILED + ) { + break; + } + + if (Date.now() > deadline) { + throw new ApiError( + 'Timeout waiting for summary generation', + 500, + Errors.GeneralErrorCodes.POLLING_TIMEOUT, + ); + } + + await new Promise((resolve) => setTimeout(resolve, 1000)); // Poll every second + } + + // Validate the summary result + validateSummaryResult(result); + + // Stream the content to get the summary + let summary = ''; + if (result.status === TextBasedAnswerStatus.SUCCEEDED) { + const stream = await wrenAIAdaptor.streamTextBasedAnswer(task.queryId); + + // Collect the streamed content + const streamPromise = new Promise((resolve, reject) => { + stream.on('data', (chunk) => { + const chunkString = chunk.toString('utf-8'); + const match = chunkString.match(/data: {"message":"([\s\S]*?)"}/); + if (match && match[1]) { + summary += match[1]; + } + }); + + stream.on('end', () => { + resolve(); + }); + + stream.on('error', (error) => { + reject(error); + }); + + // Handle client disconnect + req.on('close', () => { + stream.destroy(); + reject(new Error('Client disconnected')); + }); + }); + + await streamPromise; + } + + // Return the summary with ID and threadId + await respondWith({ + res, + statusCode: 200, + responsePayload: { + summary, + threadId: newThreadId, + }, + projectId: project.id, + apiType: ApiType.GENERATE_SUMMARY, + startTime, + requestPayload: req.body, + threadId: newThreadId, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: project?.id, + apiType: ApiType.GENERATE_SUMMARY, + requestPayload: req.body, + threadId, + headers: req.headers as Record, + startTime, + logger, + }); + } +} From 6a59430b44e51f78f56af56aa886ae2a91af0993 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Mon, 23 Jun 2025 18:24:04 +0800 Subject: [PATCH 02/16] add ask api --- docker/.env.example | 2 +- .../repositories/apiHistoryRepository.ts | 1 + wren-ui/src/apollo/server/utils/apiUtils.ts | 103 ++++++- wren-ui/src/pages/api/v1/ask.ts | 251 ++++++++++++++++++ wren-ui/src/pages/api/v1/generate_sql.ts | 83 +----- wren-ui/src/pages/api/v1/generate_summary.ts | 25 +- 6 files changed, 363 insertions(+), 102 deletions(-) create mode 100644 wren-ui/src/pages/api/v1/ask.ts diff --git a/docker/.env.example b/docker/.env.example index 7e511d4a15..fe9a96e9c5 100644 --- a/docker/.env.example +++ b/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 diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index 6f90908592..c1ce91729a 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -7,6 +7,7 @@ export enum ApiType { RUN_SQL = 'RUN_SQL', GENERATE_VEGA_CHART = 'GENERATE_VEGA_CHART', GENERATE_SUMMARY = 'GENERATE_SUMMARY', + ASK = 'ASK', } export interface ApiHistory { diff --git a/wren-ui/src/apollo/server/utils/apiUtils.ts b/wren-ui/src/apollo/server/utils/apiUtils.ts index 03ec365589..ee4c3253d1 100644 --- a/wren-ui/src/apollo/server/utils/apiUtils.ts +++ b/wren-ui/src/apollo/server/utils/apiUtils.ts @@ -1,11 +1,112 @@ 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 = {}; + + // 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, + })); +}; + /** * Common error class for API endpoints */ diff --git a/wren-ui/src/pages/api/v1/ask.ts b/wren-ui/src/pages/api/v1/ask.ts new file mode 100644 index 0000000000..bf7173f72f --- /dev/null +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -0,0 +1,251 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import * as Errors from '@/apollo/server/utils/error'; +import { v4 as uuidv4 } from 'uuid'; +import { + ApiError, + respondWith, + handleApiError, + MAX_WAIT_TIME, + isAskResultFinished, + validateAskResult, + validateSummaryResult, + transformHistoryInput, +} from '@/apollo/server/utils/apiUtils'; +import { + AskResult, + WrenAILanguage, + TextBasedAnswerInput, + TextBasedAnswerResult, + TextBasedAnswerStatus, +} from '@/apollo/server/models/adaptor'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('API_ASK'); +logger.level = 'debug'; + +const { + apiHistoryRepository, + projectService, + deployService, + wrenAIAdaptor, + queryService, +} = components; + +interface AskRequest { + question: string; + sampleSize?: number; + language?: string; + threadId?: string; +} + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const { question, sampleSize, language, threadId } = req.body as AskRequest; + const startTime = Date.now(); + let project; + + try { + project = await projectService.getCurrentProject(); + + // Only allow POST method + if (req.method !== 'POST') { + throw new ApiError('Method not allowed', 405); + } + + // Input validation + if (!question) { + throw new ApiError('Question is required', 400); + } + + // Get current project's last deployment + const lastDeploy = await deployService.getLastDeployment(project.id); + if (!lastDeploy) { + throw new ApiError( + 'No deployment found, please deploy your project first', + 400, + Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, + ); + } + + // Create a new thread if it's a new question + const newThreadId = threadId || uuidv4(); + + // Get conversation history if threadId is provided + const histories = threadId + ? await apiHistoryRepository.findAllBy({ threadId }) + : undefined; + + // Step 1: Generate SQL + const askTask = await wrenAIAdaptor.ask({ + query: question, + deployId: lastDeploy.hash, + histories: transformHistoryInput(histories) as any, + configurations: { + language: + language || WrenAILanguage[project.language] || WrenAILanguage.EN, + }, + }); + + // Poll for the SQL generation result + const deadline = Date.now() + MAX_WAIT_TIME; + let askResult: AskResult; + while (true) { + askResult = await wrenAIAdaptor.getAskResult(askTask.queryId); + if (isAskResultFinished(askResult)) { + break; + } + + if (Date.now() > deadline) { + throw new ApiError( + 'Timeout waiting for SQL generation', + 500, + Errors.GeneralErrorCodes.POLLING_TIMEOUT, + ); + } + + await new Promise((resolve) => setTimeout(resolve, 1000)); // Poll every second + } + + // Validate the AI result + validateAskResult(askResult, askTask.queryId); + + // Get the generated SQL + const sql = askResult.response?.[0]?.sql; + if (!sql) { + throw new ApiError('No SQL generated', 400); + } + + // Step 2: Execute SQL to get data + let sqlData; + try { + const queryResult = await queryService.preview(sql, { + project, + limit: sampleSize || 500, + manifest: lastDeploy.manifest, + modelingOnly: false, + }); + sqlData = queryResult; + } catch (queryError) { + throw new ApiError( + queryError.message || 'Error executing SQL query', + 400, + Errors.GeneralErrorCodes.INVALID_SQL_ERROR, + ); + } + + // Step 3: Generate summary using text-based answer + const textBasedAnswerInput: TextBasedAnswerInput = { + query: question, + sql, + sqlData, + threadId: newThreadId, + configurations: { + language: + language || WrenAILanguage[project.language] || WrenAILanguage.EN, + }, + }; + + // Start the summary generation task + const summaryTask = + await wrenAIAdaptor.createTextBasedAnswer(textBasedAnswerInput); + + if (!summaryTask || !summaryTask.queryId) { + throw new ApiError('Failed to start summary generation task', 500); + } + + // Poll for the summary result + let summaryResult: TextBasedAnswerResult; + while (true) { + summaryResult = await wrenAIAdaptor.getTextBasedAnswerResult( + summaryTask.queryId, + ); + if ( + summaryResult.status === TextBasedAnswerStatus.SUCCEEDED || + summaryResult.status === TextBasedAnswerStatus.FAILED + ) { + break; + } + + if (Date.now() > deadline) { + throw new ApiError( + 'Timeout waiting for summary generation', + 500, + Errors.GeneralErrorCodes.POLLING_TIMEOUT, + ); + } + + await new Promise((resolve) => setTimeout(resolve, 1000)); // Poll every second + } + + // Validate the summary result + validateSummaryResult(summaryResult); + + // Step 4: Stream the content to get the summary + let summary = ''; + if (summaryResult.status === TextBasedAnswerStatus.SUCCEEDED) { + const stream = await wrenAIAdaptor.streamTextBasedAnswer( + summaryTask.queryId, + ); + + // Collect the streamed content + const streamPromise = new Promise((resolve, reject) => { + stream.on('data', (chunk) => { + const chunkString = chunk.toString('utf-8'); + const match = chunkString.match(/data: {"message":"([\s\S]*?)"}/); + if (match && match[1]) { + summary += match[1]; + } + }); + + stream.on('end', () => { + resolve(); + }); + + stream.on('error', (error) => { + reject(error); + }); + + // Handle client disconnect + req.on('close', () => { + stream.destroy(); + reject(new Error('Client disconnected')); + }); + }); + + await streamPromise; + } + + // Return the combined result + await respondWith({ + res, + statusCode: 200, + responsePayload: { + sql, + summary, + threadId: newThreadId, + }, + projectId: project.id, + apiType: ApiType.ASK, + startTime, + requestPayload: req.body, + threadId: newThreadId, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: project?.id, + apiType: ApiType.ASK, + requestPayload: req.body, + threadId, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/generate_sql.ts b/wren-ui/src/pages/api/v1/generate_sql.ts index 8c7553958b..fa5b3859e1 100644 --- a/wren-ui/src/pages/api/v1/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/generate_sql.ts @@ -1,13 +1,7 @@ import { NextApiRequest, NextApiResponse } from 'next'; import { components } from '@/common'; -import { ApiType, ApiHistory } from '@server/repositories/apiHistoryRepository'; -import { - AskResult, - AskResultStatus, - AskResultType, - WrenAILanguage, - WrenAIError, -} from '@/apollo/server/models/adaptor'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { AskResult, WrenAILanguage } from '@/apollo/server/models/adaptor'; import * as Errors from '@/apollo/server/utils/error'; import { getLogger } from '@server/utils'; import { v4 as uuidv4 } from 'uuid'; @@ -15,6 +9,10 @@ import { ApiError, respondWith, handleApiError, + MAX_WAIT_TIME, + isAskResultFinished, + validateAskResult, + transformHistoryInput, } from '@/apollo/server/utils/apiUtils'; import { DataSourceName } from '@server/types'; @@ -37,75 +35,6 @@ interface GenerateSqlRequest { returnSqlDialect?: boolean; } -const MAX_WAIT_TIME = 1000 * 60 * 3; // 3 minutes - -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 - */ -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 = {}; - - // 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 }, - ); - } -}; - -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, - })); -}; - export default async function handler( req: NextApiRequest, res: NextApiResponse, diff --git a/wren-ui/src/pages/api/v1/generate_summary.ts b/wren-ui/src/pages/api/v1/generate_summary.ts index 2b0626fbcf..2ca538ce51 100644 --- a/wren-ui/src/pages/api/v1/generate_summary.ts +++ b/wren-ui/src/pages/api/v1/generate_summary.ts @@ -7,6 +7,8 @@ import { ApiError, respondWith, handleApiError, + MAX_WAIT_TIME, + validateSummaryResult, } from '@/apollo/server/utils/apiUtils'; import { TextBasedAnswerInput, @@ -22,29 +24,6 @@ logger.level = 'debug'; const { projectService, wrenAIAdaptor, deployService, queryService } = components; -const MAX_WAIT_TIME = 1000 * 60 * 3; // 3 minutes - -/** - * 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 - */ -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); - } -}; - interface GenerateSummaryRequest { question: string; sql: string; From 3ec215457cbf03272bda86509e8b14f2210f8dee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Wed, 25 Jun 2025 01:56:28 +0800 Subject: [PATCH 03/16] add instructions, sql_pairs API --- .../repositories/apiHistoryRepository.ts | 30 ++- wren-ui/src/apollo/server/schema.ts | 5 + .../server/services/instructionService.ts | 6 + wren-ui/src/apollo/server/utils/apiUtils.ts | 70 ++++++ .../api/v1/knowledge/instructions/[id].ts | 184 ++++++++++++++++ .../api/v1/knowledge/instructions/index.ts | 207 ++++++++++++++++++ .../pages/api/v1/knowledge/sql_pairs/[id].ts | 167 ++++++++++++++ .../pages/api/v1/knowledge/sql_pairs/index.ts | 138 ++++++++++++ 8 files changed, 806 insertions(+), 1 deletion(-) create mode 100644 wren-ui/src/pages/api/v1/knowledge/instructions/[id].ts create mode 100644 wren-ui/src/pages/api/v1/knowledge/instructions/index.ts create mode 100644 wren-ui/src/pages/api/v1/knowledge/sql_pairs/[id].ts create mode 100644 wren-ui/src/pages/api/v1/knowledge/sql_pairs/index.ts diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index c1ce91729a..98c8062fd5 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -1,4 +1,10 @@ -import { camelCase, isPlainObject, mapKeys, mapValues } from 'lodash'; +import { + camelCase, + isPlainObject, + mapKeys, + mapValues, + snakeCase, +} from 'lodash'; import { BaseRepository, IBasicRepository } from './baseRepository'; import { Knex } from 'knex'; @@ -8,6 +14,14 @@ export enum ApiType { 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', } export interface ApiHistory { @@ -149,6 +163,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 */ diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 86b3089956..87fca75072 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -9,6 +9,11 @@ export const typeDefs = gql` RUN_SQL GENERATE_VEGA_CHART GENERATE_SUMMARY + ASK + GET_INSTRUCTIONS + CREATE_INSTRUCTION + UPDATE_INSTRUCTION + DELETE_INSTRUCTION } input ApiHistoryFilterInput { diff --git a/wren-ui/src/apollo/server/services/instructionService.ts b/wren-ui/src/apollo/server/services/instructionService.ts index 51e1070a6c..37112ad5f4 100644 --- a/wren-ui/src/apollo/server/services/instructionService.ts +++ b/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; + getInstruction(id: number): Promise; createInstruction(instruction: InstructionInput): Promise; createInstructions(instructions: InstructionInput[]): Promise; updateInstruction(instruction: UpdateInstructionInput): Promise; @@ -30,10 +31,15 @@ export class InstructionService implements IInstructionService { this.instructionRepository = instructionRepository; this.wrenAIAdaptor = wrenAIAdaptor; } + public async getInstructions(projectId: number): Promise { return this.instructionRepository.findAllBy({ projectId }); } + public async getInstruction(id: number): Promise { + return this.instructionRepository.findOneBy({ id }); + } + public async createInstruction( input: InstructionInput, ): Promise { diff --git a/wren-ui/src/apollo/server/utils/apiUtils.ts b/wren-ui/src/apollo/server/utils/apiUtils.ts index ee4c3253d1..5c84817403 100644 --- a/wren-ui/src/apollo/server/utils/apiUtils.ts +++ b/wren-ui/src/apollo/server/utils/apiUtils.ts @@ -107,6 +107,37 @@ export const transformHistoryInput = (histories: ApiHistory[]) => { })); }; +/** + * 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 */ @@ -172,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; + headers?: Record; +}) => { + 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 */ diff --git a/wren-ui/src/pages/api/v1/knowledge/instructions/[id].ts b/wren-ui/src/pages/api/v1/knowledge/instructions/[id].ts new file mode 100644 index 0000000000..0fd54578c3 --- /dev/null +++ b/wren-ui/src/pages/api/v1/knowledge/instructions/[id].ts @@ -0,0 +1,184 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + ApiError, + respondWithSimple, + handleApiError, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('API_INSTRUCTION_BY_ID'); +logger.level = 'debug'; + +const { projectService, instructionService } = components; + +/** + * Instructions API - Supports two types of instructions: + * + * 1. Global Instructions (isGlobal: true) + * - Apply to every query that Wren AI generates + * - Ideal for setting consistent standards, enforcing business rules + * - Should NOT include questions field + * + * 2. Question-Matching Instructions (isGlobal: false or undefined) + * - Applied only when user's question matches certain patterns + * - Ideal for guiding how Wren AI handles specific business concepts + * - MUST include questions array with at least one question + */ +interface UpdateInstructionRequest { + instruction?: string; + questions?: string[]; + isGlobal?: boolean; +} + +/** + * Validate instruction ID from request query + */ +const validateInstructionId = (id: any): number => { + if (!id || typeof id !== 'string') { + throw new ApiError('Instruction ID is required', 400); + } + + const instructionId = parseInt(id, 10); + if (isNaN(instructionId)) { + throw new ApiError('Invalid instruction ID', 400); + } + + return instructionId; +}; + +/** + * Handle PUT request - update an existing instruction + */ +const handleUpdateInstruction = async ( + req: NextApiRequest, + res: NextApiResponse, + project: any, + startTime: number, +) => { + const { id } = req.query; + const instructionId = validateInstructionId(id); + + const { instruction, questions, isGlobal } = + req.body as UpdateInstructionRequest; + + // Get the original instruction + const existingInstruction = + await instructionService.getInstruction(instructionId); + + if (!existingInstruction) { + throw new ApiError('Instruction not found', 404); + } + + // Merge original with update payload + const mergedInstruction = { + instruction: instruction ?? existingInstruction.instruction, + questions: questions ?? existingInstruction.questions, + isGlobal: isGlobal ?? existingInstruction.isDefault, + }; + + // If isGlobal is true, set questions to empty array + if (mergedInstruction.isGlobal === true) { + mergedInstruction.questions = []; + } + + // Update the instruction + const updatedInstruction = await instructionService.updateInstruction({ + id: instructionId, + instruction: mergedInstruction.instruction, + questions: mergedInstruction.questions, + isDefault: mergedInstruction.isGlobal, + projectId: project.id, + }); + + // Return the updated instruction directly + const isGlobalValue = + typeof updatedInstruction.isDefault === 'boolean' + ? updatedInstruction.isDefault + : Boolean(updatedInstruction.isDefault); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { + id: updatedInstruction.id, + instruction: updatedInstruction.instruction, + questions: updatedInstruction.questions, + isGlobal: isGlobalValue, + }, + projectId: project.id, + apiType: ApiType.UPDATE_INSTRUCTION, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); +}; + +/** + * Handle DELETE request - delete an instruction + */ +const handleDeleteInstruction = async ( + req: NextApiRequest, + res: NextApiResponse, + project: any, + startTime: number, +) => { + const { id } = req.query; + const instructionId = validateInstructionId(id); + + // Delete the instruction + await instructionService.deleteInstruction(instructionId, project.id); + + // Return 204 No Content with no payload + await respondWithSimple({ + res, + statusCode: 204, + responsePayload: {}, + projectId: project.id, + apiType: ApiType.DELETE_INSTRUCTION, + startTime, + requestPayload: { id: instructionId }, + headers: req.headers as Record, + }); +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + let project; + + try { + project = await projectService.getCurrentProject(); + + // Handle PUT method - update instruction + if (req.method === 'PUT') { + await handleUpdateInstruction(req, res, project, startTime); + return; + } + + // Handle DELETE method - delete instruction + if (req.method === 'DELETE') { + await handleDeleteInstruction(req, res, project, startTime); + return; + } + + // Method not allowed + throw new ApiError('Method not allowed', 405); + } catch (error) { + await handleApiError({ + error, + res, + projectId: project?.id, + apiType: + req.method === 'PUT' + ? ApiType.UPDATE_INSTRUCTION + : ApiType.DELETE_INSTRUCTION, + requestPayload: req.method === 'PUT' ? req.body : { id: req.query.id }, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/knowledge/instructions/index.ts b/wren-ui/src/pages/api/v1/knowledge/instructions/index.ts new file mode 100644 index 0000000000..aa33653e91 --- /dev/null +++ b/wren-ui/src/pages/api/v1/knowledge/instructions/index.ts @@ -0,0 +1,207 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + ApiError, + respondWithSimple, + handleApiError, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { isNil } from 'lodash'; + +const logger = getLogger('API_INSTRUCTIONS'); +logger.level = 'debug'; + +const { projectService, instructionService } = components; + +/** + * Instructions API - Supports two types of instructions: + * + * 1. Global Instructions (isGlobal: true) + * - Apply to every query that Wren AI generates + * - Ideal for setting consistent standards, enforcing business rules + * - Should NOT include questions field + * + * 2. Question-Matching Instructions (isGlobal: false or undefined) + * - Applied only when user's question matches certain patterns + * - Ideal for guiding how Wren AI handles specific business concepts + * - MUST include questions array with at least one question + */ +interface CreateInstructionRequest { + instruction: string; + questions?: string[]; + isGlobal?: boolean; +} + +/** + * Handle GET request - list all instructions for the current project + */ +const handleGetInstructions = async ( + req: NextApiRequest, + res: NextApiResponse, + project: any, + startTime: number, +) => { + // Get all instructions for the current project + const instructions = ( + (await instructionService.getInstructions(project.id)) || [] + ).map((instruction) => { + const isGlobalValue = + typeof instruction.isDefault === 'boolean' + ? instruction.isDefault + : Boolean(instruction.isDefault); + return { + id: instruction.id, + instruction: instruction.instruction, + questions: instruction.questions, + isGlobal: isGlobalValue, + }; + }); + + // Return the instructions array directly + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: instructions, + projectId: project.id, + apiType: ApiType.GET_INSTRUCTIONS, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); +}; + +/** + * Handle POST request - create a new instruction + */ +const handleCreateInstruction = async ( + req: NextApiRequest, + res: NextApiResponse, + project: any, + startTime: number, +) => { + const { instruction, questions, isGlobal } = + req.body as CreateInstructionRequest; + + // Input validation + if (!instruction) { + throw new ApiError('Instruction is required', 400); + } + + if (instruction.length > 1000) { + throw new ApiError('Instruction is too long (max 1000 characters)', 400); + } + + if (isNil(isGlobal) && isNil(questions)) { + throw new ApiError('isGlobal or questions is required', 400); + } + + // Validate instruction type and fields + if (isGlobal === true) { + // Global instruction - questions should not be provided + if (questions && questions.length > 0) { + throw new ApiError( + 'Global instructions should not include questions. Questions are only for question-matching instructions.', + 400, + ); + } + } else { + // Question-matching instruction - questions are required + if (!questions || !Array.isArray(questions) || questions.length === 0) { + throw new ApiError( + 'Question-matching instructions require at least one question', + 400, + ); + } + + // Validate each question + questions.forEach((question, index) => { + if ( + !question || + typeof question !== 'string' || + question.trim().length === 0 + ) { + throw new ApiError( + `Question at index ${index} is required and cannot be empty`, + 400, + ); + } + if (question.length > 500) { + throw new ApiError( + `Question at index ${index} is too long (max 500 characters)`, + 400, + ); + } + }); + } + + // Create the instruction + const newInstruction = await instructionService.createInstruction({ + instruction, + questions: questions || [], + isDefault: isGlobal === true, + projectId: project.id, + }); + + // Return the created instruction directly + const isGlobalValue = + typeof newInstruction.isDefault === 'boolean' + ? newInstruction.isDefault + : Boolean(newInstruction.isDefault); + await respondWithSimple({ + res, + statusCode: 201, + responsePayload: { + id: newInstruction.id, + instruction: newInstruction.instruction, + questions: newInstruction.questions, + isGlobal: isGlobalValue, + }, + projectId: project.id, + apiType: ApiType.CREATE_INSTRUCTION, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + let project; + + try { + project = await projectService.getCurrentProject(); + + // Handle GET method - list instructions + if (req.method === 'GET') { + await handleGetInstructions(req, res, project, startTime); + return; + } + + // Handle POST method - create instruction + if (req.method === 'POST') { + await handleCreateInstruction(req, res, project, startTime); + return; + } + + // Method not allowed + throw new ApiError('Method not allowed', 405); + } catch (error) { + await handleApiError({ + error, + res, + projectId: project?.id, + apiType: + req.method === 'GET' + ? ApiType.GET_INSTRUCTIONS + : ApiType.CREATE_INSTRUCTION, + requestPayload: req.method === 'GET' ? {} : req.body, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/knowledge/sql_pairs/[id].ts b/wren-ui/src/pages/api/v1/knowledge/sql_pairs/[id].ts new file mode 100644 index 0000000000..46b0d0a7ea --- /dev/null +++ b/wren-ui/src/pages/api/v1/knowledge/sql_pairs/[id].ts @@ -0,0 +1,167 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + ApiError, + respondWithSimple, + handleApiError, + validateSql, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('API_SQL_PAIR_BY_ID'); +logger.level = 'debug'; + +const { projectService, sqlPairService, deployService, queryService } = + components; + +/** + * SQL Pairs API - Manages SQL query and question pairs for knowledge base + */ +interface UpdateSqlPairRequest { + sql?: string; + question?: string; +} + +/** + * Validate SQL pair ID from request query + */ +const validateSqlPairId = (id: any): number => { + if (!id || typeof id !== 'string') { + throw new ApiError('SQL pair ID is required', 400); + } + + const sqlPairId = parseInt(id, 10); + if (isNaN(sqlPairId)) { + throw new ApiError('Invalid SQL pair ID', 400); + } + + return sqlPairId; +}; + +/** + * Handle PUT request - update an existing SQL pair + */ +const handleUpdateSqlPair = async ( + req: NextApiRequest, + res: NextApiResponse, + project: any, + startTime: number, +) => { + const { id } = req.query; + const sqlPairId = validateSqlPairId(id); + + const { sql, question } = req.body as UpdateSqlPairRequest; + + // Input validation for provided fields + if (sql !== undefined) { + if (!sql) { + throw new ApiError('SQL cannot be empty', 400); + } + if (sql.length > 10000) { + throw new ApiError('SQL is too long (max 10000 characters)', 400); + } + // Validate SQL syntax and compatibility + await validateSql(sql, project, deployService, queryService); + } + + if (question !== undefined) { + if (!question) { + throw new ApiError('Question cannot be empty', 400); + } + if (question.length > 1000) { + throw new ApiError('Question is too long (max 1000 characters)', 400); + } + } + + // Update the SQL pair + const updatedSqlPair = await sqlPairService.editSqlPair( + project.id, + sqlPairId, + { + sql, + question, + }, + ); + + // Return the updated SQL pair directly + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: updatedSqlPair, + projectId: project.id, + apiType: ApiType.UPDATE_SQL_PAIR, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); +}; + +/** + * Handle DELETE request - delete a SQL pair + */ +const handleDeleteSqlPair = async ( + req: NextApiRequest, + res: NextApiResponse, + project: any, + startTime: number, +) => { + const { id } = req.query; + const sqlPairId = validateSqlPairId(id); + + // Delete the SQL pair + await sqlPairService.deleteSqlPair(project.id, sqlPairId); + + // Return 204 No Content with no payload + await respondWithSimple({ + res, + statusCode: 204, + responsePayload: {}, + projectId: project.id, + apiType: ApiType.DELETE_SQL_PAIR, + startTime, + requestPayload: { id: sqlPairId }, + headers: req.headers as Record, + }); +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + let project; + + try { + project = await projectService.getCurrentProject(); + + // Handle PUT method - update SQL pair + if (req.method === 'PUT') { + await handleUpdateSqlPair(req, res, project, startTime); + return; + } + + // Handle DELETE method - delete SQL pair + if (req.method === 'DELETE') { + await handleDeleteSqlPair(req, res, project, startTime); + return; + } + + // Method not allowed + throw new ApiError('Method not allowed', 405); + } catch (error) { + await handleApiError({ + error, + res, + projectId: project?.id, + apiType: + req.method === 'PUT' + ? ApiType.UPDATE_SQL_PAIR + : ApiType.DELETE_SQL_PAIR, + requestPayload: req.method === 'PUT' ? req.body : { id: req.query.id }, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/knowledge/sql_pairs/index.ts b/wren-ui/src/pages/api/v1/knowledge/sql_pairs/index.ts new file mode 100644 index 0000000000..55de946f11 --- /dev/null +++ b/wren-ui/src/pages/api/v1/knowledge/sql_pairs/index.ts @@ -0,0 +1,138 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + ApiError, + respondWithSimple, + handleApiError, + validateSql, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('API_SQL_PAIRS'); +logger.level = 'debug'; + +const { projectService, sqlPairService, deployService, queryService } = + components; + +/** + * SQL Pairs API - Manages SQL query and question pairs for knowledge base + */ +interface CreateSqlPairRequest { + sql: string; + question: string; +} + +/** + * Handle GET request - list all SQL pairs for the current project + */ +const handleGetSqlPairs = async ( + req: NextApiRequest, + res: NextApiResponse, + project: any, + startTime: number, +) => { + // Get all SQL pairs for the current project + const sqlPairs = await sqlPairService.getProjectSqlPairs(project.id); + + // Return the SQL pairs array directly + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: sqlPairs, + projectId: project.id, + apiType: ApiType.GET_SQL_PAIRS, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); +}; + +/** + * Handle POST request - create a new SQL pair + */ +const handleCreateSqlPair = async ( + req: NextApiRequest, + res: NextApiResponse, + project: any, + startTime: number, +) => { + const { sql, question } = req.body as CreateSqlPairRequest; + + // Input validation + if (!sql) { + throw new ApiError('SQL is required', 400); + } + + if (!question) { + throw new ApiError('Question is required', 400); + } + + if (sql.length > 10000) { + throw new ApiError('SQL is too long (max 10000 characters)', 400); + } + + if (question.length > 1000) { + throw new ApiError('Question is too long (max 1000 characters)', 400); + } + + // Validate SQL syntax and compatibility + await validateSql(sql, project, deployService, queryService); + + // Create the SQL pair + const newSqlPair = await sqlPairService.createSqlPair(project.id, { + sql, + question, + }); + + // Return the created SQL pair directly + await respondWithSimple({ + res, + statusCode: 201, + responsePayload: newSqlPair, + projectId: project.id, + apiType: ApiType.CREATE_SQL_PAIR, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + let project; + + try { + project = await projectService.getCurrentProject(); + + // Handle GET method - list SQL pairs + if (req.method === 'GET') { + await handleGetSqlPairs(req, res, project, startTime); + return; + } + + // Handle POST method - create SQL pair + if (req.method === 'POST') { + await handleCreateSqlPair(req, res, project, startTime); + return; + } + + // Method not allowed + throw new ApiError('Method not allowed', 405); + } catch (error) { + await handleApiError({ + error, + res, + projectId: project?.id, + apiType: + req.method === 'GET' ? ApiType.GET_SQL_PAIRS : ApiType.CREATE_SQL_PAIR, + requestPayload: req.method === 'GET' ? {} : req.body, + headers: req.headers as Record, + startTime, + logger, + }); + } +} From 5e36459bd0a26a27aa5fa426696b01fa1ea810f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Wed, 25 Jun 2025 02:48:24 +0800 Subject: [PATCH 04/16] add async ask api --- wren-ui/src/pages/api/v1/async/ask.ts | 465 ++++++++++++++++++++++++++ 1 file changed, 465 insertions(+) create mode 100644 wren-ui/src/pages/api/v1/async/ask.ts diff --git a/wren-ui/src/pages/api/v1/async/ask.ts b/wren-ui/src/pages/api/v1/async/ask.ts new file mode 100644 index 0000000000..8259a787ad --- /dev/null +++ b/wren-ui/src/pages/api/v1/async/ask.ts @@ -0,0 +1,465 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import * as Errors from '@/apollo/server/utils/error'; +import { v4 as uuidv4 } from 'uuid'; +import { + ApiError, + MAX_WAIT_TIME, + isAskResultFinished, + validateAskResult, + validateSummaryResult, + transformHistoryInput, +} from '@/apollo/server/utils/apiUtils'; +import { + AskResult, + AskResultStatus, + WrenAILanguage, + TextBasedAnswerInput, + TextBasedAnswerResult, + TextBasedAnswerStatus, +} from '@/apollo/server/models/adaptor'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('API_ASYNC_ASK'); +logger.level = 'debug'; + +const { + apiHistoryRepository, + projectService, + deployService, + wrenAIAdaptor, + queryService, +} = components; + +interface AsyncAskRequest { + question: string; + sampleSize?: number; + language?: string; + threadId?: string; +} + +interface StreamEvent { + type: 'state' | 'content' | 'error' | 'complete'; + data: any; + timestamp: number; +} + +/** + * Send SSE event to client + */ +const sendSSEEvent = (res: NextApiResponse, event: StreamEvent) => { + const eventData = `data: ${JSON.stringify(event)}\n\n`; + res.write(eventData); +}; + +/** + * Send state update to client + */ +const sendStateUpdate = ( + res: NextApiResponse, + state: string, + message?: string, + data?: any, +) => { + sendSSEEvent(res, { + type: 'state', + data: { + state, + message, + data, + }, + timestamp: Date.now(), + }); +}; + +/** + * Send content update to client + */ +const sendContentUpdate = (res: NextApiResponse, content: any) => { + sendSSEEvent(res, { + type: 'content', + data: content, + timestamp: Date.now(), + }); +}; + +/** + * Send error to client + */ +const sendError = (res: NextApiResponse, error: string, code?: string) => { + sendSSEEvent(res, { + type: 'error', + data: { + error, + code, + }, + timestamp: Date.now(), + }); +}; + +/** + * Send completion event to client + */ +const sendComplete = (res: NextApiResponse, result: any) => { + sendSSEEvent(res, { + type: 'complete', + data: result, + timestamp: Date.now(), + }); +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const { question, sampleSize, language, threadId } = + req.body as AsyncAskRequest; + const startTime = Date.now(); + let project; + + try { + project = await projectService.getCurrentProject(); + + // Only allow POST method + if (req.method !== 'POST') { + throw new ApiError('Method not allowed', 405); + } + + // Input validation + if (!question) { + throw new ApiError('Question is required', 400); + } + + // Set up SSE headers + res.setHeader('Content-Type', 'text/event-stream'); + res.setHeader('Cache-Control', 'no-cache, no-transform'); + res.setHeader('Connection', 'keep-alive'); + res.flushHeaders(); + + // Send initial connection event + sendStateUpdate(res, 'connected', 'Stream connected successfully'); + + // Get current project's last deployment + sendStateUpdate(res, 'validating', 'Validating project deployment'); + const lastDeploy = await deployService.getLastDeployment(project.id); + if (!lastDeploy) { + sendError( + res, + 'No deployment found, please deploy your project first', + Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, + ); + res.end(); + return; + } + + // Create a new thread if it's a new question + const newThreadId = threadId || uuidv4(); + sendStateUpdate(res, 'preparing', 'Preparing conversation context', { + threadId: newThreadId, + }); + + // Get conversation history if threadId is provided + const histories = threadId + ? await apiHistoryRepository.findAllBy({ threadId }) + : undefined; + + // Step 1: Generate SQL + sendStateUpdate(res, 'generating_sql', 'Generating SQL query'); + const askTask = await wrenAIAdaptor.ask({ + query: question, + deployId: lastDeploy.hash, + histories: transformHistoryInput(histories) as any, + configurations: { + language: + language || WrenAILanguage[project.language] || WrenAILanguage.EN, + }, + }); + + // Poll for the SQL generation result + const deadline = Date.now() + MAX_WAIT_TIME; + let askResult: AskResult; + let pollCount = 0; + let previousStatus: AskResultStatus | null = null; + + while (true) { + askResult = await wrenAIAdaptor.getAskResult(askTask.queryId); + + // Send status change updates when AskResultStatus changes + if (askResult.status !== previousStatus) { + const statusMessage = getStatusMessage(askResult.status); + sendStateUpdate(res, 'ask_status_change', statusMessage, { + status: askResult.status, + pollCount: pollCount + 1, + rephrasedQuestion: askResult.rephrasedQuestion, + intentReasoning: askResult.intentReasoning, + sqlGenerationReasoning: askResult.sqlGenerationReasoning, + retrievedTables: askResult.retrievedTables, + invalidSql: askResult.invalidSql, + traceId: askResult.traceId, + }); + previousStatus = askResult.status; + } + + if (isAskResultFinished(askResult)) { + break; + } + + if (Date.now() > deadline) { + sendError( + res, + 'Timeout waiting for SQL generation', + Errors.GeneralErrorCodes.POLLING_TIMEOUT, + ); + res.end(); + return; + } + + await new Promise((resolve) => setTimeout(resolve, 1000)); // Poll every second + pollCount++; + } + + // Validate the AI result + try { + validateAskResult(askResult, askTask.queryId); + } catch (error) { + sendError(res, error.message, error.code); + res.end(); + return; + } + + // Get the generated SQL + const sql = askResult.response?.[0]?.sql; + if (!sql) { + sendError(res, 'No SQL generated'); + res.end(); + return; + } + + // Send SQL generation complete + sendStateUpdate(res, 'sql_generated', 'SQL query generated successfully'); + sendContentUpdate(res, { sql }); + + // Step 2: Execute SQL to get data + sendStateUpdate(res, 'executing_sql', 'Executing SQL query to fetch data'); + let sqlData; + try { + const queryResult = await queryService.preview(sql, { + project, + limit: sampleSize || 500, + manifest: lastDeploy.manifest, + modelingOnly: false, + }); + sqlData = queryResult; + sendStateUpdate(res, 'sql_executed', 'SQL query executed successfully'); + } catch (queryError) { + sendError( + res, + queryError.message || 'Error executing SQL query', + Errors.GeneralErrorCodes.INVALID_SQL_ERROR, + ); + res.end(); + return; + } + + // Step 3: Generate summary using text-based answer + sendStateUpdate(res, 'generating_summary', 'Generating summary from data'); + const textBasedAnswerInput: TextBasedAnswerInput = { + query: question, + sql, + sqlData, + threadId: newThreadId, + configurations: { + language: + language || WrenAILanguage[project.language] || WrenAILanguage.EN, + }, + }; + + // Start the summary generation task + const summaryTask = + await wrenAIAdaptor.createTextBasedAnswer(textBasedAnswerInput); + + if (!summaryTask || !summaryTask.queryId) { + sendError(res, 'Failed to start summary generation task'); + res.end(); + return; + } + + // Poll for the summary result + let summaryResult: TextBasedAnswerResult; + pollCount = 0; + + while (true) { + summaryResult = await wrenAIAdaptor.getTextBasedAnswerResult( + summaryTask.queryId, + ); + + // Send polling status updates + if (pollCount % 3 === 0) { + // Send update every 3 polls + sendStateUpdate( + res, + 'polling_summary', + 'Waiting for summary generation', + { + pollCount: Math.floor(pollCount / 3) + 1, + }, + ); + } + + if ( + summaryResult.status === TextBasedAnswerStatus.SUCCEEDED || + summaryResult.status === TextBasedAnswerStatus.FAILED + ) { + break; + } + + if (Date.now() > deadline) { + sendError( + res, + 'Timeout waiting for summary generation', + Errors.GeneralErrorCodes.POLLING_TIMEOUT, + ); + res.end(); + return; + } + + await new Promise((resolve) => setTimeout(resolve, 1000)); // Poll every second + pollCount++; + } + + // Validate the summary result + try { + validateSummaryResult(summaryResult); + } catch (error) { + sendError(res, error.message, error.code); + res.end(); + return; + } + + // Step 4: Stream the content to get the summary + let summary = ''; + if (summaryResult.status === TextBasedAnswerStatus.SUCCEEDED) { + sendStateUpdate(res, 'streaming_summary', 'Streaming summary content'); + + const stream = await wrenAIAdaptor.streamTextBasedAnswer( + summaryTask.queryId, + ); + + // Stream the content in real-time + const streamPromise = new Promise((resolve, reject) => { + stream.on('data', (chunk) => { + const chunkString = chunk.toString('utf-8'); + const match = chunkString.match(/data: {"message":"([\s\S]*?)"}/); + if (match && match[1]) { + summary += match[1]; + // Send incremental content updates + sendContentUpdate(res, { + summary: match[1], + summaryComplete: false, + }); + } + }); + + stream.on('end', () => { + resolve(); + }); + + stream.on('error', (error) => { + reject(error); + }); + + // Handle client disconnect + req.on('close', () => { + stream.destroy(); + reject(new Error('Client disconnected')); + }); + }); + + try { + await streamPromise; + sendStateUpdate( + res, + 'summary_complete', + 'Summary generation completed', + ); + } catch (_streamError) { + sendError(res, 'Error streaming summary content'); + res.end(); + return; + } + } + + // Send final completion event + const finalResult = { + sql, + summary, + threadId: newThreadId, + }; + + sendComplete(res, finalResult); + + // Log the API call + await apiHistoryRepository.createOne({ + projectId: project.id, + apiType: ApiType.ASK, + threadId: newThreadId, + headers: req.headers as Record, + requestPayload: req.body, + responsePayload: finalResult, + statusCode: 200, + durationMs: Date.now() - startTime, + }); + + res.end(); + } catch (error) { + logger.error('Error in async ask API:', error); + + // Send error event to client + const errorMessage = + error instanceof ApiError ? error.message : 'Internal server error'; + const errorCode = error instanceof ApiError ? error.code : undefined; + + sendError(res, errorMessage, errorCode); + + // Log the error + await apiHistoryRepository.createOne({ + projectId: project?.id || 0, + apiType: ApiType.ASK, + threadId, + headers: req.headers as Record, + requestPayload: req.body, + responsePayload: { error: errorMessage }, + statusCode: error instanceof ApiError ? error.statusCode : 500, + durationMs: Date.now() - startTime, + }); + + res.end(); + } +} + +/** + * Get human-readable message for AskResultStatus + */ +const getStatusMessage = (status: AskResultStatus): string => { + switch (status) { + case AskResultStatus.UNDERSTANDING: + return 'Understanding your question'; + case AskResultStatus.SEARCHING: + return 'Searching for relevant data'; + case AskResultStatus.PLANNING: + return 'Planning the SQL generation approach'; + case AskResultStatus.GENERATING: + return 'Generating SQL query'; + case AskResultStatus.CORRECTING: + return 'Correcting and improving SQL'; + case AskResultStatus.FINISHED: + return 'SQL generation completed'; + case AskResultStatus.FAILED: + return 'SQL generation failed'; + case AskResultStatus.STOPPED: + return 'SQL generation stopped'; + default: + return 'Processing your request'; + } +}; From 6c455b042eb3ef351c77de2d1c3bade7e596f2d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Wed, 25 Jun 2025 18:55:40 +0800 Subject: [PATCH 05/16] ask async api --- wren-ui/src/apollo/server/utils/index.ts | 1 + wren-ui/src/apollo/server/utils/sseTypes.ts | 98 ++++++ wren-ui/src/pages/api/v1/async/ask.ts | 328 ++++++++++---------- 3 files changed, 270 insertions(+), 157 deletions(-) create mode 100644 wren-ui/src/apollo/server/utils/sseTypes.ts diff --git a/wren-ui/src/apollo/server/utils/index.ts b/wren-ui/src/apollo/server/utils/index.ts index 8db524806d..ba31bd8f88 100644 --- a/wren-ui/src/apollo/server/utils/index.ts +++ b/wren-ui/src/apollo/server/utils/index.ts @@ -6,3 +6,4 @@ export * from './docker'; export * from './model'; export * from './helper'; export * from './regex'; +export * from './sseTypes'; diff --git a/wren-ui/src/apollo/server/utils/sseTypes.ts b/wren-ui/src/apollo/server/utils/sseTypes.ts new file mode 100644 index 0000000000..7c2015d65a --- /dev/null +++ b/wren-ui/src/apollo/server/utils/sseTypes.ts @@ -0,0 +1,98 @@ +// Server-Sent Events (SSE) types for real-time streaming APIs +export enum EventType { + MESSAGE_START = 'message_start', + MESSAGE_STOP = 'message_stop', + STATE = 'state', + CONTENT_BLOCK_START = 'content_block_start', + CONTENT_BLOCK_DELTA = 'content_block_delta', + CONTENT_BLOCK_STOP = 'content_block_stop', + ERROR = 'error', +} + +export enum StateType { + SQL_GENERATION_START = 'sql_generation_start', + SQL_GENERATION_UNDERSTANDING = 'sql_generation_understanding', + SQL_GENERATION_SEARCHING = 'sql_generation_searching', + SQL_GENERATION_PLANNING = 'sql_generation_planning', + SQL_GENERATION_GENERATING = 'sql_generation_generating', + SQL_GENERATION_CORRECTING = 'sql_generation_correcting', + SQL_GENERATION_FINISHED = 'sql_generation_finished', + SQL_GENERATION_FAILED = 'sql_generation_failed', + SQL_GENERATION_STOPPED = 'sql_generation_stopped', + SQL_GENERATION_END = 'sql_generation_end', + SQL_EXECUTION_START = 'sql_execution_start', + SQL_EXECUTION_END = 'sql_execution_end', +} + +export enum ContentBlockContentType { + SUMMARY_GENERATION = 'summary_generation', +} + +// Interfaces for request and events +export interface AsyncAskRequest { + question: string; + sampleSize?: number; + language?: string; + threadId?: string; +} + +export interface BaseEvent { + timestamp: number; +} + +export interface MessageStartEvent extends BaseEvent { + type: EventType.MESSAGE_START; +} + +export interface MessageStopEvent extends BaseEvent { + type: EventType.MESSAGE_STOP; + data: { + threadId: string; + duration: number; + }; +} + +export interface StateEvent extends BaseEvent { + type: EventType.STATE; + data: { + state: StateType; + [key: string]: any; + }; +} + +export interface ContentBlockStartEvent extends BaseEvent { + type: EventType.CONTENT_BLOCK_START; + content_block: { + type: 'text'; + name: ContentBlockContentType; + }; +} + +export interface ContentBlockDeltaEvent extends BaseEvent { + type: EventType.CONTENT_BLOCK_DELTA; + delta: { + type: 'text_delta'; + text: string; + }; +} + +export interface ContentBlockStopEvent extends BaseEvent { + type: EventType.CONTENT_BLOCK_STOP; +} + +export interface ErrorEvent extends BaseEvent { + type: EventType.ERROR; + data: { + error: string; + code?: string; + }; +} + +export type StreamEvent = + | MessageStartEvent + | MessageStopEvent + | StateEvent + | ContentBlockStartEvent + | ContentBlockDeltaEvent + | ContentBlockStopEvent + | ErrorEvent; diff --git a/wren-ui/src/pages/api/v1/async/ask.ts b/wren-ui/src/pages/api/v1/async/ask.ts index 8259a787ad..a55962a3f6 100644 --- a/wren-ui/src/pages/api/v1/async/ask.ts +++ b/wren-ui/src/pages/api/v1/async/ask.ts @@ -20,6 +20,20 @@ import { TextBasedAnswerStatus, } from '@/apollo/server/models/adaptor'; import { getLogger } from '@server/utils'; +import { + EventType, + StateType, + ContentBlockContentType, + AsyncAskRequest, + StreamEvent, + StateEvent, + ErrorEvent, + MessageStartEvent, + MessageStopEvent, + ContentBlockStartEvent, + ContentBlockDeltaEvent, + ContentBlockStopEvent, +} from '@/apollo/server/utils'; const logger = getLogger('API_ASYNC_ASK'); logger.level = 'debug'; @@ -32,19 +46,6 @@ const { queryService, } = components; -interface AsyncAskRequest { - question: string; - sampleSize?: number; - language?: string; - threadId?: string; -} - -interface StreamEvent { - type: 'state' | 'content' | 'error' | 'complete'; - data: any; - timestamp: number; -} - /** * Send SSE event to client */ @@ -53,60 +54,138 @@ const sendSSEEvent = (res: NextApiResponse, event: StreamEvent) => { res.write(eventData); }; +/** + * Send message start event to client + */ +const sendMessageStart = (res: NextApiResponse) => { + const messageStartEvent: MessageStartEvent = { + type: EventType.MESSAGE_START, + timestamp: Date.now(), + }; + sendSSEEvent(res, messageStartEvent); +}; + +/** + * Send message stop event to client + */ +const sendMessageStop = ( + res: NextApiResponse, + threadId: string, + duration: number, +) => { + const messageStopEvent: MessageStopEvent = { + type: EventType.MESSAGE_STOP, + data: { + threadId, + duration, + }, + timestamp: Date.now(), + }; + sendSSEEvent(res, messageStopEvent); +}; + /** * Send state update to client */ const sendStateUpdate = ( res: NextApiResponse, - state: string, - message?: string, + state: StateType, data?: any, ) => { - sendSSEEvent(res, { - type: 'state', + const stateEvent: StateEvent = { + type: EventType.STATE, data: { state, - message, - data, + ...data, + }, + timestamp: Date.now(), + }; + sendSSEEvent(res, stateEvent); +}; + +/** + * Send content block start event to client + */ +const sendContentBlockStart = ( + res: NextApiResponse, + name: ContentBlockContentType, +) => { + const contentBlockStartEvent: ContentBlockStartEvent = { + type: EventType.CONTENT_BLOCK_START, + content_block: { + type: 'text', + name, + }, + timestamp: Date.now(), + }; + sendSSEEvent(res, contentBlockStartEvent); +}; + +/** + * Send content block delta event to client + */ +const sendContentBlockDelta = (res: NextApiResponse, text: string) => { + const contentBlockDeltaEvent: ContentBlockDeltaEvent = { + type: EventType.CONTENT_BLOCK_DELTA, + delta: { + type: 'text_delta', + text, }, timestamp: Date.now(), - }); + }; + sendSSEEvent(res, contentBlockDeltaEvent); }; /** - * Send content update to client + * Send content block stop event to client */ -const sendContentUpdate = (res: NextApiResponse, content: any) => { - sendSSEEvent(res, { - type: 'content', - data: content, +const sendContentBlockStop = (res: NextApiResponse) => { + const contentBlockStopEvent: ContentBlockStopEvent = { + type: EventType.CONTENT_BLOCK_STOP, timestamp: Date.now(), - }); + }; + sendSSEEvent(res, contentBlockStopEvent); }; /** * Send error to client */ const sendError = (res: NextApiResponse, error: string, code?: string) => { - sendSSEEvent(res, { - type: 'error', + const errorEvent: ErrorEvent = { + type: EventType.ERROR, data: { error, code, }, timestamp: Date.now(), - }); + }; + sendSSEEvent(res, errorEvent); }; /** - * Send completion event to client + * Transform AskResultStatus to descriptive SQL generation state */ -const sendComplete = (res: NextApiResponse, result: any) => { - sendSSEEvent(res, { - type: 'complete', - data: result, - timestamp: Date.now(), - }); +const getSqlGenerationState = (status: AskResultStatus): StateType => { + switch (status) { + case AskResultStatus.UNDERSTANDING: + return StateType.SQL_GENERATION_UNDERSTANDING; + case AskResultStatus.SEARCHING: + return StateType.SQL_GENERATION_SEARCHING; + case AskResultStatus.PLANNING: + return StateType.SQL_GENERATION_PLANNING; + case AskResultStatus.GENERATING: + return StateType.SQL_GENERATION_GENERATING; + case AskResultStatus.CORRECTING: + return StateType.SQL_GENERATION_CORRECTING; + case AskResultStatus.FINISHED: + return StateType.SQL_GENERATION_FINISHED; + case AskResultStatus.FAILED: + return StateType.SQL_GENERATION_FAILED; + case AskResultStatus.STOPPED: + return StateType.SQL_GENERATION_STOPPED; + default: + return StateType.SQL_GENERATION_UNDERSTANDING; + } }; export default async function handler( @@ -137,11 +216,10 @@ export default async function handler( res.setHeader('Connection', 'keep-alive'); res.flushHeaders(); - // Send initial connection event - sendStateUpdate(res, 'connected', 'Stream connected successfully'); + // Send message start event + sendMessageStart(res); // Get current project's last deployment - sendStateUpdate(res, 'validating', 'Validating project deployment'); const lastDeploy = await deployService.getLastDeployment(project.id); if (!lastDeploy) { sendError( @@ -155,9 +233,6 @@ export default async function handler( // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); - sendStateUpdate(res, 'preparing', 'Preparing conversation context', { - threadId: newThreadId, - }); // Get conversation history if threadId is provided const histories = threadId @@ -165,7 +240,12 @@ export default async function handler( : undefined; // Step 1: Generate SQL - sendStateUpdate(res, 'generating_sql', 'Generating SQL query'); + sendStateUpdate(res, StateType.SQL_GENERATION_START, { + question, + threadId: newThreadId, + language: + language || WrenAILanguage[project.language] || WrenAILanguage.EN, + }); const askTask = await wrenAIAdaptor.ask({ query: question, deployId: lastDeploy.hash, @@ -187,9 +267,8 @@ export default async function handler( // Send status change updates when AskResultStatus changes if (askResult.status !== previousStatus) { - const statusMessage = getStatusMessage(askResult.status); - sendStateUpdate(res, 'ask_status_change', statusMessage, { - status: askResult.status, + const sqlGenerationState = getSqlGenerationState(askResult.status); + sendStateUpdate(res, sqlGenerationState, { pollCount: pollCount + 1, rephrasedQuestion: askResult.rephrasedQuestion, intentReasoning: askResult.intentReasoning, @@ -201,32 +280,26 @@ export default async function handler( previousStatus = askResult.status; } + pollCount++; + + // Check if the result is finished if (isAskResultFinished(askResult)) { break; } + // Check if we've exceeded the maximum wait time if (Date.now() > deadline) { - sendError( - res, - 'Timeout waiting for SQL generation', - Errors.GeneralErrorCodes.POLLING_TIMEOUT, - ); + sendError(res, 'Request timeout', 'TIMEOUT'); res.end(); return; } - await new Promise((resolve) => setTimeout(resolve, 1000)); // Poll every second - pollCount++; + // Wait before polling again + await new Promise((resolve) => setTimeout(resolve, 1000)); } - // Validate the AI result - try { - validateAskResult(askResult, askTask.queryId); - } catch (error) { - sendError(res, error.message, error.code); - res.end(); - return; - } + // Validate the ask result + validateAskResult(askResult, askTask.queryId); // Get the generated SQL const sql = askResult.response?.[0]?.sql; @@ -236,12 +309,11 @@ export default async function handler( return; } - // Send SQL generation complete - sendStateUpdate(res, 'sql_generated', 'SQL query generated successfully'); - sendContentUpdate(res, { sql }); + // Send SQL generation end with the SQL content + sendStateUpdate(res, StateType.SQL_GENERATION_END, { sql }); // Step 2: Execute SQL to get data - sendStateUpdate(res, 'executing_sql', 'Executing SQL query to fetch data'); + sendStateUpdate(res, StateType.SQL_EXECUTION_START, { sql }); let sqlData; try { const queryResult = await queryService.preview(sql, { @@ -251,19 +323,18 @@ export default async function handler( modelingOnly: false, }); sqlData = queryResult; - sendStateUpdate(res, 'sql_executed', 'SQL query executed successfully'); + sendStateUpdate(res, StateType.SQL_EXECUTION_END); } catch (queryError) { sendError( res, - queryError.message || 'Error executing SQL query', - Errors.GeneralErrorCodes.INVALID_SQL_ERROR, + `SQL execution failed: ${queryError.message || 'Unknown error'}`, + 'SQL_EXECUTION_ERROR', ); res.end(); return; } // Step 3: Generate summary using text-based answer - sendStateUpdate(res, 'generating_summary', 'Generating summary from data'); const textBasedAnswerInput: TextBasedAnswerInput = { query: question, sql, @@ -285,28 +356,15 @@ export default async function handler( return; } - // Poll for the summary result + // Poll for the summary generation result + const summaryDeadline = Date.now() + MAX_WAIT_TIME; let summaryResult: TextBasedAnswerResult; - pollCount = 0; while (true) { summaryResult = await wrenAIAdaptor.getTextBasedAnswerResult( summaryTask.queryId, ); - // Send polling status updates - if (pollCount % 3 === 0) { - // Send update every 3 polls - sendStateUpdate( - res, - 'polling_summary', - 'Waiting for summary generation', - { - pollCount: Math.floor(pollCount / 3) + 1, - }, - ); - } - if ( summaryResult.status === TextBasedAnswerStatus.SUCCEEDED || summaryResult.status === TextBasedAnswerStatus.FAILED @@ -314,33 +372,25 @@ export default async function handler( break; } - if (Date.now() > deadline) { - sendError( - res, - 'Timeout waiting for summary generation', - Errors.GeneralErrorCodes.POLLING_TIMEOUT, - ); + // Check if we've exceeded the maximum wait time + if (Date.now() > summaryDeadline) { + sendError(res, 'Summary generation timeout', 'TIMEOUT'); res.end(); return; } - await new Promise((resolve) => setTimeout(resolve, 1000)); // Poll every second - pollCount++; + // Wait before polling again + await new Promise((resolve) => setTimeout(resolve, 1000)); } // Validate the summary result - try { - validateSummaryResult(summaryResult); - } catch (error) { - sendError(res, error.message, error.code); - res.end(); - return; - } + validateSummaryResult(summaryResult); // Step 4: Stream the content to get the summary let summary = ''; if (summaryResult.status === TextBasedAnswerStatus.SUCCEEDED) { - sendStateUpdate(res, 'streaming_summary', 'Streaming summary content'); + // Send content block start + sendContentBlockStart(res, ContentBlockContentType.SUMMARY_GENERATION); const stream = await wrenAIAdaptor.streamTextBasedAnswer( summaryTask.queryId, @@ -354,10 +404,7 @@ export default async function handler( if (match && match[1]) { summary += match[1]; // Send incremental content updates - sendContentUpdate(res, { - summary: match[1], - summaryComplete: false, - }); + sendContentBlockDelta(res, match[1]); } }); @@ -378,11 +425,8 @@ export default async function handler( try { await streamPromise; - sendStateUpdate( - res, - 'summary_complete', - 'Summary generation completed', - ); + // Send content block stop + sendContentBlockStop(res); } catch (_streamError) { sendError(res, 'Error streaming summary content'); res.end(); @@ -390,23 +434,17 @@ export default async function handler( } } - // Send final completion event - const finalResult = { - sql, - summary, - threadId: newThreadId, - }; - - sendComplete(res, finalResult); + // Send message end event + sendMessageStop(res, newThreadId, Date.now() - startTime); // Log the API call await apiHistoryRepository.createOne({ + id: uuidv4(), projectId: project.id, apiType: ApiType.ASK, threadId: newThreadId, - headers: req.headers as Record, - requestPayload: req.body, - responsePayload: finalResult, + requestPayload: { question, sampleSize, language }, + responsePayload: { sql, summary }, statusCode: 200, durationMs: Date.now() - startTime, }); @@ -415,51 +453,27 @@ export default async function handler( } catch (error) { logger.error('Error in async ask API:', error); - // Send error event to client - const errorMessage = - error instanceof ApiError ? error.message : 'Internal server error'; - const errorCode = error instanceof ApiError ? error.code : undefined; - - sendError(res, errorMessage, errorCode); + // Send message end event even on error + sendMessageStop(res, threadId || uuidv4(), Date.now() - startTime); // Log the error await apiHistoryRepository.createOne({ + id: uuidv4(), projectId: project?.id || 0, apiType: ApiType.ASK, - threadId, - headers: req.headers as Record, - requestPayload: req.body, - responsePayload: { error: errorMessage }, - statusCode: error instanceof ApiError ? error.statusCode : 500, + threadId: threadId || uuidv4(), + requestPayload: { question, sampleSize, language }, + responsePayload: { + error: error instanceof Error ? error.message : String(error), + }, + statusCode: 500, durationMs: Date.now() - startTime, }); + sendError( + res, + error instanceof Error ? error.message : 'Internal server error', + ); res.end(); } } - -/** - * Get human-readable message for AskResultStatus - */ -const getStatusMessage = (status: AskResultStatus): string => { - switch (status) { - case AskResultStatus.UNDERSTANDING: - return 'Understanding your question'; - case AskResultStatus.SEARCHING: - return 'Searching for relevant data'; - case AskResultStatus.PLANNING: - return 'Planning the SQL generation approach'; - case AskResultStatus.GENERATING: - return 'Generating SQL query'; - case AskResultStatus.CORRECTING: - return 'Correcting and improving SQL'; - case AskResultStatus.FINISHED: - return 'SQL generation completed'; - case AskResultStatus.FAILED: - return 'SQL generation failed'; - case AskResultStatus.STOPPED: - return 'SQL generation stopped'; - default: - return 'Processing your request'; - } -}; From 7718c19ec99b279b95c9ddb1e192437cffec599e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Thu, 26 Jun 2025 04:27:00 +0800 Subject: [PATCH 06/16] add stream explanation to ask async api --- wren-ui/src/apollo/server/utils/error.ts | 7 + wren-ui/src/apollo/server/utils/sseTypes.ts | 1 + wren-ui/src/pages/api/v1/async/ask.ts | 161 +++++++++++++++----- 3 files changed, 135 insertions(+), 34 deletions(-) diff --git a/wren-ui/src/apollo/server/utils/error.ts b/wren-ui/src/apollo/server/utils/error.ts index db0faef2fe..941754c6f0 100644 --- a/wren-ui/src/apollo/server/utils/error.ts +++ b/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 = ( diff --git a/wren-ui/src/apollo/server/utils/sseTypes.ts b/wren-ui/src/apollo/server/utils/sseTypes.ts index 7c2015d65a..5107eba118 100644 --- a/wren-ui/src/apollo/server/utils/sseTypes.ts +++ b/wren-ui/src/apollo/server/utils/sseTypes.ts @@ -26,6 +26,7 @@ export enum StateType { export enum ContentBlockContentType { SUMMARY_GENERATION = 'summary_generation', + EXPLANATION = 'explanation', } // Interfaces for request and events diff --git a/wren-ui/src/pages/api/v1/async/ask.ts b/wren-ui/src/pages/api/v1/async/ask.ts index a55962a3f6..388bb056ec 100644 --- a/wren-ui/src/pages/api/v1/async/ask.ts +++ b/wren-ui/src/pages/api/v1/async/ask.ts @@ -7,7 +7,6 @@ import { ApiError, MAX_WAIT_TIME, isAskResultFinished, - validateAskResult, validateSummaryResult, transformHistoryInput, } from '@/apollo/server/utils/apiUtils'; @@ -18,6 +17,8 @@ import { TextBasedAnswerInput, TextBasedAnswerResult, TextBasedAnswerStatus, + WrenAIError, + AskResultType, } from '@/apollo/server/models/adaptor'; import { getLogger } from '@server/utils'; import { @@ -188,6 +189,16 @@ const getSqlGenerationState = (status: AskResultStatus): StateType => { } }; +const endStream = ( + res: NextApiResponse, + threadId: string, + startTime: number, +) => { + // Send message stop event + sendMessageStop(res, threadId || uuidv4(), Date.now() - startTime); + res.end(); +}; + export default async function handler( req: NextApiRequest, res: NextApiResponse, @@ -222,13 +233,11 @@ export default async function handler( // Get current project's last deployment const lastDeploy = await deployService.getLastDeployment(project.id); if (!lastDeploy) { - sendError( - res, + throw new ApiError( 'No deployment found, please deploy your project first', + 400, Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, ); - res.end(); - return; } // Create a new thread if it's a new question @@ -289,9 +298,11 @@ export default async function handler( // Check if we've exceeded the maximum wait time if (Date.now() > deadline) { - sendError(res, 'Request timeout', 'TIMEOUT'); - res.end(); - return; + throw new ApiError( + 'Request timeout', + 400, + Errors.GeneralErrorCodes.POLLING_TIMEOUT, + ); } // Wait before polling again @@ -299,14 +310,98 @@ export default async function handler( } // Validate the ask result - validateAskResult(askResult, askTask.queryId); + // Check for error in result + if (askResult.error) { + const errorMessage = + (askResult.error as WrenAIError).message || 'Unknown error'; + const additionalData: Record = {}; + + // Include invalid SQL if available + if (askResult.invalidSql) { + additionalData.invalidSql = askResult.invalidSql; + } + + throw new ApiError( + errorMessage, + 400, + askResult.error.code, + additionalData, + ); + } + + // Check for general type response + // Stream the content to client + if (askResult.type === AskResultType.GENERAL) { + // Send content block start for explanation + sendContentBlockStart(res, ContentBlockContentType.EXPLANATION); + + const stream = await wrenAIAdaptor.getAskStreamingResult(askTask.queryId); + + // Stream the content in real-time + const streamPromise = new Promise((resolve, reject) => { + stream.on('data', (chunk) => { + const chunkString = chunk.toString('utf-8'); + const match = chunkString.match(/data: {"message":"([\s\S]*?)"}/); + if (match && match[1]) { + // Send incremental content updates + sendContentBlockDelta(res, match[1]); + } + }); + + stream.on('end', () => { + resolve(); + }); + + stream.on('error', (error) => { + reject(error); + }); + + // Handle client disconnect + req.on('close', () => { + stream.destroy(); + reject(new Error('Client disconnected')); + }); + }); + + try { + await streamPromise; + // Send content block stop + sendContentBlockStop(res); + } catch (_streamError) { + throw new ApiError( + 'Error streaming explanation content', + 400, + Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + ); + } + + // Log the API call and end the stream + await apiHistoryRepository.createOne({ + id: uuidv4(), + projectId: project.id, + apiType: ApiType.ASK, + threadId: newThreadId, + requestPayload: { question, sampleSize, language }, + responsePayload: { + type: 'general', + explanation: 'Streamed explanation', + }, + statusCode: 200, + durationMs: Date.now() - startTime, + }); + + endStream(res, newThreadId, startTime); + return; + } // Get the generated SQL const sql = askResult.response?.[0]?.sql; if (!sql) { - sendError(res, 'No SQL generated'); - res.end(); - return; + throw new ApiError( + 'No SQL generated', + 400, + Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + ); } // Send SQL generation end with the SQL content @@ -325,13 +420,11 @@ export default async function handler( sqlData = queryResult; sendStateUpdate(res, StateType.SQL_EXECUTION_END); } catch (queryError) { - sendError( - res, + throw new ApiError( `SQL execution failed: ${queryError.message || 'Unknown error'}`, - 'SQL_EXECUTION_ERROR', + 400, + Errors.GeneralErrorCodes.SQL_EXECUTION_ERROR, ); - res.end(); - return; } // Step 3: Generate summary using text-based answer @@ -351,9 +444,11 @@ export default async function handler( await wrenAIAdaptor.createTextBasedAnswer(textBasedAnswerInput); if (!summaryTask || !summaryTask.queryId) { - sendError(res, 'Failed to start summary generation task'); - res.end(); - return; + throw new ApiError( + 'Failed to start summary generation task', + 400, + Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + ); } // Poll for the summary generation result @@ -374,9 +469,11 @@ export default async function handler( // Check if we've exceeded the maximum wait time if (Date.now() > summaryDeadline) { - sendError(res, 'Summary generation timeout', 'TIMEOUT'); - res.end(); - return; + throw new ApiError( + 'Summary generation timeout', + 400, + Errors.GeneralErrorCodes.POLLING_TIMEOUT, + ); } // Wait before polling again @@ -428,15 +525,14 @@ export default async function handler( // Send content block stop sendContentBlockStop(res); } catch (_streamError) { - sendError(res, 'Error streaming summary content'); - res.end(); - return; + throw new ApiError( + 'Error streaming summary content', + 400, + Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + ); } } - // Send message end event - sendMessageStop(res, newThreadId, Date.now() - startTime); - // Log the API call await apiHistoryRepository.createOne({ id: uuidv4(), @@ -449,13 +545,10 @@ export default async function handler( durationMs: Date.now() - startTime, }); - res.end(); + endStream(res, newThreadId, startTime); } catch (error) { logger.error('Error in async ask API:', error); - // Send message end event even on error - sendMessageStop(res, threadId || uuidv4(), Date.now() - startTime); - // Log the error await apiHistoryRepository.createOne({ id: uuidv4(), @@ -474,6 +567,6 @@ export default async function handler( res, error instanceof Error ? error.message : 'Internal server error', ); - res.end(); + endStream(res, threadId || uuidv4(), startTime); } } From f219e6fe43f1e235b0f78fdd3ff514721d00dc7f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Thu, 26 Jun 2025 04:37:42 +0800 Subject: [PATCH 07/16] add generate_sql async api --- wren-ui/src/apollo/server/utils/index.ts | 1 + wren-ui/src/apollo/server/utils/sseUtils.ts | 126 +++++++++++ wren-ui/src/pages/api/v1/async/ask.ts | 119 +---------- .../src/pages/api/v1/async/generate_sql.ts | 197 ++++++++++++++++++ 4 files changed, 330 insertions(+), 113 deletions(-) create mode 100644 wren-ui/src/apollo/server/utils/sseUtils.ts create mode 100644 wren-ui/src/pages/api/v1/async/generate_sql.ts diff --git a/wren-ui/src/apollo/server/utils/index.ts b/wren-ui/src/apollo/server/utils/index.ts index ba31bd8f88..f9628a9dc3 100644 --- a/wren-ui/src/apollo/server/utils/index.ts +++ b/wren-ui/src/apollo/server/utils/index.ts @@ -7,3 +7,4 @@ export * from './model'; export * from './helper'; export * from './regex'; export * from './sseTypes'; +export * from './sseUtils'; diff --git a/wren-ui/src/apollo/server/utils/sseUtils.ts b/wren-ui/src/apollo/server/utils/sseUtils.ts new file mode 100644 index 0000000000..d4f7f37bbf --- /dev/null +++ b/wren-ui/src/apollo/server/utils/sseUtils.ts @@ -0,0 +1,126 @@ +import { NextApiResponse } from 'next'; +import { AskResultStatus } from '@/apollo/server/models/adaptor'; +import { + EventType, + StateType, + StreamEvent, + StateEvent, + ErrorEvent, + MessageStartEvent, + MessageStopEvent, +} from './sseTypes'; + +/** + * Send SSE event to client + */ +export const sendSSEEvent = (res: NextApiResponse, event: StreamEvent) => { + const eventData = `data: ${JSON.stringify(event)}\n\n`; + res.write(eventData); +}; + +/** + * Send message start event to client + */ +export const sendMessageStart = (res: NextApiResponse) => { + const messageStartEvent: MessageStartEvent = { + type: EventType.MESSAGE_START, + timestamp: Date.now(), + }; + sendSSEEvent(res, messageStartEvent); +}; + +/** + * Send message stop event to client + */ +export const sendMessageStop = ( + res: NextApiResponse, + threadId: string, + duration: number, +) => { + const messageStopEvent: MessageStopEvent = { + type: EventType.MESSAGE_STOP, + data: { + threadId, + duration, + }, + timestamp: Date.now(), + }; + sendSSEEvent(res, messageStopEvent); +}; + +/** + * Send state update to client + */ +export const sendStateUpdate = ( + res: NextApiResponse, + state: StateType, + data?: any, +) => { + const stateEvent: StateEvent = { + type: EventType.STATE, + data: { + state, + ...data, + }, + timestamp: Date.now(), + }; + sendSSEEvent(res, stateEvent); +}; + +/** + * Send error to client + */ +export const sendError = ( + res: NextApiResponse, + error: string, + code?: string, +) => { + const errorEvent: ErrorEvent = { + type: EventType.ERROR, + data: { + error, + code, + }, + timestamp: Date.now(), + }; + sendSSEEvent(res, errorEvent); +}; + +/** + * Transform AskResultStatus to descriptive SQL generation state + */ +export const getSqlGenerationState = (status: AskResultStatus): StateType => { + switch (status) { + case AskResultStatus.UNDERSTANDING: + return StateType.SQL_GENERATION_UNDERSTANDING; + case AskResultStatus.SEARCHING: + return StateType.SQL_GENERATION_SEARCHING; + case AskResultStatus.PLANNING: + return StateType.SQL_GENERATION_PLANNING; + case AskResultStatus.GENERATING: + return StateType.SQL_GENERATION_GENERATING; + case AskResultStatus.CORRECTING: + return StateType.SQL_GENERATION_CORRECTING; + case AskResultStatus.FINISHED: + return StateType.SQL_GENERATION_FINISHED; + case AskResultStatus.FAILED: + return StateType.SQL_GENERATION_FAILED; + case AskResultStatus.STOPPED: + return StateType.SQL_GENERATION_STOPPED; + default: + return StateType.SQL_GENERATION_UNDERSTANDING; + } +}; + +/** + * End the SSE stream with message stop event + */ +export const endStream = ( + res: NextApiResponse, + threadId: string, + startTime: number, +) => { + // Send message stop event + sendMessageStop(res, threadId, Date.now() - startTime); + res.end(); +}; diff --git a/wren-ui/src/pages/api/v1/async/ask.ts b/wren-ui/src/pages/api/v1/async/ask.ts index 388bb056ec..5bc9e5a0ff 100644 --- a/wren-ui/src/pages/api/v1/async/ask.ts +++ b/wren-ui/src/pages/api/v1/async/ask.ts @@ -26,14 +26,15 @@ import { StateType, ContentBlockContentType, AsyncAskRequest, - StreamEvent, - StateEvent, - ErrorEvent, - MessageStartEvent, - MessageStopEvent, ContentBlockStartEvent, ContentBlockDeltaEvent, ContentBlockStopEvent, + sendSSEEvent, + sendMessageStart, + sendStateUpdate, + sendError, + getSqlGenerationState, + endStream, } from '@/apollo/server/utils'; const logger = getLogger('API_ASYNC_ASK'); @@ -47,63 +48,6 @@ const { queryService, } = components; -/** - * Send SSE event to client - */ -const sendSSEEvent = (res: NextApiResponse, event: StreamEvent) => { - const eventData = `data: ${JSON.stringify(event)}\n\n`; - res.write(eventData); -}; - -/** - * Send message start event to client - */ -const sendMessageStart = (res: NextApiResponse) => { - const messageStartEvent: MessageStartEvent = { - type: EventType.MESSAGE_START, - timestamp: Date.now(), - }; - sendSSEEvent(res, messageStartEvent); -}; - -/** - * Send message stop event to client - */ -const sendMessageStop = ( - res: NextApiResponse, - threadId: string, - duration: number, -) => { - const messageStopEvent: MessageStopEvent = { - type: EventType.MESSAGE_STOP, - data: { - threadId, - duration, - }, - timestamp: Date.now(), - }; - sendSSEEvent(res, messageStopEvent); -}; - -/** - * Send state update to client - */ -const sendStateUpdate = ( - res: NextApiResponse, - state: StateType, - data?: any, -) => { - const stateEvent: StateEvent = { - type: EventType.STATE, - data: { - state, - ...data, - }, - timestamp: Date.now(), - }; - sendSSEEvent(res, stateEvent); -}; - /** * Send content block start event to client */ @@ -148,57 +92,6 @@ const sendContentBlockStop = (res: NextApiResponse) => { sendSSEEvent(res, contentBlockStopEvent); }; -/** - * Send error to client - */ -const sendError = (res: NextApiResponse, error: string, code?: string) => { - const errorEvent: ErrorEvent = { - type: EventType.ERROR, - data: { - error, - code, - }, - timestamp: Date.now(), - }; - sendSSEEvent(res, errorEvent); -}; - -/** - * Transform AskResultStatus to descriptive SQL generation state - */ -const getSqlGenerationState = (status: AskResultStatus): StateType => { - switch (status) { - case AskResultStatus.UNDERSTANDING: - return StateType.SQL_GENERATION_UNDERSTANDING; - case AskResultStatus.SEARCHING: - return StateType.SQL_GENERATION_SEARCHING; - case AskResultStatus.PLANNING: - return StateType.SQL_GENERATION_PLANNING; - case AskResultStatus.GENERATING: - return StateType.SQL_GENERATION_GENERATING; - case AskResultStatus.CORRECTING: - return StateType.SQL_GENERATION_CORRECTING; - case AskResultStatus.FINISHED: - return StateType.SQL_GENERATION_FINISHED; - case AskResultStatus.FAILED: - return StateType.SQL_GENERATION_FAILED; - case AskResultStatus.STOPPED: - return StateType.SQL_GENERATION_STOPPED; - default: - return StateType.SQL_GENERATION_UNDERSTANDING; - } -}; - -const endStream = ( - res: NextApiResponse, - threadId: string, - startTime: number, -) => { - // Send message stop event - sendMessageStop(res, threadId || uuidv4(), Date.now() - startTime); - res.end(); -}; - export default async function handler( req: NextApiRequest, res: NextApiResponse, diff --git a/wren-ui/src/pages/api/v1/async/generate_sql.ts b/wren-ui/src/pages/api/v1/async/generate_sql.ts new file mode 100644 index 0000000000..bbe768ff9b --- /dev/null +++ b/wren-ui/src/pages/api/v1/async/generate_sql.ts @@ -0,0 +1,197 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import * as Errors from '@/apollo/server/utils/error'; +import { v4 as uuidv4 } from 'uuid'; +import { + ApiError, + MAX_WAIT_TIME, + isAskResultFinished, + transformHistoryInput, + validateAskResult, +} from '@/apollo/server/utils/apiUtils'; +import { + AskResult, + AskResultStatus, + WrenAILanguage, +} from '@/apollo/server/models/adaptor'; +import { getLogger } from '@server/utils'; +import { + StateType, + AsyncAskRequest, + sendMessageStart, + sendStateUpdate, + sendError, + getSqlGenerationState, + endStream, +} from '@/apollo/server/utils'; + +const logger = getLogger('API_ASYNC_GENERATE_SQL'); +logger.level = 'debug'; + +const { apiHistoryRepository, projectService, deployService, wrenAIAdaptor } = + components; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const { question, language, threadId } = req.body as AsyncAskRequest; + const startTime = Date.now(); + let project; + + try { + project = await projectService.getCurrentProject(); + + // Only allow POST method + if (req.method !== 'POST') { + throw new ApiError('Method not allowed', 405); + } + + // Input validation + if (!question) { + throw new ApiError('Question is required', 400); + } + + // Set up SSE headers + res.setHeader('Content-Type', 'text/event-stream'); + res.setHeader('Cache-Control', 'no-cache, no-transform'); + res.setHeader('Connection', 'keep-alive'); + res.flushHeaders(); + + // Create a new thread if it's a new question + const newThreadId = threadId || uuidv4(); + + // Send message start event + sendMessageStart(res); + + // Get current project's last deployment + const lastDeploy = await deployService.getLastDeployment(project.id); + if (!lastDeploy) { + throw new ApiError( + 'No deployment found, please deploy your project first', + 400, + Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, + ); + } + + // Get conversation history if threadId is provided + const histories = threadId + ? await apiHistoryRepository.findAllBy({ threadId }) + : undefined; + + // Step 1: Generate SQL + sendStateUpdate(res, StateType.SQL_GENERATION_START, { + question, + threadId: newThreadId, + language: + language || WrenAILanguage[project.language] || WrenAILanguage.EN, + }); + + const askTask = await wrenAIAdaptor.ask({ + query: question, + deployId: lastDeploy.hash, + histories: transformHistoryInput(histories) as any, + configurations: { + language: + language || WrenAILanguage[project.language] || WrenAILanguage.EN, + }, + }); + + // Poll for the SQL generation result + const deadline = Date.now() + MAX_WAIT_TIME; + let askResult: AskResult; + let pollCount = 0; + let previousStatus: AskResultStatus | null = null; + + while (true) { + askResult = await wrenAIAdaptor.getAskResult(askTask.queryId); + + // Send status change updates when AskResultStatus changes + if (askResult.status !== previousStatus) { + const sqlGenerationState = getSqlGenerationState(askResult.status); + sendStateUpdate(res, sqlGenerationState, { + pollCount: pollCount + 1, + rephrasedQuestion: askResult.rephrasedQuestion, + intentReasoning: askResult.intentReasoning, + sqlGenerationReasoning: askResult.sqlGenerationReasoning, + retrievedTables: askResult.retrievedTables, + invalidSql: askResult.invalidSql, + traceId: askResult.traceId, + }); + previousStatus = askResult.status; + } + + pollCount++; + + // Check if the result is finished + if (isAskResultFinished(askResult)) { + break; + } + + // Check if we've exceeded the maximum wait time + if (Date.now() > deadline) { + throw new ApiError( + 'Request timeout', + 400, + Errors.GeneralErrorCodes.POLLING_TIMEOUT, + ); + } + + // Wait before polling again + await new Promise((resolve) => setTimeout(resolve, 1000)); + } + + // Validate the result + validateAskResult(askResult, askTask.queryId); + + // Get the generated SQL + const sql = askResult.response?.[0]?.sql; + if (!sql) { + throw new ApiError( + 'No SQL generated', + 400, + Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + ); + } + + // Send SQL generation end with the SQL content + sendStateUpdate(res, StateType.SQL_GENERATION_END, { sql }); + + // Log the API call + await apiHistoryRepository.createOne({ + id: uuidv4(), + projectId: project.id, + apiType: ApiType.GENERATE_SQL, + threadId: newThreadId, + requestPayload: { question, language }, + responsePayload: { sql }, + statusCode: 200, + durationMs: Date.now() - startTime, + }); + + endStream(res, newThreadId, startTime); + } catch (error) { + logger.error('Error in async generate SQL API:', error); + + // Log the error + await apiHistoryRepository.createOne({ + id: uuidv4(), + projectId: project?.id || 0, + apiType: ApiType.GENERATE_SQL, + threadId: threadId || uuidv4(), + requestPayload: { question, language }, + responsePayload: { + error: error instanceof Error ? error.message : String(error), + }, + statusCode: 500, + durationMs: Date.now() - startTime, + }); + + sendError( + res, + error instanceof Error ? error.message : 'Internal server error', + ); + endStream(res, threadId || uuidv4(), startTime); + } +} From 07ba5fa7cecb2bf2506d943cd78d6567748f8667 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Thu, 26 Jun 2025 11:57:24 +0800 Subject: [PATCH 08/16] fix ask to answer general question --- wren-ui/src/pages/api/v1/ask.ts | 74 ++++++++++++++++++++++++++++++++- 1 file changed, 72 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/pages/api/v1/ask.ts b/wren-ui/src/pages/api/v1/ask.ts index bf7173f72f..41586f46c1 100644 --- a/wren-ui/src/pages/api/v1/ask.ts +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -9,7 +9,6 @@ import { handleApiError, MAX_WAIT_TIME, isAskResultFinished, - validateAskResult, validateSummaryResult, transformHistoryInput, } from '@/apollo/server/utils/apiUtils'; @@ -19,6 +18,8 @@ import { TextBasedAnswerInput, TextBasedAnswerResult, TextBasedAnswerStatus, + AskResultType, + WrenAIError, } from '@/apollo/server/models/adaptor'; import { getLogger } from '@server/utils'; @@ -111,7 +112,76 @@ export default async function handler( } // Validate the AI result - validateAskResult(askResult, askTask.queryId); + // Check for error in result + if (askResult.error) { + const errorMessage = + (askResult.error as WrenAIError).message || 'Unknown error'; + const additionalData: Record = {}; + + // Include invalid SQL if available + if (askResult.invalidSql) { + additionalData.invalidSql = askResult.invalidSql; + } + + throw new ApiError( + errorMessage, + 400, + askResult.error.code, + additionalData, + ); + } + + // Check for general type response (explanation streaming) + if (askResult.type === AskResultType.GENERAL) { + // Stream the explanation content + let explanation = ''; + const stream = await wrenAIAdaptor.getAskStreamingResult(askTask.queryId); + + // Collect the streamed content + const streamPromise = new Promise((resolve, reject) => { + stream.on('data', (chunk) => { + const chunkString = chunk.toString('utf-8'); + const match = chunkString.match(/data: {"message":"([\s\S]*?)"}/); + if (match && match[1]) { + explanation += match[1]; + } + }); + + stream.on('end', () => { + resolve(); + }); + + stream.on('error', (error) => { + reject(error); + }); + + // Handle client disconnect + req.on('close', () => { + stream.destroy(); + reject(new Error('Client disconnected')); + }); + }); + + await streamPromise; + + // Return the explanation result + await respondWith({ + res, + statusCode: 200, + responsePayload: { + type: 'NON_SQL_QUERY', + explanation, + threadId: newThreadId, + }, + projectId: project.id, + apiType: ApiType.ASK, + startTime, + requestPayload: req.body, + threadId: newThreadId, + headers: req.headers as Record, + }); + return; + } // Get the generated SQL const sql = askResult.response?.[0]?.sql; From 7e2deec2790aff0c29c47d1ae2426a657fb02519 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Thu, 26 Jun 2025 14:07:06 +0800 Subject: [PATCH 09/16] add models api --- wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index 98c8062fd5..3871fbb520 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -22,6 +22,7 @@ export enum ApiType { CREATE_SQL_PAIR = 'CREATE_SQL_PAIR', UPDATE_SQL_PAIR = 'UPDATE_SQL_PAIR', DELETE_SQL_PAIR = 'DELETE_SQL_PAIR', + GET_MODELS = 'GET_MODELS', } export interface ApiHistory { From 10b584cf1c10a0600c1348016671c809ebc4db01 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Thu, 26 Jun 2025 14:07:19 +0800 Subject: [PATCH 10/16] models api --- wren-ui/src/pages/api/v1/models.ts | 76 ++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 wren-ui/src/pages/api/v1/models.ts diff --git a/wren-ui/src/pages/api/v1/models.ts b/wren-ui/src/pages/api/v1/models.ts new file mode 100644 index 0000000000..0f86e41a38 --- /dev/null +++ b/wren-ui/src/pages/api/v1/models.ts @@ -0,0 +1,76 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { components } from '@/common'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import * as Errors from '@/apollo/server/utils/error'; +import { + ApiError, + respondWithSimple, + handleApiError, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('API_MODELS'); +logger.level = 'debug'; + +const { projectService, deployService } = components; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + let project; + + try { + project = await projectService.getCurrentProject(); + + // Only allow GET method + if (req.method !== 'GET') { + throw new ApiError('Method not allowed', 405); + } + + // Get current project's last deployment + const lastDeploy = await deployService.getLastDeployment(project.id); + if (!lastDeploy) { + throw new ApiError( + 'No deployment found, please deploy your project first', + 400, + Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, + ); + } + + // Get the MDL from the deployment manifest + const mdl = lastDeploy.manifest as any; + + // Extract models, views, and relationships from the MDL with defaults + const models = mdl?.models || []; + const views = mdl?.views || []; + const relationships = mdl?.relationships || []; + + // Return the restructured response + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { + hash: lastDeploy.hash, + models, + relationships, + views, + }, + projectId: project.id, + apiType: ApiType.GET_MODELS, + startTime, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: project?.id, + apiType: ApiType.GET_MODELS, + headers: req.headers as Record, + startTime, + logger, + }); + } +} From f9dec0a837fdc110b470c0fa7f31af1347b794c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Thu, 26 Jun 2025 15:10:31 +0800 Subject: [PATCH 11/16] fix issues --- .../src/apollo/client/graphql/__types__.ts | 25 ++++++++++++++----- .../client/graphql/dashboard.generated.ts | 2 +- .../repositories/apiHistoryRepository.ts | 2 ++ wren-ui/src/apollo/server/schema.ts | 7 ++++++ wren-ui/src/pages/api-management/history.tsx | 4 ++- wren-ui/src/pages/api/v1/async/ask.ts | 9 ++++--- .../src/pages/api/v1/async/generate_sql.ts | 6 +++-- 7 files changed, 42 insertions(+), 13 deletions(-) diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index 3dfb3608d1..2bc93258d8 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/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; day?: Maybe; - frequency: ScheduleFrequencyEnum; - hour: Scalars['Int']; - minute: Scalars['Int']; + frequency?: Maybe; + hour?: Maybe; + minute?: Maybe; timezone?: Maybe; }; @@ -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; name: Scalars['String']; nextScheduledAt?: Maybe; - schedule: DashboardSchedule; + schedule?: Maybe; }; export type DetailedModel = { diff --git a/wren-ui/src/apollo/client/graphql/dashboard.generated.ts b/wren-ui/src/apollo/client/graphql/dashboard.generated.ts index 141a1d852e..1014a11ee8 100644 --- a/wren-ui/src/apollo/client/graphql/dashboard.generated.ts +++ b/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 { diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index 3871fbb520..36a77f4bc6 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -23,6 +23,8 @@ export enum ApiType { 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 { diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 87fca75072..e70c311b9f 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -14,6 +14,13 @@ export const typeDefs = gql` 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 { diff --git a/wren-ui/src/pages/api-management/history.tsx b/wren-ui/src/pages/api-management/history.tsx index 9c5824c178..42fd0b7681 100644 --- a/wren-ui/src/pages/api-management/history.tsx +++ b/wren-ui/src/pages/api-management/history.tsx @@ -99,7 +99,9 @@ export default function APIHistory() { ); } return ( -
{payload.question || payload.sql || '-'}
+
+ {payload?.question || payload?.sql || '-'} +
); }, }, diff --git a/wren-ui/src/pages/api/v1/async/ask.ts b/wren-ui/src/pages/api/v1/async/ask.ts index 5bc9e5a0ff..b2ca1aa8b9 100644 --- a/wren-ui/src/pages/api/v1/async/ask.ts +++ b/wren-ui/src/pages/api/v1/async/ask.ts @@ -272,8 +272,9 @@ export default async function handler( await apiHistoryRepository.createOne({ id: uuidv4(), projectId: project.id, - apiType: ApiType.ASK, + apiType: ApiType.ASYNC_ASK, threadId: newThreadId, + headers: req.headers as Record, requestPayload: { question, sampleSize, language }, responsePayload: { type: 'general', @@ -430,8 +431,9 @@ export default async function handler( await apiHistoryRepository.createOne({ id: uuidv4(), projectId: project.id, - apiType: ApiType.ASK, + apiType: ApiType.ASYNC_ASK, threadId: newThreadId, + headers: req.headers as Record, requestPayload: { question, sampleSize, language }, responsePayload: { sql, summary }, statusCode: 200, @@ -446,8 +448,9 @@ export default async function handler( await apiHistoryRepository.createOne({ id: uuidv4(), projectId: project?.id || 0, - apiType: ApiType.ASK, + apiType: ApiType.ASYNC_ASK, threadId: threadId || uuidv4(), + headers: req.headers as Record, requestPayload: { question, sampleSize, language }, responsePayload: { error: error instanceof Error ? error.message : String(error), diff --git a/wren-ui/src/pages/api/v1/async/generate_sql.ts b/wren-ui/src/pages/api/v1/async/generate_sql.ts index bbe768ff9b..4c09d1c60e 100644 --- a/wren-ui/src/pages/api/v1/async/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/async/generate_sql.ts @@ -162,8 +162,9 @@ export default async function handler( await apiHistoryRepository.createOne({ id: uuidv4(), projectId: project.id, - apiType: ApiType.GENERATE_SQL, + apiType: ApiType.ASYNC_GENERATE_SQL, threadId: newThreadId, + headers: req.headers as Record, requestPayload: { question, language }, responsePayload: { sql }, statusCode: 200, @@ -178,8 +179,9 @@ export default async function handler( await apiHistoryRepository.createOne({ id: uuidv4(), projectId: project?.id || 0, - apiType: ApiType.GENERATE_SQL, + apiType: ApiType.ASYNC_GENERATE_SQL, threadId: threadId || uuidv4(), + headers: req.headers as Record, requestPayload: { question, language }, responsePayload: { error: error instanceof Error ? error.message : String(error), From abec81d1f58fd6230a5a0fd235395c094fc126e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Thu, 26 Jun 2025 19:00:05 +0800 Subject: [PATCH 12/16] enhance history --- wren-ui/src/apollo/server/utils/apiUtils.ts | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/wren-ui/src/apollo/server/utils/apiUtils.ts b/wren-ui/src/apollo/server/utils/apiUtils.ts index 5c84817403..813fff0e4d 100644 --- a/wren-ui/src/apollo/server/utils/apiUtils.ts +++ b/wren-ui/src/apollo/server/utils/apiUtils.ts @@ -96,10 +96,20 @@ export const transformHistoryInput = (histories: ApiHistory[]) => { if (!histories) { return []; } + + const validApiTypes = [ + ApiType.GENERATE_SQL, + ApiType.ASK, + ApiType.ASYNC_GENERATE_SQL, + ApiType.ASYNC_ASK, + ]; + return histories .filter( (history) => - history.responsePayload?.sql && history.requestPayload?.question, + validApiTypes.includes(history.apiType) && + history.responsePayload?.sql && + history.requestPayload?.question, ) .map((history) => ({ question: history.requestPayload?.question, From c0fe8af2a749f54ef2d2550407ebea41603ac2e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Fri, 27 Jun 2025 02:10:25 +0800 Subject: [PATCH 13/16] fix some issues --- .../server/resolvers/apiHistoryResolver.ts | 6 ++++++ wren-ui/src/apollo/server/utils/sseTypes.ts | 2 +- wren-ui/src/apollo/server/utils/sseUtils.ts | 2 ++ wren-ui/src/pages/api/v1/ask.ts | 2 +- wren-ui/src/pages/api/v1/async/ask.ts | 16 +++++++++++----- wren-ui/src/pages/api/v1/async/generate_sql.ts | 6 ++++-- 6 files changed, 25 insertions(+), 9 deletions(-) diff --git a/wren-ui/src/apollo/server/resolvers/apiHistoryResolver.ts b/wren-ui/src/apollo/server/resolvers/apiHistoryResolver.ts index 8fcfc525f7..589850b3a3 100644 --- a/wren-ui/src/apollo/server/resolvers/apiHistoryResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/apiHistoryResolver.ts @@ -147,6 +147,12 @@ export class ApiHistoryResolver { }, responsePayload: (apiHistory: ApiHistory) => { if (!apiHistory.responsePayload) return null; + + // If the response payload is an array, return it as is + if (Array.isArray(apiHistory.responsePayload)) + return apiHistory.responsePayload; + + // Otherwise, sanitize the response payload return sanitizeResponsePayload( apiHistory.responsePayload, apiHistory.apiType, diff --git a/wren-ui/src/apollo/server/utils/sseTypes.ts b/wren-ui/src/apollo/server/utils/sseTypes.ts index 5107eba118..654047b631 100644 --- a/wren-ui/src/apollo/server/utils/sseTypes.ts +++ b/wren-ui/src/apollo/server/utils/sseTypes.ts @@ -19,7 +19,7 @@ export enum StateType { SQL_GENERATION_FINISHED = 'sql_generation_finished', SQL_GENERATION_FAILED = 'sql_generation_failed', SQL_GENERATION_STOPPED = 'sql_generation_stopped', - SQL_GENERATION_END = 'sql_generation_end', + SQL_GENERATION_SUCCESS = 'sql_generation_success', SQL_EXECUTION_START = 'sql_execution_start', SQL_EXECUTION_END = 'sql_execution_end', } diff --git a/wren-ui/src/apollo/server/utils/sseUtils.ts b/wren-ui/src/apollo/server/utils/sseUtils.ts index d4f7f37bbf..36dd96936d 100644 --- a/wren-ui/src/apollo/server/utils/sseUtils.ts +++ b/wren-ui/src/apollo/server/utils/sseUtils.ts @@ -74,12 +74,14 @@ export const sendError = ( res: NextApiResponse, error: string, code?: string, + additionalData?: Record, ) => { const errorEvent: ErrorEvent = { type: EventType.ERROR, data: { error, code, + ...additionalData, }, timestamp: Date.now(), }; diff --git a/wren-ui/src/pages/api/v1/ask.ts b/wren-ui/src/pages/api/v1/ask.ts index 41586f46c1..d2e98d0eb2 100644 --- a/wren-ui/src/pages/api/v1/ask.ts +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -126,7 +126,7 @@ export default async function handler( throw new ApiError( errorMessage, 400, - askResult.error.code, + askResult.error?.code || Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, additionalData, ); } diff --git a/wren-ui/src/pages/api/v1/async/ask.ts b/wren-ui/src/pages/api/v1/async/ask.ts index b2ca1aa8b9..f360789da8 100644 --- a/wren-ui/src/pages/api/v1/async/ask.ts +++ b/wren-ui/src/pages/api/v1/async/ask.ts @@ -231,12 +231,14 @@ export default async function handler( const stream = await wrenAIAdaptor.getAskStreamingResult(askTask.queryId); // Stream the content in real-time + let explanation = ''; const streamPromise = new Promise((resolve, reject) => { stream.on('data', (chunk) => { const chunkString = chunk.toString('utf-8'); const match = chunkString.match(/data: {"message":"([\s\S]*?)"}/); if (match && match[1]) { // Send incremental content updates + explanation += match[1]; sendContentBlockDelta(res, match[1]); } }); @@ -277,8 +279,7 @@ export default async function handler( headers: req.headers as Record, requestPayload: { question, sampleSize, language }, responsePayload: { - type: 'general', - explanation: 'Streamed explanation', + explanation, }, statusCode: 200, durationMs: Date.now() - startTime, @@ -298,8 +299,8 @@ export default async function handler( ); } - // Send SQL generation end with the SQL content - sendStateUpdate(res, StateType.SQL_GENERATION_END, { sql }); + // Send SQL generation success with the SQL content + sendStateUpdate(res, StateType.SQL_GENERATION_SUCCESS, { sql }); // Step 2: Execute SQL to get data sendStateUpdate(res, StateType.SQL_EXECUTION_START, { sql }); @@ -435,7 +436,10 @@ export default async function handler( threadId: newThreadId, headers: req.headers as Record, requestPayload: { question, sampleSize, language }, - responsePayload: { sql, summary }, + responsePayload: { + sql, + summary, + }, statusCode: 200, durationMs: Date.now() - startTime, }); @@ -462,6 +466,8 @@ export default async function handler( sendError( res, error instanceof Error ? error.message : 'Internal server error', + error.code || Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + error.additionalData, ); endStream(res, threadId || uuidv4(), startTime); } diff --git a/wren-ui/src/pages/api/v1/async/generate_sql.ts b/wren-ui/src/pages/api/v1/async/generate_sql.ts index 4c09d1c60e..7589299073 100644 --- a/wren-ui/src/pages/api/v1/async/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/async/generate_sql.ts @@ -155,8 +155,8 @@ export default async function handler( ); } - // Send SQL generation end with the SQL content - sendStateUpdate(res, StateType.SQL_GENERATION_END, { sql }); + // Send SQL generation success with the SQL content + sendStateUpdate(res, StateType.SQL_GENERATION_SUCCESS, { sql }); // Log the API call await apiHistoryRepository.createOne({ @@ -193,6 +193,8 @@ export default async function handler( sendError( res, error instanceof Error ? error.message : 'Internal server error', + error.code || Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + error.additionalData, ); endStream(res, threadId || uuidv4(), startTime); } From 28118f9d66fbef2a419592b565ec6d940d1f207e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Fri, 27 Jun 2025 02:12:31 +0800 Subject: [PATCH 14/16] fix ui --- wren-ui/src/pages/api-management/history.tsx | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/wren-ui/src/pages/api-management/history.tsx b/wren-ui/src/pages/api-management/history.tsx index 42fd0b7681..b0468c0e9b 100644 --- a/wren-ui/src/pages/api-management/history.tsx +++ b/wren-ui/src/pages/api-management/history.tsx @@ -70,8 +70,12 @@ export default function APIHistory() { width: 100, render: (status: number) => { const icon = - status === 200 ? : ; - const color = status === 200 ? 'success' : 'error'; + status >= 200 && status < 300 ? ( + + ) : ( + + ); + const color = status >= 200 && status < 300 ? 'success' : 'error'; return ( {status} From 6d3cbf972a05fc4181946e1902e55fa30378422a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Fri, 27 Jun 2025 02:14:54 +0800 Subject: [PATCH 15/16] change back env example --- docker/.env.example | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/.env.example b/docker/.env.example index fe9a96e9c5..7e511d4a15 100644 --- a/docker/.env.example +++ b/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=custom-instruction-v2 +WREN_AI_SERVICE_VERSION=0.23.2 IBIS_SERVER_VERSION=0.16.1 WREN_UI_VERSION=0.28.0 WREN_BOOTSTRAP_VERSION=0.1.5 From c1d1dc044b3225661054691fab5ab2f97b806022 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?william=20chang=28=E5=BC=B5=E4=BB=B2=E5=A8=81=29?= Date: Fri, 27 Jun 2025 02:22:48 +0800 Subject: [PATCH 16/16] rename to stream --- wren-ui/src/apollo/client/graphql/__types__.ts | 4 ++-- .../apollo/server/repositories/apiHistoryRepository.ts | 4 ++-- wren-ui/src/apollo/server/schema.ts | 4 ++-- wren-ui/src/apollo/server/utils/apiUtils.ts | 4 ++-- wren-ui/src/pages/api/v1/{async => stream}/ask.ts | 10 +++++----- .../src/pages/api/v1/{async => stream}/generate_sql.ts | 8 ++++---- 6 files changed, 17 insertions(+), 17 deletions(-) rename wren-ui/src/pages/api/v1/{async => stream}/ask.ts (98%) rename wren-ui/src/pages/api/v1/{async => stream}/generate_sql.ts (96%) diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index 2bc93258d8..198d9adf21 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/wren-ui/src/apollo/client/graphql/__types__.ts @@ -78,8 +78,6 @@ 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', @@ -91,6 +89,8 @@ export enum ApiType { GET_MODELS = 'GET_MODELS', GET_SQL_PAIRS = 'GET_SQL_PAIRS', RUN_SQL = 'RUN_SQL', + STREAM_ASK = 'STREAM_ASK', + STREAM_GENERATE_SQL = 'STREAM_GENERATE_SQL', UPDATE_INSTRUCTION = 'UPDATE_INSTRUCTION', UPDATE_SQL_PAIR = 'UPDATE_SQL_PAIR' } diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index 36a77f4bc6..7de5e5c097 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -23,8 +23,8 @@ export enum ApiType { 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', + STREAM_ASK = 'STREAM_ASK', + STREAM_GENERATE_SQL = 'STREAM_GENERATE_SQL', } export interface ApiHistory { diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index e70c311b9f..6174baed74 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -19,8 +19,8 @@ export const typeDefs = gql` DELETE_SQL_PAIR GET_SQL_PAIRS GET_MODELS - ASYNC_ASK - ASYNC_GENERATE_SQL + STREAM_ASK + STREAM_GENERATE_SQL } input ApiHistoryFilterInput { diff --git a/wren-ui/src/apollo/server/utils/apiUtils.ts b/wren-ui/src/apollo/server/utils/apiUtils.ts index 813fff0e4d..a30b1ab372 100644 --- a/wren-ui/src/apollo/server/utils/apiUtils.ts +++ b/wren-ui/src/apollo/server/utils/apiUtils.ts @@ -100,8 +100,8 @@ export const transformHistoryInput = (histories: ApiHistory[]) => { const validApiTypes = [ ApiType.GENERATE_SQL, ApiType.ASK, - ApiType.ASYNC_GENERATE_SQL, - ApiType.ASYNC_ASK, + ApiType.STREAM_GENERATE_SQL, + ApiType.STREAM_ASK, ]; return histories diff --git a/wren-ui/src/pages/api/v1/async/ask.ts b/wren-ui/src/pages/api/v1/stream/ask.ts similarity index 98% rename from wren-ui/src/pages/api/v1/async/ask.ts rename to wren-ui/src/pages/api/v1/stream/ask.ts index f360789da8..12e53eb4db 100644 --- a/wren-ui/src/pages/api/v1/async/ask.ts +++ b/wren-ui/src/pages/api/v1/stream/ask.ts @@ -37,7 +37,7 @@ import { endStream, } from '@/apollo/server/utils'; -const logger = getLogger('API_ASYNC_ASK'); +const logger = getLogger('API_STREAM_ASK'); logger.level = 'debug'; const { @@ -274,7 +274,7 @@ export default async function handler( await apiHistoryRepository.createOne({ id: uuidv4(), projectId: project.id, - apiType: ApiType.ASYNC_ASK, + apiType: ApiType.STREAM_ASK, threadId: newThreadId, headers: req.headers as Record, requestPayload: { question, sampleSize, language }, @@ -432,7 +432,7 @@ export default async function handler( await apiHistoryRepository.createOne({ id: uuidv4(), projectId: project.id, - apiType: ApiType.ASYNC_ASK, + apiType: ApiType.STREAM_ASK, threadId: newThreadId, headers: req.headers as Record, requestPayload: { question, sampleSize, language }, @@ -446,13 +446,13 @@ export default async function handler( endStream(res, newThreadId, startTime); } catch (error) { - logger.error('Error in async ask API:', error); + logger.error('Error in stream ask API:', error); // Log the error await apiHistoryRepository.createOne({ id: uuidv4(), projectId: project?.id || 0, - apiType: ApiType.ASYNC_ASK, + apiType: ApiType.STREAM_ASK, threadId: threadId || uuidv4(), headers: req.headers as Record, requestPayload: { question, sampleSize, language }, diff --git a/wren-ui/src/pages/api/v1/async/generate_sql.ts b/wren-ui/src/pages/api/v1/stream/generate_sql.ts similarity index 96% rename from wren-ui/src/pages/api/v1/async/generate_sql.ts rename to wren-ui/src/pages/api/v1/stream/generate_sql.ts index 7589299073..6102592657 100644 --- a/wren-ui/src/pages/api/v1/async/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/stream/generate_sql.ts @@ -26,7 +26,7 @@ import { endStream, } from '@/apollo/server/utils'; -const logger = getLogger('API_ASYNC_GENERATE_SQL'); +const logger = getLogger('API_STREAM_GENERATE_SQL'); logger.level = 'debug'; const { apiHistoryRepository, projectService, deployService, wrenAIAdaptor } = @@ -162,7 +162,7 @@ export default async function handler( await apiHistoryRepository.createOne({ id: uuidv4(), projectId: project.id, - apiType: ApiType.ASYNC_GENERATE_SQL, + apiType: ApiType.STREAM_GENERATE_SQL, threadId: newThreadId, headers: req.headers as Record, requestPayload: { question, language }, @@ -173,13 +173,13 @@ export default async function handler( endStream(res, newThreadId, startTime); } catch (error) { - logger.error('Error in async generate SQL API:', error); + logger.error('Error in stream generate SQL API:', error); // Log the error await apiHistoryRepository.createOne({ id: uuidv4(), projectId: project?.id || 0, - apiType: ApiType.ASYNC_GENERATE_SQL, + apiType: ApiType.STREAM_GENERATE_SQL, threadId: threadId || uuidv4(), headers: req.headers as Record, requestPayload: { question, language },