# Context-Augmented Code Generation Using Programming Knowledge Graphs

arXiv: https://arxiv.org/abs/2410.18251 · HTML: https://mlsanity.com/paper/2410.18251/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 52 references that are in the map.
- First posted on arXiv 9 Oct 2024.

## Builds on (most cited first, 28 of 28)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Retrieval-Augmented Generation for Large Language Models: A Survey](https://mlsanity.com/paper/2312.10997/) (2023)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [CodeBERT: A Pre-Trained Model for Programming and Natural Languages](https://mlsanity.com/paper/2002.08155/) (2020)
- [When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories](https://mlsanity.com/paper/2212.10511/) (2022)
- [Atlas: Few-shot Learning with Retrieval Augmented Language Models](https://mlsanity.com/paper/2208.03299/) (2022)
- [StarCoder: may the source be with you!](https://mlsanity.com/paper/2305.06161/) (2023)
- [In-Context Retrieval-Augmented Language Models](https://mlsanity.com/paper/2302.00083/) (2023)
- [GraphCodeBERT: Pre-training Code Representations with Data Flow](https://mlsanity.com/paper/2009.08366/) (2020)
- [Active Retrieval Augmented Generation](https://mlsanity.com/paper/2305.06983/) (2023)
- [StarCoder 2 and The Stack v2: The Next Generation](https://mlsanity.com/paper/2402.19173/) (2024)
- [RepoCoder: Repository-Level Code Completion Through Iterative Retrieval and Generation](https://mlsanity.com/paper/2303.12570/) (2023)
- [CodeT5+: Open Code Large Language Models for Code Understanding and Generation](https://mlsanity.com/paper/2305.07922/) (2023)
- [DeepSeek-Coder-V2: Breaking the Barrier of Closed-Source Models in Code Intelligence](https://mlsanity.com/paper/2406.11931/) (2024)
- [HippoRAG: Neurobiologically Inspired Long-Term Memory for Large Language Models](https://mlsanity.com/paper/2405.14831/) (2024)
- [DocPrompting: Generating Code by Retrieving the Docs](https://mlsanity.com/paper/2207.05987/) (2022)
- [Retrieval Augmented Code Generation and Summarization](https://mlsanity.com/paper/2108.11601/) (2021)
- [CodeRAG-Bench: Can Retrieval Augment Code Generation?](https://mlsanity.com/paper/2406.14497/) (2024)
- [RepoGraph: Enhancing AI Software Engineering with Repository-level Code Graph](https://mlsanity.com/paper/2410.14684/) (2024)
- [Reliable, Adaptable, and Attributable Language Models with Retrieval](https://mlsanity.com/paper/2403.03187/) (2024)
- [CodexGraph: Bridging Large Language Models and Code Repositories via Code Graph Databases](https://mlsanity.com/paper/2408.03910/) (2024)
- [EVOR: Evolving Retrieval for Code Generation](https://mlsanity.com/paper/2402.12317/) (2024)

## Built on by (most cited first, 4 of 4)

- [A Survey of Context Engineering for Large Language Models](https://mlsanity.com/paper/2507.13334/) (2025)
- [Retrieval-Augmented Code Generation: A Survey with Focus on Repository-Level Approaches](https://mlsanity.com/paper/2510.04905/) (2025)
- [HCAG: Hierarchical Abstraction and Retrieval-Augmented Generation on Theoretical Repositories with LLMs](https://mlsanity.com/paper/2603.20299/) (2026)
- [Model-Driven Quantum Code Generation Using Large Language Models and Retrieval-Augmented Generation](https://mlsanity.com/paper/2508.21097/) (2025)

## Abstract (arXiv)

Large Language Models (LLMs) and Code-LLMs (CLLMs) have significantly improved code generation, but, they frequently face difficulties when dealing with challenging and complex problems. Retrieval-Augmented Generation (RAG) addresses this issue by retrieving and integrating external knowledge at the inference time. However, retrieval models often fail to find most relevant context, and generation models, with limited context capacity, can hallucinate when given irrelevant data. We present a novel framework that leverages a Programming Knowledge Graph (PKG) to semantically represent and retrieve code. This approach enables fine-grained code retrieval by focusing on the most relevant segments while reducing irrelevant context through a tree-pruning technique. PKG is coupled with a re-ranking mechanism to reduce even more hallucinations by selectively integrating non-RAG solutions. We propose two retrieval approaches-block-wise and function-wise-based on the PKG, optimizing context granularity. Evaluations on the HumanEval and MBPP benchmarks show our method improves pass@1 accuracy by up to 20%, and outperforms state-of-the-art models by up to 34% on MBPP. Our contributions include PKG-based retrieval, tree pruning to enhance retrieval precision, a re-ranking method for robust solution selection and a Fill-in-the-Middle (FIM) enhancer module for automatic code augmentation with relevant comments and docstrings.