Merge pull request #1625 from arc53/handle-bad-tool-names

Handle bad tool names

application/api/user/routes.py

@@ -19,7 +19,7 @@ from application.core.settings import settings
 from application.extensions import api
 from application.tools.tool_manager import ToolManager
 from application.tts.google_tts import GoogleTTS
-from application.utils import check_required_fields
+from application.utils import check_required_fields, validate_function_name
 from application.vectorstore.vector_creator import VectorCreator
 
 mongo = MongoDB.get_client()
@@ -1932,6 +1932,16 @@ class UpdateTool(Resource):
             if "actions" in data:
                 update_data["actions"] = data["actions"]
             if "config" in data:
+                if "actions" in data["config"]:
+                    for action_name in list(data["config"]["actions"].keys()):
+                        if not validate_function_name(action_name):
+                            return make_response(
+                                jsonify({
+                                    "success": False,
+                                    "message": f"Invalid function name '{action_name}'. Function names must match pattern '^[a-zA-Z0-9_-]+$'.",
+                                    "param": "tools[].function.name"
+                                }), 400
+                            )
                 update_data["config"] = data["config"]
             if "status" in data:
                 update_data["status"] = data["status"]

application/tools/agent.py

@@ -52,6 +52,10 @@ class Agent:
                 },
             }
             for tool_id, tool in tools_dict.items()
+            if (
+                (tool["name"] == "api_tool" and "actions" in tool.get("config", {}))
+                or (tool["name"] != "api_tool" and "actions" in tool)
+            )
             for action in (
                 tool["config"]["actions"].values()
                 if tool["name"] == "api_tool"

application/utils.py

@@ -1,6 +1,7 @@
 import tiktoken
 import hashlib
 from flask import jsonify, make_response
+import re
 
 
 _encoding = None
@@ -95,3 +96,9 @@ def limit_chat_history(history, max_token_limit=None, gpt_model="docsgpt"):
                 break
 
     return trimmed_history
+
+def validate_function_name(function_name):
+    """Validates if a function name matches the allowed pattern."""
+    if not re.match(r"^[a-zA-Z0-9_-]+$", function_name):
+        return False
+    return True

frontend/src/modals/AddActionModal.tsx

@@ -1,21 +1,40 @@
-import React from 'react';
+import React, { useState } from 'react';
 import { useTranslation } from 'react-i18next';
 
 import Exit from '../assets/exit.svg';
 import Input from '../components/Input';
 import { ActiveState } from '../models/misc';
 
+const isValidFunctionName = (name: string): boolean => {
+  const pattern = /^[a-zA-Z0-9_-]+$/;
+  return pattern.test(name);
+};
+
+interface AddActionModalProps {
+  modalState: ActiveState;
+  setModalState: (state: ActiveState) => void;
+  handleSubmit: (actionName: string) => void;
+}
+
 export default function AddActionModal({
   modalState,
   setModalState,
   handleSubmit,
-}: {
-  modalState: ActiveState;
-  setModalState: (state: ActiveState) => void;
-  handleSubmit: (actionName: string) => void;
-}) {
+}: AddActionModalProps) {
   const { t } = useTranslation();
   const [actionName, setActionName] = React.useState('');
+  const [functionNameError, setFunctionNameError] = useState<boolean>(false); // New error state
+
+  const handleAddAction = () => {
+    if (!isValidFunctionName(actionName)) {
+      setFunctionNameError(true); // Set error state if invalid
+      return;
+    }
+    setFunctionNameError(false); // Clear error state if valid
+    handleSubmit(actionName);
+    setModalState('INACTIVE');
+  };
+
   return (
     <div
       className={`${
@@ -46,14 +65,21 @@ export default function AddActionModal({
                 onChange={(e) => setActionName(e.target.value)}
                 borderVariant="thin"
                 placeholder={'Enter name'}
-              ></Input>
+              />
+              <p className="mt-1 text-gray-500 text-xs">
+                Use only letters, numbers, underscores, and hyphens (e.g.,
+                `get_user_data`, `send-report`).
+              </p>
+              {functionNameError && (
+                <p className="mt-1 text-red-500 text-xs">
+                  Invalid function name format. Use only letters, numbers,
+                  underscores, and hyphens.
+                </p>
+              )}
             </div>
             <div className="mt-8 flex flex-row-reverse gap-1 px-3">
               <button
-                onClick={() => {
-                  handleSubmit(actionName);
-                  setModalState('INACTIVE');
-                }}
+                onClick={handleAddAction}
                 className="rounded-3xl bg-purple-30 px-5 py-2 text-sm text-white transition-all hover:bg-[#6F3FD1]"
               >
                 Add

frontend/src/modals/AddToolModal.tsx

@@ -13,11 +13,13 @@ export default function AddToolModal({
   modalState,
   setModalState,
   getUserTools,
+  onToolAdded,
 }: {
   message: string;
   modalState: ActiveState;
   setModalState: (state: ActiveState) => void;
   getUserTools: () => void;
+  onToolAdded: (toolId: string) => void;
 }) {
   const [availableTools, setAvailableTools] = React.useState<
     AvailableToolType[]
@@ -59,9 +61,20 @@ export default function AddToolModal({
         })
         .then((res) => {
           if (res.status === 200) {
-            getUserTools();
-            setModalState('INACTIVE');
+            return res.json();
+          } else {
+            throw new Error(
+              `Failed to create tool, status code: ${res.status}`,
+            );
           }
+        })
+        .then((data) => {
+          getUserTools();
+          setModalState('INACTIVE');
+          onToolAdded(data.id);
+        })
+        .catch((error) => {
+          console.error('Failed to create tool:', error);
         });
     } else {
       setModalState('INACTIVE');

frontend/src/modals/ChunkModal.tsx

@@ -184,7 +184,13 @@ export default function ChunkModal({
           message="Are you sure you want to delete this chunk?"
           modalState={deleteModal}
           setModalState={setDeleteModal}
-          handleSubmit={handleDelete ? handleDelete : () => {}}
+          handleSubmit={
+            handleDelete
+              ? handleDelete
+              : () => {
+                  /* no-op */
+                }
+          }
           submitLabel="Delete"
         />
       </div>

frontend/src/settings/Tools.tsx

@@ -58,9 +58,27 @@ export default function Tools() {
     getUserTools();
   };
 
+  const handleToolAdded = (toolId: string) => {
+    userService
+      .getUserTools()
+      .then((res) => res.json())
+      .then((data) => {
+        const newTool = data.tools.find(
+          (tool: UserToolType) => tool.id === toolId,
+        );
+        if (newTool) {
+          setSelectedTool(newTool);
+        } else {
+          console.error('Newly added tool not found');
+        }
+      })
+      .catch((error) => console.error('Error fetching tools:', error));
+  };
+
   React.useEffect(() => {
     getUserTools();
   }, []);
+
   return (
     <div>
       {selectedTool ? (
@@ -185,6 +203,7 @@ export default function Tools() {
             modalState={addToolModalState}
             setModalState={setAddToolModalState}
             getUserTools={getUserTools}
+            onToolAdded={handleToolAdded}
           />
         </div>
       )}