Additional cleanup for context items

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

Author: worksofliam

src/aiProviders/context.ts

@@ -4,6 +4,7 @@ import { JobManager } from "../config";
 import Schemas, { AllSQLTypes, SQLType } from "../database/schemas";
 import Statement from "../database/statement";
 import { DB2_SYSTEM_PROMPT } from "./prompts";
+import { Db2ContextItems } from "./prompt";
 
 export function canTalkToDb() {
   return JobManager.getSelection() !== undefined;
@@ -23,12 +24,6 @@ interface MarkdownRef {
   SCHMEA?: string;
 }
 
-interface ContextDefinition {
-  id: string;
-  type: string;
-  content: any;
-}
-
 /**
  * Builds a semantic representation of a database schema by fetching all objects
  * associated with the schema, cleaning them, and filtering out unnecessary properties.
@@ -167,7 +162,7 @@ function splitUpUserInput(input: string): string[] {
  *
  * @param {string} input - A string that may contain table references.
  */
-export async function getSqlContextItems(input: string): Promise<{items: ContextDefinition[], refs: ResolvedSqlObject[]}> {
+export async function getSqlContextItems(input: string): Promise<{items: Db2ContextItems[], refs: ResolvedSqlObject[]}> {
   // Parse all SCHEMA.TABLE references first
   const tokens = splitUpUserInput(input);
 
@@ -209,8 +204,8 @@ export async function getSqlContextItems(input: string): Promise<{items: Context
   };
 }
 
-export async function getContentItemsForRefs(allObjects: ResolvedSqlObject[]): Promise<ContextDefinition[]> {
-  const items: (ContextDefinition|undefined)[] = await Promise.all(
+export async function getContentItemsForRefs(allObjects: ResolvedSqlObject[]): Promise<Db2ContextItems[]> {
+  const items: (Db2ContextItems|undefined)[] = await Promise.all(
     allObjects.map(async (o) => {
       try {
         if (o.sqlType === `SCHEMA`) {
@@ -221,9 +216,9 @@ export async function getContentItemsForRefs(allObjects: ResolvedSqlObject[]): P
           const content = await Schemas.generateSQL(o.schema, o.name, o.sqlType);
 
           return {
-            id: o.name,
-            type: o.sqlType,
-            content: content,
+            name: `${o.sqlType.toLowerCase()} definition for ${o.name}`,
+            description: `${o.sqlType} definition`,
+            content: content
           };
         }
 

src/aiProviders/copilot/index.ts

@@ -79,6 +79,7 @@ export function activateChat(context: vscode.ExtensionContext) {
           });
 
           messages.push(...contextItems.context.map(c => {
+            // Passing in the same breaks the request?
             return vscode.LanguageModelChatMessage.Assistant(c.content);
           }));
 

src/aiProviders/prompt.ts

@@ -15,7 +15,6 @@ export interface Db2ContextItems {
   name: string;
   description: string;
   content: string;
-  specific?: "copilot"|"continue";
 }
 
 export interface BuildResult {
@@ -57,13 +56,7 @@ export async function getContextItems(input: string, options: PromptOptions = {}
     progress(`Finding objects to work with...`);
     const context = await getSqlContextItems(input);
 
-    for (const sqlObj of context.items) {
-      contextItems.push({
-        name: `${sqlObj.type.toLowerCase()} definition for ${sqlObj.id}`,
-        content: sqlObj.content,
-        description: `${sqlObj.type} definition`,
-      });
-    }
+    contextItems.push(...context.items);
 
     if (context.refs.filter(r => r.sqlType === `TABLE`).length >= 2) {
       const randomIndexA = Math.floor(Math.random() * context.refs.length);
@@ -85,20 +78,14 @@ export async function getContextItems(input: string, options: PromptOptions = {}
       const relatedObjects = await Schemas.getRelatedObjects(ref);
       const contentItems = await getContentItemsForRefs(relatedObjects);
 
+      contextItems.push(...contentItems);
+
       if (relatedObjects.length === 1) {
         followUps.push(`How is ${prettyNameRef} related to ${Statement.prettyName(relatedObjects[0].name)}?`);
       } else if (ref.sqlType === `TABLE`) {
         followUps.push(`What are some objects related to that table?`);
       }
 
-      for (const sqlObj of contentItems) {
-        contextItems.push({
-          name: `${sqlObj.type.toLowerCase()} definition for ${sqlObj.id}`,
-          content: sqlObj.content,
-          description: `${sqlObj.type} definition`,
-        });
-      }
-
     } else if (context.refs.length > 1) {
       const randomRef = context.refs[Math.floor(Math.random() * context.refs.length)];
       const prettyNameRef = Statement.prettyName(randomRef.name);