Feat/assistant refinements 3 (supabase#31015)

* limit local storage messages

* add tool for getting database functions

* reduce loading jitter in assistant

Author: SaxonF

apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx

@@ -336,7 +336,7 @@ export const AIAssistant = ({
                 )}
               </div>
             </div>
-            {!includeSchemaMetadata && (
+            {!includeSchemaMetadata && selectedOrganization && (
               <Admonition
                 type="default"
                 title="Project metadata is not shared"
@@ -345,7 +345,7 @@ export const AIAssistant = ({
                     ? 'Your organization has the HIPAA addon and will not send any project metadata with your prompts.'
                     : 'The Assistant can improve the quality of the answers if you send project metadata along with your prompts. Opt into sending anonymous data to share your schema and table definitions.'
                 }
-                className="border-0 border-b rounded-none"
+                className="border-0 border-b rounded-none bg-background"
               >
                 {!hasHipaaAddon && (
                   <Button

apps/studio/data/database-functions/database-functions-query.ts

@@ -16,16 +16,20 @@ const pgMetaFunctionsList = pgMeta.functions.list()
 
 export async function getDatabaseFunctions(
   { projectRef, connectionString }: DatabaseFunctionsVariables,
-  signal?: AbortSignal
+  signal?: AbortSignal,
+  headersInit?: HeadersInit
 ) {
+  let headers = new Headers(headersInit)
+
   const { result } = await executeSql(
     {
       projectRef,
       connectionString,
       sql: pgMetaFunctionsList.sql,
       queryKey: ['database-functions'],
     },
-    signal
+    signal,
+    headers
   )
 
   return result as DatabaseFunction[]

apps/studio/pages/api/ai/sql/tools.ts

@@ -6,6 +6,7 @@ import { getDatabasePolicies } from 'data/database-policies/database-policies-qu
 import { getEntityDefinitionsSql } from 'data/database/entity-definitions-query'
 import { executeSql } from 'data/sql/execute-sql-query'
 import { processSql, renderSupabaseJs } from '@supabase/sql-to-rest'
+import { getDatabaseFunctions } from 'data/database-functions/database-functions-query'
 
 export const getTools = ({
   projectRef,
@@ -342,5 +343,49 @@ export const getTools = ({
         `
       },
     }),
+    getFunctions: tool({
+      description: 'Get database functions for one or more schemas',
+      parameters: z.object({
+        schemas: z.array(z.string()).describe('The schema names to get the functions for'),
+      }),
+      execute: async ({ schemas }) => {
+        try {
+          const data = includeSchemaMetadata
+            ? await getDatabaseFunctions(
+                {
+                  projectRef,
+                  connectionString,
+                },
+                undefined,
+                {
+                  'Content-Type': 'application/json',
+                  ...(authorization && { Authorization: authorization }),
+                }
+              )
+            : []
+
+          // Filter functions by requested schemas
+          const filteredFunctions = data.filter((func) => schemas.includes(func.schema))
+
+          const formattedFunctions = filteredFunctions
+            .map(
+              (func) => `
+          Function Name: "${func.name}"
+          Schema: ${func.schema}
+          Arguments: ${func.argument_types}
+          Return Type: ${func.return_type}
+          Language: ${func.language}
+          Definition: ${func.definition}
+        `
+            )
+            .join('\n')
+
+          return formattedFunctions
+        } catch (error) {
+          console.error('Failed to fetch functions:', error)
+          return `Failed to fetch functions: ${error}`
+        }
+      },
+    }),
   }
 }

apps/studio/state/app-state.ts

@@ -222,10 +222,12 @@ export const appState = proxy({
 // Set up localStorage subscription
 if (typeof window !== 'undefined') {
   subscribe(appState, () => {
-    localStorage.setItem(
-      LOCAL_STORAGE_KEYS.AI_ASSISTANT_STATE,
-      JSON.stringify(appState.aiAssistantPanel)
-    )
+    const state = {
+      ...appState.aiAssistantPanel,
+      // limit to 20 messages so as to not overflow the context window
+      messages: appState.aiAssistantPanel.messages?.slice(-20),
+    }
+    localStorage.setItem(LOCAL_STORAGE_KEYS.AI_ASSISTANT_STATE, JSON.stringify(state))
   })
 }