Skip to content

Commit

Permalink
use same context backend as chatai
Browse files Browse the repository at this point in the history
  • Loading branch information
@ajshedivy
ajshedivy committed Oct 14, 2024
1 parent 38d26c2 commit 4fa956f
Show file tree
Hide file tree
Showing 5 changed files with 32 additions and 153 deletions.
4 changes: 2 additions & 2 deletions package-lock.json
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 5 additions & 0 deletions package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1090,6 +1090,11 @@
"command": "vscode-db2i.self.displayDetails",
"when": "view == vscode-db2i.self.nodes && viewItem == selfCodeNode",
"group": "navigation"
},
{
"command": "vscode-db2i.self.explainSelf",
"when": "view == vscode-db2i.self.nodes && viewItem == selfCodeNode",
"group": "navigation"
}
],
"editor/title": [
Expand Down
160 changes: 10 additions & 150 deletions src/aiProviders/continue/continueContextProvider.ts
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
import * as vscode from "vscode";
import { JobManager } from "../../config";
import { JobInfo, SQLJobManager } from "../../connection/manager";
import { JobInfo } from "../../connection/manager";
import Statement from "../../database/statement";
import { table } from "console";
import { selfCodesResultsView } from "../../views/jobManager/selfCodes/selfCodesResultsView";
import { SelfCodeNode } from "../../views/jobManager/selfCodes/nodes";
import { SQLJob } from "../../connection/sqlJob";
import { findPossibleTables } from "../../chat/context";
import { ContextItem, ContextProviderDescription, ContextProviderExtras, ContextSubmenuItem, IContextProvider, LoadSubmenuItemsArgs } from "../../..";

const db2ContextProviderDesc: ContextProviderDescription = {
title: "db2i",
Expand All @@ -14,8 +13,6 @@ const db2ContextProviderDesc: ContextProviderDescription = {
type: "normal",
};

type TableRefs = { [key: string]: TableColumn[]}

/**
* - Get Existing Db2 connection from vscode
* - Use connection to build query for database to add additional context
Expand All @@ -36,114 +33,6 @@ export class db2ContextProvider implements IContextProvider {
return currentJob?.job.options.libraries[0] || "QGPL";
};

async findPossibleTables(schema: string, words: string[]) {
// let refSchema: string = "";
// let table: string = "";
// words.forEach(item => {
// if (item.includes(`.`)) {
// refSchema = item.split(`.`)[0];
// table = item.split(`.`)[1];
// }
// })

words = words.map((word) =>
word.replace(/[.,\/#!?$%\^&\*;:{}=\-_`~()]/g, "")
);

// Add extra words for words with S at the end, to ignore possible plurals
words.forEach((item) => {
if (item.endsWith(`s`)) {
words.push(item.slice(0, -1));
}
});

const validWords = words
.filter((item) => item.length > 2 && !item.includes(`'`))
.map((item) => `'${Statement.delimName(item, true)}'`);

const objectFindStatement = [
`SELECT `,
` column.TABLE_NAME,`,
` column.COLUMN_NAME,`,
` key.CONSTRAINT_NAME,`,
` column.DATA_TYPE, `,
` column.CHARACTER_MAXIMUM_LENGTH,`,
` column.NUMERIC_SCALE, `,
` column.NUMERIC_PRECISION,`,
` column.IS_NULLABLE, `,
// ` column.HAS_DEFAULT, `,
// ` column.COLUMN_DEFAULT, `,
` column.COLUMN_TEXT, `,
` column.IS_IDENTITY`,
`FROM QSYS2.SYSCOLUMNS2 as column`,
`LEFT JOIN QSYS2.syskeycst as key`,
` on `,
` column.table_schema = key.table_schema and`,
` column.table_name = key.table_name and`,
` column.column_name = key.column_name`,
`WHERE column.TABLE_SCHEMA = '${schema}'`,
...[
words.length > 0
? `AND column.TABLE_NAME in (${validWords.join(`, `)})`
: ``,
],
`ORDER BY column.ORDINAL_POSITION`,
].join(` `);

// TODO
const result: TableColumn[] = await JobManager.runSQL(objectFindStatement);

const tables: TableRefs = {};

for (const row of result) {
if (!tables[row.TABLE_NAME]) {
tables[row.TABLE_NAME] = [];
}

tables[row.TABLE_NAME].push(row);
}

return tables;
}

async findAllTables(schema: string) {
const objectFindStatement = [
`SELECT `,
` column.TABLE_NAME,`,
` column.COLUMN_NAME,`,
` key.CONSTRAINT_NAME,`,
` column.DATA_TYPE, `,
` column.CHARACTER_MAXIMUM_LENGTH,`,
` column.NUMERIC_SCALE, `,
` column.NUMERIC_PRECISION,`,
` column.IS_NULLABLE, `,
// ` column.HAS_DEFAULT, `,
// ` column.COLUMN_DEFAULT, `,
` column.COLUMN_TEXT, `,
` column.IS_IDENTITY`,
`FROM QSYS2.SYSCOLUMNS2 as column`,
`LEFT JOIN QSYS2.syskeycst as key`,
` on `,
` column.table_schema = key.table_schema and`,
` column.table_name = key.table_name and`,
` column.column_name = key.column_name`,
`WHERE column.TABLE_SCHEMA = '${schema}'`,
].join(` `);

const result: TableColumn[] = await JobManager.runSQL(objectFindStatement);

const tables: TableRefs = {};

for (const row of result) {
if (!tables[row.TABLE_NAME]) {
tables[row.TABLE_NAME] = [];
}

tables[row.TABLE_NAME].push(row);
}

return tables;
}
async getSelfCodes(selected: JobInfo): Promise<SelfCodeNode[]|undefined> {
const current_job = selected.job.id;
const content = `SELECT
Expand All @@ -154,7 +43,7 @@ export class db2ContextProvider implements IContextProvider {
order by logged_time desc`;

try {
const result = await selected.job.query<SelfCodeNode>(content).run(10000);
const result = await selected.job.query<SelfCodeNode>(content).execute(10000);
if (result.success) {
const data: SelfCodeNode[] = result.data.map((row) => ({
...row,
Expand Down Expand Up @@ -202,19 +91,21 @@ export class db2ContextProvider implements IContextProvider {
return contextItems;
default:
// const contextItems: ContextItem[] = [];
const tableRefs = await this.findPossibleTables(
const tableRefs = await findPossibleTables(
null,
schema,
fullInput.split(` `)
);
for (const table of Object.keys(tableRefs)) {
const columnData: TableColumn[] = tableRefs[table];
const tableSchema = columnData.length > 0 ? columnData[0].TABLE_SCHEMA : null;

// create context item
let prompt = `Db2 for i schema ${schema} table ${table}\n`;
let prompt = `Db2 for i schema ${tableSchema} table ${table}\n`;
prompt += `Column Info: ${JSON.stringify(columnData)}\n\n`;

contextItems.push({
name: `${job.name}-${schema}-${table}`,
name: `${job.name}-${tableSchema}-${table}`,
description: `Schema and table information or ${table}`,
content: prompt,
});
Expand Down Expand Up @@ -255,41 +146,10 @@ export class db2ContextProvider implements IContextProvider {
}
}

class MyCustomProvider implements IContextProvider {

get description(): ContextProviderDescription {
return {
title: "custom",
displayTitle: "Custom",
description: "Custom description",
type: "normal",
};
}

async getContextItems(
query: string,
extras: ContextProviderExtras
): Promise<ContextItem[]> {
return [
{
name: "Custom",
description: "Custom description",
content: "Custom content",
},
];
}

async loadSubmenuItems(
args: LoadSubmenuItemsArgs
): Promise<ContextSubmenuItem[]> {
return [];
}
}

export async function registerContinueProvider() {
const provider = new db2ContextProvider();
const continueID = `Continue.continue`;
const continueEx = vscode.extensions.getExtension(continueID);
const continueAPI = continueEx?.exports;
continueAPI?.registerCustomContextProvider(provider);
}
}
5 changes: 4 additions & 1 deletion src/chat/context.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ export type TableRefs = { [key: string]: TableColumn[] };
export async function getTableMetaData(schema: string, tableName: string): Promise<TableColumn[]> {
const objectFindStatement = [
`SELECT `,
` column.TABLE_SCHEMA,`,
` column.TABLE_NAME,`,
` column.COLUMN_NAME,`,
` key.CONSTRAINT_NAME,`,
Expand Down Expand Up @@ -50,7 +51,9 @@ export async function parsePromptForRefs(stream: vscode.ChatResponseStream, prom
const [schema, table] = word.split(`.`);
const cleanedTable = table.replace(/[,\/#!?$%\^&\*;:{}=\-_`~()]/g, "");
if (schema && cleanedTable) {
stream.progress(`looking up information for ${schema}.${cleanedTable}`)
if (stream !== null) {
stream.progress(`looking up information for ${schema}.${cleanedTable}`)
}
const data = await getTableMetaData(schema, cleanedTable);
tables[cleanedTable] = tables[cleanedTable] || [];
tables[cleanedTable].push(...data);
Expand Down
11 changes: 11 additions & 0 deletions src/views/jobManager/selfCodes/contributes.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,12 @@
"category": "Db2 for i",
"icon": "$(info)"
},
{
"command": "vscode-db2i.self.explainSelf",
"title": "Explain SELF Code with continue",
"category": "Db2 for i",
"icon": "$(debug-alt)"
},
{
"command": "vscode-db2i.self.help",
"title": "Open SELF Documentation",
Expand Down Expand Up @@ -147,6 +153,11 @@
"command": "vscode-db2i.self.displayDetails",
"when": "view == vscode-db2i.self.nodes && viewItem == selfCodeNode",
"group": "navigation"
},
{
"command": "vscode-db2i.self.explainSelf",
"when": "view == vscode-db2i.self.nodes && viewItem == selfCodeNode",
"group": "navigation"
}
]
}
Expand Down

0 comments on commit 4fa956f

Please sign in to comment.