fix chat history, add set-schema command

codefori · Nov 4, 2024 · f51cd44 · f51cd44
1 parent fef79ff
commit f51cd44
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 14 deletions.
diff --git a/package.json b/package.json
@@ -58,6 +58,10 @@
           {
             "name": "activity",
             "description": "Summarize the activity on the system"
+          },
+          {
+            "name": "set-schema",
+            "description": "Set the current schema"
           }
         ]
       }

diff --git a/src/chat/chat.ts b/src/chat/chat.ts
@@ -10,6 +10,7 @@ import {
   refsToMarkdown,
 } from "./context";
 import { chatRequest } from "./send";
+import { JobManager } from "../config";
 
 const CHAT_ID = `vscode-db2i.chat`;
 
@@ -36,7 +37,7 @@ export function activateChat(context: vscode.ExtensionContext) {
     let messages: vscode.LanguageModelChatMessage[];
 
     if (canTalkToDb()) {
-      const usingSchema = getDefaultSchema();
+      let usingSchema = getDefaultSchema();
 
       switch (request.command) {
         case `activity`:
@@ -58,7 +59,20 @@ export function activateChat(context: vscode.ExtensionContext) {
           await streamModelResponse(messages, stream, token);
 
           return { metadata: { command: "activity" } };
-
+        case `set-schema`:
+          stream.progress(`Setting Current Schema for SQL Job`);
+          const newSchema = request.prompt.split(' ')[0];
+          if (newSchema) {
+            const curJob = JobManager.getSelection()
+            if (curJob) {
+              const result = await curJob.job.setCurrentSchema(newSchema);
+              if (result) {
+                stream.progress(`Set Current Schema: ${newSchema}✅`);
+                usingSchema = newSchema;
+              }
+            }
+            return;
+          } 
         default:
           stream.progress(
             `Getting information from ${Statement.prettyName(usingSchema)}...`
@@ -82,18 +96,20 @@ export function activateChat(context: vscode.ExtensionContext) {
           ];
 
           if (context.history.length > 0) {
-            messages.push(...context.history.map(h => {
-              if ('prompt' in h) {
-                return vscode.LanguageModelChatMessage.Assistant(h.prompt);
-              } else {
-                return vscode.LanguageModelChatMessage.Assistant(
-                  h.response.filter(r => 'value' in r.value).map(r => r.value.value).join(`\n\n`)
-                );
-              }
-            }));
-
-            messages = messages.filter(m => m.content.trim().length > 0);
-          }
+            const historyMessages = context.history.map(h => {
+                if ('prompt' in h) {
+                    return vscode.LanguageModelChatMessage.Assistant(h.prompt);
+                } else {
+                    const responseContent = h.response
+                        .filter(r => r.value instanceof vscode.MarkdownString)
+                        .map(r => (r.value as vscode.MarkdownString).value)
+                        .join('\n\n');
+                    return vscode.LanguageModelChatMessage.Assistant(responseContent);
+                }
+            });
+
+            messages.push(...historyMessages);
+        }
 
           if (Object.keys(refs).length > 0) {
             messages.push(

diff --git a/src/connection/sqlJob.ts b/src/connection/sqlJob.ts
@@ -159,6 +159,25 @@ export class OldSQLJob extends SQLJob {
     return connectResult;
   }
 
+  async setCurrentSchema(schema: string): Promise<QueryResult<any>> {
+    if (schema) {
+      const upperSchema = schema.toUpperCase();
+      const result = await this.execute(`set current schema = ?`, {parameters: [upperSchema]});
+      if (result.success) {
+        const config = getInstance().getConfig();
+        const currentLibrary = config.currentLibrary.toUpperCase();
+
+        if (upperSchema !== currentLibrary) {
+          config.currentLibrary = upperSchema;
+          await getInstance().setConfig(config);
+        }
+      }
+
+      return result;
+
+    }
+  }
+
   async setSelfState(code: SelfValue) {
     try {
       const query: string = `SET SYSIBMADM.SELFCODES = '${code}'`