Skip to content

Commit

Permalink
added(agent): chat/import to migrate historical chats to agent stor…
Browse files Browse the repository at this point in the history
…age for webview (#2108)

Adds a migration to import local chat history into the agent to be used
in the webview.

Tie in for sourcegraph/cody#5304

## Test plan
Tested manually. Need to add a unit test.

---------

Co-authored-by: Mikołaj Kondratek <mik.kondratek@gmail.com>
  • Loading branch information
jamesmcnamara and mkondratek authored Aug 27, 2024
1 parent ea5155f commit 2f5acfd
Show file tree
Hide file tree
Showing 6 changed files with 208 additions and 19 deletions.
1 change: 1 addition & 0 deletions src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,7 @@ private constructor(
ignore = ClientCapabilities.IgnoreEnum.Enabled,
untitledDocuments = ClientCapabilities.UntitledDocumentsEnum.Enabled,
codeActions = ClientCapabilities.CodeActionsEnum.Enabled,
globalState = ClientCapabilities.GlobalStateEnum.`Server-managed`,
webview = ClientCapabilities.WebviewEnum.Native,
webviewNativeConfig =
WebviewNativeConfigParams(
Expand Down
19 changes: 3 additions & 16 deletions src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt
Original file line number Diff line number Diff line change
Expand Up @@ -16,22 +16,7 @@ import com.sourcegraph.cody.agent.protocol.InlineEditParams
import com.sourcegraph.cody.agent.protocol.NetworkRequest
import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument
import com.sourcegraph.cody.agent.protocol.TelemetryEvent
import com.sourcegraph.cody.agent.protocol_generated.Chat_ModelsParams
import com.sourcegraph.cody.agent.protocol_generated.Chat_ModelsResult
import com.sourcegraph.cody.agent.protocol_generated.ClientInfo
import com.sourcegraph.cody.agent.protocol_generated.CodeActions_ProvideParams
import com.sourcegraph.cody.agent.protocol_generated.CodeActions_ProvideResult
import com.sourcegraph.cody.agent.protocol_generated.CodeActions_TriggerParams
import com.sourcegraph.cody.agent.protocol_generated.Diagnostics_PublishParams
import com.sourcegraph.cody.agent.protocol_generated.EditTask
import com.sourcegraph.cody.agent.protocol_generated.EditTask_AcceptParams
import com.sourcegraph.cody.agent.protocol_generated.EditTask_CancelParams
import com.sourcegraph.cody.agent.protocol_generated.EditTask_GetTaskDetailsParams
import com.sourcegraph.cody.agent.protocol_generated.EditTask_RetryParams
import com.sourcegraph.cody.agent.protocol_generated.EditTask_UndoParams
import com.sourcegraph.cody.agent.protocol_generated.ExtensionConfiguration
import com.sourcegraph.cody.agent.protocol_generated.Null
import com.sourcegraph.cody.agent.protocol_generated.ServerInfo
import com.sourcegraph.cody.agent.protocol_generated.*
import java.util.concurrent.CompletableFuture
import org.eclipse.lsp4j.jsonrpc.services.JsonNotification
import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
Expand Down Expand Up @@ -166,6 +151,8 @@ interface _LegacyAgentServer {

@JsonRequest("chat/export") fun chatExport(): CompletableFuture<List<ChatHistoryResponse>>

@JsonRequest("chat/import") fun chat_import(params: Chat_ImportParams): CompletableFuture<Null?>

@JsonRequest("ignore/test")
fun ignoreTest(params: IgnoreTestParams): CompletableFuture<IgnoreTestResponse>

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
package com.sourcegraph.cody.config.migration

import com.intellij.openapi.project.Project
import com.sourcegraph.cody.agent.CodyAgentService
import com.sourcegraph.cody.agent.protocol_generated.Chat_ImportParams
import com.sourcegraph.cody.agent.protocol_generated.SerializedChatInteraction
import com.sourcegraph.cody.agent.protocol_generated.SerializedChatMessage
import com.sourcegraph.cody.agent.protocol_generated.SerializedChatTranscript
import com.sourcegraph.cody.config.CodyAccount
import com.sourcegraph.cody.config.CodyAuthenticationManager
import com.sourcegraph.cody.history.HistoryService
import com.sourcegraph.cody.history.state.ChatState
import com.sourcegraph.cody.history.state.MessageState

// Copies chat history from locally stored jetbrains state to the cody agent
// so historical chats can be viewed in the cody webview
object ChatHistoryMigration {
fun migrate(project: Project) {
CodyAgentService.withAgent(project) { agent ->
val chats =
CodyAuthenticationManager.getInstance(project).getAccounts().associateWith { account ->
(HistoryService.getInstance(project).getChatHistoryFor(account.id) ?: listOf())
}
val history = toChatInput(chats)

agent.server.chat_import(Chat_ImportParams(history = history, merge = true)).get()
}
}

fun toChatInput(
chats: Map<CodyAccount, List<ChatState>>
): Map<String, Map<String, SerializedChatTranscript>> {
return chats
.map { (account, chats) ->
val serializedChats = chats.mapNotNull(::toSerializedChatTranscript)
val byId = serializedChats.associateBy { it.id }

"${account.server.url}-${account.name}" to byId
}
.toMap()
}

private fun toSerializedChatTranscript(chat: ChatState): SerializedChatTranscript? {
return SerializedChatTranscript(
id = chat.internalId ?: return null,
lastInteractionTimestamp = chat.updatedAt ?: return null,
interactions = toSerializedInteractions(chat.messages, chat.llm?.model),
)
}
}

private fun toSerializedInteractions(
messages: List<MessageState>,
model: String?
): List<SerializedChatInteraction> {
fun toChatMessage(message: MessageState): SerializedChatMessage? {
return SerializedChatMessage(
text = message.text ?: return null,
model = model,
speaker =
when (message.speaker) {
MessageState.SpeakerState.HUMAN -> SerializedChatMessage.SpeakerEnum.Human
MessageState.SpeakerState.ASSISTANT -> SerializedChatMessage.SpeakerEnum.Assistant
null -> return null
},
)
}

fun toInteraction(pair: List<SerializedChatMessage?>): SerializedChatInteraction? {
val (human, assistant) = pair.getOrNull(0) to pair.getOrNull(1)
if (human == null ||
human.speaker != SerializedChatMessage.SpeakerEnum.Human ||
assistant?.speaker != SerializedChatMessage.SpeakerEnum.Assistant) {
return null
}
return SerializedChatInteraction(humanMessage = human, assistantMessage = assistant)
}

return messages.map(::toChatMessage).chunked(2).mapNotNull(::toInteraction)
}
Original file line number Diff line number Diff line change
Expand Up @@ -54,16 +54,19 @@ class SettingsMigration : Activity {
organiseChatsByAccount(project)
}
RunOnceUtil.runOnceForApp("CodyApplicationSettingsMigration") { migrateApplicationSettings() }
RunOnceUtil.runOnceForApp("ToggleCodyToolWindowAfterMigration") {
RunOnceUtil.runOnceForProject(project, "ToggleCodyToolWindowAfterMigration") {
toggleCodyToolbarWindow(project)
}
RunOnceUtil.runOnceForApp("CodyAccountsIdsRefresh") { refreshAccountsIds(project) }
RunOnceUtil.runOnceForApp("CodyAssignOrphanedChatsToActiveAccount") {
RunOnceUtil.runOnceForProject(project, "CodyAccountsIdsRefresh") { refreshAccountsIds(project) }
RunOnceUtil.runOnceForProject(project, "CodyAssignOrphanedChatsToActiveAccount") {
migrateOrphanedChatsToActiveAccount(project)
}

DeprecatedChatLlmMigration.migrate(project)
ChatTagsLlmMigration.migrate(project)
RunOnceUtil.runOnceForProject(project, "CodyMigrateChatHistory") {
ChatHistoryMigration.migrate(project)
}
}

private fun migrateOrphanedChatsToActiveAccount(project: Project) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,9 @@ class HistoryService(private val project: Project) :
fun findActiveAccountChat(internalId: String): ChatState? =
getActiveAccountHistory()?.chats?.find { it.internalId == internalId }

@Synchronized
fun getChatHistoryFor(accountId: String): List<ChatState>? = findEntry(accountId)?.chats

private fun findEntry(accountId: String): AccountData? =
state.accountData.find { it.accountId == accountId }

Expand Down
115 changes: 115 additions & 0 deletions src/test/kotlin/com/sourcegraph/cody/config/SettingsMigrationTest.kt
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,10 @@ import com.intellij.testFramework.fixtures.BasePlatformTestCase
import com.intellij.testFramework.registerServiceInstance
import com.sourcegraph.cody.agent.protocol_generated.Model
import com.sourcegraph.cody.agent.protocol_generated.ModelContextWindow
import com.sourcegraph.cody.agent.protocol_generated.SerializedChatInteraction
import com.sourcegraph.cody.agent.protocol_generated.SerializedChatMessage
import com.sourcegraph.cody.agent.protocol_generated.SerializedChatTranscript
import com.sourcegraph.cody.config.migration.ChatHistoryMigration
import com.sourcegraph.cody.config.migration.ChatTagsLlmMigration
import com.sourcegraph.cody.config.migration.DeprecatedChatLlmMigration
import com.sourcegraph.cody.config.migration.SettingsMigration
Expand All @@ -16,6 +20,7 @@ import com.sourcegraph.cody.history.state.LLMState
import com.sourcegraph.cody.history.state.MessageState
import com.sourcegraph.cody.history.state.MessageState.SpeakerState
import com.sourcegraph.cody.history.state.RemoteRepositoryState
import java.time.LocalDateTime
import kotlin.test.assertContains

class SettingsMigrationTest : BasePlatformTestCase() {
Expand Down Expand Up @@ -364,4 +369,114 @@ class SettingsMigrationTest : BasePlatformTestCase() {
}
}
}

fun `test toChatInput`() {
val account1 =
CodyAccount(name = "account1", server = SourcegraphServerPath("https://sourcegraph.com"))
val account2 =
CodyAccount(name = "account2", server = SourcegraphServerPath("https://sourcegraph.com"))

val chat1 =
ChatState("chat1").apply {
updatedAt = LocalDateTime.now().toString()
messages =
mutableListOf(
MessageState().apply {
text = "Hello"
speaker = SpeakerState.HUMAN
},
MessageState().apply {
text = "Hi there!"
speaker = SpeakerState.ASSISTANT
})
llm =
LLMState.fromChatModel(
Model(
id = "model1",
usage = listOf("chat"),
contextWindow = ModelContextWindow(0, 0, null),
clientSideConfig = null,
provider = "Anthropic",
title = "Claude",
tags = emptyList(),
modelRef = null))
}

val chat2 =
ChatState("chat2").apply {
updatedAt = LocalDateTime.now().minusDays(1).toString()
messages =
mutableListOf(
MessageState().apply {
text = "What's up?"
speaker = SpeakerState.HUMAN
},
MessageState().apply {
text = "Not much."
speaker = SpeakerState.ASSISTANT
})
llm =
LLMState.fromChatModel(
Model(
id = "model2",
usage = listOf("chat"),
contextWindow = ModelContextWindow(0, 0, null),
clientSideConfig = null,
provider = "Anthropic",
title = "Claude",
tags = emptyList(),
modelRef = null))
}

val chats = mapOf(account1 to listOf(chat1), account2 to listOf(chat2))

val result = ChatHistoryMigration.toChatInput(chats)

val expectedResult =
mapOf(
"https://sourcegraph.com-account1" to
mapOf(
chat1.internalId to
SerializedChatTranscript(
id = chat1.internalId!!,
lastInteractionTimestamp = chat1.updatedAt!!,
interactions =
listOf(
SerializedChatInteraction(
humanMessage =
SerializedChatMessage(
text = "Hello",
model = "model1",
speaker = SerializedChatMessage.SpeakerEnum.Human),
assistantMessage =
SerializedChatMessage(
text = "Hi there!",
model = "model1",
speaker =
SerializedChatMessage.SpeakerEnum
.Assistant))))),
"https://sourcegraph.com-account2" to
mapOf(
chat2.internalId to
SerializedChatTranscript(
id = chat2.internalId!!,
lastInteractionTimestamp = chat2.updatedAt!!,
interactions =
listOf(
SerializedChatInteraction(
humanMessage =
SerializedChatMessage(
text = "What's up?",
model = "model2",
speaker = SerializedChatMessage.SpeakerEnum.Human),
assistantMessage =
SerializedChatMessage(
text = "Not much.",
model = "model2",
speaker =
SerializedChatMessage.SpeakerEnum
.Assistant))))))

assertEquals(expectedResult, result)
}
}

0 comments on commit 2f5acfd

Please sign in to comment.