Merge branch 'codefori:main' into fix/token_tests

janfh · web-flow · commit 2101878d2654 · 2025-01-15T07:56:28.000+01:00
diff --git a/src/aiProviders/context.ts b/src/aiProviders/context.ts
@@ -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,`,
@@ -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);
 
@@ -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> {
diff --git 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 } from "../context";
+import { canTalkToDb, findPossibleTables, refsToMarkdown } from "../context";
 import {
   ContextItem,
   ContextProviderDescription,
@@ -135,7 +135,7 @@ export class db2ContextProvider implements IContextProvider {
                 content: prompt,
               });
             }
-  
+
             return contextItems;
           default:
             // const contextItems: ContextItem[] = [];
@@ -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) {
diff --git a/src/aiProviders/copilot/index.ts b/src/aiProviders/copilot/index.ts
@@ -5,6 +5,7 @@ import {
   findPossibleTables,
   getCurrentSchema,
   getSystemStatus,
+  refsToMarkdown,
 } from "../context";
 import { JobManager } from "../../config";
 
@@ -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.`
@@ -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...`);
diff --git a/src/connection/manager.ts b/src/connection/manager.ts
@@ -30,7 +30,7 @@ export class SQLJobManager {
       const config = instance.getConfig();
 
       const newJob = predefinedJob || (new OldSQLJob({
-        libraries: [config.currentLibrary, ...config.libraryList],
+        libraries: [config.currentLibrary, ...config.libraryList.filter((item) => item != config.currentLibrary)],
         naming: `system`,
         "full open": false,
         "transaction isolation": "none",
@@ -161,4 +161,4 @@ export class SQLJobManager {
   static getSelfDefault(): SelfValue {
     return Configuration.get<SelfValue>(`jobSelfDefault`) || `*NONE`;
   }
-}
+}
