Skip to content

Commit

Permalink
cleanup list tables provider
Browse files Browse the repository at this point in the history
  • Loading branch information
@ajshedivy
ajshedivy committed Jan 10, 2025
1 parent e8b8aca commit fc05435
Show file tree
Hide file tree
Showing 5 changed files with 83 additions and 59 deletions.
20 changes: 19 additions & 1 deletion src/aiProviders/context.ts
View file
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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[] = [];
Expand All @@ -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);
Expand Down
13 changes: 2 additions & 11 deletions src/aiProviders/continue/continueContextProvider.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;
}
Expand Down
102 changes: 57 additions & 45 deletions src/aiProviders/continue/listTablesContextProvider.ts
View file
Original file line number Diff line number Diff line change
@@ -1,38 +1,47 @@
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 {
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 { JobManager } from "../../config";
import { JobInfo } from "../../connection/manager";
import Schemas from "../../database/schemas";
import Table from "../../database/table";
import { findPossibleTables } from "../context";
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"
}
type: "submenu",
};

export let listDb2TableContextProvider: Boolean = false;
let provider: ListDb2iTables = undefined;

class ListDb2iTables implements IContextProvider {
private schema: string;

constructor(schema: string) {
this.schema = schema;
}

get description(): ContextProviderDescription {
return listDb2Table;
}

getCurrentJob(): JobInfo {
const currentJob: JobInfo = JobManager.getSelection();
return currentJob;
setCurrentSchema(schema: string) {
this.schema = schema;
}

private getDefaultSchema = (): string => {
const currentJob: JobInfo = this.getCurrentJob();
return currentJob?.job.options.libraries[0] || "QGPL";
};

async getColumnInfoForAllTables(schema: string) {
const items: TableColumn[] = await Table.getItems(schema);

Expand All @@ -48,42 +57,47 @@ class ListDb2iTables implements IContextProvider {
query: string,
extras: ContextProviderExtras
): Promise<ContextItem[]> {
let contextitems: ContextItem[] = [];
const schema = this.getDefaultSchema();
if (query.toUpperCase() === schema.toUpperCase()) {
const tableInfo = await this.getColumnInfoForAllTables(schema);
contextitems.push({
name: `Info for all tables in ${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)}`,
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,
schema,
this.schema,
query.split(` `)
);
contextitems.push({
name: `${query}`,
content: `Db2 for i table Assistant: The following information is based on the ${query} table within the ${schema} schema. Utilize the provided schema and table metadata to assist the user:\n${JSON.stringify(tableInfo)}`,
description: "table metadata",
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;
return contextItems;
}

async loadSubmenuItems(
args: LoadSubmenuItemsArgs
): Promise<ContextSubmenuItem[]> {
const schema = this.getDefaultSchema();
const tables: BasicSQLObject[] = await Schemas.getObjects(schema, [
const tables: BasicSQLObject[] = await Schemas.getObjects(this.schema, [
`tables`,
]);

const schemaSubmenuItem: ContextSubmenuItem = {
id: schema,
title: schema,
description: `All table info in schema: ${schema}`,
id: this.schema,
title: this.schema,
description: `All table info in schema: ${this.schema}`,
};

const tableSubmenuItems: ContextSubmenuItem[] = tables.map((table) => ({
Expand All @@ -96,32 +110,30 @@ class ListDb2iTables implements IContextProvider {
}
}

export async function registerDb2iTablesProvider() {
const provider = new ListDb2iTables();
export async function registerDb2iTablesProvider(schema: string) {
const continueID = `Continue.continue`;
const continueEx = vscode.extensions.getExtension(continueID);
if (continueEx) {
if (!continueEx.isActive) {
await continueEx.activate();
}

const continueAPI = continueEx?.exports;
if (listDb2TableContextProvider) {

if (provider) {
provider.setCurrentSchema(schema);
// HACK: re register context provider work around
// 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);
console.error("Error saving Continue config file:", err);
return;
}
vscode.window.showInformationMessage('Updated @Db2-Tables!');
});
} else {
const continueAPI = continueEx?.exports;
provider = new ListDb2iTables(schema);
continueAPI?.registerCustomContextProvider(provider);
listDb2TableContextProvider = true;
}
}
}
}
2 changes: 1 addition & 1 deletion src/extension.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ export function activate(context: vscode.ExtensionContext): Db2i {
exampleBrowser.refresh();
selfCodesView.setRefreshEnabled(Configuration.get(`jobSelfViewAutoRefresh`) || false);
// register list tables
registerDb2iTablesProvider();
registerDb2iTablesProvider(JobManager.getSelection().job.options.libraries[0]);
if (devMode && runTests) {
runTests();
}
Expand Down
5 changes: 4 additions & 1 deletion src/views/jobManager/jobManagerView.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,8 @@ export class JobManagerView implements TreeDataProvider<any> {
await selected.job.connect();

// re register tables provider with potential new Schema
await registerDb2iTablesProvider();
const schema = selected.job.options.libraries[0];
await registerDb2iTablesProvider(schema);
} catch (e) {
window.showErrorMessage(`Failed to start new job with updated properties.`);
}
Expand Down Expand Up @@ -280,6 +281,8 @@ export class JobManagerView implements TreeDataProvider<any> {
vscode.commands.registerCommand(selectJobCommand, async (selectedName: string) => {
if (selectedName) {
await JobManager.setSelection(selectedName);
const schema = JobManager.getSelection().job.options.libraries[0];
await registerDb2iTablesProvider(schema);
this.refresh();
}
}),
Expand Down

0 comments on commit fc05435

Please sign in to comment.