LLM API 配置说明
此文档用于详细说明 AIKG 项目中涉及的 LLM API 配置。
1. 环境变量配置 (API Key)
我们通过环境变量来设置不同大语言模型(LLM)服务的 API Key 和服务地址(Endpoint)。这样可以保持敏感信息(如 API Key)的私密性,并方便在不同环境中切换。
支持的服务及对应的环境变量如下:
# VLLM (https://github.com/vllm-project/vllm)
export AIKG_VLLM_API_BASE=http://localhost:8000/v1
# Ollama (https://ollama.com/)
export AIKG_OLLAMA_API_BASE=http://localhost:11434
# 硅基流动 (https://www.siliconflow.cn/)
export AIKG_SILICONFLOW_API_KEY=sk-xxx
# DeepSeek (https://www.deepseek.com/)
export AIKG_DEEPSEEK_API_KEY=sk-xxx
# 火山引擎 (https://www.volcengine.com/)
export AIKG_HUOSHAN_API_KEY=0cbf8bxxxxxx
# Moonshot (https://www.moonshot.cn/)
export AIKG_MOONSHOT_API_KEY=sk-xxx
# 智谱大模型 (https://www.bigmodel.cn/)
export AIKG_ZHIPU_API_KEY=sk-xxx
2. LLM 模型配置 (llm_config.yaml)
此文件定义了所有可供 AIKG 使用的底层 LLM 模型。每个模型都有一个唯一的名称,并包含调用该模型所需的参数。
文件路径: aikg/python/ai_kernel_generator/core/llm/llm_config.yaml
功能:
- 注册模型: 将新的 LLM 模型接入 AIKG 框架。
- 参数预设: 为每个模型预设必要的调用参数。
通用参数:
api_base: API 基础 URLmodel: 模型名称max_tokens: 最大生成 token 数temperature: 温度参数,控制随机性top_p: 核采样参数,控制多样性frequency_penalty: 频率惩罚,控制重复presence_penalty: 存在惩罚,控制主题重复
使用方式:
在 llm_config.yaml 中添加一个新模型配置后,可以通过 create_model("my_model_name") 来直接调用该模型。
3. 任务编排方案配置 (xxx_custom_plan.yaml)
任务编排方案配置文件用于编排和组织一个完整的算子生成任务中,每个子任务(Agent)具体使用哪个在 llm_config.yaml 中定义的模型。
默认配置文件目录: aikg/python/ai_kernel_generator/config/
功能:
- 任务编排: 为
designer、coder、conductor等 Agent 指派不同的 LLM 模型。 - 灵活组合: 可以创建多个配置文件,以应对不同场景(如本地 vLLM 与云端 API 混合)。
- 默认方案: 按 DSL 提供默认方案,例如
default_triton_cuda_config.yaml或default_triton_ascend_config.yaml。
示例(coder-only,本地 vLLM):vllm_triton_ascend_coderonly_config.yaml(Ascend)或 vllm_triton_cuda_coderonly_config.yaml(CUDA),用于为 coder-only 流程配置统一的本地 vLLM 模型。
# 模型预设配置
agent_model_config:
designer: vllm_deepseek_v31_default
coder: vllm_deepseek_v31_default
conductor: vllm_deepseek_v31_default
api_generator: vllm_deepseek_v31_default
feature_extractor: vllm_deepseek_v31_default
# 日志配置
log_dir: "~/aikg_logs"
# Workflow 配置
workflow_config_path: "config/coder_only_workflow.yaml"
# 文档目录
docs_dir:
designer: "resources/docs/triton_ascend_docs" # Ascend 版本;CUDA 版本对应 triton_cuda_docs
coder: "resources/docs/triton_ascend_docs"
# 性能分析
profile_settings:
run_times: 50
warmup_times: 5
# 验证配置
verify_timeout: 600
使用方式:
在代码中,通过 load_config() 加载配置。
# 按 DSL 加载默认方案
config = load_config(dsl="triton_ascend", backend="ascend") # 或使用 "triton_cuda" 用于 CUDA 后端
# 或加载指定的配置文件
config = load_config(config_path="python/ai_kernel_generator/config/vllm_triton_ascend_coderonly_config.yaml")
# config = load_config(config_path="python/ai_kernel_generator/config/vllm_triton_cuda_coderonly_config.yaml")
# 在任务中使用该配置
task = Task(
# ...
config=config,
)
通过这种方式,您可以灵活地为不同的任务流程配置和切换底层的大语言模型。
4. 全局环境变量设置
您可以通过高优先级环境变量直接指定 LLM API。
export AIKG_BASE_URL="https://api.example.com/v1"
export AIKG_MODEL_NAME="your-model-name"
export AIKG_API_KEY="your-api-key"
export AIKG_MODEL_ENABLE_THINK="enabled" # 可选,启用thinking模式
其中 AIKG_MODEL_ENABLE_THINK 是可选参数,用于启用模型的 thinking 模式:
- DeepSeek 风格:设置为
"enabled"或"disabled" - GLM 风格:设置为
"true"或"false"
如果不设置此环境变量,则不启用 thinking 模式。
5. Embedding 模型环境变量设置
您可以通过高优先级环境变量直接指定 Embedding API,用于 RAG 检索等功能。
export AIKG_EMBEDDING_BASE_URL="https://api.siliconflow.cn/v1"
export AIKG_EMBEDDING_MODEL_NAME="Qwen/Qwen3-Embedding-8B"
export AIKG_EMBEDDING_API_KEY="your-api-key"
参数说明:
AIKG_EMBEDDING_BASE_URL:Embedding API 的基础 URLAIKG_EMBEDDING_MODEL_NAME:模型名称,支持两种格式:- 预设名:如
sflow_qwen3_embedding_8b(会自动从llm_config.yaml读取实际模型名) - 实际模型名:如
Qwen/Qwen3-Embedding-8B(直接使用)
- 预设名:如
AIKG_EMBEDDING_API_KEY:API 密钥
优先级:
- 环境变量(三者都设置时生效)
- 本地 HuggingFace 模型(
~/.aikg/text2vec-large-chinese)
使用示例:
# 使用硅流平台的 Embedding 模型
AIKG_EMBEDDING_BASE_URL="https://api.siliconflow.cn/v1" \
AIKG_EMBEDDING_MODEL_NAME="Qwen/Qwen3-Embedding-8B" \
AIKG_EMBEDDING_API_KEY="sk-xxx" \
python tools/run_single_adaptive_search.py config.yaml
可用的预设名(在 llm_config.yaml 中定义):
sflow_qwen3_embedding_8b:Qwen3-Embedding-8B(1024 维)sflow_qwen3_embedding_4b:Qwen3-Embedding-4B(768 维)sflow_qwen3_embedding_0.6b:Qwen3-Embedding-0.6B(512 维)