feat: add max query row count to mitigate LLM abuse

gregnr · gregnr · commit 64ec8fc8ed41 · 2024-08-13T21:43:22.000-05: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
@@ -39,7 +39,8 @@ export async function POST(req: Request) {
       - 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).
 
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.
