Add configuration setting for schema definition context in Continue and Copilot

Author: ajshedivy

package.json

@@ -101,6 +101,17 @@
           }
         }
       },
+      {
+        "id": "vscode-db2i.ai",
+        "title": "AI integrations",
+        "properties": {
+          "vscode-db2i.ai.useSchemaDefinition": {
+            "type": "boolean",
+            "description": "Provide Schema definition as additional context in Continue and Copilot",
+            "default": true
+          }
+        }
+      },
       {
         "id": "vscode-db2i.resultsets",
         "title": "Viewing Data",

src/aiProviders/context.ts

@@ -4,6 +4,7 @@ import { JobManager } from "../config";
 import Schemas, { AllSQLTypes, SQLType } from "../database/schemas";
 import Statement from "../database/statement";
 import { DB2_SYSTEM_PROMPT } from "./continue/prompts";
+import Configuration from "../configuration";
 
 export function canTalkToDb() {
   return JobManager.getSelection() !== undefined;
@@ -36,7 +37,7 @@ interface ContextDefinition {
  * @param schema - The name of the database schema to process.
  * @returns A promise that resolves to an array of filtered objects representing the schema.
  */
-export async function buildSchemaDefinition(schema: string) {
+export async function buildSchemaDefinition(schema: string): Promise<Partial<BasicSQLObject>[] | undefined> {
   /**
    * Cleans an object by removing properties with undefined, null, or empty string values.
    * If the value is an object (but not an array), it cleans that object as well.
@@ -108,8 +109,14 @@ export async function buildSchemaDefinition(schema: string) {
     return data.map(filterBasicSQLObject);
   }
 
-  const compressedData = filterBasicSQLObjects(allInfo);
-  return compressedData;
+  const useSchemaDef: boolean = Configuration.get<boolean>(`ai.useSchemaDefinition`);
+  if (useSchemaDef) {
+    const compressedData = filterBasicSQLObjects(allInfo);
+    return compressedData;
+  }
+
+  return undefined;
+
 }
 
 /**
@@ -150,8 +157,10 @@ export async function generateTableDefinition(schema: string, input: string[]) {
     `tables`,
   ]);
 
-  const filteredTables = validWords.filter((word) =>
-    allTables.some((table) => table.name == word)
+  const filteredTables = Array.from(
+    new Set(
+      validWords.filter((word) => allTables.some((table) => table.name == word))
+    )
   );
 
   await Promise.all(

src/aiProviders/continue/continueContextProvider.ts

@@ -143,11 +143,13 @@ export class db2ContextProvider implements IContextProvider {
 
       // 1. SCHEMA Definiton (semantic)
       const schemaSemantic = await buildSchemaDefinition(schema);
-      contextItems.push({
-        name: `SCHEMA Definition`,
-        description: `${schema} definition`,
-        content: JSON.stringify(schemaSemantic),
-      });
+      if (schemaSemantic) {
+        contextItems.push({
+          name: `SCHEMA Definition`,
+          description: `${schema} definition`,
+          content: JSON.stringify(schemaSemantic),
+        });
+      }
       try {
         switch (true) {
           case fullInput.includes(`*SELF`) || query?.includes(`*SELF`):

src/aiProviders/continue/listTablesContextProvider.ts

@@ -70,11 +70,13 @@ class ListDb2iTables implements IContextProvider {
     let contextItems: ContextItem[] = [];
     if (query.toUpperCase() === this.schema.toUpperCase()) {
       const schemaSemantic = await buildSchemaDefinition(this.schema);
-      contextItems.push({
-        name: `SCHEMA Definition`,
-        description: `${this.schema} definition`,
-        content: JSON.stringify(schemaSemantic),
-      });
+      if (schemaSemantic) {
+        contextItems.push({
+          name: `SCHEMA Definition`,
+          description: `${this.schema} definition`,
+          content: JSON.stringify(schemaSemantic),
+        });
+      }
     } else {
       const tablesRefs = await generateTableDefinition(
         this.schema,

src/aiProviders/copilot/index.ts

@@ -63,11 +63,13 @@ export function activateChat(context: vscode.ExtensionContext) {
       // 1. SCHEMA Definiton (semantic)
       let usingSchema = getCurrentSchema();
       const schemaSemantic = await buildSchemaDefinition(usingSchema);
-      messages.push(
-        vscode.LanguageModelChatMessage.Assistant(
-          JSON.stringify(schemaSemantic)
-        )
-      );
+      if (schemaSemantic) {
+        messages.push(
+          vscode.LanguageModelChatMessage.Assistant(
+            JSON.stringify(schemaSemantic)
+          )
+        );
+      }
 
       switch (request.command) {
         case `activity`:

src/contributes.json

@@ -57,6 +57,17 @@
             ]
           }
         }
+      },
+      {
+        "id": "vscode-db2i.ai",
+        "title": "AI integrations",
+        "properties": {
+          "vscode-db2i.ai.useSchemaDefinition": {
+            "type": "boolean",
+            "description": "Provide Schema definition as additional context in Continue and Copilot",
+            "default": true
+          }
+        }
       }
     ],
     "viewsContainers": {