konnosDashboard
generate-v4.ts255 lines · main
1import pgMeta from '@supabase/pg-meta'
2import type { JwtPayload } from '@supabase/supabase-js'
3import { safeValidateUIMessages } from 'ai'
4import { IS_PLATFORM } from 'common'
5import type { NextApiRequest, NextApiResponse } from 'next'
6import z from 'zod'
7
8import { executeSql } from '@/data/sql/execute-sql-query'
9import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
10import { getOrgAIDetails, getProjectAIDetails } from '@/lib/ai/ai-details'
11import { isTracingAllowed } from '@/lib/ai/braintrust-logger'
12import { generateAssistantResponse } from '@/lib/ai/generate-assistant-response'
13import { getModel } from '@/lib/ai/model'
14import {
15 DEFAULT_ASSISTANT_ADVANCE_MODEL_ID,
16 DEFAULT_ASSISTANT_BASE_MODEL_ID,
17 getAssistantModelEntry,
18 isAssistantBaseModelId,
19 isKnownAssistantModelId,
20 type AssistantModelId,
21} from '@/lib/ai/model.utils'
22import { getTools } from '@/lib/ai/tools'
23import apiWrapper from '@/lib/api/apiWrapper'
24import { executeQuery } from '@/lib/api/self-hosted/query'
25import { getURL } from '@/lib/helpers'
26
27export const maxDuration = 120
28
29export const config = {
30 api: {
31 bodyParser: {
32 sizeLimit: '5mb',
33 },
34 },
35}
36
37async function handler(req: NextApiRequest, res: NextApiResponse, claims?: JwtPayload) {
38 const { method } = req
39
40 switch (method) {
41 case 'POST':
42 return handlePost(req, res, claims)
43 default:
44 res.setHeader('Allow', ['POST'])
45 res.status(405).json({
46 data: null,
47 error: { message: `Method ${method} Not Allowed` },
48 })
49 }
50}
51
52const wrapper = (req: NextApiRequest, res: NextApiResponse) =>
53 apiWrapper(req, res, handler, { withAuth: true })
54
55export default wrapper
56
57const requestBodySchema = z.object({
58 messages: z.array(z.any()),
59 projectRef: z.string(),
60 connectionString: z.string(),
61 schema: z.string().optional(),
62 table: z.string().optional(),
63 chatId: z.string().optional(),
64 chatName: z.string().optional(),
65 orgSlug: z.string().optional(),
66 model: z.string().optional(),
67})
68
69async function handlePost(req: NextApiRequest, res: NextApiResponse, claims?: JwtPayload) {
70 const authorization = req.headers.authorization
71 const accessToken = authorization?.replace('Bearer ', '')
72
73 if (IS_PLATFORM && !accessToken) {
74 return res.status(401).json({ error: 'Authorization token is required' })
75 }
76
77 const userId = claims?.sub
78
79 const body = typeof req.body === 'string' ? JSON.parse(req.body) : req.body
80 const { data, error: parseError } = requestBodySchema.safeParse(body)
81
82 if (parseError) {
83 return res.status(400).json({ error: 'Invalid request body', issues: parseError.issues })
84 }
85
86 const {
87 messages: rawMessages,
88 projectRef,
89 connectionString,
90 orgSlug,
91 chatId,
92 chatName,
93 model: rawRequestedModel,
94 } = data
95
96 const requestedModel: AssistantModelId | undefined =
97 rawRequestedModel && isKnownAssistantModelId(rawRequestedModel) ? rawRequestedModel : undefined
98
99 const messagesValidation = await safeValidateUIMessages({
100 messages: rawMessages,
101 })
102 if (!messagesValidation.success) {
103 return res.status(400).json({
104 error: 'Invalid request body',
105 message: messagesValidation.error.message,
106 })
107 }
108 const messages = messagesValidation.data
109
110 let aiOptInLevel: AiOptInLevel = 'disabled'
111 let hasAccessToAdvanceModel = false
112 let orgHasHipaaAddon: boolean | undefined
113 let projectIsSensitive: boolean | undefined
114 let projectRegion: string | undefined
115 let orgId: number | undefined
116 let planId: string | undefined
117
118 if (!IS_PLATFORM) {
119 aiOptInLevel = 'schema'
120 hasAccessToAdvanceModel = true
121 }
122
123 if (IS_PLATFORM && orgSlug && authorization && projectRef) {
124 try {
125 const [orgDetails, projectDetails] = await Promise.all([
126 getOrgAIDetails({ orgSlug, authorization }),
127 getProjectAIDetails({ projectRef, authorization }),
128 ])
129
130 aiOptInLevel = orgDetails.aiOptInLevel
131 hasAccessToAdvanceModel = orgDetails.hasAccessToAdvanceModel
132 orgHasHipaaAddon = orgDetails.hasHipaaAddon
133 orgId = orgDetails.orgId
134 planId = orgDetails.planId
135 projectIsSensitive = projectDetails.isSensitive
136 projectRegion = projectDetails.region
137 } catch (error) {
138 return res.status(400).json({
139 error: 'There was an error fetching your organization details',
140 })
141 }
142 }
143
144 const envThrottled = process.env.IS_THROTTLED !== 'false'
145
146 let effectiveModel: AssistantModelId = requestedModel ?? DEFAULT_ASSISTANT_ADVANCE_MODEL_ID
147 if (!hasAccessToAdvanceModel || (envThrottled && !isAssistantBaseModelId(effectiveModel))) {
148 effectiveModel = DEFAULT_ASSISTANT_BASE_MODEL_ID
149 }
150
151 const {
152 modelParams,
153 error: modelError,
154 systemProviderOptions,
155 } = await getModel({
156 provider: 'openai',
157 modelEntry: getAssistantModelEntry(effectiveModel),
158 })
159
160 if (modelError) {
161 return res.status(500).json({ error: modelError.message })
162 }
163
164 try {
165 const abortController = new AbortController()
166 req.on('close', () => abortController.abort())
167 req.on('aborted', () => abortController.abort())
168
169 const tools = await getTools({
170 projectRef,
171 connectionString,
172 authorization,
173 aiOptInLevel,
174 accessToken,
175 baseUrl: getURL(),
176 })
177
178 // Get a list of all schemas to add to context
179 const getSchemas = async (): Promise<string> => {
180 const pgMetaSchemasList = pgMeta.schemas.list()
181 type Schemas = z.infer<(typeof pgMetaSchemasList)['zod']>
182
183 const { result: schemas } = await executeSql<Schemas>(
184 {
185 projectRef,
186 connectionString,
187 sql: pgMetaSchemasList.sql,
188 },
189 undefined,
190 {
191 'Content-Type': 'application/json',
192 ...(authorization && { Authorization: authorization }),
193 },
194 IS_PLATFORM ? undefined : executeQuery
195 )
196
197 return schemas?.length > 0
198 ? `The available database schema names are: ${JSON.stringify(schemas)}`
199 : "You don't have access to any schemas."
200 }
201
202 const result = await generateAssistantResponse({
203 messages,
204 ...modelParams,
205 tools,
206 aiOptInLevel,
207 getSchemas: aiOptInLevel !== 'disabled' ? getSchemas : undefined,
208 projectRef,
209 chatId,
210 chatName,
211 allowTracing: isTracingAllowed({
212 orgHasHipaaAddon,
213 projectIsSensitive,
214 projectRegion,
215 }),
216 userId,
217 orgId,
218 planId,
219 requestedModel,
220 systemProviderOptions,
221 abortSignal: abortController.signal,
222 onSpanCreated: (spanId) => {
223 res.setHeader('x-braintrust-span-id', spanId)
224 },
225 })
226
227 result.pipeUIMessageStreamToResponse(res, {
228 sendReasoning: true,
229 headers: { 'Content-Encoding': 'none' },
230 onError: (error) => {
231 console.error('Assistant stream error:', error)
232
233 if (error == null) {
234 return 'unknown error'
235 }
236
237 if (typeof error === 'string') {
238 return error
239 }
240
241 if (error instanceof Error) {
242 return error.message
243 }
244
245 return JSON.stringify(error)
246 },
247 })
248 } catch (error) {
249 console.error('Error in handlePost:', error)
250 if (error instanceof Error) {
251 return res.status(500).json({ message: error.message })
252 }
253 return res.status(500).json({ message: 'An unexpected error occurred.' })
254 }
255}