Ver código fonte

fix: return types of builtin tools

Yeuoly 7 meses atrás
pai
commit
8cc4c109d0

+ 3 - 2
api/core/tools/builtin_tool/providers/code/tools/simple_code.py

@@ -1,3 +1,4 @@
+from collections.abc import Generator
 from typing import Any, Optional
 
 from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
@@ -13,7 +14,7 @@ class SimpleCode(BuiltinTool):
         conversation_id: Optional[str] = None,
         app_id: Optional[str] = None,
         message_id: Optional[str] = None,
-    ) -> ToolInvokeMessage | list[ToolInvokeMessage]:
+    ) -> Generator[ToolInvokeMessage, None, None]:
         """
         invoke simple code
         """
@@ -26,4 +27,4 @@ class SimpleCode(BuiltinTool):
 
         result = CodeExecutor.execute_code(language, "", code)
 
-        return self.create_text_message(result)
+        yield self.create_text_message(result)

+ 9 - 6
api/core/tools/builtin_tool/providers/time/tools/current_time.py

@@ -1,5 +1,6 @@
+from collections.abc import Generator
 from datetime import UTC, datetime
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from pytz import timezone as pytz_timezone
 
@@ -15,7 +16,7 @@ class CurrentTimeTool(BuiltinTool):
         conversation_id: Optional[str] = None,
         app_id: Optional[str] = None,
         message_id: Optional[str] = None,
-    ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
+    ) -> Generator[ToolInvokeMessage, None, None]:
         """
         invoke tools
         """
@@ -23,10 +24,12 @@ class CurrentTimeTool(BuiltinTool):
         tz = tool_parameters.get("timezone", "UTC")
         fm = tool_parameters.get("format") or "%Y-%m-%d %H:%M:%S %Z"
         if tz == "UTC":
-            return self.create_text_message(f"{datetime.now(UTC).strftime(fm)}")
+            yield self.create_text_message(f"{datetime.now(UTC).strftime(fm)}")
+            return
 
         try:
             tz = pytz_timezone(tz)
-        except:
-            return self.create_text_message(f"Invalid timezone: {tz}")
-        return self.create_text_message(f"{datetime.now(tz).strftime(fm)}")
+        except Exception:
+            yield self.create_text_message(f"Invalid timezone: {tz}")
+            return
+        yield self.create_text_message(f"{datetime.now(tz).strftime(fm)}")

+ 6 - 4
api/core/tools/builtin_tool/providers/time/tools/localtime_to_timestamp.py

@@ -1,5 +1,6 @@
+from collections.abc import Generator
 from datetime import datetime
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import pytz
 
@@ -16,7 +17,7 @@ class LocaltimeToTimestampTool(BuiltinTool):
         conversation_id: Optional[str] = None,
         app_id: Optional[str] = None,
         message_id: Optional[str] = None,
-    ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
+    ) -> Generator[ToolInvokeMessage, None, None]:
         """
         Convert localtime to timestamp
         """
@@ -28,9 +29,10 @@ class LocaltimeToTimestampTool(BuiltinTool):
 
         timestamp = self.localtime_to_timestamp(localtime, time_format, timezone)
         if not timestamp:
-            return self.create_text_message(f"Invalid localtime: {localtime}")
+            yield self.create_text_message(f"Invalid localtime: {localtime}")
+            return
 
-        return self.create_text_message(f"{timestamp}")
+        yield self.create_text_message(f"{timestamp}")
 
     @staticmethod
     def localtime_to_timestamp(localtime: str, time_format: str, local_tz=None) -> int | None:

+ 6 - 4
api/core/tools/builtin_tool/providers/time/tools/timestamp_to_localtime.py

@@ -1,5 +1,6 @@
+from collections.abc import Generator
 from datetime import datetime
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import pytz
 
@@ -16,7 +17,7 @@ class TimestampToLocaltimeTool(BuiltinTool):
         conversation_id: Optional[str] = None,
         app_id: Optional[str] = None,
         message_id: Optional[str] = None,
-    ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
+    ) -> Generator[ToolInvokeMessage, None, None]:
         """
         Convert timestamp to localtime
         """
@@ -28,11 +29,12 @@ class TimestampToLocaltimeTool(BuiltinTool):
 
         locatime = self.timestamp_to_localtime(timestamp, timezone)
         if not locatime:
-            return self.create_text_message(f"Invalid timestamp: {timestamp}")
+            yield self.create_text_message(f"Invalid timestamp: {timestamp}")
+            return
 
         localtime_format = locatime.strftime(time_format)
 
-        return self.create_text_message(f"{localtime_format}")
+        yield self.create_text_message(f"{localtime_format}")
 
     @staticmethod
     def timestamp_to_localtime(timestamp: int, local_tz=None) -> datetime | None:

+ 6 - 4
api/core/tools/builtin_tool/providers/time/tools/timezone_conversion.py

@@ -1,5 +1,6 @@
+from collections.abc import Generator
 from datetime import datetime
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import pytz
 
@@ -16,7 +17,7 @@ class TimezoneConversionTool(BuiltinTool):
         conversation_id: Optional[str] = None,
         app_id: Optional[str] = None,
         message_id: Optional[str] = None,
-    ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
+    ) -> Generator[ToolInvokeMessage, None, None]:
         """
         Convert time to equivalent time zone
         """
@@ -25,11 +26,12 @@ class TimezoneConversionTool(BuiltinTool):
         target_timezone = tool_parameters.get("target_timezone", "Asia/Tokyo")
         target_time = self.timezone_convert(current_time, current_timezone, target_timezone)
         if not target_time:
-            return self.create_text_message(
+            yield self.create_text_message(
                 f"Invalid datatime and timezone: {current_time},{current_timezone},{target_timezone}"
             )
+            return
 
-        return self.create_text_message(f"{target_time}")
+        yield self.create_text_message(f"{target_time}")
 
     @staticmethod
     def timezone_convert(current_time: str, source_timezone: str, target_timezone: str) -> str:

+ 6 - 4
api/core/tools/builtin_tool/providers/time/tools/weekday.py

@@ -1,6 +1,7 @@
 import calendar
+from collections.abc import Generator
 from datetime import datetime
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from core.tools.builtin_tool.tool import BuiltinTool
 from core.tools.entities.tool_entities import ToolInvokeMessage
@@ -14,7 +15,7 @@ class WeekdayTool(BuiltinTool):
         conversation_id: Optional[str] = None,
         app_id: Optional[str] = None,
         message_id: Optional[str] = None,
-    ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
+    ) -> Generator[ToolInvokeMessage, None, None]:
         """
         Calculate the day of the week for a given date
         """
@@ -26,12 +27,13 @@ class WeekdayTool(BuiltinTool):
 
         date_obj = self.convert_datetime(year, month, day)
         if not date_obj:
-            return self.create_text_message(f"Invalid date: Year {year}, Month {month}, Day {day}.")
+            yield self.create_text_message(f"Invalid date: Year {year}, Month {month}, Day {day}.")
+            return
 
         weekday_name = calendar.day_name[date_obj.weekday()]
         month_name = calendar.month_name[month]
         readable_date = f"{month_name} {date_obj.day}, {date_obj.year}"
-        return self.create_text_message(f"{readable_date} is {weekday_name}.")
+        yield self.create_text_message(f"{readable_date} is {weekday_name}.")
 
     @staticmethod
     def convert_datetime(year, month, day) -> datetime | None: