TLDR: DeepCodeSeek is a novel multi-stage retrieval pipeline designed to improve real-time API retrieval for AI-assisted code generation in enterprise environments like ServiceNow. It uses a Knowledge Graph, enriched indexing with JSDoc, LLM-powered code expansion, and a compact, optimized 0.6B reranker. The system achieves 87.86% top-40 retrieval accuracy and its 0.6B reranker outperforms an 8B model with 2.5x reduced latency, making it highly effective for context-aware code completion.
In the rapidly evolving world of AI-assisted code generation, large language models (LLMs) are becoming indispensable tools for developers. However, their effectiveness, especially in specialized enterprise environments like ServiceNow, heavily relies on the quality of the context they receive. This is particularly true for custom APIs known as Script Includes, which are reusable JavaScript components central to encapsulating complex business logic.
A new research paper introduces DeepCodeSeek, a groundbreaking solution designed to enhance real-time API retrieval for context-aware code generation. The paper, titled “DeepCodeSeek: Real-Time API Retrieval for Context-Aware Code Generation,” addresses the critical challenge of accurately retrieving relevant Script Includes from partial developer code without explicit queries. Traditional search methods often fall short, failing to grasp nuanced developer intent or the complex hierarchical relationships within platforms like ServiceNow.
DeepCodeSeek proposes a novel multi-stage retrieval pipeline that significantly improves context relevance for LLMs. The core idea is to expand the code and index to predict necessary APIs, directly enabling high-quality, end-to-end code generation for auto-completion and AI agent applications. The researchers also tackled the issue of API leaks in existing code-to-code benchmark datasets by creating a new dataset from real-world ServiceNow Script Includes, which captures the difficulty of deciphering API usage intent from code.
Also Read:
- KG-R1: A Unified Agent for Efficient and Adaptable Knowledge Graph Reasoning
- Automating Expert Knowledge: How AI Generates Telecom Troubleshooting Data for LLMs
How DeepCodeSeek Works: A Multi-Stage Approach
The system is built on several key components:
- Knowledge Graph for Search Space Reduction: It uses a Knowledge Graph built from platform metadata to pre-filter and significantly narrow down the potential candidates before the main retrieval stage. This reduces the search space by approximately 59%.
- Enriched Indexing: Instead of indexing raw code, DeepCodeSeek creates a structured index. Methods within a Script Include are grouped under their parent namespace and enriched with metadata and structured JSDoc (JavaScript Documentation) summaries. These JSDoc summaries proved more effective than raw code for retrieval due to their structured and focused representation of API functionality.
- LLM-Powered Code Expansion: To overcome the lack of context in partial developer code, an LLM is used at runtime to generate more descriptive and effective queries. By analyzing the partial code, the LLM infers developer intent and produces a more complete code expansion, leading to more accurate retrieval results.
- Reranking: An initial retrieval stage might find the correct Script Include but not necessarily at the top of the list. DeepCodeSeek employs a reranking stage using a cross-encoder or LLM reranker to prioritize the most relevant candidates, moving them from lower relevance positions (e.g., top-40) to higher ones (e.g., top-5). This ensures the downstream code generation model receives a small, highly relevant set of options.
- Post-Training Optimization: A comprehensive post-training pipeline optimizes compact reranker models through synthetic dataset generation, supervised fine-tuning, and reinforcement learning. This allows smaller models to achieve performance comparable to much larger models while maintaining significantly reduced latency.
The evaluation metrics for DeepCodeSeek are impressive, showing an 87.86% top-40 retrieval accuracy. This performance is more than double that of the BM25 baseline, effectively providing the critical context and APIs needed for successful code generation. Furthermore, the optimized 0.6B reranker developed through this process outperforms a much larger 8B model, achieving 68.58% top-5 accuracy compared to the 8B model’s 66.10%, all while maintaining 2.5 times reduced latency. This efficiency is crucial for real-time predictions in production environments.
The research highlights that the “Hypothetical Code Generation” method consistently outperforms other approaches, demonstrating the power of using LLMs to generate richer, code-based queries. The paper also details how each component, from knowledge graph filtering to enhanced indexing and LLM reranking, contributes significantly to the overall accuracy and efficiency of the system.
DeepCodeSeek represents a significant advancement in real-time API retrieval for enterprise code completion. By combining advanced retrieval techniques with optimized compact reranker models, it offers a robust solution that addresses the nuances of enterprise-specific code without the computational overhead of larger models. This work was conducted by Esakkivel Esakkiraja, Denis Akhiyarov, Aditya Shanmugham, and Chitra Ganapathy from ServiceNow, Inc. You can read the full research paper here.
