Structure-Aware Dependency Retrieval for Repository-Level Code Completion via Graph Attention

Abstract

When Large Language Models (LLMs) write code inside an existing repository, the quality of their output depends heavily on the context they are given. The difficulty is that a function’s docstring rarely mentions the base class, callee, or field the function actually depends on. Lexical and dense retrievers both treat code as flat text, so they tend to miss dependencies that are close in the code structure but share little wording with the query. We address this by building a knowledge graph for each repository. It extends the CodexGraph schema with containment at the repository and directory level, and with call edges that cross file boundaries, resolved through the import graph. On top of this graph, we train a small graph neural network that reranks the top K candidates from a dense retriever by examining each candidate together with its immediate neighbours, using attention that is conditioned on the query. We train and evaluate it on a new benchmark built from CodeSearchNet: 68,685 queries across 294 repositories, with graded relevance labels derived from the code’s syntax trees and no repository shared between training and test. We measure retrieval quality directly and leave end-to-end completion to future work. Over the strongest baseline, which expands BM25 results by one hop in the same graph and reranks them by embedding similarity, our reranker improves NDCG@5 from 12.76 to 17.04, a 33% relative gain that is statistically significant under a paired bootstrap test. Access to the graph alone is not enough, since that baseline barely improves on plain BM25. Ablations show that the gain comes from message passing over the dependency graph rather than from the training labels, with USES edges (calls and variable accesses, resolved across files through imports) carrying most of the signal.