Skip to content

Commit a5f5138

Browse files
yangbobo2021yangbobo2021
committed
feat: Enhance variable context and type handling
- Improve findAccessibleVariables for better scope traversal - Refactor buildSymbolsContext for more detailed type information - Add helper functions for type validation and text limitation
1 parent 9279760 commit a5f5138

File tree

2 files changed

+141
-45
lines changed

2 files changed

+141
-45
lines changed

‎src/main/kotlin/ai/devchat/common/IDEUtils.kt‎

Lines changed: 61 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,9 @@ import kotlinx.coroutines.runBlocking
1818
import java.util.concurrent.CompletableFuture
1919
import java.util.concurrent.CountDownLatch
2020
import kotlin.system.measureTimeMillis
21-
21+
import com.intellij.psi.util.PsiTreeUtil
22+
import com.intellij.codeInsight.navigation.actions.GotoTypeDeclarationAction
23+
import com.intellij.openapi.fileEditor.FileEditorManager
2224

2325
object IDEUtils {
2426
fun <T> runInEdtAndGet(block: () -> T): T {
@@ -127,18 +129,65 @@ object IDEUtils {
127129
)
128130

129131
fun PsiElement.findAccessibleVariables(): Sequence<SymbolTypeDeclaration> {
130-
val projectFileIndex = ProjectFileIndex.getInstance(this.project)
131-
return generateSequence(this.parent) { it.parent }
132-
.takeWhile { it !is PsiFile }
133-
.flatMap { it.children.asSequence().filterIsInstance<PsiNameIdentifierOwner>() }
134-
.plus(this.containingFile.children.asSequence().filterIsInstance<PsiNameIdentifierOwner>())
135-
.filter { !it.name.isNullOrEmpty() && it.nameIdentifier != null }
136-
.mapNotNull {
137-
val typeDeclaration = it.getTypeDeclaration() ?: return@mapNotNull null
138-
val virtualFile = typeDeclaration.containingFile.virtualFile ?: return@mapNotNull null
139-
val isProjectContent = projectFileIndex.isInContent(virtualFile)
140-
SymbolTypeDeclaration(it, CodeNode(typeDeclaration, isProjectContent))
132+
val projectFileIndex = ProjectFileIndex.getInstance(project)
133+
134+
return sequence {
135+
// 从当前元素开始，向上遍历所有父元素
136+
var currentScope: PsiElement? = this@findAccessibleVariables
137+
while (currentScope != null && currentScope !is PsiFile) {
138+
// 在当前作用域中搜索变量声明
139+
val variablesInScope = PsiTreeUtil.findChildrenOfAnyType(
140+
currentScope,
141+
false,
142+
PsiNameIdentifierOwner::class.java
143+
)
144+
145+
for (variable in variablesInScope) {
146+
if (isLikelyVariable(variable) && !variable.name.isNullOrEmpty() && variable.nameIdentifier != null) {
147+
val typeDeclaration = getTypeElement(variable)
148+
if (typeDeclaration != null) {
149+
val virtualFile = variable.containingFile?.virtualFile
150+
if (virtualFile != null) {
151+
val isProjectContent = projectFileIndex.isInContent(virtualFile)
152+
yield(SymbolTypeDeclaration(variable, CodeNode(typeDeclaration, isProjectContent)))
153+
}
154+
}
155+
}
156+
}
157+
158+
// 移动到父作用域
159+
currentScope = currentScope.parent
141160
}
161+
162+
// 最后检查文件级别的变量
163+
yieldAll(this@findAccessibleVariables.containingFile.children
164+
.asSequence()
165+
.filterIsInstance<PsiNameIdentifierOwner>()
166+
.filter { isLikelyVariable(it) && !it.name.isNullOrEmpty() && it.nameIdentifier != null }
167+
.mapNotNull { variable ->
168+
val typeElement = getTypeElement(variable) ?: return@mapNotNull null
169+
val virtualFile = variable.containingFile?.virtualFile ?: return@mapNotNull null
170+
val isProjectContent = projectFileIndex.isInContent(virtualFile)
171+
SymbolTypeDeclaration(variable, CodeNode(typeElement, isProjectContent))
172+
})
173+
}
174+
}
175+
176+
// 辅助函数，用于判断一个元素是否可能是变量
177+
private fun isLikelyVariable(element: PsiElement): Boolean {
178+
val elementClass = element.javaClass.simpleName
179+
return elementClass.contains("Variable", ignoreCase = true) ||
180+
elementClass.contains("Parameter", ignoreCase = true) ||
181+
elementClass.contains("Field", ignoreCase = true)
182+
}
183+
184+
// 辅助函数，用于获取变量的类型元素
185+
private fun getTypeElement(element: PsiElement): PsiElement? {
186+
val project = element.project
187+
val editor = FileEditorManager.getInstance(project).selectedTextEditor ?: return null
188+
val offset = element.textOffset
189+
190+
return GotoTypeDeclarationAction.findSymbolType(editor, offset)
142191
}
143192

144193
fun PsiElement.foldTextOfLevel(foldingLevel: Int = 1): String {

‎src/main/kotlin/ai/devchat/plugin/completion/agent/ContextBuilder.kt‎

Lines changed: 80 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -116,16 +116,6 @@ class ContextBuilder(val file: PsiFile, val offset: Int) {
116116
}.lastOrNull()?.first?.last ?: content.length
117117
tokenCount += suffixTokens
118118

119-
val debugPrefixStart = maxOf(0, offset - 100)
120-
val debugSuffixEnd = minOf(content.length, offset + 100)
121-
Log.info("Debug: Offset 前 100 个字节文本:")
122-
Log.info(content.substring(debugPrefixStart, offset))
123-
Log.info("\n--- Offset 位置 ---\n")
124-
Log.info("Debug: Offset 后 100 个字节文本:")
125-
Log.info(content.substring(offset, debugSuffixEnd))
126-
127-
128-
129119
return Pair(
130120
content.substring(prefixStart, offset),
131121
content.substring(offset, suffixEnd)
@@ -146,7 +136,7 @@ class ContextBuilder(val file: PsiFile, val offset: Int) {
146136
?.takeWhile(::checkAndUpdateTokenCount)
147137
?.joinToString(separator = "") {snippet ->
148138
val commentedContent = snippet.content.lines()
149-
.joinToString(LINE_SEPARATOR.toString()) { "$commentPrefix $it" }
139+
.joinToString(LINE_SEPARATOR.toString())
150140
"$commentPrefix Function call definition:\n\n" +
151141
"$commentPrefix <filename>${snippet.filepath}\n\n" +
152142
"$commentPrefix <definition>\n$commentedContent\n\n\n\n"
@@ -156,29 +146,86 @@ class ContextBuilder(val file: PsiFile, val offset: Int) {
156146

157147
private fun buildSymbolsContext(): String {
158148
return runInEdtAndGet {
159-
file.findElementAt(offset)
160-
?.findAccessibleVariables()
161-
?.filter { it.typeDeclaration.element.containingFile.virtualFile.path != filepath }
162-
?.map {
163-
val typeElement = it.typeDeclaration.element
164-
it.symbol.name to CodeSnippet(
165-
typeElement.containingFile.virtualFile.path,
166-
if (it.typeDeclaration.isProjectContent) {
167-
typeElement.foldTextOfLevel(2)
149+
Log.info("Starting buildSymbolsContext")
150+
val element = file.findElementAt(offset)
151+
Log.info("Found element at offset: ${element?.text}")
152+
153+
val variables = element?.findAccessibleVariables() ?: emptySequence()
154+
val variablesCount = variables.count()
155+
Log.info("Found $variablesCount accessible variables")
156+
157+
val processedTypes = mutableSetOf<String>()
158+
val result = StringBuilder()
159+
160+
variables
161+
.onEach { Log.info("Processing variable: ${it.symbol.name}") }
162+
.forEach { variable ->
163+
val typeElement = variable.typeDeclaration.element
164+
val isLocalType = typeElement.containingFile.virtualFile.path == filepath
165+
val typeText = limitTypeText(typeElement.text)
166+
Log.info("Variable ${variable.symbol.name} type: ${typeText}")
167+
Log.info("Is local type: $isLocalType")
168+
169+
val typeFilePath = typeElement.containingFile.virtualFile.path
170+
Log.info("Actual type file: $typeFilePath")
171+
172+
// 如果typeFilePath表示了系统库的定义，那么不应该添加到上下文中，例如string的定义。
173+
if (isValidTypePath(typeFilePath)) {
174+
val typeKey = "${typeElement.text}:$typeFilePath"
175+
if (!processedTypes.contains(typeKey)) {
176+
processedTypes.add(typeKey)
177+
178+
val snippet = CodeSnippet(
179+
typeFilePath,
180+
if (variable.typeDeclaration.isProjectContent) {
181+
typeElement.foldTextOfLevel(2)
182+
} else {
183+
typeElement.text.lines().first() + "..."
184+
}
185+
)
186+
187+
if (checkAndUpdateTokenCount(snippet)) {
188+
Log.info("Adding context for type: ${typeText}")
189+
val commentedContent = snippet.content.lines()
190+
.joinToString(LINE_SEPARATOR.toString())
191+
result.append("$commentPrefix Symbol type definition:\n\n")
192+
.append("$commentPrefix <symbol>${variable.symbol.name}\n\n")
193+
.append("$commentPrefix <filename>${snippet.filepath}\n\n")
194+
.append("$commentPrefix <definition>\n$commentedContent\n\n\n\n")
195+
} else {
196+
Log.info("Skipping type ${variable.symbol.name} due to token limit")
197+
return@runInEdtAndGet result.toString()
198+
}
168199
} else {
169-
typeElement.text.lines().first() + "..."
200+
Log.info("Skipping duplicate type: ${typeText}")
170201
}
171-
)
202+
}
172203
}
173-
?.takeWhile { checkAndUpdateTokenCount(it.second) }
174-
?.joinToString(separator = "") {(name, snippet) ->
175-
val commentedContent = snippet.content.lines()
176-
.joinToString(LINE_SEPARATOR.toString()) { "$commentPrefix $it" }
177-
"$commentPrefix Symbol type definition:\n\n" +
178-
"$commentPrefix <symbol>${name}\n\n" +
179-
"$commentPrefix <filename>${snippet.filepath}\n\n" +
180-
"$commentPrefix <definition>\n$commentedContent\n\n\n\n"
181-
} ?: ""
204+
205+
Log.info("buildSymbolsContext result length: ${result.length}")
206+
result.toString()
207+
}
208+
}
209+
210+
private fun isValidTypePath(path: String): Boolean {
211+
// 这里需要根据具体的项目结构和依赖管理方式来实现
212+
// 例如，可以检查路径是否在项目目录下，或者是否在已知的第三方依赖目录下
213+
// 以下是一个简单的示例实现
214+
val projectPath = file.project.basePath ?: return false
215+
return path.startsWith(projectPath)
216+
//return path.startsWith(projectPath) || path.contains("/.gradle/") || path.contains("/build/")
217+
}
218+
219+
// 新增的辅助函数，用于限制类型文本的输出长度
220+
private fun limitTypeText(text: String, maxLines: Int = 5): String {
221+
val lines = text.lines()
222+
return when {
223+
lines.size <= maxLines -> text
224+
else -> {
225+
val firstLines = lines.take(maxLines / 2)
226+
val lastLines = lines.takeLast(maxLines / 2)
227+
(firstLines + "..." + lastLines).joinToString("\n")
228+
}
182229
}
183230
}
184231

@@ -192,7 +239,7 @@ class ContextBuilder(val file: PsiFile, val offset: Int) {
192239
.takeWhile(::checkAndUpdateTokenCount)
193240
.joinToString(separator = "") {snippet ->
194241
val commentedContent = snippet.content.lines()
195-
.joinToString(LINE_SEPARATOR.toString()) { "$commentPrefix $it" }
242+
.joinToString(LINE_SEPARATOR.toString())
196243
"$commentPrefix Recently open file:\n\n" +
197244
"$commentPrefix <filename>${snippet.filepath}\n\n" +
198245
"$commentedContent\n\n\n\n"
@@ -211,7 +258,7 @@ class ContextBuilder(val file: PsiFile, val offset: Int) {
211258
// similarBlockContext,
212259
// gitDiffContext,
213260
).joinToString("")
214-
// Log.info("Extras completion context:\n$extras")
261+
215262
return if (!model.isNullOrEmpty() && model.contains("deepseek"))
216263
"<｜fim▁begin｜>$extras$commentPrefix<filename>$filepath\n\n$prefix<｜fim▁hole｜>$suffix<｜fim▁end｜>"
217264
else

0 commit comments

Comments
 (0)