From 8b37cfc306457b8ded6c26f19e7d72bec2378732 Mon Sep 17 00:00:00 2001
From: yetone <yetoneful@gmail.com>
Date: Wed, 12 Mar 2025 19:57:17 +0800
Subject: [PATCH] feat: system_prompt can be a function (#1571)

---
 lua/avante/config.lua | 2 +-
 lua/avante/llm.lua    | 8 ++++++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/lua/avante/config.lua b/lua/avante/config.lua
index 07c7f2c..609892a 100644
--- a/lua/avante/config.lua
+++ b/lua/avante/config.lua
@@ -17,7 +17,6 @@ local Utils = require("avante.utils")
 ---@class avante.CoreConfig: avante.Config
 local M = {}
 ---@class avante.Config
----@field custom_tools AvanteLLMToolPublic[]
 M._defaults = {
   debug = false,
   ---@alias ProviderName "claude" | "openai" | "azure" | "gemini" | "vertex" | "cohere" | "copilot" | "bedrock" | "ollama" | string
@@ -34,6 +33,7 @@ M._defaults = {
   -- For most providers that we support we will determine this automatically.
   -- If you wish to use a given implementation, then you can override it here.
   tokenizer = "tiktoken",
+  ---@type string | (fun(): string) | nil
   system_prompt = nil,
   rag_service = {
     enabled = false, -- Enables the rag service, requires OPENAI_API_KEY to be set
diff --git a/lua/avante/llm.lua b/lua/avante/llm.lua
index 48c5844..1c1e19f 100644
--- a/lua/avante/llm.lua
+++ b/lua/avante/llm.lua
@@ -128,8 +128,12 @@ function M.generate_prompts(opts)
 
   local system_prompt = Path.prompts.render_mode(mode, template_opts)
 
-  if Config.system_prompt ~= nil and Config.system_prompt ~= "" and Config.system_prompt ~= "null" then
-    system_prompt = system_prompt .. "\n\n" .. Config.system_prompt
+  if Config.system_prompt ~= nil then
+    local custom_system_prompt = Config.system_prompt
+    if type(custom_system_prompt) == "function" then custom_system_prompt = custom_system_prompt() end
+    if custom_system_prompt ~= nil and custom_system_prompt ~= "" and custom_system_prompt ~= "null" then
+      system_prompt = system_prompt .. "\n\n" .. custom_system_prompt
+    end
   end
 
   ---@type AvanteLLMMessage[]