From 48ae1fed78f42d6072f5badbc01f28600678ad39 Mon Sep 17 00:00:00 2001
From: yetone <yetoneful@gmail.com>
Date: Fri, 21 Feb 2025 13:04:02 +0800
Subject: [PATCH] feat: python llm tool (#1334)

---
 lua/avante/llm_tools.lua              | 53 +++++++++++++++++++++++++++
 lua/avante/templates/base.avanterules |  3 +-
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/lua/avante/llm_tools.lua b/lua/avante/llm_tools.lua
index a37b953..0511f8b 100644
--- a/lua/avante/llm_tools.lua
+++ b/lua/avante/llm_tools.lua
@@ -533,6 +533,27 @@ function M.git_commit(opts, on_log)
   return true, nil
 end
 
+---@param opts { code: string, rel_path: string }
+---@param on_log? fun(log: string): nil
+---@return string|nil result
+---@return string|nil error
+function M.python(opts, on_log)
+  local abs_path = get_abs_path(opts.rel_path)
+  if not has_permission_to_access(abs_path) then return nil, "No permission to access path: " .. abs_path end
+  if not Path:new(abs_path):exists() then return nil, "Path not found: " .. abs_path end
+  if on_log then on_log("cwd: " .. abs_path) end
+  if on_log then on_log("code: " .. opts.code) end
+  ---change cwd to abs_path
+  local old_cwd = vim.fn.getcwd()
+  vim.fn.chdir(abs_path)
+  local output = vim.fn.system({ "python", "-c", opts.code })
+  local exit_code = vim.v.shell_error
+  vim.fn.chdir(old_cwd)
+  if exit_code ~= 0 then return nil, "Error: " .. output end
+  Utils.debug("output", output)
+  return output, nil
+end
+
 ---@class AvanteLLMTool
 ---@field name string
 ---@field description string
@@ -558,6 +579,38 @@ end
 
 ---@type AvanteLLMTool[]
 M.tools = {
+  {
+    name = "python",
+    description = "Run python code",
+    param = {
+      type = "table",
+      fields = {
+        {
+          name = "code",
+          description = "Python code to run",
+          type = "string",
+        },
+        {
+          name = "rel_path",
+          description = "Relative path to the directory, as cwd",
+          type = "string",
+        },
+      },
+    },
+    returns = {
+      {
+        name = "result",
+        description = "Python output",
+        type = "string",
+      },
+      {
+        name = "error",
+        description = "Error message if the python code failed",
+        type = "string",
+        optional = true,
+      },
+    },
+  },
   {
     name = "git_diff",
     description = "Get git diff for generating commit message",
diff --git a/lua/avante/templates/base.avanterules b/lua/avante/templates/base.avanterules
index c96886e..e67d85c 100644
--- a/lua/avante/templates/base.avanterules
+++ b/lua/avante/templates/base.avanterules
@@ -20,7 +20,8 @@ Tools Usage Guide:
   - When attempting to modify a file that is not in the context, please first use the `list_files` tool and `search_files` tool to check if the file you want to modify exists, then use the `read_file` tool to read the file content. Don't modify blindly!
   - When generating files, first use `list_files` tool to read the directory structure, don't generate blindly!
   - When creating files, first check if the directory exists. If it doesn't exist, create the directory before creating the file.
-  - After `web_search`, if you don't get detailed enough information, do not continue use `web_search`, just continue using the `fetch` tool to get more information you need from the links in the search results.
+  - After `web_search` tool returns, if you don't get detailed enough information, do not continue use `web_search` tool, just continue using the `fetch` tool to get more information you need from the links in the search results.
+  - For any mathematical calculation problems, please prioritize using the `python` tool to solve them. Please try to avoid mathematical symbols in the return value of the `python` tool for mathematical problems and directly output human-readable results, because large models don't understand mathematical symbols, they only understand human natural language.
 
 {% if system_info -%}
 Use the appropriate shell based on the user's system info: