forked from dbeaver/dbeaver
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
dbeaver/pro#3911 AI and SQL control commands (dbeaver#36653)
* dbeaver/pro#3911 AI and SQL control commands * dbeaver/pro#3911 AI and SQL control commands * dbeaver/pro#3741 System agent service * dbeaver/pro#3911 AI and SQL control commands * dbeaver/pro#3911 AI and SQL control commands * dbeaver/pro#3911 AI and SQL control commands * dbeaver/pro#3911 SQL output formatting * dbeaver/pro#3911 AI error messages * dbeaver/pro#3911 AI command fixes * dbeaver/pro#3911 Support custom context * dbeaver/pro#3911 Disable datasource check in headless products --------- Co-authored-by: kseniaguzeeva <112612526+kseniaguzeeva@users.noreply.github.com>
- Loading branch information
1 parent
31790fa
commit d6580fc
Showing
21 changed files
with
394 additions
and
104 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
41 changes: 41 additions & 0 deletions
41
.../org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/commands/AIOutputSeverity.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,41 @@ | ||
| /* | ||
| * DBeaver - Universal Database Manager | ||
| * Copyright (C) 2010-2024 DBeaver Corp and others | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package org.jkiss.dbeaver.model.ai.commands; | ||
|
|
||
| import org.jkiss.code.NotNull; | ||
| import org.jkiss.dbeaver.model.exec.output.DBCOutputSeverity; | ||
|
|
||
| enum AIOutputSeverity implements DBCOutputSeverity { | ||
| PROMPT("AI"); | ||
|
|
||
| private final String name; | ||
|
|
||
| AIOutputSeverity(@NotNull String name) { | ||
| this.name = name; | ||
| } | ||
|
|
||
| @NotNull | ||
| @Override | ||
| public String getName() { | ||
| return name; | ||
| } | ||
|
|
||
| @Override | ||
| public boolean isForced() { | ||
| return true; | ||
| } | ||
| } |
128 changes: 128 additions & 0 deletions
128
plugins/org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/commands/SQLCommandAI.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,128 @@ | ||
| /* | ||
| * DBeaver - Universal Database Manager | ||
| * Copyright (C) 2010-2024 DBeaver Corp and others | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package org.jkiss.dbeaver.model.ai.commands; | ||
|
|
||
| import org.jkiss.code.NotNull; | ||
| import org.jkiss.dbeaver.DBException; | ||
| import org.jkiss.dbeaver.Log; | ||
| import org.jkiss.dbeaver.model.DBPDataSourceContainer; | ||
| import org.jkiss.dbeaver.model.ai.*; | ||
| import org.jkiss.dbeaver.model.ai.completion.*; | ||
| import org.jkiss.dbeaver.model.ai.format.IAIFormatter; | ||
| import org.jkiss.dbeaver.model.logical.DBSLogicalDataSource; | ||
| import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor; | ||
| import org.jkiss.dbeaver.model.sql.*; | ||
| import org.jkiss.dbeaver.runtime.DBWorkbench; | ||
| import org.jkiss.utils.CommonUtils; | ||
|
|
||
| import java.util.Arrays; | ||
| import java.util.List; | ||
| import java.util.stream.Collectors; | ||
|
|
||
| /** | ||
| * Control command handler | ||
| */ | ||
| public class SQLCommandAI implements SQLControlCommandHandler { | ||
|
|
||
| private static final Log log = Log.getLog(SQLCommandAI.class); | ||
|
|
||
| @NotNull | ||
| @Override | ||
| public SQLControlResult handleCommand(@NotNull DBRProgressMonitor monitor, @NotNull SQLControlCommand command, @NotNull SQLScriptContext scriptContext) throws DBException { | ||
| if (command.getDataSource() == null) { | ||
| throw new DBException("Not connected to database"); | ||
| } | ||
| AISettings aiSettings = AISettingsRegistry.getInstance().getSettings(); | ||
| if (aiSettings.isAiDisabled()) { | ||
| throw new DBException("AI services are disabled"); | ||
| } | ||
| DAICompletionEngine<?> engine = AIEngineRegistry.getInstance().getCompletionEngine( | ||
| aiSettings.getActiveEngine()); | ||
|
|
||
| String prompt = command.getParameter(); | ||
| if (CommonUtils.isEmptyTrimmed(prompt)) { | ||
| throw new DBException("Empty AI prompt"); | ||
| } | ||
|
|
||
| IAIFormatter formatter = AIFormatterRegistry.getInstance().getFormatter(AIConstants.CORE_FORMATTER); | ||
|
|
||
| final DBSLogicalDataSource dataSource = new DBSLogicalDataSource( | ||
| command.getDataSourceContainer(), "AI logical wrapper", null); | ||
|
|
||
| DBPDataSourceContainer dataSourceContainer = dataSource.getDataSourceContainer(); | ||
| DAICompletionSettings completionSettings = new DAICompletionSettings(dataSourceContainer); | ||
| if (!DBWorkbench.getPlatform().getApplication().isHeadlessMode() && !completionSettings.isMetaTransferConfirmed()) { | ||
| if (DBWorkbench.getPlatformUI().confirmAction("Do you confirm AI usage", | ||
| "Do you confirm AI usage for '" + dataSourceContainer.getName() + "'?" | ||
| )) { | ||
| completionSettings.setMetaTransferConfirmed(true); | ||
| completionSettings.saveSettings(); | ||
| } else { | ||
| throw new DBException("AI services restricted for '" + dataSourceContainer.getName() + "'"); | ||
| } | ||
| } | ||
| DAICompletionScope scope = completionSettings.getScope(); | ||
| DAICompletionContext.Builder contextBuilder = new DAICompletionContext.Builder() | ||
| .setScope(scope) | ||
| .setDataSource(dataSource) | ||
| .setExecutionContext(scriptContext.getExecutionContext()); | ||
| if (scope == DAICompletionScope.CUSTOM) { | ||
| contextBuilder.setCustomEntities( | ||
| AITextUtils.loadCustomEntities( | ||
| monitor, | ||
| command.getDataSource(), | ||
| Arrays.stream(completionSettings.getCustomObjectIds()).collect(Collectors.toSet())) | ||
| ); | ||
| } | ||
| final DAICompletionContext aiContext = contextBuilder.build(); | ||
|
|
||
| DAICompletionSession aiSession = new DAICompletionSession(); | ||
| aiSession.add(new DAICompletionMessage(DAICompletionMessage.Role.USER, prompt)); | ||
|
|
||
| List<DAICompletionResponse> responses = engine.performSessionCompletion( | ||
| monitor, | ||
| aiContext, | ||
| aiSession, | ||
| formatter, | ||
| true); | ||
|
|
||
| DAICompletionResponse response = responses.get(0); | ||
| MessageChunk[] messageChunks = AITextUtils.splitIntoChunks( | ||
| CommonUtils.notEmpty(response.getResultCompletion())); | ||
|
|
||
| String finalSQL = null; | ||
| StringBuilder messages = new StringBuilder(); | ||
| for (MessageChunk chunk : messageChunks) { | ||
| if (chunk instanceof MessageChunk.Code code) { | ||
| finalSQL = code.text(); | ||
| } else if (chunk instanceof MessageChunk.Text text) { | ||
| messages.append(text.text()); | ||
| } | ||
| } | ||
| if (finalSQL == null) { | ||
| if (!messages.isEmpty()) { | ||
| throw new DBException(messages.toString()); | ||
| } | ||
| throw new DBException("Empty AI completion for '" + prompt + "'"); | ||
| } | ||
|
|
||
| scriptContext.getOutputWriter().println(AIOutputSeverity.PROMPT, prompt + " ==> " + finalSQL + "\n"); | ||
| return SQLControlResult.transform( | ||
| new SQLQuery(command.getDataSource(), finalSQL)); | ||
| } | ||
|
|
||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
48 changes: 48 additions & 0 deletions
48
plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLControlResult.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| /* | ||
| * DBeaver - Universal Database Manager | ||
| * Copyright (C) 2010-2024 DBeaver Corp and others | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package org.jkiss.dbeaver.model.sql; | ||
|
|
||
| /** | ||
| * Control command result. | ||
| * | ||
| * It may finish with no extra information or with parameters: | ||
| * - message: will be shown in UI | ||
| * - error: execution error will be shown in UI | ||
| */ | ||
| public class SQLControlResult { | ||
|
|
||
| public static SQLControlResult success() { | ||
| return new SQLControlResult(); | ||
| } | ||
|
|
||
| public static SQLControlResult transform(SQLScriptElement element) { | ||
| return new SQLControlResult(element); | ||
| } | ||
|
|
||
| private SQLScriptElement transformed; | ||
|
|
||
| private SQLControlResult() { | ||
| } | ||
|
|
||
| private SQLControlResult(SQLScriptElement transformed) { | ||
| this.transformed = transformed; | ||
| } | ||
|
|
||
| public SQLScriptElement getTransformed() { | ||
| return transformed; | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.