Source-linked AI summary
Prompt-based Code Completion via Multi-Retrieval Augmented Generation
Hanzhuo Tan, Qi Luo, Ling Jiang, Zizheng Zhan, Jing Li, Haotian Zhang, Yuqun Zhang
TL;DR
Code-completion RAG systems use a single encoder-defined perspective, which can miss the diversity of code semantics, while LLMs can struggle with complex or unfamiliar code. ProCC combines prompt-based retrieval perspectives with adaptive selection, outperforming the state-of-the-art technique by 8.6% on the open-source benchmark and 10.1% on the private-domain benchmark in Exact Match.
Problem
Existing RAG techniques rely on external encoders and a single lexical perspective, limiting representation of code’s complexity and diversity; LLMs can also struggle with complex or unfamiliar code.
Source: §1 · p. 2Method
ProCC prompts an LLM for multiple code perspectives and uses a contextual multi-armed bandit to select a retrieval perspective based on code similarity.
Source: §1 · p. 2Results
8.6% on the open-source benchmark and 10.1% on the private-domain benchmark in Exact Match over the state-of-the-art technique; ProCC also yields a 5.6% improvement over the studied fine-tuned model.
Source: §1 · p. 2Takeaways & Limitations
ProCC provides a more comprehensive and adjustable encoding strategy than rigid prior representations without requiring extra encoders.
Source: §1 · p. 2External validity is bounded by the benchmarks and techniques studied, although the evaluation includes industry data and averages results over five runs.
Source: §5 · p. 10
Abstract
from arXiv · showhide
Automated code completion, aiming at generating subsequent tokens from unfinished code, has been significantly benefited from recent progress in pre-trained Large Language Models (LLMs). However, these models often suffer from coherence issues and hallucinations when dealing with complex code logic or extrapolating beyond their training data. Existing Retrieval Augmented Generation (RAG) techniques partially address these issues by retrieving relevant code with a separate encoding model where the retrieved snippet serves as contextual reference for code completion. However, their retrieval scope is subject to a singular perspective defined by the encoding model, which largely overlooks the complexity and diversity inherent in code semantics. To address this limitation, we propose ProCC, a code completion framework leveraging prompt engineering and the contextual multi-armed bandits algorithm to flexibly incorporate and adapt to multiple perspectives of code. ProCC first employs a prompt-based multi-retriever system which crafts prompt templates to elicit LLM knowledge to understand code semantics with multiple retrieval perspectives. Then, it adopts the adaptive retrieval selection algorithm to incorporate code similarity into the decision-making process to determine the most suitable retrieval perspective for the LLM to complete the code. Experimental results demonstrate that ProCC outperforms state-of-the-art code completion technique by 8.6% on our collected open-source benchmark suite and 10.1% on the private-domain benchmark suite collected from a billion-user e-commerce company in terms of Exact Match. ProCC also allows augmenting fine-tuned techniques in a plug-and-play manner, yielding 5.6% improvement over our studied fine-tuned model.
1 INTRODUCTION
ProCC addresses limitations in code completion and retrieval by combining prompt-based perspectives on code with adaptive retrieval selection. It outperforms the state-of-the-art technique on open-source and private-domain benchmarks and improves a studied fine-tuned model.
Background
- Conventional program-analysis approaches can be limited in semantic, line-level completion, while statistical methods require costly data collection, annotation, and training to adapt across domains.These limitations constrain their applicability to real-world development.Source: §1 · p. 1
Motivation
- LLMs can generate incoherent or hallucinated code for complex logic or beyond their training data, while fine-tuning requires substantial time and computational resources.RAG offers a more feasible alternative by retrieving relevant information to support generation without fine-tuning.Source: §1 · p. 2
- Existing RAG techniques rely on external encoders and a single lexical perspective, limiting their ability to represent the complexity and diversity of code semantics.Encoding code from additional perspectives with predefined models would require substantially more training.Source: §1 · p. 2
Technique
- ProCC uses prompts to retrieve from lexical, hypothetical-line, and code-summary perspectives, then selects a perspective using code similarity and a contextual multi-armed bandit.The selection is conditioned on each incomplete code snippet; the system avoids requiring additional encoding models.Source: §1 · p. 2
Evaluation
- 8.6% on the open-source benchmark and 10.1% on the private-domain benchmark in Exact Match, while ProCC yields a 5.6% improvement over the studied fine-tuned model.The evaluation also reports that prompt-based retrievers are robust across templates and comparable to external encoders.Source: §1 · pp. 2–3
2 BACKGROUND AND MOTIVATION
The section introduces code completion with prefix and suffix context, then explains how retrieval augments generation. It motivates ProCC with cases where the most useful retrieval perspective depends on the incomplete code.
2.1 Code Completion
- Fill-in-the-Middle completion conditions missing tokens on both prefix and suffix; Code Llama reaches 85.6% pass@1 on infilling benchmarks.The section notes that access to suffix code is critical and practical for code completion.Source: §2.1 · p. 3
- LLMs may produce incoherent or hallucinated code when handling complex logic or generalizing beyond their training data.These limitations motivate accessing external memory or knowledge through retrieval.Source: §2.1 · p. 3, §2.2 · p. 3
2.2 Retrieval-Augmented Generation
- RAG encodes code snippets and incomplete code to retrieve context for generation; the section focuses on per-output RAG, which retrieves once before generating the sequence.Per-token RAG retrieves for each generated token, increasing retrieval and encoding demands and limiting integration with black-box systems.Source: §2.2 · p. 3
- Within RAG-sequence, ReACC uses a dual-encoder retriever, while RepoCoder iteratively retrieves using recently generated content.Both are cited as existing code-completion approaches.Source: §2.2 · p. 3
2.3 Motivation
- Different completion contexts favor different retrievals: Figure 1 contrasts helpful and misleading hints, motivating flexible selection across code perspectives.The examples show one retriever helping in one scenario while another is more useful in a different scenario.Source: §2.3 · p. 4, Figure 1 · p. 4
- ProCC uses prompts to elicit multiple interpretations of unfinished code without training extra encoders, then adaptively selects retrieved results for that code.The framework’s three perspectives are lexical semantics, a hypothetical line, and code summarization.Source: §2.3 · p. 4, Figure 2 · p. 5
3 APPROACH
ProCC uses prompted retrievals to represent code from multiple semantic perspectives, then adapts which perspective it selects for each incomplete snippet. Its approach combines three prompt-based retrievers with a contextual bandit selection algorithm.
3.1 Overview
- ProCC retrieves code through lexical, hypothetical-line, and summarization perspectives, then uses a contextual bandit to select one retrieval for completion.The selected retrieval is combined with the unfinished code as input to the completion model.Source: §3.1 · p. 5
3.2 Prompt-based Multi-Retriever System
- Prompted LLM representations avoid separately trained encoding models and can be extended to semantic dimensions beyond the three implemented perspectives.The framework derives representations from output hidden states.Source: §3.2 · pp. 5–6
- Lexical retrieval averages the last-layer hidden states for a prompt containing the unfinished code.The prompt is “Embedding the following code snippets: [code]”.Source: §3.2 · p. 6
- Hypothetical-line retrieval generates a possible missing line from the surrounding prefix and suffix and uses its embedding to search the database.This representation can include variable names, types, and function signatures absent from the unfinished snippet’s lexical context.Source: §3.2 · p. 6
- Code-summarization retrieval represents a snippet through an LLM-generated natural-language summary, emphasizing functional similarity over surface syntax.The method averages the summary embeddings to form the code representation.Source: §3.2 · p. 6
3.3 Adaptive Retrieval Selection Algorithm
- LinUCB treats the three retrieval perspectives as arms and selects one using retriever similarity scores and Jaccard similarity as context.Its reward is Exact Match, and its decisions adapt to the context of each incomplete snippet.Source: §3.3 · p. 6
- The selection algorithm updates LinUCB parameters from feature vectors and rewards, then passes the selected retrieval with unfinished code to the completion model.The feature vector includes Jaccard and cosine similarities between incomplete code and retrieved results.Source: §3.3 · p. 7
4 EVALUATION
ProCC is evaluated against code-completion baselines on open-source and private-domain benchmarks, with analyses of its retrieval components and use alongside fine-tuned models. Results report improvements on both benchmark suites and from augmenting a fine-tuned model.
4.1 Experiment Setup
- Evaluation uses StarCoder and Code Llama across 20 GitHub and 58 private repositories, with EM and ES metrics and BM25, ReACC, and RepoCoder baselines.The test sets contain 3,317 open-source and 3,074 private-domain instances.Source: §4.1.1 · p. 7, §4.1.3 · p. 7, §4.1.4 · p. 7
4.2 Results and Analysis
- 8.6% EM improvement over RepoCoder (50.31 to 54.66) on Code Llama’s open-source results; StarCoder gains 6.6% EM (46.11 to 49.14).The evaluation also reports an average 14.1% improvement for Code Llama over the vanilla model.Source: §4.2.1 · pp. 7–8, Table 1 · p. 8
- For Code Llama on the open-source suite, lexical, hypothetical-line, and code-summary retrievers score 50.44, 51.61, and 51.40 EM, respectively, versus RepoCoder’s 50.31.Each prompt template also produces exclusive correct completions; Prompt No.3 uniquely retrieves correct contextual hints for 65 samples.Source: §4.2.1 · p. 8, §4.2.2 · p. 8
- 10.1% EM improvement over BM25 (41.84 to 46.07) on the private-domain benchmark; ProCC improves 36.8% over the baseline Code Llama model (33.67 to 46.07).The private benchmark uses domain-specific code that the paper says falls outside the LLMs’ training corpus.Source: Table 3 · p. 8, §4.2.1 · p. 8
- Prompt combinations No.1+3, No.1+5, and No.3+5 score 52.43, 52.94, and 53.51 EM, while six perspectives add only 0.22 EM over three prompts.The three-prompt combination scores 54.66 EM; the paper notes that overlapping information may limit gains from adding perspectives.Source: Table 5 · p. 9, §4.2.2 · p. 9
- LinUCB improves EM by 3.05 over the best single retriever; maximum similarity and logistic regression improve it by 0.76 and 1.6, while union reaches 51.97 EM.The union’s gain is 0.36 over the best single perspective’s 51.61 EM.Source: Table 6 · p. 9, §4.2.2 · p. 9
- Applying ProCC to fine-tuned models adds 3.41 EM on the open-source suite and 4.94 EM on the private-domain suite.The paper reports these gains over the fine-tuned model.Source: §4.2.3 · p. 10
- Fine-tuning Code Llama takes about 9.5 hours on eight A100 GPUs, whereas ProCC needs no training and takes about 0.5 seconds for embedding and search on one A100.The paper characterizes ProCC as more resource-efficient to deploy than fine-tuning.Source: §4.2.3 · p. 10
5 THREATS TO VALIDITY
The paper identifies internal, external, and construct validity threats, with safeguards including code review, industry data, repeated runs, and established metrics.
5 THREATS TO VALIDITY
- The internal validity threat is potential implementation bugs; the authors obtained original source code, matched published hyperparameters, and reviewed their scripts for correctness.Source: §5 · p. 10
- External validity is mainly constrained by the benchmarks and techniques studied; the authors included industry data unknown to LLMs and averaged results over five runs.Source: §5 · p. 10
- The construct validity threat concerns evaluation metrics; the study uses Exact Match and Edit Similarity, following previous work.Source: §5 · p. 10
6 RELATED WORK
Related work traces code completion from statistical and pretrained language models to retrieval-augmented methods. ProCC builds on retrieval-based completion while introducing multiple prompted semantic perspectives and adaptive selection.
Language Model for Code Completion
- Modern code LMs use autoregressive prediction, while Code Llama's fill-in-the-middle objective resembles incomplete code contexts and provides an inductive bias for infilling [4] [56].Source: §6 · p. 10
Retrieval Augmented Code Completion
- RAG retrieves relevant information using similarity measures or embeddings, then appends it to the input to guide generation; code-completion adaptations include code-to-code retrieval and iterative retrieval.Source: §6 · pp. 10–11
- Other retrieval work incorporates in-domain code databases with Bayesian inference [60] or retrieves cross-file dependencies, with benchmarks and a cross-file context finder supporting that setting.Source: §6 · p. 11
- ProCC combines prompt engineering and contextual multi-armed bandits to adaptively incorporate multiple code perspectives, and its evaluation reports gains over existing RAG-based completion techniques.Source: §6 · p. 11
7 CONCLUSION
ProCC combines prompt-based multi-perspective retrieval with adaptive selection informed by code similarity. Evaluations on open-source and private-domain repositories report Exact Match gains over the state-of-the-art technique.
7 CONCLUSION
- ProCC combines a prompt-based multi-retriever system with an adaptive retrieval selection algorithm that uses code similarity to choose a retrieval perspective.Source: §7 · p. 11
- The framework is evaluated across open-source and private-domain repositories.Source: §7 · p. 11
- 8.6% and 10.1% are ProCC's reported Exact Match gains over the state-of-the-art technique on the open-source and private-domain benchmarks, respectively.Source: §7 · p. 11
8 DATA AVAILABILITY
The private-domain dataset remains undisclosed under the company's deployment and privacy policy; other materials are available through the repository.
8 DATA AVAILABILITY
- The dataset contains private-domain code from the e-commerce company where the artifact was deployed.Source: §8 · p. 11
- The private-domain dataset is not disclosed because of the company's privacy protection policy.Source: §8 · p. 11
- The repository provides the artifact's source code and the open-source dataset among the other available materials.Source: §8 · p. 11