diff --git a/src/aiProviders/context.ts b/src/aiProviders/context.ts index bd5d9907..cff091d0 100644 --- a/src/aiProviders/context.ts +++ b/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 { const sqlStatment = `SELECT * FROM TABLE(QSYS2.SYSTEM_STATUS(RESET_STATISTICS=>'YES',DETAILED_INFO=>'ALL')) X`; const result = await JobManager.runSQL(sqlStatment, undefined); diff --git a/src/aiProviders/continue/continueContextProvider.ts b/src/aiProviders/continue/continueContextProvider.ts index c4f80500..0e958b26 100644 --- a/src/aiProviders/continue/continueContextProvider.ts +++ b/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; } diff --git a/src/aiProviders/continue/listTablesContextProvider.ts b/src/aiProviders/continue/listTablesContextProvider.ts new file mode 100644 index 00000000..47969a14 --- /dev/null +++ b/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 { + 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 { + 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); + } + } +} diff --git a/src/database/table.ts b/src/database/table.ts index 06517510..27972458 100644 --- a/src/database/table.ts +++ b/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} */ - static async getItems(schema: string, name: string): Promise { + static async getItems(schema: string, table?: string): Promise { + 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}); } /** diff --git a/src/extension.ts b/src/extension.ts index b85eb065..219b0a39 100644 --- a/src/extension.ts +++ b/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) }; diff --git a/src/views/jobManager/jobManagerView.ts b/src/views/jobManager/jobManagerView.ts index 41c91a08..c8cbbb68 100644 --- a/src/views/jobManager/jobManagerView.ts +++ b/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 { 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) {