摘要:软件开发者在编写代码时, 常常会参考项目中实现了相似功能的代码. 代码生成模型在生成代码时也具有类似特点, 会以输入中给出的代码上下文信息作为参考. 基于检索增强的代码补全技术与这一思想类似, 该技术从检索库中检索到外部代码作为额外信息, 对生成模型起到提示的作用, 从而生成目标代码. 现有的基于检索增强的代码补全方法将输入代码和检索结果直接拼接到一起作为生成模型的输入, 这种方法带来了一个风险, 即检索到的代码片段可能并不能对模型起到提示作用, 反而有可能会误导模型, 导致生成的代码结果不准确. 此外, 由于无论检索到的外部代码是否与输入代码完全相关, 都会被与输入代码拼接起来输入到模型, 这导致该方法的效果在很大程度上依赖于代码检索阶段的准确性. 如果检索阶段不能返回可用的代码片段, 那么后续的代码补全效果可能也会受到影响. 首先, 针对现有的代码补全方法中的检索增强策略进行了经验研究, 通过定性和定量实验分析检索增强的各个阶段对于代码补全效果的影响, 在经验研究中重点识别了代码粒度、代码检索方法、代码后处理方法这3种影响检索增强效果的因素. 接着, 基于经验研究的结论设计改进方法, 提出一种通过分阶段优化代码检索策略来改进检索增强的代码补全方法MAGIC (multi-stage optimization for retrieval augmented code completion), 设计了代码切分、二次检索精排、模板提示生成等改进策略, 可以有效地提升检索增强对代码补全模型的辅助生成作用, 并减少模型在代码生成阶段受到的噪声干扰, 提升生成代码的质量. 最后, 在Java代码数据集上的实验结果表明: 与现有的基于检索增强的代码补全方法相比, 该方法在编辑相似度和完全匹配指标上分别提升了6.76个百分点和7.81个百分点. 与6B参数量的代码大模型相比, 该方法能够在节省94.5%的显存和73.8%的推理时间的前提下, 在编辑相似度和完全匹配指标上分别提升了5.62个百分点和4.66个百分点.