Skip to content

Commit 95f216d

Browse files
uc4w6cuc4w6c
authored
Merge pull request #19 from uc4w6c/feat/code-generator
feat: add Code Generater
2 parents 491f58a + 8cf25d7 commit 95f216d

File tree

5 files changed

+247
-1
lines changed

5 files changed

+247
-1
lines changed

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ Bedrock Assistant is an IntelliJ plugin designed to enhance your development pro
2424
## Current Features
2525

2626
- **Analyze Code**: Gain deep insights into your code with our AI-driven analysis tools.
27+
- **Generate Code**: Simply provide instructions, and AI will generate the corresponding code for you.
2728

2829
Stay tuned as we continue to expand the capabilities of Bedrock Assistant, bringing more powerful features to your fingertips.
2930

‎gradle.properties‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ pluginGroup = com.github.uc4w6c.bedrockassistant
44
pluginName = BedrockAssistant
55
pluginRepositoryUrl = https://github.com/uc4w6c/BedrockAssistant
66
# SemVer format -> https://semver.org
7-
pluginVersion = 0.1.3
7+
pluginVersion = 0.2.0
88

99
# Supported build number ranges and IntelliJ Platform versions -> https://plugins.jetbrains.com/docs/intellij/build-number-ranges.html
1010
pluginSinceBuild = 242

‎src/main/java/com/github/uc4w6c/bedrockassistant/action/CodeGenerateAction.java‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
package com.github.uc4w6c.bedrockassistant.action;
2+
3+
import com.github.uc4w6c.bedrockassistant.domain.AwsCredentials;
4+
import com.github.uc4w6c.bedrockassistant.domain.Message;
5+
import com.github.uc4w6c.bedrockassistant.domain.MessageRole;
6+
import com.github.uc4w6c.bedrockassistant.exceptions.BedrockException;
7+
import com.github.uc4w6c.bedrockassistant.exceptions.TokenException;
8+
import com.github.uc4w6c.bedrockassistant.helper.TokenHelper;
9+
import com.github.uc4w6c.bedrockassistant.popup.MfaPopupComponent;
10+
import com.github.uc4w6c.bedrockassistant.popup.PromptPopupComponent;
11+
import com.github.uc4w6c.bedrockassistant.repository.BedrockRepository;
12+
import com.github.uc4w6c.bedrockassistant.service.BedrockAssistantToolWindowService;
13+
import com.github.uc4w6c.bedrockassistant.state.BedrockAssistantCacheState;
14+
import com.github.uc4w6c.bedrockassistant.state.BedrockAssistantState;
15+
import com.github.uc4w6c.bedrockassistant.state.TokenManager;
16+
import com.github.uc4w6c.bedrockassistant.window.BedrockAssistantToolWindow;
17+
import com.intellij.notification.NotificationGroupManager;
18+
import com.intellij.notification.NotificationType;
19+
import com.intellij.openapi.actionSystem.AnAction;
20+
import com.intellij.openapi.actionSystem.AnActionEvent;
21+
import com.intellij.openapi.actionSystem.CommonDataKeys;
22+
import com.intellij.openapi.application.ApplicationManager;
23+
import com.intellij.openapi.command.WriteCommandAction;
24+
import com.intellij.openapi.editor.CaretModel;
25+
import com.intellij.openapi.editor.Document;
26+
import com.intellij.openapi.editor.Editor;
27+
import com.intellij.openapi.fileEditor.FileDocumentManager;
28+
import com.intellij.openapi.project.Project;
29+
import com.intellij.openapi.vfs.VirtualFile;
30+
import com.intellij.psi.PsiDocumentManager;
31+
import org.jetbrains.annotations.NotNull;
32+
33+
import java.util.List;
34+
import java.util.Optional;
35+
import java.util.function.Consumer;
36+
import java.util.function.Function;
37+
import java.util.regex.Matcher;
38+
import java.util.regex.Pattern;
39+
40+
public class CodeGenerateAction extends AnAction {
41+
private final String GENERATE_CODE_PROMPT = """
42+
The file name is %s.
43+
44+
Here is the surrounding code context:
45+
```
46+
%s
47+
```
48+
49+
Write the new code on line %d.
50+
51+
Generate the code according to the instructions below.
52+
53+
Return only the generated code inside triple backticks. Do not include any explanations, comments, or additional text.
54+
55+
%s
56+
""";
57+
58+
@Override
59+
public void actionPerformed(@NotNull AnActionEvent anActionEvent) {
60+
Project project = anActionEvent.getProject();
61+
if (project == null) {
62+
return;
63+
}
64+
65+
Editor editor = anActionEvent.getData(CommonDataKeys.EDITOR);
66+
if (editor == null) {
67+
return;
68+
}
69+
70+
Consumer<String> submitActionListener = (promptText) -> {
71+
BedrockAssistantToolWindow bedrockAssistantToolWindow = BedrockAssistantToolWindowService.getInstance(project)
72+
.getBedrockAssistantToolWindow();
73+
74+
if (bedrockAssistantToolWindow == null) {
75+
throw new RuntimeException();
76+
}
77+
BedrockAssistantState.State state = BedrockAssistantState.getInstance().getState();
78+
79+
TokenHelper tokenHelper = new TokenHelper();
80+
Optional<AwsCredentials> optionalAwsCredentials = tokenHelper.getCredentials(project, state.profile);
81+
82+
if (optionalAwsCredentials.isEmpty()) {
83+
return;
84+
}
85+
86+
VirtualFile file = FileDocumentManager.getInstance().getFile(editor.getDocument());
87+
if (file == null) {
88+
return;
89+
}
90+
91+
String fileName = file.getName();
92+
93+
Document document = editor.getDocument();
94+
95+
int totalLines = document.getLineCount();
96+
CaretModel caretModel = editor.getCaretModel();
97+
int currentLine = document.getLineNumber(caretModel.getOffset());
98+
99+
int startLine = Math.max(0, currentLine - 100);
100+
int endLine = Math.min(totalLines - 1, currentLine + 100);
101+
102+
int selectedLineIndex = currentLine - startLine + 1;
103+
104+
int startOffset = document.getLineStartOffset(startLine);
105+
int endOffset = document.getLineEndOffset(endLine);
106+
107+
String code = document.getText().substring(startOffset, endOffset);
108+
109+
String context = String.format(GENERATE_CODE_PROMPT
110+
,fileName
111+
,code
112+
,selectedLineIndex
113+
,promptText);
114+
List<Message> messages = List.of(new Message(MessageRole.USER, context));
115+
116+
try {
117+
BedrockRepository bedrockRepository = new BedrockRepository();
118+
String response = bedrockRepository.get(optionalAwsCredentials.get(), state.region, messages);
119+
120+
Function<String, String> extractCodeBlock = (responseText) -> {
121+
Pattern pattern = Pattern.compile("```(.*?)```", Pattern.DOTALL);
122+
Matcher matcher = pattern.matcher(responseText);
123+
124+
if (matcher.find()) {
125+
return matcher.group(1).trim();
126+
}
127+
return responseText;
128+
};
129+
130+
int caretOffset = caretModel.getOffset();
131+
int lineNumber = document.getLineNumber(caretOffset);
132+
133+
if (lineNumber < 0) {
134+
return;
135+
}
136+
137+
int lineEndOffset = document.getLineEndOffset(lineNumber);
138+
139+
WriteCommandAction.runWriteCommandAction(project, () ->
140+
document.insertString(lineEndOffset, extractCodeBlock.apply(response)));
141+
142+
PsiDocumentManager.getInstance(project).commitDocument(document);
143+
} catch (BedrockException e) {
144+
NotificationGroupManager.getInstance()
145+
.getNotificationGroup("BedrockAssistantBalloon")
146+
.createNotification(String.format("""
147+
Bedrock call failed.
148+
%s
149+
""", e.getMessage()),
150+
NotificationType.ERROR)
151+
.notify(project);
152+
}
153+
};
154+
155+
ApplicationManager.getApplication().invokeLater(() -> {
156+
new PromptPopupComponent.PromptPopupComponentBuilder()
157+
.submitActionListener(submitActionListener)
158+
.build();
159+
});
160+
}
161+
}

‎src/main/java/com/github/uc4w6c/bedrockassistant/popup/PromptPopupComponent.java‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package com.github.uc4w6c.bedrockassistant.popup;
2+
3+
import com.intellij.openapi.ui.popup.JBPopup;
4+
import com.intellij.openapi.ui.popup.JBPopupFactory;
5+
import com.intellij.ui.components.JBTextField;
6+
7+
import javax.swing.*;
8+
import javax.swing.border.EmptyBorder;
9+
import java.awt.*;
10+
import java.util.Objects;
11+
import java.util.function.Consumer;
12+
13+
public class PromptPopupComponent extends JPanel {
14+
private JBTextField promptTextfield = new JBTextField();
15+
private JBPopup popup;
16+
17+
PromptPopupComponent(Consumer<String> submitActionListener) {
18+
super(new BorderLayout());
19+
this.setBorder(new EmptyBorder(10, 10, 10, 10));
20+
this.setEnabled(true);
21+
22+
promptTextfield.setBorder(BorderFactory.createLineBorder(Color.GRAY, 2));
23+
promptTextfield.requestFocus();
24+
promptTextfield.setColumns(80);
25+
this.add(promptTextfield, BorderLayout.CENTER);
26+
27+
JPanel buttonPanel = new JPanel(new FlowLayout(FlowLayout.RIGHT));
28+
29+
JButton submitButton = new JButton("submit");
30+
submitButton.addActionListener(actionEvent -> {
31+
submitActionListener.accept(promptTextfield.getText());
32+
popup.cancel();
33+
});
34+
buttonPanel.add(submitButton);
35+
36+
this.add(buttonPanel, BorderLayout.SOUTH);
37+
}
38+
39+
public JComponent getFocusableComponent() {
40+
return promptTextfield;
41+
}
42+
43+
public void setPopup(JBPopup popup) {
44+
this.popup = popup;
45+
}
46+
47+
public static class PromptPopupComponentBuilder {
48+
private Consumer<String> submitActionListener;
49+
50+
public PromptPopupComponentBuilder submitActionListener(Consumer<String> submitActionListener) {
51+
this.submitActionListener = submitActionListener;
52+
return this;
53+
}
54+
55+
public PromptPopupComponent build() {
56+
if (Objects.isNull(this.submitActionListener)) {
57+
throw new IllegalArgumentException();
58+
}
59+
60+
PromptPopupComponent promptPopupComponent = new PromptPopupComponent(submitActionListener);
61+
62+
JBPopup popup = JBPopupFactory.getInstance()
63+
.createComponentPopupBuilder(promptPopupComponent, promptPopupComponent.getFocusableComponent())
64+
.setResizable(true)
65+
.setMovable(true)
66+
.setFocusable(true)
67+
.setRequestFocus(true)
68+
.createPopup();
69+
promptPopupComponent.setPopup(popup);
70+
71+
popup.showInFocusCenter();
72+
73+
return promptPopupComponent;
74+
}
75+
}
76+
}

‎src/main/resources/META-INF/plugin.xml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,14 @@
3636
>
3737
<add-to-group group-id="EditorPopupMenu" anchor="first"/>
3838
</action>
39+
<action id="CodeGenerateAction"
40+
class="com.github.uc4w6c.bedrockassistant.action.CodeGenerateAction"
41+
text="Generate code"
42+
description="Generate code using Bedrock."
43+
>
44+
<keyboard-shortcut keymap="$default" first-keystroke="control shift L"/>
45+
<add-to-group group-id="EditorPopupMenu" anchor="first"/>
46+
</action>
3947
</actions>
4048

4149
<!--

0 commit comments

Comments
 (0)