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 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)

View File

@ -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)}")

View File

@ -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:

View File

@ -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:

View File

@ -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:

View File

@ -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: