TLDR: SpareCodeSearch is a novel approach that uses keyword-based search to efficiently retrieve relevant code context for AI code completion, especially in resource-constrained environments like Integrated Development Environments (IDEs). Unlike computationally intensive semantic search methods, SpareCodeSearch leverages tools like Zoekt for fast indexing and retrieval. It demonstrated strong performance in the Code Context Competition, achieving top ranks in Kotlin and Python tracks, proving that a lightweight, keyword-driven strategy can effectively enhance Code Language Models without requiring substantial GPU resources.
Retrieval-Augmented Generation (RAG) frameworks are becoming increasingly popular for enhancing Code Language Models (CLMs). These frameworks work by adding a module that retrieves relevant context, which is then used to construct more informed input prompts for the CLMs. This approach significantly improves the quality of generated code.
However, a common challenge with many existing RAG solutions is their reliance on semantic search. Semantic search, while powerful, demands substantial computational resources, particularly GPUs, for training and hosting the embedded models. This makes them impractical for integration into lightweight applications, such as AI-based code completion features within Integrated Development Environments (IDEs), where resources are often constrained.
A new solution, called SpareCodeSearch, addresses this very issue. Developed with a philosophy centered on simplicity and efficiency, SpareCodeSearch demonstrates that keyword-based search is not only sufficient but also highly effective for retrieving relevant and useful code context from large codebases, all without the need for extensive GPU resources. This approach is particularly well-suited for automated code completion tasks, emphasizing lightweight deployment, fast retrieval times, easy integration with existing CLMs, and extensibility across multiple programming languages.
The core of SpareCodeSearch’s effectiveness lies in its use of Zoekt, an open-source, fast, and scalable code search engine. Zoekt is designed to handle vast codebases efficiently and has been adopted by major platforms like Sourcegraph and GitLab for their indexed code search functionalities. Its architecture, featuring separate index and web servers, makes it ideal for microservices deployment, allowing for seamless integration and customization with other systems.
The system operates in two main phases: Offline Indexing and Online Retrieval.
Offline Indexing with Zoekt
In the offline phase, Zoekt builds an index from code datasets. For instance, in the Code Context Competition, public Kotlin and Python datasets were indexed. The Zoekt indexer processes each data point, looking at its revisions, and creates an index shard – a compressed, keyword-search-optimized representation of the codebase. This process utilizes tools like ctags to extract and save symbols. The indexing is remarkably efficient; for example, indexing 400 Kotlin revisions took only about 18 minutes on a Macbook M3 Air with modest resources (1 CPU, 8GB RAM).
Online Context Retrieval
Once indexing is complete and the Zoekt web server is running, the online retrieval phase begins. This phase involves two key modules: the Zoekt Query Generator and the Post-processor of search results.
The Query Generator is designed to create a comprehensive search space. For every code completion point, it manufactures up to 19 different Zoekt query candidates. These queries are constructed using keyword terms extracted from the ‘diff string’ associated with the completion point – essentially, the changes made in a code revision. Tree-sitter is used to extract meaningful identifiers, ranging from function and class names to broader navigation expressions and all identifiers within the diff string.
Queries are formulated using various strategies from the official Zoekt query language, including exact matching, reduced term queries, regex for fuzzy searches, and OR logic to expand the search boundary. Importantly, not all 19 query variations are sent simultaneously. The Query Generator iterates through them, sending queries until a non-empty search result is found, optimizing retrieval time. The system also incorporates fallback, timeout, and retry mechanisms to ensure robust operation.
A significant feature is ‘Cross-shard searching’, which allows queries to span multiple code revisions within the same repository. This capability proved to be highly effective, yielding significantly higher ‘hit rates’ (meaning at least one successful search result) compared to single-shard queries. This mirrors how developers often search across different branches and commits to find relevant examples.
The Post-processor module then takes the numerous results returned by Zoekt, which include metadata like line numbers and file paths. It ranks these results using Zoekt’s relevance scores and fetches the actual code snippets. To manage the context window for CLMs, it uses tokenizers (like Jetbrains Mellum’s) to determine if a file is too large. If so, only relevant snippets or their merged versions are returned. It dynamically adjusts the context window based on token constraints, ensuring that the selected top-k ranked files or snippets fit within specified budgets for both individual files and the overall context.
Impressive Findings and Performance
SpareCodeSearch demonstrated strong performance in the Code Context Competition. In the Public Phase, it ranked first for Kotlin with a chRF score of 0.7125 and second for Python with 0.6152. These optimistic results are largely attributed to the high hit rate achieved by Cross-shard queries (97.5% for Kotlin and 96.7% for Python), which ensured more relevant contexts were retrieved.
The Private Phase results showed a similar trend, with SpareCodeSearch again outperforming competitors in the Kotlin track (0.748 chRF) and securing second place in Python (0.725 chRF), only slightly behind a semantic search solution. Crucially, the entire solution proved to be lightweight and efficient, running on a standard laptop with minimal resources, making it highly suitable for integration into existing in-IDE CLM-based code completion systems.
Also Read:
- Contextual Information Improves AI’s Code Understanding
- MM-ASSIST: An AI-Powered Tool for Streamlined Code Refactoring
Future Directions
While SpareCodeSearch has shown great promise, the authors acknowledge areas for future work, including further investigating the correlation between hit percentage and code completion quality, and exploring automated query optimization techniques. The project is open-sourced on GitHub, inviting community contributions to further improve and extend its capabilities. You can find more details about the research paper here: SpareCodeSearch: Searching for Code Context When You Have No Spare GPU.
