AlignCoder: Aligning Retrieval with Target Intent for Repository-Level Code Completion

TL;DR

AlignCoder addresses repository-level code completion by tackling the misalignment between unfinished queries and target code and by enabling retrieval to exploit inference information. It introduces a query enhancement mechanism that generates multiple candidate completions to bridge semantic gaps, and trains an AlignRetriever via reinforcement learning to leverage the enhanced queries for more accurate retrieval. Across CrossCodeEval and RepoEval benchmarks and five backbone LLMs, AlignCoder achieves substantial improvements (e.g., up to $18.1\%$ EM on CrossCodeEval Python) and demonstrates strong generalizability across languages. The approach combines base and dependency code snippets to form a rich retrieval codebase, a coarse-to-fine retrieval pipeline with multiple sampling, and an RL-based retriever objective, contributing a practical advancement for repository-aware code completion.

Abstract

Repository-level code completion remains a challenging task for existing code large language models (code LLMs) due to their limited understanding of repository-specific context and domain knowledge. While retrieval-augmented generation (RAG) approaches have shown promise by retrieving relevant code snippets as cross-file context, they suffer from two fundamental problems: misalignment between the query and the target code in the retrieval process, and the inability of existing retrieval methods to effectively utilize the inference information. To address these challenges, we propose AlignCoder, a repository-level code completion framework that introduces a query enhancement mechanism and a reinforcement learning based retriever training method. Our approach generates multiple candidate completions to construct an enhanced query that bridges the semantic gap between the initial query and the target code. Additionally, we employ reinforcement learning to train an AlignRetriever that learns to leverage inference information in the enhanced query for more accurate retrieval. We evaluate AlignCoder on two widely-used benchmarks (CrossCodeEval and RepoEval) across five backbone code LLMs, demonstrating an 18.1% improvement in EM score compared to baselines on the CrossCodeEval benchmark. The results show that our framework achieves superior performance and exhibits high generalizability across various code LLMs and programming languages.

Figures (6)

• Figure 1: Motivating example. Limitations of prior works: ReACC and RLCoder use the unfinished code as the query, which may miss key tokens in the ground truth. RepoCoder may cause chain errors and has potential efficiency issue.
• Figure 2: EM@k performance trends on CrossCodeEval and RepoEval with varying k parameters. The trends indicate that multiple sampling strategy provides a higher probability of producing a correct completion compared tosingle sampling.
• Figure 3: Overview of AlignCoder. Our approach utilizes the inference capability of the LLM to sample multiple candidate completions and learn to understand the inference information for accurate retrieval for repository-level code completion.
• Figure 4: Example: the initial query and the enhanced query obtained with the query enhancement mechanism. Task id: project_cc_python/74.
• Figure 5: Temperature and top-$p$ sampling parameter effects on AlignCoder performance.