|
| 1 | +import type { AppUIMessage } from '~/ai-tools' |
| 2 | +import { anthropic } from '@ai-sdk/anthropic' |
| 3 | +import { google } from '@ai-sdk/google' |
| 4 | +import { DatabaseType } from '@conar/shared/enums/database-type' |
| 5 | +import { streamToEventIterator } from '@orpc/server' |
| 6 | +import { convertToModelMessages, smoothStream, stepCountIs, streamText } from 'ai' |
| 7 | +import { createRetryable } from 'ai-retry' |
| 8 | +import { type } from 'arktype' |
| 9 | +import { consola } from 'consola' |
| 10 | +import { v7 } from 'uuid' |
| 11 | +import { tools } from '~/ai-tools' |
| 12 | +import { withPosthog } from '~/lib/posthog' |
| 13 | +import { orpc, requireSubscriptionMiddleware } from '~/orpc' |
| 14 | + |
| 15 | +const model = createRetryable({ |
| 16 | + model: anthropic('claude-sonnet-4-5'), |
| 17 | + retries: [ |
| 18 | + anthropic('claude-opus-4-5'), |
| 19 | + google('gemini-2.5-pro'), |
| 20 | + ], |
| 21 | +}) |
| 22 | + |
| 23 | +function handleError(error: unknown) { |
| 24 | + if (typeof error === 'object' && (error as { type?: string }).type === 'overloaded_error') { |
| 25 | + return 'Sorry, I was unable to generate a response due to high load. Please try again later.' |
| 26 | + } |
| 27 | + if (typeof error === 'object' && (error as { message?: string }).message?.includes('prompt is too long')) { |
| 28 | + return 'Sorry, I was unable to generate a response. Currently I cannot handle larger chats like yours. Please create a new chat.' |
| 29 | + } |
| 30 | + return 'Sorry, I was unable to generate a response due to an error. Please try again.' |
| 31 | +} |
| 32 | + |
| 33 | +export const chat = orpc |
| 34 | + .use(requireSubscriptionMiddleware) |
| 35 | + .use(async ({ context, next }) => { |
| 36 | + context.setHeader('Transfer-Encoding', 'chunked') |
| 37 | + context.setHeader('Connection', 'keep-alive') |
| 38 | + |
| 39 | + return next() |
| 40 | + }) |
| 41 | + .input(type({ |
| 42 | + id: 'string.uuid.v7', |
| 43 | + type: type.valueOf(DatabaseType), |
| 44 | + context: 'string', |
| 45 | + createdAt: 'Date', |
| 46 | + updatedAt: 'Date', |
| 47 | + messages: 'object[]' as type.cast<AppUIMessage[]>, |
| 48 | + })) |
| 49 | + .handler(async ({ input, context, signal }) => { |
| 50 | + consola.info('messages', JSON.stringify(input.messages.map(message => ({ |
| 51 | + id: message.id, |
| 52 | + chatId: input.id, |
| 53 | + role: message.role, |
| 54 | + partsCount: message.parts.length, |
| 55 | + })), null, 2)) |
| 56 | + |
| 57 | + const result = streamText({ |
| 58 | + messages: [ |
| 59 | + { |
| 60 | + role: 'system', |
| 61 | + content: [ |
| 62 | + `You are an SQL tool that generates valid SQL code for ${input.type} database.`, |
| 63 | + 'You can use several tools to improve response.', |
| 64 | + 'You can generate select queries using the tools to get data directly from the database.', |
| 65 | + 'You can also search the web for information when the user asks about external resources, provides URLs, or needs current information beyond the database schema.', |
| 66 | + '', |
| 67 | + 'Requirements:', |
| 68 | + `- Ensure the SQL is 100% valid and optimized for ${input.type} database`, |
| 69 | + '- Use proper table and column names exactly as provided in the context', |
| 70 | + '- Use 2 spaces for indentation and consistent formatting', |
| 71 | + '- Consider performance implications for complex queries', |
| 72 | + '- The SQL code will be executed directly in a production database editor', |
| 73 | + '- Generate SQL query only for the provided schemas, tables, columns and enums', |
| 74 | + '- Answer in markdown and paste the SQL code in a code block, each query in a separate code block, do not use headings', |
| 75 | + '- Answer in the same language as the user\'s message', |
| 76 | + '- Use quotes for table and column names to prevent SQL errors with case sensitivity', |
| 77 | + '- If a user asks to change specific lines generate SQL only for the lines, not for whole SQL', |
| 78 | + '', |
| 79 | + 'Additional information:', |
| 80 | + `- Current date and time: ${new Date().toISOString()}`, |
| 81 | + '', |
| 82 | + 'You can use the following tools to help you generate the SQL code:', |
| 83 | + `- ${Object.entries(tools).map(([tool, { description }]) => `${tool}: ${description}`).join('\n')}`, |
| 84 | + '', |
| 85 | + 'User provided context:', |
| 86 | + input.context, |
| 87 | + ].join('\n'), |
| 88 | + }, |
| 89 | + ...(await convertToModelMessages(input.messages)), |
| 90 | + ], |
| 91 | + stopWhen: stepCountIs(Number.POSITIVE_INFINITY), |
| 92 | + abortSignal: signal, |
| 93 | + model: withPosthog(model, { |
| 94 | + chatId: input.id, |
| 95 | + userId: context.user.id, |
| 96 | + }), |
| 97 | + experimental_transform: smoothStream(), |
| 98 | + tools, |
| 99 | + }) |
| 100 | + |
| 101 | + const stream = result.toUIMessageStream({ |
| 102 | + originalMessages: input.messages, |
| 103 | + generateMessageId: () => v7(), |
| 104 | + sendSources: true, |
| 105 | + onFinish: async (result) => { |
| 106 | + consola.info('stream finished', JSON.stringify({ |
| 107 | + ...result.responseMessage, |
| 108 | + parts: result.responseMessage.parts.map(part => part.type), |
| 109 | + }, null, 2)) |
| 110 | + }, |
| 111 | + onError: (error) => { |
| 112 | + consola.error('error toUIMessageStream onError', error) |
| 113 | + |
| 114 | + return handleError(error) |
| 115 | + }, |
| 116 | + }) |
| 117 | + |
| 118 | + return streamToEventIterator(stream) |
| 119 | + }) |
0 commit comments