Merge branch 'feat/upload-db' into feat/deploy-db-to-supabase

jgoux · jgoux · commit 92735e5947ca · 2024-08-14T17:06:59.000+02:00
diff --git a/apps/postgres-new/app/api/chat/route.ts b/apps/postgres-new/app/api/chat/route.ts
@@ -1,7 +1,7 @@
 import { openai } from '@ai-sdk/openai'
 import { ToolInvocation, convertToCoreMessages, streamText } from 'ai'
 import { codeBlock } from 'common-tags'
-import { convertToCoreTools, tools } from '~/lib/tools'
+import { convertToCoreTools, maxRowLimit, tools } from '~/lib/tools'
 
 // Allow streaming responses up to 30 seconds
 export const maxDuration = 30
@@ -33,17 +33,19 @@ export async function POST(req: Request) {
       - For primary keys, always use "id bigint primary key generated always as identity" (not serial)
       - Prefer 'text' over 'varchar'
       - Keep explanations brief but helpful
+      - Don't repeat yourself after creating the table
 
       When creating sample data:
       - Make the data realistic, including joined data
       - Check for existing records/conflicts in the table
 
-      When querying data, limit to 5 by default.
+      When querying data, limit to 5 by default. The maximum number of rows you're allowed to fetch is ${maxRowLimit} (to protect AI from token abuse).
+      If the user needs to fetch more than ${maxRowLimit} rows at once, they can export the query as a CSV.
 
       When performing FTS, always use 'simple' (languages aren't available).
 
       When importing CSVs try to solve the problem yourself (eg. use a generic text column, then refine)
-      vs. asking the user to change the CSV.
+      vs. asking the user to change the CSV. No need to select rows after importing.
       
       You also know math. All math equations and expressions must be written in KaTex and must be wrapped in double dollar \`$$\`:
         - Inline: $$\\sqrt{26}$$
@@ -56,7 +58,7 @@ export async function POST(req: Request) {
 
       Feel free to suggest corrections for suspected typos.
     `,
-    model: openai('gpt-4o-2024-05-13'),
+    model: openai('gpt-4o-2024-08-06'),
     messages: convertToCoreMessages(messages),
     tools: convertToCoreTools(tools),
   })
diff --git a/apps/postgres-new/lib/hooks.ts b/apps/postgres-new/lib/hooks.ts
@@ -20,7 +20,7 @@ import { useTablesQuery } from '~/data/tables/tables-query'
 import { embed } from './embed'
 import { loadFile, saveFile } from './files'
 import { SmoothScroller } from './smooth-scroller'
-import { OnToolCall } from './tools'
+import { maxRowLimit, OnToolCall } from './tools'
 
 export function useDebounce<T>(value: T, delay: number) {
   const [debouncedValue, setDebouncedValue] = useState(value)
@@ -265,7 +265,17 @@ export function useOnToolCall(databaseId: string) {
 
             const results = await db.exec(sql)
 
-            // Truncate vector columns due to their large size
+            const oversizedResult = results.find((result) => result.rows.length > maxRowLimit)
+
+            // We have a max row count in place to mitigate LLM token abuse
+            if (oversizedResult) {
+              return {
+                success: false,
+                error: `Query produced ${oversizedResult.rows.length} rows but the max allowed limit is ${maxRowLimit}. Rerun the query with a limit of ${maxRowLimit}.`,
+              }
+            }
+
+            // Truncate vector columns due to their large size (display purposes only)
             const filteredResults = results.map((result) => {
               const vectorFields = result.fields.filter(
                 (field) => field.dataTypeID === vectorDataTypeId
diff --git a/apps/postgres-new/lib/tools.ts b/apps/postgres-new/lib/tools.ts
@@ -16,6 +16,11 @@ function result<T extends z.ZodTypeAny>(schema: T) {
   return z.union([z.intersection(successResultSchema, schema), errorResultSchema])
 }
 
+/**
+ * The maximum SQL result row limit to prevent overloading LLM.
+ */
+export const maxRowLimit = 100
+
 /**
  * Central location for all LLM tools including their
  * description, arg schema, and result schema.
