Recent advances in retrieval-augmented generation (RAG) have initiated a new era in repository-level code completion. However, the invariable use of retrieval in existing methods exposes issues in both efficiency and robustness, with a large proportion of the retrieved contexts proving unhelpful or harmful to code language models (code LMs). To tackle the challenges, this paper proposes a selective RAG framework where retrieval is avoided when unnecessary. To power this framework, we design a self-supervised learning approach that enables a code LM to accurately self-evaluate whether retrieval can improve its output quality and robustly leverage the potentially noisy retrieved contexts. Using this LM as both the selective retrieval policy and the generation model, our framework consistently outperforms the state-of-the-art prompting with an invariable retrieval approach on diverse benchmarks including RepoEval, CrossCodeEval, and a new benchmark. Meanwhile, our selective retrieval strategy results in strong efficiency improvements by as much as 70% inference speedup without harming the performance. We demonstrate that our framework effectively accommodates different generation models, retrievers, and programming languages. These advancements position our framework as an important step towards more accurate and efficient repository-level code completion.

翻译：近期检索增强生成（RAG）技术的进展开启了仓库级代码补全的新纪元。然而，现有方法中一成不变的检索策略引发了效率与鲁棒性问题——大量检索到的上下文对代码语言模型而言不仅无益，甚至有害。为解决上述挑战，本文提出了一种选择性RAG框架，在非必要场景下主动规避检索。为支撑该框架，我们设计了一种自监督学习方法，使代码语言模型能够准确自我评估检索能否提升输出质量，并鲁棒地利用可能含噪的检索上下文。将同一语言模型同时作为选择性检索策略与生成模型，我们的框架在RepoEval、CrossCodeEval及新基准数据集上，始终优于采用固定检索策略的现有最优提示方法。与此同时，该选择性检索策略在保持性能不降的前提下，实现了高达70%的推理加速。实验证明，该框架可有效适配不同的生成模型、检索器及编程语言。这些进展标志着我们的框架朝着更精准、更高效的仓库级代码补全迈出了重要一步。