AstraAI：面向HPC代码库的LLM、检索与AST引导辅助

我们提出 AstraAI，一个面向高性能计算（HPC）软件开发的命令行界面（CLI）编码框架。AstraAI 直接在 Linux 终端中运行，将大语言模型（LLM）与检索增强生成（RAG）以及基于抽象语法树（AST）的结构分析相结合，为复杂的科学代码库实现上下文感知的代码生成。其核心思想是构建一个高保真提示并传递给 LLM 进行推理。该提示用通过 RAG 从底层框架代码库检索到的相关代码片段，以及从 AST 分析中提取的结构上下文来增强用户请求，为模型提供关于相关函数、数据结构与整体代码组织的精确信息。该框架旨在对源代码执行有范围限定的修改，同时保持与周边代码的结构一致性。AstraAI 同时支持 Hugging Face 上的本地托管模型与可通过 American Science Cloud 访问的基于 API 的前沿模型，从而可在各类 HPC 环境中灵活部署。该系统生成的代码与现有项目结构和编程模式保持一致。我们在 AMReX——一个由美国能源部（DOE）支持的面向百亿亿次（exascale）应用的 HPC 软件基础设施——内的代表性 HPC 代码生成任务上演示了 AstraAI。