Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AstraAI: LLMs, Retrieval, and AST-Guided Assistance for HPC Codebases

Mahesh Natarajan, Xiaoye Li, Weiqun Zhang

Abstract

We present AstraAI, a command-line interface (CLI) coding framework for high-performance computing (HPC) software development. AstraAI operates directly within a Linux terminal and integrates large language models (LLMs) with Retrieval-Augmented Generation (RAG) and Abstract Syntax Tree (AST)-based structural analysis to enable context-aware code generation for complex scientific codebases. The central idea is to construct a high-fidelity prompt that is passed to the LLM for inference. This prompt augments the user request with relevant code snippets retrieved from the underlying framework codebase via RAG and structural context extracted from AST analysis, providing the model with precise information about relevant functions, data structures, and overall code organization. The framework is designed to perform scoped modifications to source code while preserving structural consistency with the surrounding code. AstraAI supports both locally hosted models from Hugging Face and API-based frontier models accessible via the American Science Cloud, enabling flexible deployment across HPC environments. The system generates code that aligns with existing project structures and programming patterns. We demonstrate AstraAI on representative HPC code generation tasks within AMReX, a DOE-supported HPC software infrastructure for exascale applications.

AstraAI: LLMs, Retrieval, and AST-Guided Assistance for HPC Codebases

Abstract

We present AstraAI, a command-line interface (CLI) coding framework for high-performance computing (HPC) software development. AstraAI operates directly within a Linux terminal and integrates large language models (LLMs) with Retrieval-Augmented Generation (RAG) and Abstract Syntax Tree (AST)-based structural analysis to enable context-aware code generation for complex scientific codebases. The central idea is to construct a high-fidelity prompt that is passed to the LLM for inference. This prompt augments the user request with relevant code snippets retrieved from the underlying framework codebase via RAG and structural context extracted from AST analysis, providing the model with precise information about relevant functions, data structures, and overall code organization. The framework is designed to perform scoped modifications to source code while preserving structural consistency with the surrounding code. AstraAI supports both locally hosted models from Hugging Face and API-based frontier models accessible via the American Science Cloud, enabling flexible deployment across HPC environments. The system generates code that aligns with existing project structures and programming patterns. We demonstrate AstraAI on representative HPC code generation tasks within AMReX, a DOE-supported HPC software infrastructure for exascale applications.

Paper Structure

This paper contains 13 sections, 5 figures, 2 tables.

Table of Contents

  1. Introduction
  2. AstraAI architecture
  3. Implementation
  4. Pre-inference steps
  5. Model access and deployment
  6. Creation of code chunks for RAG
  7. Inference pipeline
  8. Retrieval of relevant code chunks for RAG
  9. Structure extraction using AST analysis
  10. Prompt construction
  11. Model inference
  12. Evaluation
  13. Conclusions and future work

Figures (5)

  • Figure 1: End-to-end AstraAI workflow: User prompts are interpreted to extract intent, and intent-specific pipelines combine retrieval-augmented context, AST-guided constraints, and controlled code generation or editing. Tools used at each stage are highlighted in red, and the workflow is orchestrated via Python scripts.
  • Figure 2: General structure of the chunks for retrieval (top), an example code snippet in AMReX with the metadata (middle), and the code chunk in .json format (bottom).
  • Figure 3: clang-query command examples to extract structure information of the code.
  • Figure 4: General structure of the high-fidelity prompt.
  • Figure 5: An example user prompt that attempts to port a nested for loop in C++ into AMReX (top), the full prompt to the LLM (middle), and the result of the inference (bottom).