Merge pull request #305 from codefori/fix/continue-context

Author: worksofliam

src/aiProviders/context.ts

@@ -1,6 +1,7 @@
 import { JobManager } from "../config";
 import * as vscode from "vscode";
 import Statement from "../database/statement";
+import { ContextItem } from "@continuedev/core";
 
 export function canTalkToDb() {
   return JobManager.getSelection() !== undefined;
@@ -99,7 +100,7 @@ export async function findPossibleTables(stream: vscode.ChatResponseStream, sche
  * 
  * Tables with names starting with 'SYS' are skipped.
  */
-export function refsToMarkdown(refs: TableRefs) {
+export function refsToMarkdown(refs: TableRefs): MarkdownRef[] {
   const condensedResult = Object.keys(refs).length > 5;
 
   let markdownRefs: MarkdownRef[] = [];
@@ -122,6 +123,23 @@ export function refsToMarkdown(refs: TableRefs) {
   return markdownRefs;
 }
 
+export function createContinueContextItems(refs: MarkdownRef[]) {
+  const contextItems: ContextItem[] = [];
+  const job = JobManager.getSelection();
+  for (const tableRef of refs) {
+    let prompt = `Table: ${tableRef.TABLE_NAME} (Schema: ${tableRef.SCHMEA}) Column Information:\n`;
+    prompt += `Format: column_name (column_text) type(length:precision) is_identity is_nullable\n`
+    prompt += `${tableRef.COLUMN_INFO}`;
+    contextItems.push({
+      name: `${job.name}-${tableRef.SCHMEA}-${tableRef.TABLE_NAME}`,
+      description: `Column information for ${tableRef.TABLE_NAME}`,
+      content: prompt,
+    });
+  }
+
+  return contextItems;
+}
+
 export async function getSystemStatus(): Promise<string> {
   const sqlStatment = `SELECT * FROM TABLE(QSYS2.SYSTEM_STATUS(RESET_STATISTICS=>'YES',DETAILED_INFO=>'ALL')) X`;
   const result = await JobManager.runSQL(sqlStatment, undefined);

src/aiProviders/continue/continueContextProvider.ts

@@ -2,7 +2,7 @@ import * as vscode from "vscode";
 import { JobManager } from "../../config";
 import { JobInfo } from "../../connection/manager";
 import { SelfCodeNode } from "../../views/jobManager/selfCodes/nodes";
-import { canTalkToDb, findPossibleTables, refsToMarkdown } from "../context";
+import { canTalkToDb, createContinueContextItems, findPossibleTables, refsToMarkdown } from "../context";
 import {
   ContextItem,
   ContextProviderDescription,
@@ -146,16 +146,7 @@ export class db2ContextProvider implements IContextProvider {
             );
             const markdownRefs = refsToMarkdown(tableRefs);
 
-            for (const tableRef of markdownRefs) {
-              let prompt = `Table: ${tableRef.TABLE_NAME} (Schema: ${tableRef.SCHMEA}) Column Information:\n`;
-              prompt += `Format: column_name (column_text) type(length:precision) is_identity is_nullable\n`
-              prompt += `${tableRef.COLUMN_INFO}`;
-              contextItems.push({
-                name: `${job.name}-${tableRef.SCHMEA}-${tableRef.TABLE_NAME}`,
-                description: `Column information for ${tableRef.TABLE_NAME}`,
-                content: prompt,
-              });
-            }
+            contextItems.push(...createContinueContextItems(markdownRefs));
 
             return contextItems;
         }

src/aiProviders/continue/listTablesContextProvider.ts

@@ -0,0 +1,143 @@
+import {
+  ContextItem,
+  ContextProviderDescription,
+  ContextProviderExtras,
+  ContextSubmenuItem,
+  IContextProvider,
+  LoadSubmenuItemsArgs,
+} from "@continuedev/core";
+import * as fs from "fs";
+import * as os from "os";
+import * as path from "path";
+import * as vscode from "vscode";
+import Schemas from "../../database/schemas";
+import Table from "../../database/table";
+import {
+  createContinueContextItems,
+  findPossibleTables,
+  refsToMarkdown,
+} from "../context";
+
+const listDb2Table: ContextProviderDescription = {
+  title: "list Db2i Tables",
+  displayTitle: "Db2i-tables",
+  description: "Add Db2i Table info to Context",
+  type: "submenu",
+};
+
+export let provider: ListDb2iTables = undefined;
+
+class ListDb2iTables implements IContextProvider {
+  constructor(private schema: string) {
+    this.schema = schema;
+  }
+
+  get description(): ContextProviderDescription {
+    return listDb2Table;
+  }
+
+  setCurrentSchema(schema: string) {
+    this.schema = schema;
+  }
+
+  getCurrentSchema() {
+    return this.schema;
+  }
+
+  async getColumnInfoForAllTables(schema: string) {
+    const items: TableColumn[] = await Table.getItems(schema);
+
+    return items.map((column) => ({
+      table_name: column.TABLE_NAME,
+      schema: column.TABLE_SCHEMA,
+      column_name: column.COLUMN_NAME,
+      column_data_type: column.DATA_TYPE,
+    }));
+  }
+
+  async getContextItems(
+    query: string,
+    extras: ContextProviderExtras
+  ): Promise<ContextItem[]> {
+    let contextItems: ContextItem[] = [];
+    if (query.toUpperCase() === this.schema.toUpperCase()) {
+      const tableInfo = await this.getColumnInfoForAllTables(this.schema);
+      contextItems.push({
+        name: `Info for all tables in ${this.schema}`,
+        content: `Db2 for i table Assistant: The following table and column information is from the ${query} schema. Utilize the provided schema and table metadata to assist the user:\n${JSON.stringify(
+          tableInfo
+        )}`,
+        description: "table metadata",
+      });
+    } else {
+      const tableInfo = await findPossibleTables(
+        null,
+        this.schema,
+        query.split(` `)
+      );
+      const markdownRefs = refsToMarkdown(tableInfo);
+
+      // add additional context for working with Db2 for i tables
+      contextItems.push({
+        name: `Instructions`,
+        content: `Db2 for i table Assistant: The following information is based on the ${query} table within the ${this.schema} schema. Utilize the provided schema and table metadata to assist the user. Only use valid Db2 for i SQL syntax and conventions. If input is unclear ask user to clarify`,
+        description: "instructions for working with Db2 for i tables",
+      });
+
+      contextItems.push(...createContinueContextItems(markdownRefs));
+    }
+    return contextItems;
+  }
+
+  async loadSubmenuItems(
+    args: LoadSubmenuItemsArgs
+  ): Promise<ContextSubmenuItem[]> {
+    const tables: BasicSQLObject[] = await Schemas.getObjects(this.schema, [
+      `tables`,
+    ]);
+
+    const schemaSubmenuItem: ContextSubmenuItem = {
+      id: this.schema,
+      title: this.schema,
+      description: `All table info in schema: ${this.schema}`,
+    };
+
+    const tableSubmenuItems: ContextSubmenuItem[] = tables.map((table) => ({
+      id: table.name,
+      title: table.name,
+      description: `${table.schema}-${table.name}`,
+    }));
+
+    return [schemaSubmenuItem, ...tableSubmenuItems];
+  }
+}
+
+export async function registerDb2iTablesProvider(schema?: string) {
+  if (!schema) {
+    return;
+  }
+  const continueID = `Continue.continue`;
+  const continueEx = vscode.extensions.getExtension(continueID);
+  if (continueEx) {
+    if (!continueEx.isActive) {
+      await continueEx.activate();
+    }
+
+    if (provider) {
+      provider.setCurrentSchema(schema);
+      // save continue config file to trigger a config reload to update list tables provider
+      const configFile = path.join(os.homedir(), `.continue`, `config.json`);
+      const now = new Date();
+      fs.utimes(configFile, now, now, (err) => {
+        if (err) {
+          console.error("Error saving Continue config file:", err);
+          return;
+        }
+      });
+    } else {
+      const continueAPI = continueEx?.exports;
+      provider = new ListDb2iTables(schema);
+      continueAPI?.registerCustomContextProvider(provider);
+    }
+  }
+}

src/database/table.ts

@@ -5,10 +5,11 @@ import { getInstance } from "../base";
 export default class Table {
   /**
    * @param {string} schema Not user input
-   * @param {string} name Not user input
+   * @param {string} table Not user input
    * @returns {Promise<TableColumn[]>}
    */
-  static async getItems(schema: string, name: string): Promise<TableColumn[]> {
+  static async getItems(schema: string, table?: string): Promise<TableColumn[]> {
+    const params = table ? [schema, table] : [schema];
     const sql = [
       `SELECT `,
       `  column.TABLE_SCHEMA,`,
@@ -30,11 +31,14 @@ export default class Table {
       `    column.table_schema = key.table_schema and`,
       `    column.table_name = key.table_name and`,
       `    column.column_name = key.column_name`,
-      `WHERE column.TABLE_SCHEMA = ? AND column.TABLE_NAME = ?`,
+      `WHERE column.TABLE_SCHEMA = ?`,
+      ...[
+        table ? `AND column.TABLE_NAME = ?` : ``,
+      ],
       `ORDER BY column.ORDINAL_POSITION`,
     ].join(` `);
 
-    return JobManager.runSQL(sql, {parameters: [schema, name]});
+    return JobManager.runSQL(sql, {parameters: params});
   }
 
   /**

src/extension.ts

@@ -23,8 +23,8 @@ import { JobManagerView } from "./views/jobManager/jobManagerView";
 import { SelfTreeDecorationProvider, selfCodesResultsView } from "./views/jobManager/selfCodes/selfCodesResultsView";
 import { registerContinueProvider } from "./aiProviders/continue/continueContextProvider";
 import { queryHistory } from "./views/queryHistoryView";
-import { activateChat, registerCopilotProvider } from "./aiProviders/copilot";
-import { SQLStatementChecker } from "./connection/syntaxChecker";
+import { registerCopilotProvider } from "./aiProviders/copilot";
+import { registerDb2iTablesProvider } from "./aiProviders/continue/listTablesContextProvider";
 import { setCheckerAvailableContext } from "./language/providers/problemProvider";
 
 export interface Db2i {
@@ -101,6 +101,10 @@ export function activate(context: vscode.ExtensionContext): Db2i {
     onConnectOrServerInstall().then(() => {
       exampleBrowser.refresh();
       selfCodesView.setRefreshEnabled(Configuration.get(`jobSelfViewAutoRefresh`) || false);
+      // register list tables
+      const currentJob = JobManager.getSelection();
+      const currentSchema = currentJob?.job.options.libraries[0];
+      registerDb2iTablesProvider(currentSchema);
       if (devMode && runTests) {
         runTests();
       }
@@ -113,6 +117,8 @@ export function activate(context: vscode.ExtensionContext): Db2i {
   // register continue provider
   registerContinueProvider();
 
+
+
   instance.subscribe(context, `disconnected`, `db2i-disconnected`, () => ServerComponent.reset());
 
   return { sqlJobManager: JobManager, sqlJob: (options?: JDBCOptions) => new OldSQLJob(options) };

src/views/jobManager/jobManagerView.ts

@@ -12,6 +12,7 @@ import { SelfCodesQuickPickItem } from "./selfCodes/selfCodesBrowser";
 import { updateStatusBar } from "./statusBar";
 import { setCancelButtonVisibility } from "../results";
 import { JDBCOptions } from "@ibm/mapepire-js/dist/src/types";
+import { provider, registerDb2iTablesProvider } from "../../aiProviders/continue/listTablesContextProvider";
 import { sqlLanguageStatus } from "../../language/providers";
 
 const selectJobCommand = `vscode-db2i.jobManager.selectJob`;
@@ -303,10 +304,28 @@ export class JobManagerView implements TreeDataProvider<any> {
     updateStatusBar();
 
     const selectedJob = JobManager.getSelection();
+    
+    // re-register db2i tables context provider with current schema
+    const selectedSchema = selectedJob?.job.options.libraries[0];
+    const currentSchema = provider?.getCurrentSchema();
+    if (
+      provider &&
+      selectedJob &&
+      selectedSchema &&
+      currentSchema.toLowerCase() !== selectedSchema.toLowerCase()
+    ) {
+      registerDb2iTablesProvider(selectedSchema);
+    }
 
-    setCancelButtonVisibility(selectedJob && selectedJob.job.getStatus() === "busy");
+    setCancelButtonVisibility(
+      selectedJob && selectedJob.job.getStatus() === "busy"
+    );
     sqlLanguageStatus.setState(selectedJob !== undefined);
-    commands.executeCommand(`setContext`, `vscode-db2i:jobManager.hasJob`, selectedJob !== undefined);
+    commands.executeCommand(
+      `setContext`,
+      `vscode-db2i:jobManager.hasJob`,
+      selectedJob !== undefined
+    );
   }
 
   getTreeItem(element: vscode.TreeItem) {