From c9e9ce931bb5af4c3945cd4b1841c762dcc3905e Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Sat, 5 Apr 2025 04:49:07 -0600
Subject: [PATCH] refac

---
 backend/open_webui/utils/middleware.py | 10 ++++++++--
 backend/open_webui/utils/tools.py      | 21 +++++++++++++++++++--
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index 72f1f30ce..93c8a5644 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -227,7 +227,9 @@ async def chat_completion_tools_handler(
                 if isinstance(tool_result, str):
                     tool = tools[tool_function_name]
                     tool_id = tool.get("toolkit_id", "")
-                    if tool.get("citation", False) or tool.get("direct", False):
+                    if tool.get("metadata", {}).get("citation", False) or tool.get(
+                        "direct", False
+                    ):
 
                         sources.append(
                             {
@@ -267,7 +269,11 @@ async def chat_completion_tools_handler(
                             }
                         )
 
-                    if tools[tool_function_name].get("file_handler", False):
+                    if (
+                        tools[tool_function_name]
+                        .get("metadata", {})
+                        .get("file_handler", False)
+                    ):
                         skip_files = True
 
             # check if "tool_calls" in result
diff --git a/backend/open_webui/utils/tools.py b/backend/open_webui/utils/tools.py
index 84def7060..842b9a1b5 100644
--- a/backend/open_webui/utils/tools.py
+++ b/backend/open_webui/utils/tools.py
@@ -47,6 +47,20 @@ def get_tools(
     for tool_id in tool_ids:
         tools = Tools.get_tool_by_id(tool_id)
         if tools is None:
+
+            tool_dict = {
+                "spec": spec,
+                "callable": callable,
+                "toolkit_id": tool_id,
+                "pydantic_model": function_to_pydantic_model(callable),
+                # Misc info
+                "metadata": {
+                    "file_handler": hasattr(module, "file_handler")
+                    and module.file_handler,
+                    "citation": hasattr(module, "citation") and module.citation,
+                },
+            }
+
             continue
 
         module = request.app.state.TOOLS.get(tool_id, None)
@@ -97,8 +111,11 @@ def get_tools(
                 "toolkit_id": tool_id,
                 "pydantic_model": function_to_pydantic_model(callable),
                 # Misc info
-                "file_handler": hasattr(module, "file_handler") and module.file_handler,
-                "citation": hasattr(module, "citation") and module.citation,
+                "metadata": {
+                    "file_handler": hasattr(module, "file_handler")
+                    and module.file_handler,
+                    "citation": hasattr(module, "citation") and module.citation,
+                },
             }
 
             # TODO: if collision, prepend toolkit name