Skip to content

Commit

Permalink
Merge pull request #308 from codefori/update/ai-context
Browse files Browse the repository at this point in the history 
update context, use simple string list
  • Loading branch information
@worksofliam
worksofliam authored Jan 10, 2025
2 parents 8d401f6 + f9b0a3a commit f8b99fe
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 105 deletions.
111 changes: 29 additions & 82 deletions src/aiProviders/context.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,73 +14,31 @@ export function getCurrentSchema(): string {
};

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,`,
` 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 = '${Statement.delimName(schema, true)}'`,
`AND column.TABLE_NAME = '${Statement.delimName(tableName, true)}'`,
`ORDER BY column.ORDINAL_POSITION`,
].join(` `);

return await JobManager.runSQL(objectFindStatement);
}

export async function parsePromptForRefs(stream: vscode.ChatResponseStream, prompt: string[]): Promise<TableRefs> {
const tables: TableRefs = {};
for (const word of prompt) {
const [schema, table] = word.split(`.`);
const cleanedTable = table.replace(/[,\/#!?$%\^&\*;:{}=\-_`~()]/g, "");
if (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);
}
}
return tables;
interface MarkdownRef {
TABLE_NAME: string,
COLUMN_INFO?: string,
SCHMEA?: string,
}

export async function findPossibleTables(stream: vscode.ChatResponseStream, schema: string, words: string[]) {

let tables: TableRefs = {}

// parse all SCHEMA.TABLE references first
tables = await parsePromptForRefs(stream, words.filter(word => word.includes('.')));

// Parse all SCHEMA.TABLE references first
const schemaTableRefs = words.filter(word => word.includes('.'));
const justWords = words.map(word => word.replace(/[,\/#!?$%\^&\*;:{}=\-_`~()]/g, ""));

// Remove plurals from words
justWords.push(...justWords.filter(word => word.endsWith('s')).map(word => word.slice(0, -1)));

// filter prompt for possible refs to tables
// Filter prompt for possible refs to tables
const validWords = justWords
.filter(word => word.length > 2 && !word.endsWith('s') && !word.includes(`'`))
.map(word => `'${Statement.delimName(word, true)}'`);

const objectFindStatement = [
`SELECT `,
` column.TABLE_SCHEMA,`,
` column.TABLE_NAME,`,
` column.COLUMN_NAME,`,
` key.CONSTRAINT_NAME,`,
Expand All @@ -99,15 +57,18 @@ export async function findPossibleTables(stream: vscode.ChatResponseStream, sche
` 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}'`,
`WHERE column.TABLE_SCHEMA = '${Statement.delimName(schema, true)}'`,
...[
words.length > 0
? `AND column.TABLE_NAME in (${validWords.join(`, `)})`
: ``,
schemaTableRefs.length > 0
? `AND (column.TABLE_NAME in (${validWords.join(`, `)}) OR (${schemaTableRefs.map(ref => {
const [schema, table] = ref.split('.');
const cleanedTable = table.replace(/[,\/#!?$%\^&\*;:{}=\-_`~()]/g, "");
return `(column.TABLE_SCHEMA = '${Statement.delimName(schema, true)}' AND column.TABLE_NAME = '${Statement.delimName(cleanedTable, true)}')`;
}).join(' OR ')}))`
: `AND column.TABLE_NAME in (${validWords.join(`, `)})`,
],
`ORDER BY column.ORDINAL_POSITION`,
].join(` `);

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

Expand Down Expand Up @@ -141,38 +102,24 @@ export async function findPossibleTables(stream: vscode.ChatResponseStream, sche
export function refsToMarkdown(refs: TableRefs) {
const condensedResult = Object.keys(refs).length > 5;

let markdown: string[] = [];

let markdownRefs: MarkdownRef[] = [];
for (const tableName in refs) {
if (tableName.startsWith(`SYS`)) continue;

markdown.push(`# ${tableName}`, ``);

if (condensedResult) {
markdown.push(`| Column | Type | Text |`);
markdown.push(`| - | - | - |`);
} else {
markdown.push(
`| Column | Type | Nullable | Identity | Text | Constraint |`
);
markdown.push(`| - | - | - | - | - | - |`);
}
for (const column of refs[tableName]) {
if (condensedResult) {
markdown.push(
`| ${column.COLUMN_NAME} | ${column.DATA_TYPE} | ${column.COLUMN_TEXT} |`
);
} else {
markdown.push(
`| ${column.COLUMN_NAME} | ${column.DATA_TYPE} | ${column.IS_NULLABLE} | ${column.IS_IDENTITY} | ${column.COLUMN_TEXT} | ${column.CONSTRAINT_NAME} |`
);
}
}

markdown.push(``);
const curRef: MarkdownRef = {
TABLE_NAME: tableName,
SCHMEA: refs[tableName][0].TABLE_SCHEMA,
COLUMN_INFO: refs[tableName].map(column => {
const lengthPrecision = column.CHARACTER_MAXIMUM_LENGTH
? `(${column.CHARACTER_MAXIMUM_LENGTH}${column.NUMERIC_PRECISION ? `:${column.NUMERIC_PRECISION}` : ``})`
: ``;
return `${column.COLUMN_NAME}${column.COLUMN_TEXT ? ` - ${column.COLUMN_TEXT}` : ``} ${column.DATA_TYPE}${lengthPrecision} is_identity: ${column.IS_IDENTITY} is_nullable: ${column.IS_NULLABLE}`;
}).join(`\n`)
};
markdownRefs.push(curRef);
}

return markdown.join(`\n`);
return markdownRefs;
}

export async function getSystemStatus(): Promise<string> {
Expand Down
31 changes: 13 additions & 18 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 } from "../context";
import { canTalkToDb, findPossibleTables, refsToMarkdown } from "../context";
import {
ContextItem,
ContextProviderDescription,
Expand Down Expand Up @@ -135,7 +135,7 @@ export class db2ContextProvider implements IContextProvider {
content: prompt,
});
}

return contextItems;
default:
// const contextItems: ContextItem[] = [];
Expand All @@ -144,24 +144,19 @@ export class db2ContextProvider implements IContextProvider {
schema,
fullInput.split(` `)
);
for (const table of Object.keys(tableRefs)) {
const columnData: TableColumn[] = tableRefs[table];
if (columnData && columnData.length > 0) {
const tableSchema =
columnData.length > 0 ? columnData[0].TABLE_SCHEMA : null;
const markdownRefs = refsToMarkdown(tableRefs);

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

contextItems.push({
name: `${job.name}-${tableSchema}-${table}`,
description: `Schema and table information for ${table}`,
content: prompt,
});
}
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,
});
}

return contextItems;
}
} catch (error) {
Expand Down
16 changes: 11 additions & 5 deletions src/aiProviders/copilot/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import {
findPossibleTables,
getCurrentSchema,
getSystemStatus,
refsToMarkdown,
} from "../context";
import { JobManager } from "../../config";

Expand Down Expand Up @@ -89,6 +90,8 @@ export function activateChat(context: vscode.ExtensionContext) {
request.prompt.split(` `)
);

const markdownRefs = refsToMarkdown(refs);

messages = [
vscode.LanguageModelChatMessage.Assistant(
`You are a an IBM i savant speciallizing in database features in Db2 for i. Your job is to help developers write and debug their SQL along with offering SQL programming advice.`
Expand Down Expand Up @@ -118,11 +121,14 @@ export function activateChat(context: vscode.ExtensionContext) {
}

if (Object.keys(refs).length > 0) {
messages.push(
vscode.LanguageModelChatMessage.Assistant(
`Here are new table references ${JSON.stringify(refs)}`
),
);
for (const tableRef of markdownRefs) {
messages.push(
vscode.LanguageModelChatMessage.Assistant(
`Here are new table references for ${tableRef.TABLE_NAME} (Schema: ${tableRef.SCHMEA})\nFormat: column_name (column_text) type(length:precision) is_identity is_nullable\n${tableRef.COLUMN_INFO}`
),
);

}
}

stream.progress(`Building response...`);
Expand Down

0 comments on commit f8b99fe

Please sign in to comment.