Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Implement user-selected context files #367

Merged
merged 8 commits into from
Feb 6, 2024
Merged

Implement user-selected context files #367

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
68ae69e
basic context file menu
beyang Feb 1, 2024
ae7d1ec
respond to PR comments
beyang Feb 2, 2024
585905e
fix displayPath; re-register webview message listener when agent sess…
beyang Feb 3, 2024
d6c3dd1
fix list appearance
beyang Feb 3, 2024
ed300e4
refactor
beyang Feb 3, 2024
43f608d
Merge remote-tracking branch 'origin/main' into bl/wip
beyang Feb 3, 2024
2192c16
Merge remote-tracking branch 'origin/main' into bl/wip
beyang Feb 5, 2024
27dfaa6
make webview message handling consistent with other message handling
beyang Feb 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
294 changes: 220 additions & 74 deletions src/main/java/com/sourcegraph/cody/PromptPanel.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
package com.sourcegraph.cody

import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.CustomShortcutSet
import com.intellij.openapi.actionSystem.KeyboardShortcut
Expand All @@ -10,6 +9,8 @@ import com.intellij.ui.DocumentAdapter
import com.intellij.util.concurrency.annotations.RequiresEdt
import com.intellij.util.ui.JBUI
import com.intellij.util.ui.UIUtil
import com.sourcegraph.cody.agent.WebviewMessage
import com.sourcegraph.cody.agent.protocol.ContextFile
import com.sourcegraph.cody.chat.ChatSession
import com.sourcegraph.cody.chat.CodyChatMessageHistory
import com.sourcegraph.cody.chat.ui.SendButton
Expand All @@ -18,50 +19,60 @@ import com.sourcegraph.cody.vscode.CancellationToken
import java.awt.Dimension
import java.awt.event.ComponentAdapter
import java.awt.event.ComponentEvent
import java.awt.event.KeyAdapter
import java.awt.event.KeyEvent
import java.awt.event.MouseAdapter
import java.awt.event.MouseEvent
import java.io.File
import javax.swing.DefaultListModel
import javax.swing.JLayeredPane
import javax.swing.JList
import javax.swing.JScrollPane
import javax.swing.KeyStroke
import javax.swing.border.EmptyBorder
import javax.swing.event.AncestorEvent
import javax.swing.event.AncestorListener
import javax.swing.event.DocumentEvent

class PromptPanel(private val chatSession: ChatSession) : JLayeredPane() {
class PromptPanel(
private val chatSession: ChatSession,
) : JLayeredPane() {

/** View components */
private val autoGrowingTextArea = AutoGrowingTextArea(5, 9, this)
private val scrollPane = autoGrowingTextArea.scrollPane
private val textArea = autoGrowingTextArea.textArea
private val sendButton = SendButton()
private var contextFilesSelectorModel = DefaultListModel<DisplayedContextFile>()
private val contextFilesSelector = JList(contextFilesSelectorModel)
private val contextFilesScroller = JScrollPane(contextFilesSelector)

/** Externally updated state */
private val selectedContextFiles: ArrayList<ContextFile> = ArrayList()

/** Related components */
private val promptMessageHistory =
CodyChatMessageHistory(CHAT_MESSAGE_HISTORY_CAPACITY, chatSession)
private val sendButton = SendButton()

init {
/** Initialize view */
textArea.emptyText.text = "Ask a question about this code..."
scrollPane.border = EmptyBorder(JBUI.emptyInsets())
scrollPane.background = UIUtil.getPanelBackground()

sendButton.addActionListener { _ -> chatSession.sendMessage(getTextAndReset()) }
// Set initial bounds for the scrollPane (100x100) to ensure proper initialization;
// later adjusted dynamically based on component resizing in the component listener.
scrollPane.setBounds(0, 0, 100, 100)
add(scrollPane, DEFAULT_LAYER)
scrollPane.setBounds(0, 0, width, scrollPane.preferredSize.height)

val upperMessageAction: AnAction =
object : DumbAwareAction() {
override fun actionPerformed(e: AnActionEvent) {
promptMessageHistory.popUpperMessage(textArea)
}
}
val lowerMessageAction: AnAction =
object : DumbAwareAction() {
override fun actionPerformed(e: AnActionEvent) {
promptMessageHistory.popLowerMessage(textArea)
}
}
val sendMessageAction: AnAction =
object : DumbAwareAction() {
override fun actionPerformed(e: AnActionEvent) {
if (sendButton.isEnabled) {
chatSession.sendMessage(getTextAndReset())
}
}
}
contextFilesSelector.border = EmptyBorder(JBUI.emptyInsets())
add(contextFilesScroller, PALETTE_LAYER, 0)

add(sendButton, PALETTE_LAYER, 0)

preferredSize = Dimension(scrollPane.width, scrollPane.height)

/** Add listeners */
addAncestorListener(
object : AncestorListener {
override fun ancestorAdded(event: AncestorEvent?) {
Expand All @@ -73,87 +84,222 @@ class PromptPanel(private val chatSession: ChatSession) : JLayeredPane() {

override fun ancestorMoved(event: AncestorEvent?) {}
})

sendMessageAction.registerCustomShortcutSet(DEFAULT_SUBMIT_ACTION_SHORTCUT, textArea)
upperMessageAction.registerCustomShortcutSet(POP_UPPER_MESSAGE_ACTION_SHORTCUT, textArea)
lowerMessageAction.registerCustomShortcutSet(POP_LOWER_MESSAGE_ACTION_SHORTCUT, textArea)

textArea.addKeyListener(
object : KeyAdapter() {
override fun keyReleased(e: KeyEvent) {
val keyCode = e.keyCode
if (keyCode != KeyEvent.VK_UP && keyCode != KeyEvent.VK_DOWN) {}
addComponentListener(
object : ComponentAdapter() {
override fun componentResized(e: ComponentEvent?) {
// HACK
val jButtonPreferredSize = sendButton.preferredSize
sendButton.setBounds(
scrollPane.width - jButtonPreferredSize.width,
scrollPane.height - jButtonPreferredSize.height,
jButtonPreferredSize.width,
jButtonPreferredSize.height)
refreshViewLayout()
}
})

// Add user action listeners
sendButton.addActionListener { _ -> didSubmitChatMessage() }
textArea.document.addDocumentListener(
object : DocumentAdapter() {
override fun textChanged(e: DocumentEvent) {
refreshSendButton()
didUserInputChange(textArea.text)
}
})
scrollPane.border = EmptyBorder(JBUI.emptyInsets())
scrollPane.background = UIUtil.getPanelBackground()
contextFilesSelector.addMouseListener(
object : MouseAdapter() {
override fun mouseClicked(e: MouseEvent) {
contextFilesSelector.selectedIndex = contextFilesSelector.locationToIndex(e.getPoint())
didSelectContextFile()
textArea.requestFocusInWindow()
}
})
for (shortcut in listOf(ENTER, UP, DOWN, TAB)) { // key listeners
object : DumbAwareAction() {
override fun actionPerformed(e: AnActionEvent) {
didUseShortcut(shortcut)
}
}
.registerCustomShortcutSet(shortcut, textArea)
}
}

// Set initial bounds for the scrollPane (100x100) to ensure proper initialization;
// later adjusted dynamically based on component resizing in the component listener.
scrollPane.setBounds(0, 0, 100, 100)
private fun didUseShortcut(shortcut: CustomShortcutSet) {
if (contextFilesSelector.model.size > 0) {
when (shortcut) {
UP -> setSelectedContextFileIndex(-1)
DOWN -> setSelectedContextFileIndex(1)
ENTER,
TAB -> didSelectContextFile()
}
return
}
when (shortcut) {
ENTER -> if (sendButton.isEnabled) didSubmitChatMessage()
UP -> promptMessageHistory.popUpperMessage(textArea)
DOWN -> promptMessageHistory.popLowerMessage(textArea)
}
}

add(scrollPane, DEFAULT_LAYER)
/** View handlers */
private fun didSubmitChatMessage() {
val cf = findContextFiles(selectedContextFiles, textArea.text)
val text = textArea.text

add(sendButton, PALETTE_LAYER, 0)
// Reset text
promptMessageHistory.messageSent(text)
textArea.text = ""
selectedContextFiles.clear()

scrollPane.setBounds(0, 0, width, scrollPane.preferredSize.height)
chatSession.sendMessage(text, cf)
}

preferredSize = Dimension(scrollPane.width, scrollPane.height)
private fun didSelectContextFile() {
if (contextFilesSelector.selectedIndex == -1) return

addComponentListener(
object : ComponentAdapter() {
override fun componentResized(e: ComponentEvent?) {
revalidate()
val jButtonPreferredSize = sendButton.preferredSize
sendButton.setBounds(
scrollPane.width - jButtonPreferredSize.width,
scrollPane.height - jButtonPreferredSize.height,
jButtonPreferredSize.width,
jButtonPreferredSize.height)
}
})
val selected = contextFilesSelector.model.getElementAt(contextFilesSelector.selectedIndex)
this.selectedContextFiles.add(selected.contextFile)
val cfDisplayPath = selected.toString()
val expr = findAtExpressions(textArea.text).lastOrNull() ?: return

textArea.replaceRange("@${cfDisplayPath} ", expr.startIndex, expr.endIndex)

setContextFilesSelector(listOf())
refreshViewLayout()
}

private fun didUserInputChange(text: String) {
val exp = findAtExpressions(text).lastOrNull()
if (exp == null ||
exp.endIndex <
text.length) { // TODO(beyang): instead of text.length, should be current cursor index
setContextFilesSelector(listOf())
refreshViewLayout()
return
}
this.chatSession.sendWebviewMessage(
WebviewMessage(command = "getUserContext", submitType = "user", query = exp.value))
}

/** State updaters */
private fun setSelectedContextFileIndex(increment: Int) {
var newSelectedIndex =
(contextFilesSelector.selectedIndex + increment) % contextFilesSelector.model.size
if (newSelectedIndex < 0) {
newSelectedIndex += contextFilesSelector.model.size
}
contextFilesSelector.selectedIndex = newSelectedIndex
refreshViewLayout()
}

/** View updaters */
@RequiresEdt
fun refreshSendButton() {
private fun refreshViewLayout() {
// get the height of the context files list based on font height and number of context files
val contextFilesHeight = contextFilesSelector.preferredSize.height
contextFilesScroller.size = Dimension(scrollPane.width, contextFilesHeight)

val margin = 10
scrollPane.setBounds(0, contextFilesHeight, width, scrollPane.preferredSize.height + margin)
preferredSize = Dimension(scrollPane.width, scrollPane.height + contextFilesHeight)

sendButton.setLocation(
scrollPane.width - sendButton.preferredSize.width,
scrollPane.height + contextFilesSelector.height - sendButton.preferredSize.height)

revalidate()
}

@RequiresEdt
private fun refreshSendButton() {
sendButton.isEnabled =
textArea.getText().isNotEmpty() && chatSession.getCancellationToken().isDone
}

/** External prop setters */
fun registerCancellationToken(cancellationToken: CancellationToken) {
cancellationToken.onFinished {
ApplicationManager.getApplication().invokeLater { refreshSendButton() }
}
}

override fun revalidate() {
super.revalidate()

scrollPane.setBounds(0, 0, width, scrollPane.preferredSize.height)
preferredSize = Dimension(scrollPane.width, scrollPane.height)
}
@RequiresEdt
fun setContextFilesSelector(newUserContextFiles: List<ContextFile>) {
val changed = contextFilesSelectorModel.elements().toList() != newUserContextFiles
if (changed) {
val newModel = DefaultListModel<DisplayedContextFile>()
newModel.addAll(newUserContextFiles.map { f -> DisplayedContextFile(f) })
contextFilesSelector.model = newModel
contextFilesSelectorModel = newModel

private fun getTextAndReset(): String {
val text = textArea.text
promptMessageHistory.messageSent(text)
textArea.text = ""
return text
if (newUserContextFiles.isNotEmpty()) {
contextFilesSelector.selectedIndex = 0
} else {
contextFilesSelector.selectedIndex = -1
}
refreshViewLayout()
}
}

companion object {
private const val CHAT_MESSAGE_HISTORY_CAPACITY = 100
private val JUST_ENTER = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_ENTER, 0), null)
private val KEY_ENTER = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_ENTER, 0), null)
private val KEY_UP = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_UP, 0), null)
private val KEY_DOWN = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_DOWN, 0), null)
private val KEY_TAB = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_TAB, 0), null)

val ENTER = CustomShortcutSet(KEY_ENTER)
val UP = CustomShortcutSet(KEY_UP)
val DOWN = CustomShortcutSet(KEY_DOWN)
val TAB = CustomShortcutSet(KEY_TAB)
}
}

data class DisplayedContextFile(val contextFile: ContextFile) {
override fun toString(): String {
return displayPath(contextFile)
}
}

data class AtExpression(
val startIndex: Int,
val endIndex: Int,
val rawValue: String,
val value: String
)

val atExpressionPattern = """(@(?:\\\s|[^\s])+)(?:\s|$)""".toRegex()

fun findAtExpressions(text: String): List<AtExpression> {
val matches = atExpressionPattern.findAll(text)
val expressions = ArrayList<AtExpression>()
for (match in matches) {
val subMatch = match.groups.get(1)
if (subMatch != null) {
val value = subMatch.value.substring(1).replace("\\ ", " ")
expressions.add(
AtExpression(subMatch.range.first, subMatch.range.last + 1, subMatch.value, value))
}
}
return expressions
}

fun findContextFiles(contextFiles: List<ContextFile>, text: String): List<ContextFile> {
val atExpressions = findAtExpressions(text)
return contextFiles.filter { f -> atExpressions.any { it.value == displayPath(f) } }
}

// TODO(beyang): temporary displayPath implementation, should be updated to mirror what the VS Code
// plugin does
fun displayPath(contextFile: ContextFile): String {
// if the path contains more than three components, display the last three
val path = contextFile.uri.path

val UP = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_UP, 0), null)
val DOWN = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_DOWN, 0), null)
val DEFAULT_SUBMIT_ACTION_SHORTCUT = CustomShortcutSet(JUST_ENTER)
val POP_UPPER_MESSAGE_ACTION_SHORTCUT = CustomShortcutSet(UP)
val POP_LOWER_MESSAGE_ACTION_SHORTCUT = CustomShortcutSet(DOWN)
// split path on separator (OS agnostic)
val pathComponents = path.split(File.separator)
Copy link
Contributor

@pkukielka pkukielka Feb 2, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's safer to avoid string operations on paths.
More idiomatic way to do it could be:

val pathLength = path.getNameCount()
val shortenPath = Paths.get("..").
    resolve(path.subpath(pathLength - 3, pathLength))

Copy link
Member Author

@beyang beyang Feb 3, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't quite get the recommendation here—converting uri.path to a Path object seems inadvisable, since the former uses posix style path separators, whereas Path seems like it would be OS-specific. This does flag an issue where I was using File.separator to split the path string, where I should just be using / (and then replacing with File.separator in the return value).

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, I haven't noticed that we start with uri and not path, and so we have are quarantined to have well defined schema we want to preserve.
You can ignore my comment then.

if (pathComponents.size > 3) {
return "...${File.separator}${pathComponents.subList(pathComponents.size - 3, pathComponents.size).joinToString(File.separator)}"
}
return path
}
Loading
Loading