diff --git a/src/main/kotlin/ai/welltested/fluttergpt/repository/GeminiRepository.kt b/src/main/kotlin/ai/welltested/fluttergpt/repository/GeminiRepository.kt new file mode 100644 index 0000000..34e8472 --- /dev/null +++ b/src/main/kotlin/ai/welltested/fluttergpt/repository/GeminiRepository.kt @@ -0,0 +1,69 @@ +package ai.welltested.fluttergpt.repository + +import ai.welltested.fluttergpt.utilities.configManager.SecretKeyConfig +import ai.welltested.fluttergpt.utilities.configManager.SecretKeyListener +import com.google.gson.Gson +import com.intellij.openapi.application.ApplicationManager +import com.intellij.util.messages.MessageBusConnection +import java.net.HttpURLConnection +import java.net.URL + + +data class Candidate( + val content: Content, + val finishReason: String, + val index: Int +) + +data class Content( + val parts: List, + val role: String +) + +data class Part( + val text: String +) + +data class GeminiCompletionResponse( + val candidates: List +) + +class GeminiRepository { + private var apiKey: String? = null + + init { + apiKey = SecretKeyConfig.getInstance().secretKey; + val connection: MessageBusConnection = ApplicationManager.getApplication().messageBus.connect() + connection.subscribe(SecretKeyListener.SECRET_KEY_TOPIC, object : SecretKeyListener { + override fun secretKeyChanged(newSecretKey: String) { + apiKey = newSecretKey + } + }) + } + + fun getCompletion(prompt: List): String { + if (apiKey == null) { + throw Exception("API token not set, please go to extension settings to set it (read README.md for more info)") + } + val url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent?key=$apiKey" + + val connection = URL(url).openConnection() as HttpURLConnection + connection.requestMethod = "POST" + connection.setRequestProperty("Content-Type", "application/json") + connection.doOutput = true + + val gson = Gson() + val requestBody = gson.toJson(mapOf("contents" to prompt.map { mapOf("parts" to listOf(mapOf("text" to it))) })) + connection.outputStream.write(requestBody.toByteArray()) + val inputStream = connection.inputStream + val response = gson.fromJson(inputStream.reader(), GeminiCompletionResponse::class.java) + inputStream.close() + + + if(response.candidates.isNotEmpty()) { + return response.candidates[0].content.parts[0].text + } + return "Failed to generate content" + + } +} \ No newline at end of file diff --git a/src/main/kotlin/ai/welltested/fluttergpt/tools/refactor/optimizeFunction.kt b/src/main/kotlin/ai/welltested/fluttergpt/tools/refactor/optimizeFunction.kt new file mode 100644 index 0000000..4d1e3c6 --- /dev/null +++ b/src/main/kotlin/ai/welltested/fluttergpt/tools/refactor/optimizeFunction.kt @@ -0,0 +1,66 @@ +package ai.welltested.fluttergpt.tools.refactor + +import ai.welltested.fluttergpt.repository.GeminiRepository +import ai.welltested.fluttergpt.utilities.extractDartCode +import ai.welltested.fluttergpt.utilities.getMethodElement +import ai.welltested.fluttergpt.utilities.isCursorAtMethodDeclaration +import com.intellij.codeInsight.intention.PsiElementBaseIntentionAction; +import com.intellij.codeInsight.intention.IntentionAction; +import com.intellij.openapi.editor.Editor; +import com.intellij.openapi.progress.ProgressIndicator +import com.intellij.openapi.progress.Task +import com.intellij.openapi.project.Project; +import com.intellij.psi.PsiElement; + +class OptimizeFunction : PsiElementBaseIntentionAction(), IntentionAction { + + override fun isAvailable(project: Project, editor: Editor?, element: PsiElement): Boolean { + if (editor == null) return false + return isCursorAtMethodDeclaration(element, editor) + } + + override fun invoke(p0: Project, p1: Editor?, p2: PsiElement) { + + val methodElement = getMethodElement(p2) + if (methodElement == null) { + return + } + val fullMethodCode = methodElement.text + + var prompt = + "You're an expert Flutter/Dart coding assistant. Follow the instructions carefully and output response in the modified format..\\n\\n" + prompt += "Develop and optimize the following Flutter code by troubleshooting errors, fixing errors, and identifying root causes of issues. Reflect and critique your solution to ensure it meets the requirements and specifications of speed, flexibility and user friendliness.\n\n Subject Code:\n${fullMethodCode}\n\n"; + prompt += "Here is the full code for context:\n"; + if (p1 != null) { + prompt += "```" + p1.document.text + "```"; + } + prompt += "\n\n"; + prompt += "Output the optimized code in a single code block to be replaced over selected code."; + object : Task.Backgroundable(p0, "Optimizing function", false) { + override fun run(indicator: ProgressIndicator) { + + val result = GeminiRepository().getCompletion(listOf(prompt)) + val dartCode = extractDartCode(result) + if (p1 != null) { + p1.document.replaceString( + methodElement.textRange.startOffset, + methodElement.textRange.endOffset, + dartCode + ) + } + } + }.queue() + + } + + + override fun getFamilyName(): String { + return "FlutterGPT" + } + + override fun getText(): String { + return "Optimize code" + } + + +} \ No newline at end of file diff --git a/src/main/kotlin/ai/welltested/fluttergpt/utilities/codeActionsUtils.kt b/src/main/kotlin/ai/welltested/fluttergpt/utilities/codeActionsUtils.kt new file mode 100644 index 0000000..8ab7d21 --- /dev/null +++ b/src/main/kotlin/ai/welltested/fluttergpt/utilities/codeActionsUtils.kt @@ -0,0 +1,22 @@ +package ai.welltested.fluttergpt.utilities + +import com.intellij.openapi.editor.Editor +import com.intellij.psi.PsiElement +import com.intellij.psi.util.PsiTreeUtil +import com.intellij.psi.util.elementType + +fun isCursorAtMethodDeclaration(element: PsiElement, editor: Editor): Boolean { + val methodElement = PsiTreeUtil.findFirstParent(element) { it.elementType.toString() == "METHOD_DECLARATION" } + if (methodElement == null) { + return false + } + val document = editor.document + val lineNumber = document.getLineNumber(editor.caretModel.offset) + val methodLine = methodElement.textRange?.startOffset?.let { document.getLineNumber(it) } + return lineNumber == methodLine +} + +fun getMethodElement(element: PsiElement): PsiElement? { + val methodElement = PsiTreeUtil.findFirstParent(element) { it.elementType.toString() == "METHOD_DECLARATION" } + return methodElement +} \ No newline at end of file diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml index d2904e3..5f6f671 100644 --- a/src/main/resources/META-INF/plugin.xml +++ b/src/main/resources/META-INF/plugin.xml @@ -139,7 +139,12 @@ + displayType="BALLOON"/> + + + + ai.welltested.fluttergpt.tools.refactor.OptimizeFunction +