dbeaver/pro#2236 add an ability to send all messages to AI (#22556)

* dbeaver/pro#2236 add an ability to send all messages to AI * #20602 remove excessive changes --------- Co-authored-by: Matvey16 <82543000+Matvey16@users.noreply.github.com>
dbeaver · Jan 25, 2024 · 892fe88 · 892fe88
1 parent 8c8725f
commit 892fe88
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 10 deletions.
diff --git a/...beaver.model.ai/src/org/jkiss/dbeaver/model/ai/completion/AbstractAICompletionEngine.java b/...beaver.model.ai/src/org/jkiss/dbeaver/model/ai/completion/AbstractAICompletionEngine.java
@@ -16,6 +16,10 @@
  */
 package org.jkiss.dbeaver.model.ai.completion;
 
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
 import org.jkiss.code.NotNull;
 import org.jkiss.code.Nullable;
 import org.jkiss.dbeaver.DBException;
@@ -28,10 +32,6 @@
 import org.jkiss.dbeaver.model.struct.DBSObjectContainer;
 import org.jkiss.utils.CommonUtils;
 
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
-
 public abstract class AbstractAICompletionEngine<SERVICE, REQUEST> implements DAICompletionEngine<SERVICE> {
 
     @NotNull
@@ -48,19 +48,38 @@ public List<DAICompletionResponse> performQueryCompletion(
     }
 
     @NotNull
-    @Override
     public List<DAICompletionResponse> performQueryCompletion(
         @NotNull DBRProgressMonitor monitor,
         @NotNull DAICompletionContext context,
         @NotNull DAICompletionSession session,
         @NotNull IAIFormatter formatter
     ) throws DBException {
-        final String result = requestCompletion(monitor, context, session.getMessages(), formatter);
+        return performQueryCompletion(monitor, context, session, formatter, false);
+    }
+
+    @NotNull
+    @Override
+    public List<DAICompletionResponse> performQueryCompletion(
+        @NotNull DBRProgressMonitor monitor,
+        @NotNull DAICompletionContext context,
+        @NotNull DAICompletionSession session,
+        @NotNull IAIFormatter formatter,
+        boolean includeAssistantMessages
+    ) throws DBException {
+        final String result = requestCompletion(monitor, context, filterMessages(session.getMessages(), includeAssistantMessages), formatter);
         final DAICompletionResponse response = new DAICompletionResponse();
         response.setResultCompletion(result);
         return List.of(response);
     }
 
+    @NotNull
+    private static List<DAICompletionMessage> filterMessages(List<DAICompletionMessage> messages, boolean includeAssistantMessages ) {
+        if (includeAssistantMessages) {
+            return messages;
+        }
+        return messages.stream().filter(it -> DAICompletionMessage.Role.USER.equals(it.role())).toList();
+    }
+
     public abstract Map<String, SERVICE> getServiceMap();
 
     protected abstract int getMaxTokens();

diff --git a/...jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/completion/DAICompletionEngine.java b/...jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/completion/DAICompletionEngine.java
@@ -17,14 +17,14 @@
 
 package org.jkiss.dbeaver.model.ai.completion;
 
+import java.util.List;
+import java.util.Map;
+
 import org.jkiss.code.NotNull;
 import org.jkiss.dbeaver.DBException;
 import org.jkiss.dbeaver.model.ai.format.IAIFormatter;
 import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
 
-import java.util.List;
-import java.util.Map;
-
 /**
  * Completion engine
  */
@@ -50,13 +50,27 @@ List<DAICompletionResponse> performQueryCompletion(
      * Do query completion in a session
      */
     @NotNull
+    @Deprecated
     List<DAICompletionResponse> performQueryCompletion(
         @NotNull DBRProgressMonitor monitor,
         @NotNull DAICompletionContext context,
         @NotNull DAICompletionSession session,
         @NotNull IAIFormatter formatter
     ) throws DBException;
 
+    /**
+     * Do query completion in a session
+     */
+    @NotNull
+    List<DAICompletionResponse> performQueryCompletion(
+        @NotNull DBRProgressMonitor monitor,
+        @NotNull DAICompletionContext context,
+        @NotNull DAICompletionSession session,
+        @NotNull IAIFormatter formatter,
+        boolean allMessages
+    ) throws DBException;
+
+
     boolean isValidConfiguration();
 
     Map<String, SERVICE> getServiceMap();

diff --git a/...kiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/completion/DAICompletionMessage.java b/...kiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/completion/DAICompletionMessage.java
@@ -30,13 +30,15 @@ public record DAICompletionMessage(@NotNull Role role, @NotNull String content)
      */
     public enum Role {
         SYSTEM,
-        USER;
+        USER,
+        ASSISTANT;
 
         @NotNull
         public String getId() {
             return switch (this) {
                 case SYSTEM -> "system";
                 case USER -> "user";
+                case ASSISTANT -> "assistant";
             };
         }
     }