[aia-eval] LME-89 Support static datasets for code generation scenario in End-to-End evaluation

GitOrigin-RevId: 558e41cd22748746e2eaf1b1ab299b0c787905f1

plugins/evaluation-plugin/core/src/com/intellij/cce/workspace/Config.kt

@@ -62,6 +62,7 @@ data class Config private constructor(
     val language: String,
     val evaluationRoots: List<String>,
     val ignoreFileNames: Set<String>,
+    val sourceFile: String?,
   )
 
   /**

plugins/evaluation-plugin/core/src/com/intellij/cce/workspace/ConfigFactory.kt

@@ -90,7 +90,8 @@ object ConfigFactory {
         projectName ?: map.getIfExists<String>("projectName")?.handleEnv() ?: resultProjectPath.split('/').last(),
         language ?: map.getAs("language"),
         map.getAs("evaluationRoots"),
-        map.getIfExists<List<String>>("ignoreFileNames")?.toSet() ?: emptySet()
+        map.getIfExists<List<String>>("ignoreFileNames")?.toSet() ?: emptySet(),
+        map.getIfExists<String>("sourceFile"),
       )
       return
     }

plugins/evaluation-plugin/core/src/com/intellij/cce/workspace/storages/storage/ActionsSingleFileStorage.kt

@@ -6,21 +6,16 @@ import com.intellij.cce.actions.FileActions
 import com.intellij.openapi.diagnostic.logger
 import java.io.File
 import java.nio.file.Files
-import java.nio.file.Paths
+import java.nio.file.Path
 import kotlin.io.path.absolutePathString
 
-class ActionsSingleFileStorage(storageDir: String) : ActionsStorage {
+class ActionsSingleFileStorage(filePath: Path) : ActionsStorage {
   private val LOG = logger<ActionsSingleFileStorage>()
 
-  private val file: File
-
-  init {
-    val filePath = Paths.get(storageDir, "actions")
-    file = if (Files.exists(filePath)) {
-      File(filePath.absolutePathString())
-    } else {
-      Files.createFile(filePath).toFile()
-    }
+  private val file: File = if (Files.exists(filePath)) {
+    File(filePath.absolutePathString())
+  } else {
+    Files.createFile(filePath).toFile()
   }
 
   override fun saveActions(actions: FileActions) {

plugins/evaluation-plugin/core/src/com/intellij/cce/workspace/storages/storage/ActionsStorage.kt

@@ -3,6 +3,7 @@ package com.intellij.cce.workspace.storages.storage
 
 import com.intellij.cce.actions.FileActions
 import com.intellij.cce.workspace.storages.ensureDirExists
+import java.nio.file.Paths
 
 interface ActionsStorage {
   fun saveActions(actions: FileActions)
@@ -27,7 +28,7 @@ object ActionsStorageFactory {
 
     return when (type) {
       ActionsStorageType.MULTIPLY_FILES -> ActionsMultiplyFilesStorage(storageDir)
-      ActionsStorageType.SINGLE_FILE -> ActionsSingleFileStorage(storageDir)
+      ActionsStorageType.SINGLE_FILE -> ActionsSingleFileStorage(Paths.get(storageDir, "actions"))
     }
   }
 }

plugins/evaluation-plugin/intellij.evaluationPlugin.iml

@@ -24,5 +24,9 @@
     <orderEntry type="library" name="commons-text" level="project" />
     <orderEntry type="module" module-name="intellij.platform.inspect" />
     <orderEntry type="module" module-name="intellij.platform.warmup" />
+    <orderEntry type="library" name="kotlinx-coroutines-core" level="project" />
+    <orderEntry type="library" name="jetbrains-annotations" level="project" />
+    <orderEntry type="library" name="ktor-client-core" level="project" />
+    <orderEntry type="library" name="ktor-client-java" level="project" />
   </component>
 </module>

plugins/evaluation-plugin/languages/evaluationconfig/codegeneration/java_static_dataset.json

@@ -0,0 +1,33 @@
+{
+  "language": "Java",
+  "projectPath": "../../ml-eval/keycloak",
+  "outputDir": "ml-eval-code-generation-java-output",
+  "strategy": {
+    "model": "",
+    "mode": "",
+    "behaviour": "REMOVE_METHOD_BODY_ASK_FOR_BODY_DESCRIPTION_AND_CALL_FEATURE",
+    "collectContextOnly": false,
+    "userPrompts": {}
+  },
+  "actions": {
+    "evaluationRoots": [
+      "."
+    ],
+    "sourceFile": "remote:https://huggingface.co/datasets/JetBrains/eval_plugin/resolve/main/code-generation/java/keycloak_keycloak/actions"
+  },
+  "interpret": {
+    "saveLogs": true,
+    "sessionProbability": 1.0,
+    "sessionSeed": null,
+    "sessionsLimit": 15,
+    "filesLimit": 2,
+    "order": "LINEAR",
+    "trainTestSplit": 70
+  },
+  "reports": {
+    "evaluationTitle": "Default",
+    "defaultMetrics": null,
+    "sessionsFilters": [],
+    "comparisonFilters": []
+  }
+}

plugins/evaluation-plugin/src/com/intellij/cce/actions/ContextCollectionEvaluationCommand.kt

@@ -71,7 +71,7 @@ internal class ContextCollectionEvaluationCommand : CompletionEvaluationStarter.
         evaluationRootInfo,
         dataset.project,
         dataset.processor,
-        feature.name
+        feature.name,
       ) {
         override fun prepare(datasetContext: DatasetContext, progress: Progress) {
           val files = runReadAction {

plugins/evaluation-plugin/src/com/intellij/cce/actions/DatasetContext.kt

@@ -20,15 +20,27 @@ class DatasetContext(
 ) {
   private val datasetDir = Paths.get("ml-eval-datasets")
 
-  val actionsStorage: ActionsStorage by lazy {
-    check(actionWorkspace != null) { "It seems that workspace with actions wasn't configured" }
+  /*
+    EvaluationDataset should be extracted into 3 parts - steps, dataset creation and dataset.
+    this will become simpler after refactoring
+   */
+  private var _actionsStorage: ActionsStorage? = null
+  val actionsStorage: ActionsStorage
+    get(): ActionsStorage {
+      if (_actionsStorage == null) {
+        check(actionWorkspace != null) { "It seems that workspace with actions wasn't configured" }
 
-    val directory = actionWorkspace.path().resolve("actions")
-    Files.createDirectories(directory)
-    ActionsStorageFactory.create(directory.toString(), getActionsStorageTypeFromEnv())
+        val directory = actionWorkspace.path().resolve("actions")
+        Files.createDirectories(directory)
+        _actionsStorage = ActionsStorageFactory.create(directory.toString(), getActionsStorageTypeFromEnv())
+      }
+      return _actionsStorage!!
+    }
+
+  fun replaceActionsStorage(storage: ActionsStorage) {
+    _actionsStorage = storage
   }
 
-
   val errorsStorage: FileErrorsStorage = outputWorkspace.errorsStorage
 
   fun saveAdditionalStats(name: String, stats: Map<String, Any>) {

plugins/evaluation-plugin/src/com/intellij/cce/actions/DatasetRef.kt

@@ -1,5 +1,7 @@
 package com.intellij.cce.actions
 
+import com.intellij.cce.util.httpGet
+import com.intellij.cce.workspace.storages.storage.ActionsSingleFileStorage
 import java.nio.file.Files
 import java.nio.file.Path
 import java.nio.file.StandardCopyOption
@@ -13,12 +15,17 @@ sealed interface DatasetRef {
 
   companion object {
     private const val EXISTING_PROTOCOL = "existing:"
+    private const val REMOTE_PROTOCOL = "remote:"
 
     fun parse(ref: String): DatasetRef {
       if (ref.startsWith(EXISTING_PROTOCOL)) {
         return ExistingRef(ref.substring(EXISTING_PROTOCOL.length))
       }
 
+      if (ref.startsWith(REMOTE_PROTOCOL)) {
+        return RemoteFileRef(ref.substring(REMOTE_PROTOCOL.length))
+      }
+
       if (ref.contains(":")) {
         throw IllegalArgumentException("Protocol is not supported: $ref")
       }
@@ -61,4 +68,21 @@ internal data class ExistingRef(override val name: String) : DatasetRef {
       "Dataset $name does not exist: $path"
     }
   }
+}
+
+internal data class RemoteFileRef(private val url: String) : DatasetRef {
+  override val name: String = url
+    .removePrefix("https://huggingface.co/datasets/JetBrains/eval_plugin/resolve/main/")
+    .replace("/", "_")
+
+  override fun prepare(datasetContext: DatasetContext) {
+    val readToken = System.getenv("AIA_EVALUATION_DATASET_READ_TOKEN")
+    check(readToken.isNotBlank()) {
+      "Token for dataset $url should be configured"
+    }
+
+    val content = httpGet(url, readToken)
+    val path = datasetContext.path(name)
+    path.toFile().writeText(content)
+  }
 }

plugins/evaluation-plugin/src/com/intellij/cce/actions/ProjectActionsDataset.kt

@@ -20,6 +20,7 @@ import com.intellij.cce.util.text
 import com.intellij.cce.visitor.CodeFragmentBuilder
 import com.intellij.cce.workspace.Config
 import com.intellij.cce.workspace.info.FileErrorInfo
+import com.intellij.cce.workspace.storages.storage.ActionsSingleFileStorage
 import com.intellij.openapi.application.ReadAction
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.project.Project
@@ -34,31 +35,54 @@ open class ProjectActionsDataset(
   private val evaluationRootInfo: EvaluationRootInfo,
   val project: Project,
   val processor: GenerateActionsProcessor,
-  private val featureName: String
+  private val featureName: String,
 ) : EvaluationDataset {
+  private val datasetRef = config.sourceFile?.run { DatasetRef.parse(this) }
+  private var datasetRefIsHandled = false
+
   override val setupSdk: EvaluationStep? = SetupSdkStep.forLanguage(project, Language.resolve(config.language))
   override val checkSdk: EvaluationStep? = CheckProjectSdkStep(project, config.language)
 
   override val preparationDescription: String = "Generating actions by selected files"
 
   override fun prepare(datasetContext: DatasetContext, progress: Progress) {
-    val filesForEvaluation = ReadAction.compute<List<VirtualFile>, Throwable> {
-      FilesHelper.getFilesOfLanguage(project, config.evaluationRoots, config.ignoreFileNames, config.language)
+    if (datasetRef != null) {
+      ensureDataRefIsHandled(datasetContext)
+    } else {
+      val filesForEvaluation = ReadAction.compute<List<VirtualFile>, Throwable> {
+        FilesHelper.getFilesOfLanguage(project, config.evaluationRoots, config.ignoreFileNames, config.language)
+      }
+
+      generateActions(
+        datasetContext,
+        config.language,
+        filesForEvaluation,
+        evaluationRootInfo,
+        progress,
+        filesLimit = this.filesLimit ?: Int.MAX_VALUE,
+        sessionsLimit = this.sessionsLimit ?: Int.MAX_VALUE,
+      )
     }
-    generateActions(
-      datasetContext,
-      config.language,
-      filesForEvaluation,
-      evaluationRootInfo,
-      progress,
-      filesLimit = this.filesLimit ?: Int.MAX_VALUE,
-      sessionsLimit = this.sessionsLimit ?: Int.MAX_VALUE,
-    )
   }
 
-  override fun sessionCount(datasetContext: DatasetContext): Int = datasetContext.actionsStorage.computeSessionsCount()
+  private fun ensureDataRefIsHandled(datasetContext: DatasetContext) {
+    if (!datasetRefIsHandled) {
+      if (datasetRef != null) {
+       datasetRef.prepare(datasetContext)
+       val path = datasetContext.path(datasetRef.name)
+       datasetContext.replaceActionsStorage(ActionsSingleFileStorage(path))
+     }
+      datasetRefIsHandled = true
+    }
+  }
+
+  override fun sessionCount(datasetContext: DatasetContext): Int {
+    ensureDataRefIsHandled(datasetContext)
+    return datasetContext.actionsStorage.computeSessionsCount()
+  }
 
   override fun chunks(datasetContext: DatasetContext): Iterator<EvaluationDatasetChunk> {
+    ensureDataRefIsHandled(datasetContext)
     val files = datasetContext.actionsStorage.getActionFiles()
     return files.shuffled(FILES_RANDOM).asSequence().map { file ->
       val fileActions = datasetContext.actionsStorage.getActions(file)

plugins/evaluation-plugin/src/com/intellij/cce/evaluable/EvaluableFeatureBase.kt

@@ -1,6 +1,7 @@
 // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.cce.evaluable
 
+import com.intellij.cce.actions.DatasetRef
 import com.intellij.cce.actions.ProjectActionsDataset
 import com.intellij.cce.core.Language
 import com.intellij.cce.evaluation.*

plugins/evaluation-plugin/src/com/intellij/cce/util/HttpUtil.kt

@@ -0,0 +1,46 @@
+package com.intellij.cce.util
+
+import com.intellij.openapi.progress.runBlockingCancellable
+import io.ktor.client.HttpClient
+import io.ktor.client.engine.java.Java
+import io.ktor.client.plugins.HttpRequestRetry
+import io.ktor.client.plugins.HttpTimeout
+import io.ktor.client.request.get
+import io.ktor.client.request.headers
+import io.ktor.client.statement.bodyAsText
+import io.ktor.http.HttpHeaders
+import io.ktor.http.isSuccess
+
+private val httpClient: HttpClient by lazy {
+  HttpClient(Java) {
+    expectSuccess = true
+    install(HttpRequestRetry) {
+      retryOnExceptionOrServerErrors(maxRetries = 3)
+      exponentialDelay()
+    }
+    install(HttpTimeout) {
+      connectTimeoutMillis = System.getProperty("idea.connection.timeout")?.toLongOrNull() ?: 10_000
+    }
+  }
+}
+
+private suspend fun httpGetSuspend(url: String, authToken: String): String {
+  val response = httpClient.get(url) {
+    headers {
+      append(HttpHeaders.Authorization, "Bearer $authToken")
+    }
+  }
+
+  if (!response.status.isSuccess()) {
+    error("failed to access $url. status=${response.status}")
+  }
+
+  return response.bodyAsText()
+}
+
+fun httpGet(url: String, authToken: String): String {
+  //todo refac eval framework to make it work with suspend funs
+  return runBlockingCancellable {
+    httpGetSuspend(url, authToken)
+  }
+}

plugins/evaluation-plugin/test/com/intellij/cce/workspace/ConfigFactoryTest.kt

@@ -84,6 +84,45 @@ class ConfigFactoryTest {
     }
   }
 
+  @Test
+  fun `test fileDataset serialization`() {
+    deserialize(
+      """
+      {
+        "outputDir": "test_outputDir",
+        "strategy": {},
+        "fileDataset": {
+          "url": "test_url",
+          "chunkSize": 300
+        }
+      }
+      """.trimIndent()
+    ).also {
+      assertEquals("test_url", it.fileDataset?.url)
+      assertEquals(300, it.fileDataset?.chunkSize)
+    }
+  }
+
+  @Test
+  fun `test actions sourceFile serialization`() {
+    deserialize(
+      """
+      {
+        "outputDir": "test_outputDir",
+        "strategy": {},
+        "projectPath": "projectPath",
+        "language": "LANG",
+        "actions": {
+        "evaluationRoots": [],
+          "sourceFile": "test_sourceFile"
+        }
+      }
+      """.trimIndent()
+    ).also {
+      assertEquals("test_sourceFile", it.actions?.sourceFile)
+    }
+  }
+
   private fun deserialize(text: String) =
     ConfigFactory.deserialize(ConfigFactory.createGson(TestStrategySerializer), text, TestStrategySerializer)