Context-Augmented Code Generation Using Programming Knowledge Graphs

TL;DR

This work tackles the brittleness of Retrieval-Augmented Generation (RAG) for code by introducing Programming Knowledge Graphs (PKG), a dual- modal, structure-aware representation that captures code (via AST-derived function/block hierarchies) and text (via constrained JSON path-value graphs). PKG enables multi-granularity retrieval (function-level vs block-level) with tree pruning and a post-generation reranker to mitigate noise and hallucinations, improving code correctness on benchmarks like HumanEval and MBPP. Empirical results show that code-centric PKG with proper candidate selection yields substantial gains (up to ~20% pass@1 on MBPP and notable improvements on HumanEval), while text-centric PKG offers more modest, model- and topic-dependent benefits; the reranker consistently enhances robustness. The paper also discusses model-adaptive retrieval, topic-sensitive effects, error-mode shifts, and cost considerations, offering practical guidelines for deploying PKG-based retrieval in real-world coding assistants and outlining open directions such as adaptive granularity and tighter integration of execution-aware checks.

Abstract

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

Figures (12)

• Figure 1: The overview of generating PKG
• Figure 2: Overview of the retrieval process from PKG
• Figure 3: Comparison of different approaches across 10 topics using the MBPP benchmark on StarCoder2-7B
• Figure 4: Comparison of different approaches across 10 topics using the MBPP benchmark on CodeLlama-7B.
• Figure 5: Comparison of different approaches across 10 topics using the MBPP benchmark on DeepSeek-Coder-7B