feat: add allow_llm_to_see_data flag for vanna

langgenius · Oct 10, 2024 · 4f382c2 · 4f382c2
1 parent 68107fe
commit 4f382c2
Show file tree

Hide file tree

Showing 4 changed files with 102 additions and 195 deletions.
diff --git a/api/core/tools/provider/builtin/vanna/tools/vanna.py b/api/core/tools/provider/builtin/vanna/tools/vanna.py
@@ -111,9 +111,10 @@ def _invoke(
         # with "visualize" set to True (default behavior) leads to remote code execution.
         # Affected versions: <= 0.5.5
         #########################################################################################
-        generate_chart = False
-        # generate_chart = tool_parameters.get("generate_chart", True)
-        res = vn.ask(prompt, False, True, generate_chart)
+        allow_llm_to_see_data = tool_parameters.get("allow_llm_to_see_data", False)
+        res = vn.ask(
+            prompt, print_results=False, auto_train=True, visualize=False, allow_llm_to_see_data=allow_llm_to_see_data
+        )
 
         result = []
 

diff --git a/api/core/tools/provider/builtin/vanna/tools/vanna.yaml b/api/core/tools/provider/builtin/vanna/tools/vanna.yaml
@@ -200,14 +200,14 @@ parameters:
       en_US: If enabled, it will attempt to train on the metadata of that database
       zh_Hans: 是否自动从数据库获取元数据来训练
     form: form
-  - name: generate_chart
+  - name: allow_llm_to_see_data
     type: boolean
     required: false
-    default: True
+    default: false
     label:
-      en_US: Generate Charts
-      zh_Hans: 生成图表
+      en_US: Whether to allow the LLM to see the data
+      zh_Hans: 是否允许LLM查看数据
     human_description:
-      en_US: Generate Charts
-      zh_Hans: 是否生成图表
+      en_US: Whether to allow the LLM to see the data
+      zh_Hans: 是否允许LLM查看数据
     form: form