fix: return types of builtin tools

This commit is contained in:
Yeuoly 2024-12-19 01:09:15 +08:00
parent a1cdca02e3
commit 8cc4c109d0
No known key found for this signature in database
GPG Key ID: A66E7E320FB19F61
6 changed files with 36 additions and 24 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -1,6 +1,7 @@
import calendar import calendar
from collections.abc import Generator
from datetime import datetime 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.builtin_tool.tool import BuiltinTool
from core.tools.entities.tool_entities import ToolInvokeMessage from core.tools.entities.tool_entities import ToolInvokeMessage
@ -14,7 +15,7 @@ class WeekdayTool(BuiltinTool):
conversation_id: Optional[str] = None, conversation_id: Optional[str] = None,
app_id: Optional[str] = None, app_id: Optional[str] = None,
message_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 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) date_obj = self.convert_datetime(year, month, day)
if not date_obj: 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()] weekday_name = calendar.day_name[date_obj.weekday()]
month_name = calendar.month_name[month] month_name = calendar.month_name[month]
readable_date = f"{month_name} {date_obj.day}, {date_obj.year}" 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 @staticmethod
def convert_datetime(year, month, day) -> datetime | None: def convert_datetime(year, month, day) -> datetime | None: