feat: add {branch} variable (#43)

Compute the most common branch from all files that are being committed and use it to replace {branch} variable in prompt.

Author: Blarc

CHANGELOG.md

@@ -1,6 +1,9 @@
 # Changelog
 
 ## [Unreleased]
+### Added
+- Add `{branch}` variable for prompt customisation. 
+
 ### Fixed
 - Commit message generation does not respect locale.
 

src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/AICommitAction.kt

@@ -36,13 +36,19 @@ class AICommitAction : AnAction(), DumbAware {
 
         runBackgroundableTask(message("action.background"), project) {
             val diff = computeDiff(includedChanges, project)
-
             if (diff.isBlank()) {
                 sendNotification(Notification.emptyDiff())
                 return@runBackgroundableTask
             }
 
-            val prompt = AppSettings.instance.getPrompt(diff)
+            var branch = commonBranch(includedChanges, project)
+            if (branch == null) {
+                sendNotification(Notification.noCommonBranch())
+                // hardcoded fallback branch
+                branch = "main"
+            }
+
+            val prompt = AppSettings.instance.getPrompt(diff, branch)
             if (isPromptTooLarge(prompt)) {
                 sendNotification(Notification.promptTooLarge())
                 return@runBackgroundableTask
@@ -120,4 +126,11 @@ class AICommitAction : AnAction(), DumbAware {
         val encoding = registry.getEncoding(modelType.encodingType)
         return encoding.countTokens(prompt) > modelType.maxContextLength
     }
+
+    private fun commonBranch(changes: List<Change>, project: Project): String? {
+        val repositoryManager = GitRepositoryManager.getInstance(project)
+        return changes.map {
+            repositoryManager.getRepositoryForFileQuick(it.virtualFile)?.currentBranchName
+        }.groupingBy { it }.eachCount().maxByOrNull { it.value }?.key
+    }
 }

src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/notifications/Notification.kt

@@ -42,8 +42,9 @@ data class Notification(
         fun emptyDiff() = Notification(DEFAULT_TITLE, message = message("notifications.empty-diff"))
         fun promptTooLarge() = Notification(DEFAULT_TITLE, message = message("notifications.prompt-too-large"))
         fun unsuccessfulRequest(message: String) = Notification(message = message("notifications.unsuccessful-request", message))
-        fun noCommitMessage(): Notification = Notification(message = message("notifications.no-commit-message"))
-        fun unableToSaveToken(): Notification = Notification(message = message("notifications.unable-to-save-token"))
+        fun noCommitMessage() = Notification(message = message("notifications.no-commit-message"))
+        fun unableToSaveToken() = Notification(message = message("notifications.unable-to-save-token"))
+        fun noCommonBranch() = Notification(message = message("notifications.no-common-branch"))
 
     }
 

src/main/kotlin/com/github/blarc/ai/commits/intellij/plugin/settings/AppSettings.kt

@@ -31,25 +31,26 @@ class AppSettings : PersistentStateComponent<AppSettings> {
 
     var requestSupport = true
     var lastVersion: String? = null
-    var openAIHost: String = OpenAIHost.OpenAI.baseUrl
-    var openAIHosts: MutableSet<String> = mutableSetOf(OpenAIHost.OpenAI.baseUrl)
+    var openAIHost = OpenAIHost.OpenAI.baseUrl
+    var openAIHosts = mutableSetOf(OpenAIHost.OpenAI.baseUrl)
     var proxyUrl: String? = null
 
-    var prompts: MutableMap<String, Prompt> = initPrompts()
-    var currentPrompt: Prompt = prompts["basic"]!!
+    var prompts = initPrompts()
+    var currentPrompt = prompts["basic"]!!
 
-    var openAIModelId: String = "gpt-3.5-turbo"
-    var openAIModelIds: List<String> = listOf("gpt-3.5-turbo", "gpt-4")
+    var openAIModelId = "gpt-3.5-turbo"
+    var openAIModelIds = listOf("gpt-3.5-turbo", "gpt-4")
 
     companion object {
         const val SERVICE_NAME = "com.github.blarc.ai.commits.intellij.plugin.settings.AppSettings"
         val instance: AppSettings
             get() = ApplicationManager.getApplication().getService(AppSettings::class.java)
     }
 
-    fun getPrompt(diff: String): String {
+    fun getPrompt(diff: String, branch: String): String {
         var content = currentPrompt.content
         content = content.replace("{locale}", locale.displayLanguage)
+        content = content.replace("{branch}", branch)
 
         return if (content.contains("{diff}")) {
             content.replace("{diff}", diff)

src/main/resources/messages/MyBundle.properties

@@ -28,6 +28,7 @@ notifications.unknown-error=Unknown error ({0}).
 notification.group.important.name=AI Commits important
 notification.group.general.name=AI Commits general
 notifications.empty-diff=Git diff is empty.
+notifications.no-common-branch=Branch could not be found. Replacing {branch} with main.
 
 actions.do-not-ask-again = Do not ask me again.
 actions.take-me-there = Take me there.
@@ -40,7 +41,7 @@ settings.prompt.name=Name
 settings.prompt.content=Content
 validation.required=This value is required.
 validation.number=Value is not a number.
-settings.prompt.comment=You can use {locale} and {diff} variables to customise your prompt.
+settings.prompt.comment=You can use {locale}, {diff}, {branch} variables to customise your prompt.
 actions.update=Update
 actions.add=Add
 settings.prompt.edit.title=Edit Prompt