diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index 3dfb3608d1..198d9adf21 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', + 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', + STREAM_ASK = 'STREAM_ASK', + STREAM_GENERATE_SQL = 'STREAM_GENERATE_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 9e2164bc9e..7de5e5c097 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'; @@ -6,6 +12,19 @@ export enum ApiType { GENERATE_SQL = 'GENERATE_SQL', RUN_SQL = 'RUN_SQL', GENERATE_VEGA_CHART = 'GENERATE_VEGA_CHART', + GENERATE_SUMMARY = 'GENERATE_SUMMARY', + ASK = 'ASK', + GET_INSTRUCTIONS = 'GET_INSTRUCTIONS', + CREATE_INSTRUCTION = 'CREATE_INSTRUCTION', + UPDATE_INSTRUCTION = 'UPDATE_INSTRUCTION', + DELETE_INSTRUCTION = 'DELETE_INSTRUCTION', + GET_SQL_PAIRS = 'GET_SQL_PAIRS', + CREATE_SQL_PAIR = 'CREATE_SQL_PAIR', + UPDATE_SQL_PAIR = 'UPDATE_SQL_PAIR', + DELETE_SQL_PAIR = 'DELETE_SQL_PAIR', + GET_MODELS = 'GET_MODELS', + STREAM_ASK = 'STREAM_ASK', + STREAM_GENERATE_SQL = 'STREAM_GENERATE_SQL', } export interface ApiHistory { @@ -147,6 +166,20 @@ export class ApiHistoryRepository return formattedData; }; + protected override transformToDBData = (data: any) => { + if (!isPlainObject(data)) { + throw new Error('Unexpected dbdata'); + } + const transformedData = mapValues(data, (value, key) => { + if (this.jsonbColumns.includes(key)) { + return JSON.stringify(value); + } else { + return value; + } + }); + return mapKeys(transformedData, (_value, key) => snakeCase(key)); + }; + /** * Convert camelCase to snake_case for DB column names */ 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/schema.ts b/wren-ui/src/apollo/server/schema.ts index 5590055dbc..6174baed74 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -8,6 +8,19 @@ export const typeDefs = gql` GENERATE_SQL RUN_SQL GENERATE_VEGA_CHART + GENERATE_SUMMARY + ASK + GET_INSTRUCTIONS + CREATE_INSTRUCTION + UPDATE_INSTRUCTION + DELETE_INSTRUCTION + CREATE_SQL_PAIR + UPDATE_SQL_PAIR + DELETE_SQL_PAIR + GET_SQL_PAIRS + GET_MODELS + STREAM_ASK + STREAM_GENERATE_SQL } 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 03ec365589..a30b1ab372 100644 --- a/wren-ui/src/apollo/server/utils/apiUtils.ts +++ b/wren-ui/src/apollo/server/utils/apiUtils.ts @@ -1,11 +1,153 @@ 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 []; + } + + const validApiTypes = [ + ApiType.GENERATE_SQL, + ApiType.ASK, + ApiType.STREAM_GENERATE_SQL, + ApiType.STREAM_ASK, + ]; + + return histories + .filter( + (history) => + validApiTypes.includes(history.apiType) && + history.responsePayload?.sql && + history.requestPayload?.question, + ) + .map((history) => ({ + question: history.requestPayload?.question, + sql: history.responsePayload?.sql, + })); +}; + +/** + * Validates SQL syntax and compatibility with the project's manifest. + * Throws an ApiError if the SQL is invalid or cannot be previewed. + * @param sql The SQL string to validate + * @param project The project object (must have id) + * @param deployService The deployment service instance + * @param queryService The query service instance + */ +export const validateSql = async ( + sql: string, + project: any, + deployService: any, + queryService: any, +) => { + const lastDeployment = await deployService.getLastDeployment(project.id); + const manifest = lastDeployment.manifest; + try { + await queryService.preview(sql, { + manifest, + project, + dryRun: true, + }); + } catch (err: any) { + throw new ApiError( + err.message || 'Invalid SQL', + 400, + Errors.GeneralErrorCodes.INVALID_SQL_ERROR, + ); + } +}; + /** * Common error class for API endpoints */ @@ -71,6 +213,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/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/index.ts b/wren-ui/src/apollo/server/utils/index.ts index 8db524806d..f9628a9dc3 100644 --- a/wren-ui/src/apollo/server/utils/index.ts +++ b/wren-ui/src/apollo/server/utils/index.ts @@ -6,3 +6,5 @@ export * from './docker'; export * from './model'; export * from './helper'; export * from './regex'; +export * from './sseTypes'; +export * from './sseUtils'; 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..654047b631 --- /dev/null +++ b/wren-ui/src/apollo/server/utils/sseTypes.ts @@ -0,0 +1,99 @@ +// 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_SUCCESS = 'sql_generation_success', + SQL_EXECUTION_START = 'sql_execution_start', + SQL_EXECUTION_END = 'sql_execution_end', +} + +export enum ContentBlockContentType { + SUMMARY_GENERATION = 'summary_generation', + EXPLANATION = 'explanation', +} + +// 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/apollo/server/utils/sseUtils.ts b/wren-ui/src/apollo/server/utils/sseUtils.ts new file mode 100644 index 0000000000..36dd96936d --- /dev/null +++ b/wren-ui/src/apollo/server/utils/sseUtils.ts @@ -0,0 +1,128 @@ +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, + additionalData?: Record, +) => { + const errorEvent: ErrorEvent = { + type: EventType.ERROR, + data: { + error, + code, + ...additionalData, + }, + 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-management/history.tsx b/wren-ui/src/pages/api-management/history.tsx index 9c5824c178..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} @@ -99,7 +103,9 @@ export default function APIHistory() { ); } return ( -
{payload.question || payload.sql || '-'}
+
+ {payload?.question || payload?.sql || '-'} +
); }, }, 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..d2e98d0eb2 --- /dev/null +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -0,0 +1,321 @@ +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, + validateSummaryResult, + transformHistoryInput, +} from '@/apollo/server/utils/apiUtils'; +import { + AskResult, + WrenAILanguage, + TextBasedAnswerInput, + TextBasedAnswerResult, + TextBasedAnswerStatus, + AskResultType, + WrenAIError, +} 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 + // 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 || Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + 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; + 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 new file mode 100644 index 0000000000..2ca538ce51 --- /dev/null +++ b/wren-ui/src/pages/api/v1/generate_summary.ts @@ -0,0 +1,199 @@ +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, + validateSummaryResult, +} 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; + +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, + }); + } +} 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, + }); + } +} 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, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/stream/ask.ts b/wren-ui/src/pages/api/v1/stream/ask.ts new file mode 100644 index 0000000000..12e53eb4db --- /dev/null +++ b/wren-ui/src/pages/api/v1/stream/ask.ts @@ -0,0 +1,474 @@ +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, + validateSummaryResult, + transformHistoryInput, +} from '@/apollo/server/utils/apiUtils'; +import { + AskResult, + AskResultStatus, + WrenAILanguage, + TextBasedAnswerInput, + TextBasedAnswerResult, + TextBasedAnswerStatus, + WrenAIError, + AskResultType, +} from '@/apollo/server/models/adaptor'; +import { getLogger } from '@server/utils'; +import { + EventType, + StateType, + ContentBlockContentType, + AsyncAskRequest, + ContentBlockStartEvent, + ContentBlockDeltaEvent, + ContentBlockStopEvent, + sendSSEEvent, + sendMessageStart, + sendStateUpdate, + sendError, + getSqlGenerationState, + endStream, +} from '@/apollo/server/utils'; + +const logger = getLogger('API_STREAM_ASK'); +logger.level = 'debug'; + +const { + apiHistoryRepository, + projectService, + deployService, + wrenAIAdaptor, + queryService, +} = components; + +/** + * 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 block stop event to client + */ +const sendContentBlockStop = (res: NextApiResponse) => { + const contentBlockStopEvent: ContentBlockStopEvent = { + type: EventType.CONTENT_BLOCK_STOP, + timestamp: Date.now(), + }; + sendSSEEvent(res, contentBlockStopEvent); +}; + +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 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, + ); + } + + // 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 + 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 ask result + // 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 + 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]); + } + }); + + 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.STREAM_ASK, + threadId: newThreadId, + headers: req.headers as Record, + requestPayload: { question, sampleSize, language }, + responsePayload: { + explanation, + }, + statusCode: 200, + durationMs: Date.now() - startTime, + }); + + endStream(res, newThreadId, startTime); + return; + } + + // 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 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 }); + let sqlData; + try { + const queryResult = await queryService.preview(sql, { + project, + limit: sampleSize || 500, + manifest: lastDeploy.manifest, + modelingOnly: false, + }); + sqlData = queryResult; + sendStateUpdate(res, StateType.SQL_EXECUTION_END); + } catch (queryError) { + throw new ApiError( + `SQL execution failed: ${queryError.message || 'Unknown error'}`, + 400, + Errors.GeneralErrorCodes.SQL_EXECUTION_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', + 400, + Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + ); + } + + // Poll for the summary generation result + const summaryDeadline = Date.now() + MAX_WAIT_TIME; + let summaryResult: TextBasedAnswerResult; + + while (true) { + summaryResult = await wrenAIAdaptor.getTextBasedAnswerResult( + summaryTask.queryId, + ); + + if ( + summaryResult.status === TextBasedAnswerStatus.SUCCEEDED || + summaryResult.status === TextBasedAnswerStatus.FAILED + ) { + break; + } + + // Check if we've exceeded the maximum wait time + if (Date.now() > summaryDeadline) { + throw new ApiError( + 'Summary generation timeout', + 400, + Errors.GeneralErrorCodes.POLLING_TIMEOUT, + ); + } + + // Wait before polling again + await new Promise((resolve) => setTimeout(resolve, 1000)); + } + + // Validate the summary result + validateSummaryResult(summaryResult); + + // Step 4: Stream the content to get the summary + let summary = ''; + if (summaryResult.status === TextBasedAnswerStatus.SUCCEEDED) { + // Send content block start + sendContentBlockStart(res, ContentBlockContentType.SUMMARY_GENERATION); + + 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 + 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 summary content', + 400, + Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + ); + } + } + + // Log the API call + await apiHistoryRepository.createOne({ + id: uuidv4(), + projectId: project.id, + apiType: ApiType.STREAM_ASK, + threadId: newThreadId, + headers: req.headers as Record, + requestPayload: { question, sampleSize, language }, + responsePayload: { + sql, + summary, + }, + statusCode: 200, + durationMs: Date.now() - startTime, + }); + + endStream(res, newThreadId, startTime); + } catch (error) { + logger.error('Error in stream ask API:', error); + + // Log the error + await apiHistoryRepository.createOne({ + id: uuidv4(), + projectId: project?.id || 0, + apiType: ApiType.STREAM_ASK, + threadId: threadId || uuidv4(), + headers: req.headers as Record, + 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', + error.code || Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + error.additionalData, + ); + endStream(res, threadId || uuidv4(), startTime); + } +} diff --git a/wren-ui/src/pages/api/v1/stream/generate_sql.ts b/wren-ui/src/pages/api/v1/stream/generate_sql.ts new file mode 100644 index 0000000000..6102592657 --- /dev/null +++ b/wren-ui/src/pages/api/v1/stream/generate_sql.ts @@ -0,0 +1,201 @@ +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_STREAM_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 success with the SQL content + sendStateUpdate(res, StateType.SQL_GENERATION_SUCCESS, { sql }); + + // Log the API call + await apiHistoryRepository.createOne({ + id: uuidv4(), + projectId: project.id, + apiType: ApiType.STREAM_GENERATE_SQL, + threadId: newThreadId, + headers: req.headers as Record, + requestPayload: { question, language }, + responsePayload: { sql }, + statusCode: 200, + durationMs: Date.now() - startTime, + }); + + endStream(res, newThreadId, startTime); + } catch (error) { + logger.error('Error in stream generate SQL API:', error); + + // Log the error + await apiHistoryRepository.createOne({ + id: uuidv4(), + projectId: project?.id || 0, + apiType: ApiType.STREAM_GENERATE_SQL, + threadId: threadId || uuidv4(), + headers: req.headers as Record, + 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', + error.code || Errors.GeneralErrorCodes.INTERNAL_SERVER_ERROR, + error.additionalData, + ); + endStream(res, threadId || uuidv4(), startTime); + } +}