Bridging Code and Context: A Knowledge Graph-Based Repository-Level Code Generation
Presented at ICSE 2025(LLM4Code 2025), Ottawa, Canada
In the evolving landscape of software engineering, large language models (LLMs) like GPT-4 and Claude-3.5 Sonnet have emerged as revolutionary tools, capable of translating natural language queries into functional, high-quality code. These models significantly boost developer productivity by automating routine coding tasks. However, despite their impressive capabilities, LLMs still face substantial challenges, particularly when generating code within complex, real-world software repositories.
The Challenges
Today’s code generation tools often overlook critical repository-specific context and produce code in isolation, lacking awareness of the existing codebase’s architectural patterns, dependencies, and style conventions. As a result, developers must spend significant time adjusting and integrating generated code, reducing the overall productivity gains. Common issues include redundant code creation, inconsistent coding styles, and difficulties in accurately retrieving relevant context from large code repositories. These errors from highly capable models are due to lack of accurate contextual information or challenges in the retrieval of relevant information.
To address these challenges, we introduce an innovative knowledge graph-based code generation framework designed specifically to enhance the quality, contextual accuracy, and relevance of code generation at the repository level. Our approach integrates the power of knowledge graphs with hybrid code retrieval and LLM-based generation.
How the Framework Works
Figure 1: A diagrammatic representation of our three-step approach to code generation using knowledge graphs: (1) Knowledge Graph Creation: Transform the code repository into a knowledge graph, capturing connections among elements like classes, functions, etc. (2) Graph-based Code Retrieval: Create an index (knowledge base) for retrieving relevant sub-graph based on user queries, utilizes a hybrid search system (3) Code Generation with LLM: Employ a Large Language Model (LLM) to generate code, based on the retrieved sub-graph provided as context to the LLM.
1. Knowledge Graph Creation
Figure 2: Knowledge Graph Creation Process via Parsing and Component-Schema Alignment
To transform an entire code repository into a structured knowledge graph, we begin by parsing the code files using Abstract Syntax Trees (ASTs), to identify essential code components like classes, methods, functions, and attributes. These elements are structured into a graph schema (as shown in Figure 3), with nodes representing distinct entities (e.g., Files, Class, Function) and relationships capturing the interactions between these nodes (e.g., defines class, defines function, used in, has definition).
Figure 3: A sample schema that defines the node types and relations between them. Every node in the knowledge graph will be one of these types, linked by a relevant relation to another node
We further enrich this graph schema with additional metadata from documentation and comments from code files, stored as additional nodes in our knowledge graph. Additionally, we use an LLM to generate descriptions for code snippets to capture functional meaning and context from the code in the repository, and we store these as LLM-generated descriptions in the knowledge graph.
Figure 4: Index Creation (knowledge base) and Embeddings Generation
The structured data is then ingested into a Neo4j graph database. To optimize search operations, we create full-text and vector indexes and utilize an encoder model to generate embeddings for the metadata. These embeddings are then ingested into Neo4j’s vector database to create the search index.
This setup supports hybrid retrieval, enabling both structural and semantic searches on the documentation and descriptions stored within the knowledge graph.
2. Graph-Based Hybrid Code Retrieval
Figure 5: Retrieving top-k nodes via full-text and semantic search on user query
Once the knowledge graph is constructed, we adopt a hybrid retrieval approach that combines syntactic, semantic, and graph-based querying. When a user submits a natural language query, the system employs large language models (LLMs) to extract schema-specific entities and generate query embeddings. Once we have the extracted entities and the query embedding, we perform full-text and semantic searches over the knowledge graph to identify top-k relevant nodes. These nodes are then expanded into an n-hop subgraph through graph traversal, capturing the contextual relationships needed to interpret the user query accurately.
Figure 6: Node expansion to n-hop subgraph and semantic reranking to get filtered subgraph
To enhance precision, we apply semantic re-ranking to this subgraph, prioritizing nodes that best match the intent of the user’s query and filtering out less relevant elements. This pruning step is crucial for maintaining computational efficiency and ensuring that the refined subgraph passed to the LLM contains only the most contextually relevant information. As a result, the system is able to generate more accurate and context-aware code responses.
3. Contextual Code Generation with LLMs
In the final stage of code generation, the refined sub-graph is provided as context to the LLM, guiding the code generation process. We write a custom prompt to guide the LLM, which includes instructions to utilize the relationships, dependencies, and contextual information embedded within the sub-graph. The prompt directs the LLM to focus on producing code aligned with the user’s query and grounded with relevance to the existing code base.
By grounding the code generation in a semantically enriched sub-graph, with the detailed repository-specific relationships and structural context, our system ensures that the produced code aligns closely with the coding conventions, dependencies, and functionality of the base repository, significantly reducing the need for manual adjustments post-generation.
Dataset, Benchmarking, and Results
To evaluate our knowledge graph-based code generation methodology, we utilized the EvoCodeBench dataset—a benchmark tailored for realistic, repository-level code generation tasks. EvoCodeBench consists of 275 samples derived from 25 open-source repositories. Each sample task aims to generate the function body using context such as its namespace, location in the repo, and expected behavior. We only evaluate our approach for code written in the Python programming language.
| Category | Method | LLM Model | pass@1 Score |
| Our Approach | Graph-based Retrieval | GPT-4 | 32.00% |
| GPT-4o | 33.45% | ||
| Claude 3.5 Sonnet | 36.36% | ||
| Baseline (EvoCodeBench) | Local File (Infilling) | GPT-4 | 20.73% |
| GPT-3.5 | 17.82% | ||
| Local File (Completion) | GPT-4 | 17.45% | |
| GPT-3.5 | 15.64% | ||
| Without Context | GPT-4 | 7.27% | |
| GPT-3.5 | 6.55% | ||
| CodeXGraph | Graph-based (212/275 Samples) | Qwen2 | 19.34% |
| DS-Coder | 25.47% | ||
| GPT-4o | 36.02% |
Table 1: Comparative Analysis of pass@1 scores achieved by various Code Generation Methodologies on the Evocodebench Dataset
*Note: “The evaluation results above are based on experiments conducted during Nov 2024. Since then, updated releases have been made, likely bringing further improvements.”
To measure performance, we used the pass@1 metric, evaluating whether the LLM-generated function passed predefined test cases on the first try. The results were promising. Using Claude 3.5 Sonnet, our approach achieved a pass@1 score of 36.36%, narrowly outperforming CodeXGraph (36.02%), a strong baseline in the field. Importantly, while CodeXGraph evaluated only 212 samples due to environment setup issues, we addressed those limitations and ran our evaluation on all 275 samples for a more comprehensive analysis.
Compared to traditional methods, many of which ignore the surrounding context, our approach demonstrated significant improvements. For instance, OpenAI’s GPT-4 and GPT-4o, when paired with our graph-based retrieval system, achieved scores above 32%, outperforming baseline techniques that ranged from 7% to 20%.
Advantages of Our Knowledge Graph Framework
- Contextual Accuracy: Significantly improves the retrieval of relevant, context-aware code snippets.
- Consistency: Ensures code generation aligns with repository-specific coding styles and architectural patterns.
- Reduced Manual Effort: Minimizes the need for manual adjustments, thus improving developer efficiency and reducing code redundancy.
- Scalability: Capable of supporting complex queries and codebases through a robust hybrid retrieval mechanism.
Future Directions and Potential Enhancements
- Efficiency and Scalability Improvements: Optimizing retrieval algorithms to balance depth and computational overhead.
- Broader Schema Coverage: Expanding the knowledge graph schema to include additional code elements, such as decorators, variable types, and auxiliary files, to enhance contextual relevance.
- Multi-Language Support: Extending the framework to support other programming languages to broaden applicability.
- Agent-Based Refinement: Integrating iterative agent-driven refinement workflows to enhance code accuracy, efficiency, and autonomy in generation tasks.
Conclusion
Our knowledge graph-based methodology demonstrates significant promise in advancing the state-of-the-art for repository-specific code generation. By combining structural insights with advanced LLM capabilities, our framework generates more accurate, context-sensitive code, providing a strong foundation for future AI-driven software development tools.
As software development continues to evolve, integrating knowledge graphs with LLMs presents an exciting frontier, promising substantial improvements in developer productivity, code quality, and consistency across complex codebases.
