From c93a10388bee97da0bd491d6391c6b8e1412159e Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Sat, 28 Sep 2024 19:51:28 +0200
Subject: [PATCH] refac

---
 backend/open_webui/apps/ollama/main.py |  1 +
 backend/open_webui/main.py             | 66 ++++++++++++++------------
 2 files changed, 37 insertions(+), 30 deletions(-)

diff --git a/backend/open_webui/apps/ollama/main.py b/backend/open_webui/apps/ollama/main.py
index 1337fbb31..2e02d069f 100644
--- a/backend/open_webui/apps/ollama/main.py
+++ b/backend/open_webui/apps/ollama/main.py
@@ -787,6 +787,7 @@ async def generate_chat_completion(
 ):
     payload = {**form_data.model_dump(exclude_none=True)}
     log.debug(f"{payload = }")
+
     if "metadata" in payload:
         del payload["metadata"]
 
diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py
index 40fac171f..5c964c597 100644
--- a/backend/open_webui/main.py
+++ b/backend/open_webui/main.py
@@ -448,38 +448,44 @@ async def chat_completion_tools_handler(
         if not content:
             return body, {}
 
-        content = content[content.find("{") : content.rfind("}") + 1]
-        result = json.loads(content)
-
-        tool_function_name = result.get("name", None)
-        if tool_function_name not in tools:
-            return body, {}
-
-        tool_function_params = result.get("parameters", {})
-
         try:
-            tool_output = await tools[tool_function_name]["callable"](
-                **tool_function_params
-            )
+            content = content[content.find("{") : content.rfind("}") + 1]
+            if not content:
+                raise Exception("No JSON object found in the response")
+
+            result = json.loads(content)
+
+            tool_function_name = result.get("name", None)
+            if tool_function_name not in tools:
+                return body, {}
+
+            tool_function_params = result.get("parameters", {})
+
+            try:
+                tool_output = await tools[tool_function_name]["callable"](
+                    **tool_function_params
+                )
+            except Exception as e:
+                tool_output = str(e)
+
+            if tools[tool_function_name]["citation"]:
+                citations.append(
+                    {
+                        "source": {
+                            "name": f"TOOL:{tools[tool_function_name]['toolkit_id']}/{tool_function_name}"
+                        },
+                        "document": [tool_output],
+                        "metadata": [{"source": tool_function_name}],
+                    }
+                )
+            if tools[tool_function_name]["file_handler"]:
+                skip_files = True
+
+            if isinstance(tool_output, str):
+                contexts.append(tool_output)
         except Exception as e:
-            tool_output = str(e)
-
-        if tools[tool_function_name]["citation"]:
-            citations.append(
-                {
-                    "source": {
-                        "name": f"TOOL:{tools[tool_function_name]['toolkit_id']}/{tool_function_name}"
-                    },
-                    "document": [tool_output],
-                    "metadata": [{"source": tool_function_name}],
-                }
-            )
-        if tools[tool_function_name]["file_handler"]:
-            skip_files = True
-
-        if isinstance(tool_output, str):
-            contexts.append(tool_output)
-
+            log.exception(f"Error: {e}")
+            content = None
     except Exception as e:
         log.exception(f"Error: {e}")
         content = None