Fetch related items when getting info about one object

Signed-off-by: worksofliam <[email protected]>

Author: worksofliam

src/aiProviders/context.ts

@@ -167,7 +167,7 @@ function splitUpUserInput(input: string): string[] {
  *
  * @param {string} input - A string that may contain table references.
  */
-export async function getSqlContextItems(input: string): Promise<ContextDefinition[]> {
+export async function getSqlContextItems(input: string): Promise<{items: ContextDefinition[], refs: ResolvedSqlObject[]}> {
   // Parse all SCHEMA.TABLE references first
   const tokens = splitUpUserInput(input);
 
@@ -201,30 +201,39 @@ export async function getSqlContextItems(input: string): Promise<ContextDefiniti
 
   const allObjects = await Schemas.resolveObjects(possibleRefs);
 
-  const contextItems = (await Promise.all(
+  const contextItems = await getContentItemsForRefs(allObjects);
+
+  return {
+    items: contextItems,
+    refs: allObjects,
+  };
+}
+
+export async function getContentItemsForRefs(allObjects: ResolvedSqlObject[]): Promise<ContextDefinition[]> {
+  const items: (ContextDefinition|undefined)[] = await Promise.all(
     allObjects.map(async (o) => {
       try {
         if (o.sqlType === `SCHEMA`) {
           // TODO: maybe we want to include info about a schema here?
           return undefined;
-          
+
         } else {
           const content = await Schemas.generateSQL(o.schema, o.name, o.sqlType);
 
           return {
             id: o.name,
             type: o.sqlType,
             content: content,
-          }
+          };
         }
 
       } catch (e) {
         return undefined;
       }
     })
-  )).filter((item) => item !== undefined);
+  );
 
-  return contextItems;
+  return items.filter((item) => item !== undefined);
 }
 
 /**

src/aiProviders/prompt.ts

@@ -1,7 +1,9 @@
 import { JobManager } from "../config";
 import Configuration from "../configuration";
 import { JobInfo } from "../connection/manager";
-import { buildSchemaDefinition, canTalkToDb, getSqlContextItems } from "./context";
+import Schemas from "../database/schemas";
+import Statement from "../database/statement";
+import { buildSchemaDefinition, canTalkToDb, getContentItemsForRefs, getSqlContextItems } from "./context";
 import { DB2_SYSTEM_PROMPT } from "./prompts";
 
 export interface PromptOptions {
@@ -48,21 +50,39 @@ export async function buildPrompt(input: string, options: PromptOptions = {}): P
     // TODO: self?
 
     progress(`Finding objects to work with...`);
-    const refs = await getSqlContextItems(input);
+    const context = await getSqlContextItems(input);
 
     if (options.history) {
       contextItems.push(...options.history);
     }
 
-    for (const table of refs) {
+    for (const sqlObj of context.items) {
       contextItems.push({
-        name: `table definition for ${table.id}`,
-        content: table.content,
-        description: `${table.type} definition`,
+        name: `${sqlObj.type.toLowerCase()} definition for ${sqlObj.id}`,
+        content: sqlObj.content,
+        description: `${sqlObj.type} definition`,
         type: `assistant`
       });
     }
 
+    // If the user only requests one reference, then let's find related objects
+    if (context.refs.length === 1) {
+      const ref = context.refs[0];
+      progress(`Finding objects related to ${Statement.prettyName(ref.name)}...`);
+
+      const relatedObjects = await Schemas.getRelatedObjects(ref);
+      const contentItems = await getContentItemsForRefs(relatedObjects);
+
+      for (const sqlObj of contentItems) {
+        contextItems.push({
+          name: `${sqlObj.type.toLowerCase()} definition for ${sqlObj.id}`,
+          content: sqlObj.content,
+          description: `${sqlObj.type} definition`,
+          type: `assistant`
+        });
+      }
+    }
+
     if (!options.history) {
       contextItems.push({
         name: `system prompt`,

src/database/schemas.ts

@@ -46,6 +46,8 @@ function getFilterClause(againstColumn: string, filter: string, noAnd?: boolean)
   };
 }
 
+export interface ObjectReference {name: string, schema?: string};
+
 const BASE_RESOLVE_SELECT = [
   `select `,
   `OBJLONGNAME as name, `,
@@ -60,7 +62,7 @@ export default class Schemas {
   /**
    * Resolves to the following SQL types: SCHEMA, TABLE, VIEW, ALIAS, INDEX, FUNCTION and PROCEDURE
    */
-  static async resolveObjects(sqlObjects: {name: string, schema?: string}[]): Promise<ResolvedSqlObject[]> {
+  static async resolveObjects(sqlObjects: ObjectReference[]): Promise<ResolvedSqlObject[]> {
     let statements: string[] = [];
     let parameters: BasicColumnType[] = [];
 
@@ -111,7 +113,7 @@ export default class Schemas {
     const query = `${statements.join(" UNION ALL ")}`;
     const objects: any[] = await JobManager.runSQL(query, { parameters });
 
-    const resolvedObjects: ResolvedSqlObject[] = objects.map(object => ({
+    let resolvedObjects: ResolvedSqlObject[] = objects.map(object => ({
       name: object.NAME,
       schema: object.SCHEMA,
       sqlType: object.SQLTYPE
@@ -120,6 +122,27 @@ export default class Schemas {
     return resolvedObjects;
   }
 
+  static async getRelatedObjects(object: ResolvedSqlObject): Promise<ResolvedSqlObject[]> {
+    const sql = [
+      `with refs as (`,
+      `  SELECT `,
+      `    schema_name as schema, `,
+      `    sql_name as name, `,
+      `    case when sql_object_type = 'FOREIGN KEY' then 'TABLE' else sql_object_type end as type`,
+      `  FROM TABLE(SYSTOOLS.RELATED_OBJECTS(?, ?))`,
+      `)`,
+      `select * from refs `,
+      `where type in ('TABLE', 'FUNCTION', 'PROCEDURE')`,
+    ].join(` `);
+
+    const related: any[] = await JobManager.runSQL(sql, { parameters: [object.schema, object.name] });
+    return related.map(item => ({
+      name: item.NAME,
+      schema: item.SCHEMA,
+      sqlType: item.TYPE
+    }));
+  }
+
   /**
    * @param schema Not user input
    */