[evaluation-plugin] Add Feature, Invoker, and Strategy for documentation generation task

GitOrigin-RevId: c81a856e0c5658a2e800f7163aba954b0c5dfbc3
This commit is contained in:
Michele Conti
2024-05-08 14:33:56 +00:00
committed by intellij-monorepo-bot
parent ce0c710c4a
commit 881f34e404
12 changed files with 319 additions and 0 deletions

View File

@@ -27,6 +27,10 @@ interface TokenProperties {
if (simpleTokenProperties.tokenType == TypeProperty.LINE) {
return context.deserialize(json, LineProperties::class.java)
}
else if (json.asJsonObject.has("docComment")) {
return context.deserialize<DocumentationProperties>(json, DocumentationProperties::class.java)
}
return simpleTokenProperties
}
@@ -132,6 +136,15 @@ class SimpleTokenProperties private constructor(
SimpleTokenProperties(tokenType, location, this.features.apply { addAll(features) }, additional)
}
class DocumentationProperties(val docComment: String, val startOffset: Int, val endOffset: Int, val docStartOffset: Int, val docEndOffset: Int, val nameIdentifierOffset: Int) : TokenProperties {
override val tokenType: TypeProperty = TypeProperty.UNKNOWN
override val location: SymbolLocation = SymbolLocation.UNKNOWN
override fun additionalProperty(name: String): String? = null
override fun describe(): String = ""
override fun hasFeature(feature: String): Boolean = false
override fun withFeatures(features: Set<String>): TokenProperties = this
}
enum class SymbolLocation {
PROJECT, LIBRARY, UNKNOWN
}

View File

@@ -6,6 +6,7 @@
<setupSdkStep implementation="com.intellij.cce.evaluation.SetupJDKStep"/>
<lineCompletionVisitorFactory implementation="com.intellij.cce.visitor.JavaLineCompletionVisitorFactory"/>
<completionEvaluationVisitor implementation="com.intellij.cce.visitor.JavaTestGenerationVisitor"/>
<completionEvaluationVisitor implementation="com.intellij.cce.visitor.JavaDocGenerationVisitor"/>
<completionEvaluationVisitor implementation="com.intellij.cce.visitor.JavaCodeGenerationVisitor"/>
<completionEvaluationVisitor implementation="com.intellij.cce.visitor.JavaCompletionContextEvaluationVisitor"/>
</extensions>

View File

@@ -0,0 +1,56 @@
package com.intellij.cce.visitor
import com.intellij.cce.core.*
import com.intellij.cce.visitor.exceptions.PsiConverterException
import com.intellij.psi.JavaRecursiveElementVisitor
import com.intellij.psi.PsiClass
import com.intellij.psi.PsiJavaFile
import com.intellij.psi.PsiMethod
import com.intellij.psi.util.endOffset
import com.intellij.psi.util.startOffset
class JavaDocGenerationVisitor : EvaluationVisitor, JavaRecursiveElementVisitor() {
override val feature: String = "doc-generation"
override val language: Language = Language.JAVA
private var codeFragment: CodeFragment? = null
override fun getFile(): CodeFragment = codeFragment ?: throw PsiConverterException("Invoke 'accept' with visitor on PSI first")
override fun visitJavaFile(file: PsiJavaFile) {
codeFragment = CodeFragment(file.textOffset, file.textLength)
super.visitJavaFile(file)
}
override fun visitMethod(method: PsiMethod) {
val docComment = method.docComment
val nameIdentifier = method.nameIdentifier
if (docComment != null && nameIdentifier != null) {
codeFragment?.addChild(
CodeToken(
method.text,
method.textRange.startOffset,
DocumentationProperties(docComment.text, method.startOffset, method.endOffset, docComment.startOffset, docComment.endOffset, nameIdentifier.startOffset)
)
)
}
super.visitMethod(method)
}
override fun visitClass(klass: PsiClass) {
val docComment = klass.docComment
val nameIdentifier = klass.nameIdentifier
if (docComment != null && nameIdentifier != null) {
codeFragment?.addChild(
CodeToken(
klass.text,
klass.textRange.startOffset,
DocumentationProperties(docComment.text, klass.startOffset, klass.endOffset, docComment.startOffset, docComment.endOffset, nameIdentifier.startOffset)
)
)
}
super.visitClass(klass)
}
}

View File

@@ -4,5 +4,6 @@
<completionEvaluationVisitor implementation="com.intellij.cce.visitor.KotlinMultiLineEvaluationVisitor"/>
<completionEvaluationVisitor implementation="com.intellij.cce.visitor.KotlinCompletionContextEvaluationVisitor"/>
<lineCompletionVisitorFactory implementation="com.intellij.cce.visitor.KotlinLineCompletionVisitorFactory"/>
<completionEvaluationVisitor implementation="com.intellij.cce.visitor.KotlinDocGenerationVisitor"/>
</extensions>
</idea-plugin>

View File

@@ -0,0 +1,67 @@
package com.intellij.cce.visitor
import com.intellij.cce.core.*
import com.intellij.cce.visitor.exceptions.PsiConverterException
import com.intellij.psi.util.endOffset
import com.intellij.psi.util.startOffset
import org.jetbrains.kotlin.psi.*
class KotlinDocGenerationVisitor : EvaluationVisitor, KtTreeVisitorVoid() {
override val feature: String = "doc-generation"
override val language: Language = Language.KOTLIN
private var codeFragment: CodeFragment? = null
override fun getFile(): CodeFragment = codeFragment ?: throw PsiConverterException("Invoke 'accept' with visitor on PSI first")
override fun visitKtFile(file: KtFile) {
codeFragment = CodeFragment(file.textOffset, file.textLength).apply { text = file.text }
super.visitKtFile(file)
}
override fun visitNamedFunction(function: KtNamedFunction) {
val docComment = function.docComment
val nameIdentifier = function.nameIdentifier
if (docComment != null && nameIdentifier != null) {
codeFragment?.addChild(
CodeToken(
function.text,
function.textRange.startOffset,
DocumentationProperties(docComment.text, function.startOffset, function.endOffset, docComment.startOffset, docComment.endOffset, nameIdentifier.startOffset)
)
)
}
super.visitNamedFunction(function)
}
override fun visitClass(klass: KtClass) {
val docComment = klass.docComment
val nameIdentifier = klass.nameIdentifier
if (docComment != null && nameIdentifier != null) {
codeFragment?.addChild(
CodeToken(
klass.text,
klass.textRange.startOffset,
DocumentationProperties(docComment.text, klass.startOffset, klass.endOffset, docComment.startOffset, docComment.endOffset, nameIdentifier.startOffset)
)
)
}
super.visitClass(klass)
}
override fun visitObjectDeclaration(declaration: KtObjectDeclaration) {
val docComment = declaration.docComment
val nameIdentifier = declaration.nameIdentifier
if (docComment != null && nameIdentifier != null) {
codeFragment?.addChild(
CodeToken(
declaration.text,
declaration.textRange.startOffset,
DocumentationProperties(docComment.text, declaration.startOffset, declaration.endOffset, docComment.startOffset, docComment.endOffset, nameIdentifier.startOffset)
)
)
}
super.visitObjectDeclaration(declaration)
}
}

View File

@@ -5,5 +5,6 @@
<lineCompletionVisitorFactory implementation="com.intellij.cce.visitor.PythonLineCompletionVisitorFactory"/>
<setupSdkStep implementation="com.intellij.cce.evaluation.SetupPythonInterpreterStep"/>
<completionEvaluationVisitor implementation="com.intellij.cce.visitor.PythonCompletionContextEvaluationVisitor"/>
<completionEvaluationVisitor implementation="com.intellij.cce.visitor.PythonDocGenerationVisitor"/>
</extensions>
</idea-plugin>

View File

@@ -0,0 +1,52 @@
package com.intellij.cce.visitor
import com.intellij.cce.core.*
import com.intellij.cce.visitor.exceptions.PsiConverterException
import com.intellij.psi.util.endOffset
import com.intellij.psi.util.startOffset
import com.jetbrains.python.psi.*
class PythonDocGenerationVisitor : EvaluationVisitor, PyRecursiveElementVisitor() {
override val feature: String = "doc-generation"
private var codeFragment: CodeFragment? = null
override val language: Language = Language.PYTHON
override fun getFile(): CodeFragment = codeFragment
?: throw PsiConverterException("Invoke 'accept' with visitor on PSI first")
override fun visitPyFile(file: PyFile) {
codeFragment = CodeFragment(file.textOffset, file.textLength).apply { text = file.text }
super.visitPyFile(file)
}
override fun visitPyFunction(function: PyFunction) {
val docComment = function.docStringExpression
val nameIdentifier = function.nameIdentifier
if (docComment != null && nameIdentifier != null) {
codeFragment?.addChild(
CodeToken(
function.text,
function.textRange.startOffset,
DocumentationProperties(docComment.text, function.startOffset, function.endOffset, docComment.startOffset, docComment.endOffset, nameIdentifier.startOffset)
)
)
}
super.visitPyFunction(function)
}
override fun visitPyClass(klass: PyClass) {
val docComment = klass.docStringExpression
val nameIdentifier = klass.nameIdentifier
if (docComment != null && nameIdentifier != null) {
codeFragment?.addChild(
CodeToken(
klass.text,
klass.textRange.startOffset,
DocumentationProperties(docComment.text, klass.startOffset, klass.endOffset, docComment.startOffset, docComment.endOffset, nameIdentifier.startOffset)
)
)
}
super.visitPyClass(klass)
}
}

View File

@@ -43,6 +43,7 @@
<extensions defaultExtensionNs="com.intellij.cce">
<evaluableFeature implementation="com.intellij.cce.evaluable.rename.RenameFeature"/>
<evaluableFeature implementation="com.intellij.cce.evaluable.docGeneration.DocGenerationFeature"/>
<evaluableFeature implementation="com.intellij.cce.evaluable.testGeneration.TestGenerationFeature"/>
<evaluableFeature implementation="com.intellij.cce.evaluable.completion.TokenCompletionFeature"/>
<evaluableFeature implementation="com.intellij.cce.actions.ContextCollectionFeature"/>

View File

@@ -0,0 +1,37 @@
// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package com.intellij.cce.evaluable.docGeneration
import com.intellij.cce.actions.CallFeature
import com.intellij.cce.actions.DeleteRange
import com.intellij.cce.actions.MoveCaret
import com.intellij.cce.actions.PrintText
import com.intellij.cce.core.CodeFragment
import com.intellij.cce.core.CodeToken
import com.intellij.cce.core.DocumentationProperties
import com.intellij.cce.processor.GenerateActionsProcessor
class DocGenerationActionsProcessor : GenerateActionsProcessor() {
override fun process(code: CodeFragment) {
for (token in code.getChildren()) {
processToken(token as CodeToken)
}
}
private fun processToken(token: CodeToken) {
val properties = token.properties as DocumentationProperties
val docLen = properties.docComment.length
// Recompute nameIdentifierOffset after the docComment deletion
val correctNameIdentifierOffset = when {
properties.nameIdentifierOffset > properties.docStartOffset -> properties.nameIdentifierOffset - docLen
else -> properties.nameIdentifierOffset
}
addAction(DeleteRange(properties.docStartOffset, properties.docEndOffset))
addAction(MoveCaret(correctNameIdentifierOffset))
addAction(CallFeature(token.text, correctNameIdentifierOffset, properties))
addAction(MoveCaret(properties.docStartOffset))
addAction(PrintText(properties.docComment))
}
}

View File

@@ -0,0 +1,35 @@
// 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.docGeneration
import com.google.gson.JsonObject
import com.google.gson.JsonSerializationContext
import com.intellij.cce.core.Language
import com.intellij.cce.evaluable.EvaluableFeatureBase
import com.intellij.cce.evaluable.StrategySerializer
import com.intellij.cce.evaluation.EvaluationStep
import com.intellij.cce.interpreter.FeatureInvoker
import com.intellij.cce.metric.EmptyContextSessionRatio
import com.intellij.cce.metric.Metric
import com.intellij.cce.metric.SessionsCountMetric
import com.intellij.cce.processor.GenerateActionsProcessor
import com.intellij.openapi.project.Project
import java.lang.reflect.Type
class DocGenerationFeature : EvaluableFeatureBase<DocGenerationStrategy>("doc-generation") {
override fun getGenerateActionsProcessor(strategy: DocGenerationStrategy): GenerateActionsProcessor =
DocGenerationActionsProcessor()
override fun getFeatureInvoker(project: Project, language: Language, strategy: DocGenerationStrategy): FeatureInvoker =
DocGenerationInvoker(project, language, strategy)
override fun getStrategySerializer(): StrategySerializer<DocGenerationStrategy> = object : StrategySerializer<DocGenerationStrategy> {
override fun deserialize(map: Map<String, Any>, language: String): DocGenerationStrategy = DocGenerationStrategy()
override fun serialize(src: DocGenerationStrategy, typeOfSrc: Type, context: JsonSerializationContext): JsonObject = JsonObject()
}
override fun getMetrics(): List<Metric> = listOf(SessionsCountMetric(), EmptyContextSessionRatio())
override fun getEvaluationSteps(language: Language, strategy: DocGenerationStrategy): List<EvaluationStep> = emptyList()
}

View File

@@ -0,0 +1,43 @@
package com.intellij.cce.evaluable.docGeneration
import com.intellij.cce.core.*
import com.intellij.cce.evaluable.common.getEditorSafe
import com.intellij.cce.evaluation.SuggestionsProvider
import com.intellij.cce.interpreter.FeatureInvoker
import com.intellij.openapi.application.runInEdt
import com.intellij.openapi.application.runReadAction
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.project.Project
import com.intellij.psi.PsiDocumentManager
class DocGenerationInvoker(private val project: Project,
private val language: Language,
private val strategy: DocGenerationStrategy) : FeatureInvoker {
override fun callFeature(expectedText: String, offset: Int, properties: TokenProperties): Session {
val editor = runReadAction {
getEditorSafe(project)
}
runInEdt {
PsiDocumentManager.getInstance(project).commitDocument(editor.document)
}
val session = Session(offset, expectedText, expectedText.length, TokenProperties.UNKNOWN)
val lookup = getSuggestions(expectedText, editor, strategy.suggestionsProvider)
val docProperties = properties as DocumentationProperties
val res = mapOf("docComment" to docProperties.docComment)
session.addLookup(Lookup.fromExpectedText("", "", lookup.suggestions, lookup.latency, null, false, additionalInfo = lookup.additionalInfo + res, comparator = this::comparator))
return session
}
protected fun getSuggestions(expectedLine: String, editor: Editor, suggestionsProviderName: String): Lookup {
val lang = com.intellij.lang.Language.findLanguageByID(language.ideaLanguageId)
?: throw IllegalStateException("Can't find language \"${language.ideaLanguageId}\"")
val provider = SuggestionsProvider.find(project, suggestionsProviderName)
?: throw IllegalStateException("Can't find suggestions provider \"${suggestionsProviderName}\"")
return provider.getSuggestions(expectedLine, editor, lang, this::comparator)
}
override fun comparator(generated: String, expected: String): Boolean = expected == generated
}

View File

@@ -0,0 +1,12 @@
package com.intellij.cce.evaluable.docGeneration
import com.intellij.cce.evaluable.EvaluationStrategy
import com.intellij.cce.filter.EvaluationFilter
class DocGenerationStrategy(val suggestionsProvider: String = DEFAULT_PROVIDER) : EvaluationStrategy {
override val filters: Map<String, EvaluationFilter> = emptyMap()
companion object {
private const val DEFAULT_PROVIDER: String = "LLM-doc-gen"
}
}