Teaching Code LLMs to Use Autocompletion Tools in Repository-Level Code Generation

TL;DR

ToolGen addresses repository-level code generation challenges by integrating static autocompletion tools into LLM generation, using an offline trigger insertion and fine-tuning phase and an online tool-aware generation phase. It augments code with <COMP> markers at candidate dependency identifiers, enabling Jedi to propose completions that are then selected by a constrained greedy mechanism. Across 12,406 Python functions from 671 repos and 176 CoderEval tasks, ToolGen improves Dependency Coverage and Static Validity Rate substantially across CodeGPT, CodeT5, CodeLlama, while maintaining or improving standard similarity metrics and pass rates. The approach demonstrates generation efficiency and generalizability across model architectures, suggesting practical utility for repository-aware code generation.

Abstract

Code large language models (LLMs) face limitations in repository-level code generation due to their lack of awareness of repository-level dependencies (e.g., user-defined attributes), resulting in dependency errors such as undefined-variable and no-member errors. In this work, we introduce ToolGen, an approach that integrates autocompletion tools into the code LLM generation process to address these dependencies. ToolGen comprises two main phases: Trigger Insertion and Model Fine-tuning (Offline), and Tool-integrated Code Generation (Online). During the offline phase, ToolGen augments functions within a given code corpus with a special mark token, indicating positions to trigger autocompletion tools. These augmented functions, along with their corresponding docstrings, are then used to fine-tune a selected code LLM. In the online phase, ToolGen iteratively generates functions by predicting tokens step-by-step using the fine-tuned LLM. Whenever a mark token is encountered, ToolGen invokes the autocompletion tool to suggest code completions and selects the most appropriate one. We conduct comprehensive experiments to evaluate ToolGen's effectiveness in repository-level code generation. To facilitate this evaluation, we create a benchmark comprising 671 real-world code repositories and introduce two new dependency-based metrics: Dependency Coverage and Static Validity Rate. The results demonstrate that ToolGen significantly improves Dependency Coverage by 31.4% to 39.1% and Static Validity Rate by 44.9% to 57.7% across the three LLMs, while maintaining competitive or improved performance in widely recognized similarity metrics such as BLEU-4, CodeBLEU, Edit Similarity, and Exact Match. On the CoderEval dataset, ToolGen achieves improvements of 40.0% and 25.0% in Pass@1 for CodeT5 and CodeLlama, respectively.

Figures (7)

• Figure 1: Illustrative Example of LLM Prediction and Tool Completion
• Figure 2: Decoder-only Model and Encoder-Decoder Model
• Figure 3: Approach Overview of ToolGen
• Figure 4: Augmented Function
• Figure 5: Example Prefix Tree