Exploration of Structural Code Relationship Space for Context Collection

This paper presents a solution to the Context Collection Competition organized by JetBrains and Mistral AI, addressing the challenge of optimizing code completion across diverse codebases and multiple LLMs. The competition includes both Kotlin and a Python tracks, and this paper represents a Kotlin submission. The problem was approached by designing a dynamic, composable framework, enabling rapid experimentation with multiple context-collection strategies. Building on the Kotlin Psi library, a semantic tree layer was introduced to capture meaningful code associations and dynamically adjust the level of detail provided in the context, such as omitting method bodies or private methods. This approach, similar to Hierarchical Context Pruning, maximizes the signal-to-noise ratio in retrieved context. Experiments show that concise structural associations and targeted pruning consistently improve performance, though results vary across LLMs, with Codestral benefiting most. The framework’s adaptability allows systematic exploration of configuration parameters and the combination of multiple strategies to optimize performance. Limited computational resources constrained exhaustive exploration, motivating future work on automated large-scale evaluation to identify the upper bounds of achievable performance.