Context-Augmented Code Generation Using Programming Knowledge Graphs

Assessing the Computational Trade-offs and Limitations of PKG in Code Generation: Efficiency, Scalability, and Task-Specific Challenges

# Reviewer Concerns: ## 1. The PKG generation and retrieval involve multiple modules, adding complexity and cost. However, the paper lacks discussion on scalability and computational trade-offs, limiting its feasibility assessment. Thank you for your insightful observation. We will add the cost trade-off when we use PKG compared to current RAG approaches. We will explain these details in the paper by the Nov 27th in a new section in appendix (cost trade-off section): We can consider two versions of PKGs for discussion: ### **PKG with enhancer module, which aims to enhance the graph with doc-string data**: Step1: downloading dataset (143,000 QA, 280MB): negligible, a few seconds. Step2: extract python code using function analyzer: 3 minutes. Step3: Code block extraction: 25 minutes. Step4: enhanced PKG with doc_strings using LLM on one A100: 82 hours. Step5: encode PKG using voyage embeddings using API calls: 4 hours. Step6: generating neo4j graph: 33 minutes. In general for 143,000 Q&A (~500,000 nodes will be added in graph) it took around: ***87 Hours***. Total storage: ***12,530 MB*** (voyage embeddings of code blocks on Neo4j) ### **PKG without enhancer module** (Step4 is excluded, and Func-block results would be excluded in the result tables): Step1: downloading dataset (143,000 QA, 280MB): negligible, a few seconds. Step2: extract python code using function analyzer: 3 minutes. Step3: Code block extraction: 25 minutes. Step5: encode PKG using voyage embeddings using API calls: 4 hours. Step6: generating neo4j graph: 33 minutes. In general for 143,000 Q&A (~500,000 nodes will be added in graph): it took around: ***5 Hours***. Total storage: ***12,530 MB*** (voyage embeddings of code blocks on Neo4j) ### **Voyage Embedding RAG:** Step1: downloading dataset (143,000 QA, 280MB): negligible, a few seconds. Step2: encode dataset using voyage embeddings using API calls: 4 hours. In general for 143,000 Q&A: it took around: ***4 Hours.*** Total storage: ***8440 MB*** (voyage embeddings of Q&A samples) ### **BM25**: Step1: downloading dataset (143,000 QA, 280MB): negligible, a few seconds. Step2: index dataset: 44 minutes. In general for 143,000 Q&A: it took around: ***44 minutes.*** Total storage:***315mb*** (index and text data) The results in Tables 1 and 2 show that removing the func-block has minimal impact on performance. While our approach takes an additional hour to process the dataset compared to embedding-based RAG methods like Voyage Embeddings, it leads to a significant performance improvement. **graph updates**: Neo4j's semantic vector indexing ensures efficient additions, with time complexities of O(logN) for nodes and O(logM) for relationships, where 𝑁 and 𝑀 are the total nodes and relationships. **maintenance**: The graph remains useful as long as its content are not deprecated in the target programming language. ## 2. The re-ranking mechanism in the paper is used to enhance retrieval accuracy, but the decision-making process is not detailed enough. The author should provide a more comprehensive description and explanation of the method. Thank you for your consideration. Our re-ranking mechanism involves three steps: AST Validation: Ensuring syntactic correctness via Abstract Syntax Tree (AST) checks. Execution Testing: Running valid answers and excluding those with runtime errors. Embedding Similarity: Selecting the answer most similar to the query based on embeddings. Details are in Section 2.3. Let us know which parts require further clarification. ## 3. While PKG performs well overall, the paper lacks analysis of categories like string manipulation and data structures where PKG retrieval falls short. Examining these cases could uncover limitations of graph-based retrieval and identify conditions where PKG is less effective. Thank you for your comment. We agree that exploring cases where PKG retrieval does not improve results, especially in string manipulation and data structure tasks, could provide valuable insights into the limitations of graph-based retrieval and how LLMs interpret input data. In string manipulation tasks, the challenge lies in the model's focus on semantic meaning rather than string structure. Example Problem: Write a python code to convert lowercase to uppercase and vice versa: "Hello" to "hELLO" and "pYthon" to "PyTHON." Challenges: Embedding Model’s Focus on Semantics: In RAG, the embedder retrieves content based on meaning, not formatting. It may focus on "hello" as a greeting rather than the case transformation. LLM’s Tokenization and Semantic Bias: LLMs tokenize based on meaning, not formatting, making case transformation difficult as "Hello" and "hello" are treated the same. In summary, both RAG retrieval and LLM tokenization prioritize semantics over formatting, complicating string manipulation tasks and limiting PKG’s effectiveness in these cases.

New version is uploaded.

Dear Reviewer, We would like to kindly inform you that a revised version of the paper has been uploaded. In this updated version, we have extended our evaluation to include a new text-centric data source (Python tutorials). Our results demonstrate that the proposed approach effectively leverages text-centric data to retrieve more precise and relevant content, leading to measurable performance improvements in the code generation task. You can find the updated results in Section 7.2 of the Appendix.

Evaluating the Resource Trade-offs, Language Generalizability, and Performance Enhancements of PKG

# Reviewer Concerns: ## 1.Building and maintaining the Programming Knowledge Graph may be resource-intensive and require domain expertise. Thank you for your insightful observation. We will add the cost trade-off when we use PKG compared to current RAG approaches. We will explain these details in the paper by the Nov 27th in a new section in appendix (cost trade-off section): We can consider two versions of PKGs for discussion: ### **PKG with enhancer module, which aims to enhance the graph with doc-string data**: Step1: downloading dataset (143,000 QA, 280MB): negligible, a few seconds. Step2: extract python code using function analyzer: 3 minutes. Step3: Code block extraction: 25 minutes. Step4: enhanced PKG with doc_strings using LLM on one A100: 82 hours. Step5: encode PKG using voyage embeddings using API calls: 4 hours. Step6: generating neo4j graph: 33 minutes. In general for 143,000 Q&A (~500,000 nodes will be added in graph) it took around: ***87 Hours***. Total storage: ***12,530 MB*** (voyage embeddings of code blocks on Neo4j) ### **PKG without enhancer module** (Step4 is excluded, and Func-block results would be excluded in the result tables): Step1: downloading dataset (143,000 QA, 280MB): negligible, a few seconds. Step2: extract python code using function analyzer: 3 minutes. Step3: Code block extraction: 25 minutes. Step5: encode PKG using voyage embeddings using API calls: 4 hours. Step6: generating neo4j graph: 33 minutes. In general for 143,000 Q&A (~500,000 nodes will be added in graph): it took around: ***5 Hours***. Total storage: ***12,530 MB*** (voyage embeddings of code blocks on Neo4j) ### **Voyage Embedding RAG:** Step1: downloading dataset (143,000 QA, 280MB): negligible, a few seconds. Step2: encode dataset using voyage embeddings using API calls: 4 hours. In general for 143,000 Q&A: it took around: ***4 Hours.*** Total storage: ***8440 MB*** (voyage embeddings of Q&A samples) ### **BM25**: Step1: downloading dataset (143,000 QA, 280MB): negligible, a few seconds. Step2: index dataset: 44 minutes. In general for 143,000 Q&A: it took around: ***44 minutes.*** Total storage:***315mb*** (index and text data) Based on the comparisons and the results presented in Tables 1 and 2, we can conclude that removing the func-block does not significantly impact performance. In comparison to existing retrieval-augmented generation (RAG) methods, which primarily rely on embedding approaches (such as Voyage Embeddings), our approach takes an additional hour to process the selected dataset. However, this extra time results in a significant performance improvement over the standard embedding-based RAG methods. ## 2. The framework’s effectiveness may be constrained to specific programming languages (e.g. Python) of code tasks Thank you for your insightful comment. Our approach works with any programming language that supports AST extraction. For languages without a native AST library, third-party tools like tree-sitter, which supports over 25 languages, can be used. PKG can be built on any language with structural code blocks, treating each block as the smallest semantic node and retrieving the most similar nodes during inference. ## 3. The paper could benefit from a deeper exploration where PKG and retrieval mechanisms fail to improve or potentially hinder code generation quality. Thank you for the feedback. If the model requires domain expertise, the PKG should reflect that domain. For example, if the model targets a specific framework, the dataset must include it, or for project-specific code, the PKG should contain project data. Failures occur when querying a graph lacking domain knowledge. We will elaborate on this in the paper by Nov 27th. ## 4. What is the computational cost of building and updating the PKG, and how frequently does it need maintenance to remain effective? We have provided the computational costs regarding the building PKG in the **first concern**. **graph updates:** Neo4j's semantic vector indexing ensures efficient graph updates, with O(logN) complexity for adding nodes and O(logM) for relationships, where 𝑁 and 𝑀 are the total nodes and relationships. This logarithmic growth ensures scalability. **maintenance:** As long as the existing content in the graph is not deprecated in the target programming language, we could benefit from that knowledge graph. ## 5. How does the additional complexity introduced by tree pruning and re-ranking affect code generation performance, and is this overhead manageable? Thanks for pointing this out! We conducted tree pruning and re-ranking in a Neo4j environment using Cypher queries and the Graph Data Science plugin. With a graph of around 500,000 nodes, each query took under 5 seconds on an M1 chip. Despite the added complexity, the performance improvement was significant and manageable.

New version is uploaded.

Dear Reviewer, We would like to kindly inform you that a revised version of the paper has been uploaded. In this updated version, we have extended our evaluation to include a new text-centric data source (Python tutorials). Our results demonstrate that the proposed approach effectively leverages text-centric data to retrieve more precise and relevant content, leading to measurable performance improvements in the code generation task. You can find the updated results in Section 7.2 of the Appendix.

Addressing Retrieval Limitations, Dataset Structure, and Open-Domain Contexts in RAG Settings

# Reviewer Concerns: ## 1. The setting is kind of weird to me: in real-world applications, code generation is usually augmented by code documents, natural language thoughts, or similar question-solution pairs, which is to say, natural language could be used to retrieve helpful information. Instead, only focusing on code representations may be limited: Thank you for your thoughtful observation. In our initial experiments, we tried augmenting the model with a knowledge graph that included Python documents and natural language explanations. However, this approach led to worse results than the No-RAG baseline. This is because, when the goal is to generate accurate code, providing code-based context is more effective than natural language context. Natural language inputs often lead the model to focus on generating explanations or descriptions rather than precise code outputs. This observation aligns with findings from the Code-RAG benchmark paper [1], where the authors noted that Stack Overflow content, which is code-centric, yields better results than tutorial or documentation content when used as context in RAG settings. Our approach provides a technique to enable retrieving relevant code, which is suitable for real-world applications. For example, PKG can be applied on the code-base of a project or proprietary repositories, helping with the retrieval of similar code that is tailored towards a specific context. [1] Zora Zhiruo Wang, Akari Asai, Xinyan Velocity Yu, Frank F Xu, Yiqing Xie, Graham Neubig, and Daniel Fried. Coderag-bench: Can retrieval augment code generation? ## 2. Accordingly, in a not realistic setting, the experiments are not convincing enough to me. For example, on both HumanEval and MBPP, using BM25 for RAG consistently gets lower performance than no RAG. Also, with PKG the performance improves, it seems to be not a fair comparison: Thank you for your attention to this detail. The reason BM25 and Voyage perform lower than No-RAG on both HumanEval and MBPP is due to how the data is structured. In these experiments, the dataset is composed of question-answer pairs. When we apply BM25 or VoyageEmb without any post-processing, the retrieved content includes question-answer pairs. Including these in the model’s context introduces additional questions and answers, which can confuse the model and lead to hallucinated outputs. However, when we clean the dataset to contain only functions instead of full question-answer pairs (referred to as Func-BM25 in the tables 1 and 2), the performance improves. Func-BM25 outperforms No-RAG in HumanEval, indicating that retrieving only relevant function information is beneficial. ## 3. What's the advantage of PKG to normal RAG in an open retrieval setting? Thanks for your comment. We consider your concern as about retrieving from a dynamic data source such as google search. The advantage of PKG over normal RAG in an open retrieval setting is that it allows for more precise and relevant retrieval of information. In standard RAG, if we retrieve 100 answers from an open-domain source, we typically use a re-ranker to rank these answers and then select the top-n answers to add as an additional context for the model. In normal RAG we chunk data paragraph-wise or page-wise. Even though retriever tries to retrieve most similar chunks, but they might contain irrelevant data (e.g. irrelevant sentences). In contrast, PKG organizes these 100 answers into a structured Programming Knowledge Graph, where only the most relevant nodes (i.e., the specific, necessary information) are retrieved. This approach minimizes irrelevant data and focuses on retrieving content in a fine-grained manner, even if it's within larger irrelevant sections. For example, consider a small code block that closely resembles the query but is contained within a function unrelated to the query. In standard RAG, this code block would not be retrieved because the retrieval process compares the embeddings of the query against the embeddings of the entire function, rather than focusing on the individual code block. As a result, the function’s broader context may obscure the relevance of the specific code block to the query. This way, PKG provides a more targeted, context-rich augmentation that helps the model perform better. ## 4. Is max_token=512 in the experiments enough? Yes, as the benchmarks contains general python programming and as the canonical solutions are less than 512 tokens, max_new_tokens=512 is enough for experiments. The authors of CodeT5+[1] also evaluated their model with the same max length. [1] Yue Wang, Hung Le, Akhilesh Deepak Gotmare, Nghi DQ Bui, Junnan Li, and Steven CH Hoi. Codet5+: Open code large language models for code understanding and generation.

Thanks for your rebuttal (2)

Thanks for the revised paper. I still feel my concerns have not been adequately addressed. Take some of them listed below: **1. Thank you for recognizing that SO is actually not code-centric like GitHub repositories, which then fails to support your rebuttal to the first concern about the retrieval setting.** **Table 7 HumanEval** | Method | Tutorial | Docs | SO | GitHub | All | |-------------|----------|------|-------|--------|------| | **BM25** | 27.4 | 29.3 | 32.9 | 30.5 | 97.6 | | **GIST-large** | 34.8 | 26.7 | 32.3 | 32.9 | 69.1 | | **OpenAI** | 29.3 | 24.4 | 36.0 | 31.1 | 97.6 | **Table 8 ODEX** | Method | Tutorial | Docs | SO | GitHub | All | |-------------|----------|-------|-------|--------|------| | **BM25** | 13.4 | 14.1 | 11.6 | 15.9 | 16.2 | | **GIST-large** | 15.7 | 17.3 | 11.4 | 15.5 | 17.1 | | **OpenAI** | 14.1 | 15.9 | 10.9 | 16.9 | 15.3 | **Table 9 RepoEval** | Method | Tutorial | Docs | SO | GitHub | Open | L+O | |-------------|----------|-------|-------|--------|-------|------| | **BM25** | 25.2 | 23.9 | 23.6 | 25.5 | 23.6 | 31.4 | | **GIST-large** | 23.3 | 21.7 | 24.7 | 24.4 | 24.1 | 41.8 | | **OpenAI** | 24.1 | 24.1 | 23.1 | 22.8 | 24.9 | 50.9 | **2. The authors insisted in their two responses that SO (StackOverflow) yielded better scores than other resources. I'd like to refer to the original tables from Code-RAG-Bench above.** From the results, we can find that SO shows the worst performance on ODEX, the average performance on RepoEval, and the best performance on HumanEval. However, even on HumanEval, SO gets a lower score than Tutorial and GitHub. Obviously, it's not correct to claim that code-centric content is better than other resources. Furthermore, the authors misinterpreted my concerns as more context, which is not relevant to the retrieval setting concern. **3. The authors kept referring to other papers (e.g., Code-RAG-Bench) and haven't yet provided some additional experimental results.** For example, as the authors replied, "In our initial experiments, we tried augmenting the model with a knowledge graph that included Python documents and natural language explanations. However, this approach led to worse results than the No-RAG baseline." I explicitly expressed my concerns about the performance of other retrieval methods in my reply. Through the discussions, I have a better understanding of the paper, and the authors haven't addressed my concerns. Therefore, I will decrease my rating accordingly.

Exploring the Reliability, Errors, and Comparative Effectiveness of PKG-Based Contextual Retrieval in Code Generation.

# Reviewer concerns: ## 1. There still exists a gap between ideal ranker & re-ranker, what was the reason for it. Thank you for your thoughtful observation. The gap between the ideal ranker and the re-ranker arises because the ideal ranker has access to information that the re-ranker does not. The ideal ranker can always pick the correct solution from the candidate pool if one exists because it assumes that we know beforehand whether each solution is correct or incorrect. It acts as an upper bound for the re-ranker in our work. In contrast, the re-ranker operates under real-world conditions, where it does not have access to ground-truth labels. Instead, it must evaluate each solution based on available signals like Abstract Syntax Trees (ASTs), execution results, and embedding comparisons to select the best candidate. Without direct labels, the re-ranker’s choice is based on these approximations, which can lead to some incorrect selections and thus a performance gap compared to the ideal ranker. ## 2. How reliable is the result? Based on multiple generations with low temperatures? To ensure result reliability, all experiments were conducted with a temperature of 0, allowing for deterministic outputs. Additionally, all prompts used in these experiments are provided in the appendix (sections 7.1-7.4), enabling full reproducibility of the results. ## 3. Samples in MBPP & HumanEval present with misaligned or incorrect NL, under which error is that catered to? Thanks for your attention, can you clarify which examples are misaligned? We used the version of HumanEval samples available in the original Huggingface dataset. For MBPP, we use a filtered version curated by the CodeRAG-Bench authors[1], who removed samples with fewer test cases, as these were deemed less reliable. We are not sure which examples do you mean. If you provide more details, we will answer more specifically. [1] Zora Zhiruo Wang, Akari Asai, Xinyan Velocity Yu, Frank F Xu, Yiqing Xie, Graham Neubig, and Daniel Fried. Coderag-bench: Can retrieval augment code generation? ## 4. When block-PKG performs better than func-PKG then why is ranking done for considering func-PKG? Thank you for your insightful question. The re-ranking is applied to four approaches: NoRAG, BM25, Func-PKG, and Block-PKG. We include Func-PKG in the evaluation because, in certain cases, it outperforms Block-PKG. By incorporating Func-PKG, we aim to leverage the strengths of this approach where it shows an advantage. ## 5. Need to show cases of extraction with PKG providing better function than current RAG. To better illustrate how PKG outperforms current RAG approaches, we have already included results from standard RAG methods like BM25 and Voyage Embedding in our comparison. In the appendix (section 7.4), we provided examples showing the distinctions between PKG and NoRAG. We will further strengthen this section by adding more examples that compare PKG to BM25 and Voyage RAG, making PKG's advantages even clearer. ## Styling format suggestions: Thank you for your detailed feedback. We will address your comments and incorporate the necessary revisions in the updated version of the paper by November 27th.