Skip to content

Commit 378f9d1

Browse files
ajshedivyajshedivy
committed
Refactor Db2i context providers and enhance SQL prompt guidelines
1 parent 7d8efa6 commit 378f9d1

File tree

5 files changed

+197
-39
lines changed

5 files changed

+197
-39
lines changed

‎src/aiProviders/context.ts

Lines changed: 141 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ import { JobManager } from "../config";
22
import * as vscode from "vscode";
33
import Statement from "../database/statement";
44
import { ContextItem } from "@continuedev/core";
5+
import { DB2_SYSTEM_PROMPT } from "./continue/prompts";
6+
import Schemas, { AllSQLTypes } from "../database/schemas";
57

68
export function canTalkToDb() {
79
return JobManager.getSelection() !== undefined;
@@ -21,6 +23,130 @@ interface MarkdownRef {
2123
SCHMEA?: string,
2224
}
2325

26+
export async function buildSchemaSemantic(schema: string) {
27+
/**
28+
* Cleans an object by removing properties with undefined, null, or empty string values.
29+
* If the value is an object (but not an array), it cleans that object as well.
30+
* This function is meant for properties other than the top-level `schema` and `system`.
31+
*
32+
* @param obj - The object to clean.
33+
* @returns A new object with only non-empty properties.
34+
*/
35+
const allInfo: BasicSQLObject[] = await Schemas.getObjects(schema, AllSQLTypes);
36+
function cleanObject(obj: any): any {
37+
const cleaned: any = {};
38+
39+
for (const key in obj) {
40+
if (Object.prototype.hasOwnProperty.call(obj, key)) {
41+
const value = obj[key];
42+
43+
// Skip undefined, null or empty string values.
44+
if (
45+
value === undefined ||
46+
value === null ||
47+
(typeof value === "string" && value.trim() === "")
48+
) {
49+
continue;
50+
}
51+
52+
// If the property is a plain object (and not an array), clean it recursively.
53+
if (typeof value === "object" && !Array.isArray(value)) {
54+
const nested = cleanObject(value);
55+
if (Object.keys(nested).length > 0) {
56+
cleaned[key] = nested;
57+
}
58+
} else {
59+
cleaned[key] = value;
60+
}
61+
}
62+
}
63+
64+
return cleaned;
65+
}
66+
67+
/**
68+
* Filters a single BasicSQLObject by:
69+
* - Removing the top-level `schema` and `system` properties.
70+
* - Removing any property (or nested property) that is undefined, null, or an empty string.
71+
*
72+
* @param obj - The BasicSQLObject to filter.
73+
* @returns A new object with filtered properties.
74+
*/
75+
function filterBasicSQLObject(obj: BasicSQLObject): Partial<BasicSQLObject> {
76+
// Remove the top-level `schema` and `system` properties.
77+
const { schema, system, ...rest } = obj;
78+
79+
// Clean the remaining properties.
80+
return cleanObject(rest);
81+
}
82+
83+
/**
84+
* Processes an array of BasicSQLObject items.
85+
*
86+
* @param data - The array of BasicSQLObject items.
87+
* @returns A new array of filtered objects.
88+
*/
89+
function filterBasicSQLObjects(
90+
data: BasicSQLObject[]
91+
): Partial<BasicSQLObject>[] {
92+
return data.map(filterBasicSQLObject);
93+
}
94+
95+
96+
const compressedData = filterBasicSQLObjects(allInfo);
97+
return compressedData;
98+
}
99+
100+
101+
export async function generateTableDefinition(schema: string, input: string[]) {
102+
let tables: {[key: string]: string|undefined} = {}
103+
// Parse all SCHEMA.TABLE references first
104+
const schemaTableRefs = input.filter(word => word.includes('.'));
105+
const justWords = input.map(word => word.replace(/[,\/#!?$%\^&\*;:{}=\-_`~()]/g, ""));
106+
107+
// Remove plurals from words
108+
justWords.push(...justWords.filter(word => word.endsWith('s')).map(word => word.slice(0, -1)));
109+
110+
// Filter prompt for possible refs to tables
111+
const validWords = justWords
112+
.filter(word => word.length > 2 && !word.endsWith('s') && !word.includes(`'`))
113+
.map(word => `${Statement.delimName(word, true)}`);
114+
115+
const allTables: BasicSQLObject[] = await Schemas.getObjects(schema, [
116+
`tables`,
117+
]);
118+
119+
const filteredTables = validWords.filter((word) =>
120+
allTables.some((table) => table.name == word)
121+
);
122+
123+
await Promise.all(filteredTables.map(async (token) => {
124+
try {
125+
const content = await Schemas.generateSQL(schema, token, `TABLE`);
126+
if (content) {
127+
tables[token] = content;
128+
}
129+
} catch (e) {
130+
// ignore
131+
}
132+
}));
133+
134+
// check for QSYS2
135+
// for (const item of schemaTableRefs) {
136+
// const [curSchema, table] = item.split(`.`);
137+
// if (curSchema.toUpperCase() === `QSYS2`) {
138+
// try {
139+
// const content = await Schemas.getObjects
140+
// } catch (e) {
141+
// continue
142+
// }
143+
// }
144+
// }
145+
146+
return tables;
147+
148+
}
149+
24150
export async function findPossibleTables(stream: vscode.ChatResponseStream, schema: string, words: string[]) {
25151

26152
let tables: TableRefs = {}
@@ -126,15 +252,23 @@ export function refsToMarkdown(refs: TableRefs): MarkdownRef[] {
126252
export function createContinueContextItems(refs: MarkdownRef[]) {
127253
const contextItems: ContextItem[] = [];
128254
const job = JobManager.getSelection();
129-
for (const tableRef of refs) {
130-
let prompt = `Table: ${tableRef.TABLE_NAME} (Schema: ${tableRef.SCHMEA}) Column Information:\n`;
131-
prompt += `Format: column_name (column_text) type(length:precision) is_identity is_nullable\n`
132-
prompt += `${tableRef.COLUMN_INFO}`;
255+
if (refs.length === 0) {
133256
contextItems.push({
134-
name: `${job.name}-${tableRef.SCHMEA}-${tableRef.TABLE_NAME}`,
135-
description: `Column information for ${tableRef.TABLE_NAME}`,
136-
content: prompt,
257+
name: `SYSTEM PROMPT`,
258+
description: `system prompt context`,
259+
content: DB2_SYSTEM_PROMPT + `\n\nNo references found`,
137260
});
261+
} else {
262+
for (const tableRef of refs) {
263+
let prompt = `Table: ${tableRef.TABLE_NAME} (Schema: ${tableRef.SCHMEA}) Column Information:\n`;
264+
prompt += `Format: column_name (column_text) type(length:precision) is_identity is_nullable\n`
265+
prompt += `${tableRef.COLUMN_INFO}`;
266+
contextItems.push({
267+
name: `${job.name}-${tableRef.SCHMEA}-${tableRef.TABLE_NAME}`,
268+
description: `Column information for ${tableRef.TABLE_NAME}`,
269+
content: DB2_SYSTEM_PROMPT + prompt,
270+
});
271+
}
138272
}
139273

140274
return contextItems;

‎src/aiProviders/continue/continueContextProvider.ts

Lines changed: 22 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import * as vscode from "vscode";
22
import { JobManager } from "../../config";
33
import { JobInfo } from "../../connection/manager";
44
import { SelfCodeNode } from "../../views/jobManager/selfCodes/nodes";
5-
import { canTalkToDb, createContinueContextItems, findPossibleTables, refsToMarkdown } from "../context";
5+
import { buildSchemaSemantic, canTalkToDb, createContinueContextItems, findPossibleTables, generateTableDefinition, refsToMarkdown } from "../context";
66
import {
77
ContextItem,
88
ContextProviderDescription,
@@ -12,6 +12,7 @@ import {
1212
LoadSubmenuItemsArgs,
1313
} from "@continuedev/core";
1414
import { DB2_SELF_PROMPT, DB2_SYSTEM_PROMPT } from "./prompts";
15+
import { table } from "console";
1516

1617
export let isContinueActive = false;
1718

@@ -111,10 +112,11 @@ export class db2ContextProvider implements IContextProvider {
111112
const job: JobInfo = this.getCurrentJob();
112113
const schema = this.getDefaultSchema();
113114
const fullInput = extras.fullInput;
115+
const schemaSemantic = await buildSchemaSemantic(schema);
114116
contextItems.push({
115-
name: `SYSTEM PROMPT`,
116-
description: `system prompt context`,
117-
content: DB2_SYSTEM_PROMPT,
117+
name: `SCHEMA Semantic`,
118+
description: `${schema} definition`,
119+
content: JSON.stringify(schemaSemantic),
118120
});
119121
try {
120122
switch (true) {
@@ -139,14 +141,23 @@ export class db2ContextProvider implements IContextProvider {
139141
return contextItems;
140142
default:
141143
// const contextItems: ContextItem[] = [];
142-
const tableRefs = await findPossibleTables(
143-
null,
144-
schema,
145-
fullInput.split(` `)
146-
);
147-
const markdownRefs = refsToMarkdown(tableRefs);
144+
// const tableRefs = await findPossibleTables(
145+
// null,
146+
// schema,
147+
// fullInput.split(` `)
148+
// );
149+
// const markdownRefs = refsToMarkdown(tableRefs);
148150

149-
contextItems.push(...createContinueContextItems(markdownRefs));
151+
// contextItems.push(...createContinueContextItems(markdownRefs));
152+
153+
const tablesRefs = await generateTableDefinition(schema, fullInput.split(` `));
154+
for (const table in tablesRefs) {
155+
contextItems.push({
156+
name: `table definition for ${table}`,
157+
content: tablesRefs[table],
158+
description: `Table definition`
159+
})
160+
}
150161

151162
return contextItems;
152163
}

‎src/aiProviders/continue/listTablesContextProvider.ts

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -20,20 +20,32 @@ import {
2020

2121
const listDb2Table: ContextProviderDescription = {
2222
title: "list Db2i Tables",
23-
displayTitle: "Db2i-tables",
23+
displayTitle: `Db2i-{tables}`,
2424
description: "Add Db2i Table info to Context",
2525
type: "submenu",
26+
dependsOnIndexing: true
2627
};
2728

28-
export let provider: ListDb2iTables = undefined;
29+
interface SchemaContextProvider {
30+
schema: string;
31+
provider: IContextProvider,
32+
}
33+
34+
let providers: SchemaContextProvider[] = []
2935

3036
class ListDb2iTables implements IContextProvider {
3137
constructor(private schema: string) {
3238
this.schema = schema;
3339
}
3440

3541
get description(): ContextProviderDescription {
36-
return listDb2Table;
42+
return {
43+
title: `Db2i-${this.schema}`,
44+
displayTitle: `Db2i-${this.schema}`,
45+
description: "Add Db2i Table info to Context",
46+
type: "submenu",
47+
dependsOnIndexing: true
48+
};
3749
}
3850

3951
setCurrentSchema(schema: string) {
@@ -123,21 +135,14 @@ export async function registerDb2iTablesProvider(schema?: string) {
123135
await continueEx.activate();
124136
}
125137

126-
if (provider) {
127-
provider.setCurrentSchema(schema);
128-
// save continue config file to trigger a config reload to update list tables provider
129-
const configFile = path.join(os.homedir(), `.continue`, `config.json`);
130-
const now = new Date();
131-
fs.utimes(configFile, now, now, (err) => {
132-
if (err) {
133-
console.error("Error saving Continue config file:", err);
134-
return;
135-
}
136-
});
138+
const existingProvider: SchemaContextProvider = providers.find(p => p.schema === schema);
139+
if (existingProvider !== undefined) {
140+
return;
137141
} else {
138142
const continueAPI = continueEx?.exports;
139-
provider = new ListDb2iTables(schema);
143+
let provider = new ListDb2iTables(schema);
140144
continueAPI?.registerCustomContextProvider(provider);
145+
providers.push({provider: provider, schema: schema});
141146
}
142147
}
143148
}

‎src/aiProviders/continue/prompts.ts

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,15 @@
1-
export const DB2_SYSTEM_PROMPT = `You are an expert in IBM i, specializing in database features of Db2 for i. Your role is to assist developers in writing and debugging their SQL queries, as well as providing SQL programming advice and best practices.`;
1+
export const DB2_SYSTEM_PROMPT = `You are an expert in IBM i, specializing in the database features of Db2 for i. Your role is to assist developers in writing and debugging SQL queries using only the provided table and column metadata.
2+
3+
Guidelines:
4+
- Use only the provided table and column metadata. If no metadata is available, inform the user that no references were found.
5+
- Do not generate or assume table names, column names, or SQL references that are not explicitly provided.
6+
- Ensure all SQL statements are valid for Db2 for i.
7+
- If the user’s request is ambiguous or lacks necessary details, ask clarifying questions before generating a response.
8+
- suggest relevant follow-up questions to help the user refine their request
9+
10+
Stay accurate, clear, and helpful while guiding the user through SQL query development.
11+
12+
Refernces:`;
213

314
export const DB2_SELF_PROMPT = [`Db2 for i self code errors\n`,
415
`Summarize the SELF code errors provided. The SQL Error Logging Facility (SELF) provides a mechanism that can be used to understand when SQL statements are encountering specific SQL errors or warnings. SELF is built into Db2 for i and can be enabled in specific jobs or system wide. Provide additional details about the errors and how to fix them.\n`,

‎src/views/jobManager/jobManagerView.ts

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import { SelfCodesQuickPickItem } from "./selfCodes/selfCodesBrowser";
1212
import { updateStatusBar } from "./statusBar";
1313
import { setCancelButtonVisibility } from "../results";
1414
import { JDBCOptions } from "@ibm/mapepire-js/dist/src/types";
15-
import { provider, registerDb2iTablesProvider } from "../../aiProviders/continue/listTablesContextProvider";
15+
import { registerDb2iTablesProvider } from "../../aiProviders/continue/listTablesContextProvider";
1616
import { sqlLanguageStatus } from "../../language/providers";
1717

1818
const selectJobCommand = `vscode-db2i.jobManager.selectJob`;
@@ -307,12 +307,9 @@ export class JobManagerView implements TreeDataProvider<any> {
307307

308308
// re-register db2i tables context provider with current schema
309309
const selectedSchema = selectedJob?.job.options.libraries[0];
310-
const currentSchema = provider?.getCurrentSchema();
311310
if (
312-
provider &&
313311
selectedJob &&
314-
selectedSchema &&
315-
currentSchema.toLowerCase() !== selectedSchema.toLowerCase()
312+
selectedSchema
316313
) {
317314
registerDb2iTablesProvider(selectedSchema);
318315
}

0 commit comments

Comments
 (0)