konnosDashboard
complete.ts311 lines · main
1import pgMeta, { getEntityDefinitionsSql } from '@supabase/pg-meta'
2import { generateText, ModelMessage, stepCountIs, tool } from 'ai'
3import { IS_PLATFORM } from 'common'
4import { source } from 'common-tags'
5import { NextApiRequest, NextApiResponse } from 'next'
6import z from 'zod'
7
8import { executeSql } from '@/data/sql/execute-sql-query'
9import { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
10import { getOrgAIDetails } from '@/lib/ai/ai-details'
11import { getModel } from '@/lib/ai/model'
12import { DEFAULT_COMPLETION_MODEL } from '@/lib/ai/model.utils'
13import {
14 COMPLETION_PROMPT,
15 EDGE_FUNCTION_PROMPT,
16 PG_BEST_PRACTICES,
17 SECURITY_PROMPT,
18 SQL_COMPLETION_INSTRUCTIONS,
19} from '@/lib/ai/prompts'
20import apiWrapper from '@/lib/api/apiWrapper'
21import { executeQuery } from '@/lib/api/self-hosted/query'
22
23export const maxDuration = 60
24
25const pgMetaSchemasList = pgMeta.schemas.list()
26type Schemas = z.infer<(typeof pgMetaSchemasList)['zod']>
27type EntityDefinitionRow = { data: { definitions: Array<{ id: number; sql: string }> } }
28
29type SqlFetchParams = {
30 projectRef: string
31 connectionString: string | null | undefined
32 headers: Record<string, string>
33}
34
35type SchemaListResult =
36 | { error: true }
37 | { error: false; queriedSchemas: string[]; otherSchemas: string[] }
38
39type SchemaDDLResult = { error: true } | { error: false; sqlDefinitions: string[] }
40
41async function fetchSchemas(
42 includeSchema: boolean,
43 { projectRef, connectionString, headers }: SqlFetchParams
44): Promise<{ schemas: Schemas; error: boolean }> {
45 if (!includeSchema) return { schemas: [], error: false }
46 try {
47 const { result } = await executeSql<Schemas>(
48 { projectRef, connectionString, sql: pgMetaSchemasList.sql },
49 undefined,
50 headers,
51 IS_PLATFORM ? undefined : executeQuery
52 )
53 return { schemas: result, error: false }
54 } catch {
55 return { schemas: [], error: true }
56 }
57}
58
59async function fetchSchemaDDL(
60 schemas: string[],
61 { projectRef, connectionString, headers }: SqlFetchParams
62): Promise<SchemaDDLResult> {
63 if (schemas.length === 0) return { error: false, sqlDefinitions: [] }
64 try {
65 const { result } = await executeSql<EntityDefinitionRow[]>(
66 { projectRef, connectionString, sql: getEntityDefinitionsSql({ schemas }) },
67 undefined,
68 headers,
69 IS_PLATFORM ? undefined : executeQuery
70 )
71 const definitions = result?.[0]?.data?.definitions ?? []
72 return {
73 error: false,
74 sqlDefinitions: definitions.map((d) => d.sql),
75 }
76 } catch {
77 return { error: true }
78 }
79}
80
81function buildDatabaseSchemaSection({
82 includeSchema,
83 schemaListResult,
84 schemaDDLResult,
85}: {
86 includeSchema: boolean
87 schemaListResult: SchemaListResult
88 schemaDDLResult: SchemaDDLResult
89}): string {
90 if (!includeSchema) {
91 return 'Schema context is unavailable — data opt-in is not enabled for this project.'
92 }
93 const lines: string[] = []
94
95 if (schemaListResult.error) {
96 lines.push(
97 "Unable to fetch list of available database schemas. Assume `public` schema, infer others from the user's existing code."
98 )
99 } else {
100 lines.push(`Queried schemas: ${schemaListResult.queriedSchemas.join(', ')}`)
101 if (schemaListResult.otherSchemas.length > 0)
102 lines.push(
103 `Other available schemas (use getSchemaDefinitions tool): ${schemaListResult.otherSchemas.join(', ')}`
104 )
105 }
106
107 if (schemaDDLResult.error) {
108 lines.push('Failed to fetch table definitions due to a database error.')
109 } else {
110 const defsText =
111 schemaDDLResult.sqlDefinitions.length > 0
112 ? schemaDDLResult.sqlDefinitions.join('\n\n')
113 : 'No table definitions found.'
114 lines.push(`\n${defsText}`)
115 }
116
117 return lines.join('\n')
118}
119
120const requestBodySchema = z.object({
121 completionMetadata: z.object({
122 textBeforeCursor: z.string(),
123 textAfterCursor: z.string(),
124 prompt: z.string(),
125 selection: z.string(),
126 }),
127 projectRef: z.string(),
128 connectionString: z.string().nullish(),
129 orgSlug: z.string().optional(),
130 language: z.string().optional(),
131})
132
133async function handler(req: NextApiRequest, res: NextApiResponse) {
134 if (req.method !== 'POST') {
135 return res.status(405).json({ error: `Method ${req.method} Not Allowed` })
136 }
137
138 try {
139 let body: unknown
140 try {
141 body = typeof req.body === 'string' ? JSON.parse(req.body) : req.body
142 } catch {
143 return res.status(400).json({ error: 'Malformed JSON' })
144 }
145 const { data, error: parseError } = requestBodySchema.safeParse(body)
146
147 if (parseError) {
148 return res.status(400).json({ error: 'Invalid request body', issues: parseError.issues })
149 }
150
151 const { completionMetadata, projectRef, connectionString, orgSlug, language } = data
152 const { textBeforeCursor, textAfterCursor, prompt, selection } = completionMetadata
153
154 const authorization = req.headers.authorization
155 let aiOptInLevel: AiOptInLevel = IS_PLATFORM ? 'disabled' : 'schema'
156
157 if (IS_PLATFORM && orgSlug && authorization && projectRef) {
158 const { aiOptInLevel: orgAIOptInLevel } = await getOrgAIDetails({
159 orgSlug,
160 authorization,
161 })
162 aiOptInLevel = orgAIOptInLevel
163 }
164
165 const {
166 modelParams,
167 error: modelError,
168 systemProviderOptions,
169 } = await getModel({
170 provider: 'openai',
171 modelEntry: DEFAULT_COMPLETION_MODEL,
172 })
173
174 if (modelError) {
175 return res.status(500).json({ error: modelError.message })
176 }
177
178 const headers = {
179 'Content-Type': 'application/json',
180 ...(authorization && { Authorization: authorization }),
181 }
182
183 const includeSchema = aiOptInLevel !== 'disabled'
184
185 // Fetch schema list first so we can determine which schemas to load DDL for.
186 // These are best-effort — if they fail, we proceed without DDL context.
187 const { schemas, error: schemaListError } = await fetchSchemas(includeSchema, {
188 projectRef,
189 connectionString,
190 headers,
191 })
192
193 // Always include public; also eagerly include any non-public schema whose name
194 // appears as `name.` in the cursor context. Checking against the real schema list
195 // avoids fetching DDL for table aliases or other false matches. This is robust to
196 // incomplete SQL (the user may be mid-typing, so a full parser would fail here).
197 const cursorContext = textBeforeCursor + selection + textAfterCursor
198 const lowerContext = cursorContext.toLowerCase()
199 const schemasToFetch = includeSchema
200 ? [
201 'public',
202 ...schemas
203 .filter((s) => {
204 const lower = s.name.toLowerCase()
205 return (
206 s.name !== 'public' &&
207 (lowerContext.includes(lower + '.') || lowerContext.includes(`"${lower}".`))
208 )
209 })
210 .map((s) => s.name),
211 ]
212 : []
213
214 const schemaDDLResult = await fetchSchemaDDL(schemasToFetch, {
215 projectRef,
216 connectionString,
217 headers,
218 })
219
220 // Reshape the fetched schemas and candidates into a discriminated union over error states
221 const fetchedSchemaSet = new Set(schemasToFetch)
222 const schemaListResult: SchemaListResult = schemaListError
223 ? { error: true }
224 : {
225 error: false,
226 queriedSchemas: schemasToFetch,
227 otherSchemas: schemas.filter((s) => !fetchedSchemaSet.has(s.name)).map((s) => s.name),
228 }
229
230 // Important: do not use dynamic content in the system prompt or Bedrock will not cache it
231 const system = source`
232 ${COMPLETION_PROMPT}
233 ${language === 'sql' ? SQL_COMPLETION_INSTRUCTIONS : ''}
234 ${language === 'sql' ? PG_BEST_PRACTICES : EDGE_FUNCTION_PROMPT}
235 ${SECURITY_PROMPT}
236 `
237
238 const userMessage = source`
239 ## Database Schema
240
241 ${buildDatabaseSchemaSection({ includeSchema, schemaListResult, schemaDDLResult })}
242
243 ## Code
244
245 \`\`\`${language ?? ''}
246 ${textBeforeCursor}<selection>${selection}</selection>${textAfterCursor}
247 \`\`\`
248
249 ## Instruction
250
251 ${prompt}
252 `
253
254 // Note: these must be of type `CoreMessage` to prevent AI SDK from stripping `providerOptions`
255 // https://github.com/vercel/ai/blob/81ef2511311e8af34d75e37fc8204a82e775e8c3/packages/ai/core/prompt/standardize-prompt.ts#L83-L88
256 const coreMessages: ModelMessage[] = [
257 {
258 role: 'system',
259 content: system,
260 ...(systemProviderOptions && { providerOptions: systemProviderOptions }),
261 },
262 {
263 role: 'user',
264 content: userMessage,
265 },
266 ]
267
268 const { text } = await generateText({
269 ...modelParams,
270 stopWhen: stepCountIs(5),
271 messages: coreMessages,
272 tools:
273 includeSchema && !schemaListResult.error
274 ? {
275 getSchemaDefinitions: tool({
276 description: 'Get table and column definitions for one or more schemas',
277 inputSchema: z.object({
278 schemas: z
279 .array(z.string())
280 .describe('The schema names to get the definitions for'),
281 }),
282 execute: async ({ schemas: maybeSchemas }) => {
283 const validSchemas = maybeSchemas.filter((name) =>
284 schemas.some((s) => s.name === name)
285 )
286 const result = await fetchSchemaDDL(validSchemas, {
287 projectRef,
288 connectionString,
289 headers,
290 })
291 if (result.error)
292 return 'Failed to fetch schema definitions due to a database error.'
293 if (result.sqlDefinitions.length === 0) return 'No table definitions found.'
294 return result.sqlDefinitions.join('\n\n')
295 },
296 }),
297 }
298 : undefined,
299 })
300
301 return res.status(200).json(text)
302 } catch (error) {
303 console.error('Completion error:', error)
304 return res.status(500).json({ error: 'Failed to generate completion' })
305 }
306}
307
308const wrapper = (req: NextApiRequest, res: NextApiResponse) =>
309 apiWrapper(req, res, handler, { withAuth: true })
310
311export default wrapper