konnosDashboard
functions.ts207 lines · main
1import { SchemaBuilder } from '@serafin/schema-builder'
2import { codeBlock, stripIndent } from 'common-tags'
3import { jsonrepair } from 'jsonrepair'
4import type OpenAI from 'openai'
5
6import { ContextLengthError, EmptyResponseError, EmptySqlError } from '../errors'
7
8const debugSqlSchema = SchemaBuilder.emptySchema()
9 .addString('solution', {
10 description: 'A short suggested solution for the error (as concise as possible).',
11 })
12 .addString('sql', {
13 description:
14 'The SQL rewritten to apply the solution. Includes all the original logic, but modified to fix the issue.',
15 })
16
17const generateTitleSchema = SchemaBuilder.emptySchema()
18 .addString('title', {
19 description: stripIndent`
20 The generated title for the SQL snippet (short and concise).
21 - Omit these words: 'SQL', 'Postgres', 'Query', 'Database'
22 `,
23 })
24 .addString('description', {
25 description: stripIndent`
26 The generated description for the SQL snippet.
27 `,
28 })
29
30// Reference auto-generated types for each JSON schema
31export type DebugSqlResult = typeof debugSqlSchema.T
32export type GenerateTitleResult = typeof generateTitleSchema.T
33
34// Combine the completion functions
35const completionFunctions = {
36 debugSql: {
37 name: 'debugSql',
38 description: stripIndent`
39 Debugs a Postgres SQL error. Returns the fixed SQL and a solution explaining it.
40 - Create extensions if they are missing (only for valid extensions)
41 - Suggest creating tables if they are missing
42 - Include all of the original SQL
43 - For primary keys, always use "id bigint primary key generated always as identity" (not serial)
44 - When creating tables, always add foreign key references inline
45 - Prefer 'text' over 'varchar'
46 - Prefer 'timestamp with time zone' over 'date'
47 - Use vector(384) data type for any embedding/vector related query
48 - Always use double apostrophe in SQL strings (eg. 'Night''s watch')
49 - Always use semicolons
50 `,
51 parameters: debugSqlSchema.schema as Record<string, unknown>,
52 },
53 generateTitle: {
54 name: 'generateTitle',
55 description: stripIndent`
56 Generates a short title and summarized description for a Postgres SQL snippet.
57
58 The description should describe why this table was created (eg. "Table to track todos")
59 `,
60 parameters: generateTitleSchema.schema as Record<string, unknown>,
61 },
62} satisfies Record<string, OpenAI.Chat.Completions.ChatCompletionCreateParams.Function>
63
64/**
65 * Debugs SQL errors.
66 *
67 * @returns A suggested SQL fix along with a solution explanation.
68 */
69export async function debugSql(
70 openai: OpenAI,
71 errorMessage: string,
72 sql: string,
73 entityDefinitions?: string[]
74) {
75 const hasEntityDefinitions = entityDefinitions !== undefined && entityDefinitions.length > 0
76
77 const completionMessages: OpenAI.Chat.Completions.ChatCompletionMessageParam[] = []
78
79 if (hasEntityDefinitions) {
80 completionMessages.push({
81 role: 'user',
82 content: codeBlock`
83 Here is my database schema for reference:
84 ${entityDefinitions.join('\n\n')}
85 `,
86 })
87 }
88
89 completionMessages.push(
90 {
91 role: 'user',
92 content: stripIndent`
93 Here is my current SQL:
94 ${sql}
95 `,
96 },
97 {
98 role: 'user',
99 content: stripIndent`
100 Here is the error I am getting:
101 ${errorMessage}
102 `,
103 }
104 )
105
106 try {
107 const completionResponse = await openai.chat.completions.create({
108 model: 'gpt-4o-mini-2024-07-18',
109 messages: completionMessages,
110 max_tokens: 2048,
111 temperature: 0,
112 tool_choice: {
113 type: 'function',
114 function: {
115 name: completionFunctions.debugSql.name,
116 },
117 },
118 tools: [
119 {
120 type: 'function',
121 function: completionFunctions.debugSql,
122 },
123 ],
124 stream: false,
125 })
126
127 const [firstChoice] = completionResponse.choices
128 const [firstTool] = firstChoice.message?.tool_calls ?? []
129
130 const sqlResponseString = firstTool?.function.arguments
131
132 if (!sqlResponseString) {
133 throw new EmptyResponseError()
134 }
135
136 const repairedJsonString = jsonrepair(sqlResponseString)
137
138 const result: DebugSqlResult = JSON.parse(repairedJsonString)
139
140 if (!result.sql) {
141 throw new EmptySqlError()
142 }
143
144 return result
145 } catch (error) {
146 if (error instanceof Error && 'code' in error && error.code === 'context_length_exceeded') {
147 throw new ContextLengthError()
148 }
149 throw error
150 }
151}
152
153/**
154 * Generates a snippet title based on the provided SQL.
155 *
156 * @returns A title and description for the SQL snippet.
157 */
158export async function titleSql(openai: OpenAI, sql: string) {
159 const completionMessages: OpenAI.Chat.Completions.ChatCompletionMessageParam[] = [
160 {
161 role: 'user',
162 content: sql,
163 },
164 ]
165
166 try {
167 const completionResponse = await openai.chat.completions.create({
168 model: 'gpt-4o-mini-2024-07-18',
169 messages: completionMessages,
170 max_tokens: 1024,
171 temperature: 0,
172 tool_choice: {
173 type: 'function',
174 function: {
175 name: completionFunctions.generateTitle.name,
176 },
177 },
178 tools: [
179 {
180 type: 'function',
181 function: completionFunctions.generateTitle,
182 },
183 ],
184 stream: false,
185 })
186
187 const [firstChoice] = completionResponse.choices
188 const [firstTool] = firstChoice.message?.tool_calls ?? []
189
190 const sqlResponseString = firstTool?.function.arguments
191
192 if (!sqlResponseString) {
193 throw new EmptyResponseError()
194 }
195
196 const repairedJsonString = jsonrepair(sqlResponseString)
197
198 const result: GenerateTitleResult = JSON.parse(repairedJsonString)
199
200 return result
201 } catch (error) {
202 if (error instanceof Error && 'code' in error && error.code === 'context_length_exceeded') {
203 throw new ContextLengthError()
204 }
205 throw error
206 }
207}