# Context-Augmented Code Generation Using Programming Knowledge Graphs

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 38 of its 52 references that are in the map.
- First posted on arXiv 28 Jan 2026.

## Builds on (most cited first, 30 of 38)

- [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)
- [Enhancing repository-level software repair via repository-aware knowledge graphs](https://mlsanity.com/paper/2503.21710/) (2025)
- [Provence: efficient and robust context pruning for retrieval-augmented generation](https://mlsanity.com/paper/2501.16214/) (2025)
- [Retrieval-Augmented Code Generation: A Survey with Focus on Repository-Level Approaches](https://mlsanity.com/paper/2510.04905/) (2025)

## Abstract (arXiv)

Large Language Models (LLMs) excel at code generation but struggle with complex problems. Retrieval-Augmented Generation (RAG) mitigates this issue by integrating external knowledge, yet retrieval models often miss relevant context, and generation models hallucinate with irrelevant data. We propose Programming Knowledge Graph (PKG) for semantic representation and fine-grained retrieval of code and text. Our approach enhances retrieval precision through tree pruning and mitigates hallucinations via a re-ranking mechanism that integrates non-RAG solutions. Structuring external data into finer-grained nodes improves retrieval granularity. Evaluations on HumanEval and MBPP show up to 20% pass@1 accuracy gains and a 34% improvement over baselines on MBPP. Our findings demonstrate that our proposed PKG approach along with re-ranker effectively address complex problems while maintaining minimal negative impact on solutions that are already correct without RAG. The replication package is published at https://github.com/iamshahd/ProgrammingKnowledgeGraph