ooodc/PicoBot
1
0
Fork 0
Code Pull Requests Activity

feat: 添加TimeTool以获取当前时间和计算相对时间，支持时区覆盖

Browse Source
This commit is contained in:
ooodc 2026-04-27 09:54:38 +08:00
parent ed45ec54ed
commit 37dc3385bb
5 changed files with 523 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/gateway/mod.rs
View File

@ -43,6 +43,7 @@ impl GatewayState {
session_ttl_hours,
agent_prompt_reinject_every,
show_tool_results,
config.time.timezone.clone(),
provider_config,
provider_configs,
skills,

80
src/gateway/session.rs
View File

@ -23,7 +23,7 @@ use crate::storage::{SessionRecord, SessionStore, persistent_session_id};
use crate::tools::{
BashTool, CalculatorTool, FileEditTool, FileReadTool, FileWriteTool,
HttpRequestTool, MemoryManageTool, MemorySearchTool, SchedulerManageTool, SkillListTool, SkillManageTool, ToolContext, ToolRegistry,
WebFetchTool,
TimeTool, WebFetchTool,
};
const DEFAULT_AGENT_PROMPT: &str = include_str!("default_agent_prompt.md");
@ -822,9 +822,15 @@ struct SessionManagerInner {
session_ttl: Duration,
}
fn default_tools(skills: Arc<SkillRuntime>, store: Arc<SessionStore>, known_agents: HashSet<String>) -> ToolRegistry {
fn default_tools(
skills: Arc<SkillRuntime>,
store: Arc<SessionStore>,
known_agents: HashSet<String>,
default_timezone: String,
) -> ToolRegistry {
let mut registry = ToolRegistry::new();
registry.register(CalculatorTool::new());
registry.register(TimeTool::new(default_timezone));
registry.register(FileReadTool::new());
registry.register(FileWriteTool::new());
registry.register(FileEditTool::new());
@ -963,6 +969,7 @@ impl SessionManager {
session_ttl_hours: u64,
agent_prompt_reinject_every: u64,
show_tool_results: bool,
default_timezone: String,
provider_config: LLMProviderConfig,
provider_configs: HashMap<String, LLMProviderConfig>,
skills: Arc<SkillRuntime>,
@ -985,7 +992,12 @@ impl SessionManager {
})),
provider_config,
provider_configs: Arc::new(provider_configs),
tools: Arc::new(default_tools(skills.clone(), store.clone(), known_agents)),
tools: Arc::new(default_tools(
skills.clone(),
store.clone(),
known_agents,
default_timezone,
)),
skills,
store,
agent_prompt_reinject_every,
@ -1690,7 +1702,12 @@ mod tests {
let store = Arc::new(SessionStore::in_memory().unwrap());
let (user_tx, _user_rx) = mpsc::channel(4);
let skills = Arc::new(SkillRuntime::default());
let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new()));
let tools = Arc::new(default_tools(
skills.clone(),
store.clone(),
HashSet::new(),
"Asia/Shanghai".to_string(),
));
let mut session = Session::new(
"feishu".to_string(),
test_provider_config(),
@ -1863,6 +1880,7 @@ mod tests {
4,
100,
false,
"Asia/Shanghai".to_string(),
provider_config.clone(),
HashMap::from([("default".to_string(), provider_config)]),
Arc::new(SkillRuntime::default()),
@ -1906,6 +1924,7 @@ mod tests {
4,
100,
false,
"Asia/Shanghai".to_string(),
default_provider.clone(),
HashMap::from([
("default".to_string(), default_provider),
@ -1971,6 +1990,7 @@ mod tests {
4,
100,
false,
"Asia/Shanghai".to_string(),
provider_config.clone(),
HashMap::from([("default".to_string(), provider_config)]),
Arc::new(SkillRuntime::default()),
@ -2045,6 +2065,7 @@ mod tests {
4,
100,
false,
"Asia/Shanghai".to_string(),
provider_config.clone(),
HashMap::from([("default".to_string(), provider_config)]),
Arc::new(SkillRuntime::default()),
@ -2134,6 +2155,7 @@ mod tests {
4,
100,
false,
"Asia/Shanghai".to_string(),
provider_config.clone(),
HashMap::from([("default".to_string(), provider_config)]),
Arc::new(SkillRuntime::default()),
@ -2194,6 +2216,7 @@ mod tests {
4,
100,
false,
"Asia/Shanghai".to_string(),
provider_config.clone(),
HashMap::from([("default".to_string(), provider_config)]),
Arc::new(SkillRuntime::default()),
@ -2364,7 +2387,12 @@ mod tests {
let store = Arc::new(SessionStore::in_memory().unwrap());
let (user_tx, _user_rx) = mpsc::channel(4);
let skills = Arc::new(SkillRuntime::default());
let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new()));
let tools = Arc::new(default_tools(
skills.clone(),
store.clone(),
HashSet::new(),
"Asia/Shanghai".to_string(),
));
let mut session = Session::new(
"feishu".to_string(),
test_provider_config(),
@ -2412,7 +2440,12 @@ mod tests {
let store = Arc::new(SessionStore::in_memory().unwrap());
let (user_tx, _user_rx) = mpsc::channel(4);
let skills = Arc::new(SkillRuntime::default());
let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new()));
let tools = Arc::new(default_tools(
skills.clone(),
store.clone(),
HashSet::new(),
"Asia/Shanghai".to_string(),
));
let mut session = Session::new(
"feishu".to_string(),
test_provider_config(),
@ -2439,7 +2472,12 @@ mod tests {
let store = Arc::new(SessionStore::in_memory().unwrap());
let (user_tx, _user_rx) = mpsc::channel(4);
let skills = Arc::new(SkillRuntime::default());
let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new()));
let tools = Arc::new(default_tools(
skills.clone(),
store.clone(),
HashSet::new(),
"Asia/Shanghai".to_string(),
));
let mut session = Session::new(
"feishu".to_string(),
test_provider_config(),
@ -2484,7 +2522,12 @@ mod tests {
let store = Arc::new(SessionStore::in_memory().unwrap());
let (user_tx, _user_rx) = mpsc::channel(4);
let skills = Arc::new(SkillRuntime::default());
let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new()));
let tools = Arc::new(default_tools(
skills.clone(),
store.clone(),
HashSet::new(),
"Asia/Shanghai".to_string(),
));
let mut session = Session::new(
"feishu".to_string(),
test_provider_config(),
@ -2518,7 +2561,26 @@ mod tests {
let store = Arc::new(SessionStore::in_memory().unwrap());
let (user_tx, _user_rx) = mpsc::channel(4);
let skills = Arc::new(SkillRuntime::default());
let tools = Arc::new(default_tools(skills.clone(), store.clone(), HashSet::new()));
let tools = Arc::new(default_tools(
skills.clone(),
store.clone(),
HashSet::new(),
"Asia/Shanghai".to_string(),
));
#[test]
fn test_default_tools_registers_get_time() {
let skills = Arc::new(SkillRuntime::default());
let store = Arc::new(SessionStore::in_memory().unwrap());
let tools = default_tools(
skills,
store,
HashSet::new(),
"Asia/Shanghai".to_string(),
);
assert!(tools.tool_names().iter().any(|name| name == "get_time"));
}
let mut session = Session::new(
"feishu".to_string(),
test_provider_config(),

2
src/scheduler/mod.rs
View File

@ -856,6 +856,7 @@ mod tests {
4,
100,
false,
"Asia/Shanghai".to_string(),
provider_config.clone(),
HashMap::from([("default".to_string(), provider_config)]),
Arc::new(SkillRuntime::default()),
@ -906,6 +907,7 @@ mod tests {
4,
100,
false,
"Asia/Shanghai".to_string(),
provider_config.clone(),
HashMap::from([("default".to_string(), provider_config)]),
Arc::new(SkillRuntime::default()),

2
src/tools/mod.rs
View File

@ -10,6 +10,7 @@ pub mod registry;
pub mod scheduler_manage;
pub mod schema;
pub mod skill_manage;
pub mod time;
pub mod traits;
pub mod web_fetch;
@ -25,5 +26,6 @@ pub use registry::ToolRegistry;
pub use scheduler_manage::SchedulerManageTool;
pub use schema::{CleaningStrategy, SchemaCleanr};
pub use skill_manage::{SkillListTool, SkillManageTool};
pub use time::TimeTool;
pub use traits::{Tool, ToolContext, ToolResult};
pub use web_fetch::WebFetchTool;

447
src/tools/time.rs Normal file
View File

@ -0,0 +1,447 @@
use async_trait::async_trait;
use chrono::{DateTime, Days, Duration, Months, Utc};
use serde_json::{Value, json};
use super::traits::{Tool, ToolResult};
pub struct TimeTool {
default_timezone: String,
}
impl TimeTool {
pub fn new(default_timezone: impl Into<String>) -> Self {
Self {
default_timezone: default_timezone.into(),
}
}
fn execute_at(&self, now_utc: DateTime<Utc>, args: Value) -> ToolResult {
match execute_time_request(now_utc, &self.default_timezone, args) {
Ok(output) => ToolResult {
success: true,
output,
error: None,
},
Err(error) => ToolResult {
success: false,
output: String::new(),
error: Some(error),
},
}
}
}
impl Default for TimeTool {
fn default() -> Self {
Self::new("Asia/Shanghai")
}
}
#[async_trait]
impl Tool for TimeTool {
fn name(&self) -> &str {
"get_time"
}
fn description(&self) -> &str {
"Get the current time or calculate a time before/after now in minutes, hours, days, months, or years. Use this for current time, relative time offsets, or when exact timestamps matter. If amount/unit/direction are omitted, it returns the current time in the default timezone. Examples: current time, 30 minutes later, 2 hours ago, 7 days later, 1 month ago, 1 year later."
}
fn read_only(&self) -> bool {
true
}
fn parameters_schema(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"direction": {
"type": "string",
"description": "Relative direction from the current time. Optional when only the current time is needed.",
"enum": ["past", "future", "before", "after", "ago", "later"]
},
"amount": {
"type": "integer",
"description": "How many units to move relative to now. Required when direction and unit are used.",
"minimum": 0
},
"unit": {
"type": "string",
"description": "Time unit for the relative offset. Required when direction and amount are used.",
"enum": [
"minute", "minutes",
"hour", "hours",
"day", "days",
"month", "months",
"year", "years"
]
},
"timezone": {
"type": "string",
"description": "Optional IANA timezone override, for example Asia/Shanghai or Europe/Berlin. Defaults to the PicoBot configured timezone."
}
}
})
}
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
Ok(self.execute_at(Utc::now(), args))
}
}
fn execute_time_request(
now_utc: DateTime<Utc>,
default_timezone: &str,
args: Value,
) -> Result<String, String> {
let timezone_name = args
.get("timezone")
.and_then(Value::as_str)
.unwrap_or(default_timezone);
let timezone = timezone_name.parse::<chrono_tz::Tz>().map_err(|_| {
format!(
"Invalid timezone: {timezone_name}. Expected an IANA timezone like Asia/Shanghai"
)
})?;
let now_local = now_utc.with_timezone(&timezone);
let offset = parse_offset_request(&args)?;
let (kind, result_local, description, offset_json) = match offset {
None => (
"current",
now_local,
format!(
"Current time in {timezone_name} is {}",
format_human_time(&now_local)
),
Value::Null,
),
Some(offset) => {
let result_local = apply_offset(now_local, &offset)?;
let direction_text = match offset.direction {
OffsetDirection::Past => "before",
OffsetDirection::Future => "after",
};
let unit_label = pluralized_unit(offset.unit, offset.amount);
let description = format!(
"{} {} {} current time in {timezone_name} is {}",
offset.amount,
unit_label,
direction_text,
format_human_time(&result_local)
);
(
"offset",
result_local,
description,
json!({
"direction": offset.direction.as_str(),
"amount": offset.amount,
"unit": offset.unit.as_str(),
}),
)
}
};
serde_json::to_string_pretty(&json!({
"kind": kind,
"timezone": timezone_name,
"current_time": now_local.to_rfc3339(),
"result_time": result_local.to_rfc3339(),
"unix_timestamp": result_local.timestamp(),
"offset": offset_json,
"description": description,
}))
.map_err(|error| format!("Failed to serialize time tool output: {error}"))
}
fn parse_offset_request(args: &Value) -> Result<Option<OffsetRequest>, String> {
let direction = args.get("direction").and_then(Value::as_str);
let amount = args.get("amount");
let unit = args.get("unit").and_then(Value::as_str);
if direction.is_none() && amount.is_none() && unit.is_none() {
return Ok(None);
}
let direction = direction.ok_or_else(|| {
"Missing required parameter: direction when requesting a relative time".to_string()
})?;
let amount = amount
.and_then(Value::as_u64)
.ok_or_else(|| "Missing required parameter: amount when requesting a relative time".to_string())?;
let amount = u32::try_from(amount)
.map_err(|_| "amount is too large; expected a 32-bit unsigned integer".to_string())?;
let unit = unit
.ok_or_else(|| "Missing required parameter: unit when requesting a relative time".to_string())?;
Ok(Some(OffsetRequest {
direction: OffsetDirection::parse(direction)?,
amount,
unit: TimeUnit::parse(unit)?,
}))
}
fn apply_offset(
now_local: DateTime<chrono_tz::Tz>,
offset: &OffsetRequest,
) -> Result<DateTime<chrono_tz::Tz>, String> {
match (offset.direction, offset.unit) {
(OffsetDirection::Future, TimeUnit::Minute) => Ok(now_local + Duration::minutes(i64::from(offset.amount))),
(OffsetDirection::Past, TimeUnit::Minute) => Ok(now_local - Duration::minutes(i64::from(offset.amount))),
(OffsetDirection::Future, TimeUnit::Hour) => Ok(now_local + Duration::hours(i64::from(offset.amount))),
(OffsetDirection::Past, TimeUnit::Hour) => Ok(now_local - Duration::hours(i64::from(offset.amount))),
(OffsetDirection::Future, TimeUnit::Day) => now_local
.checked_add_days(Days::new(u64::from(offset.amount)))
.ok_or_else(|| "Failed to add days to the current time".to_string()),
(OffsetDirection::Past, TimeUnit::Day) => now_local
.checked_sub_days(Days::new(u64::from(offset.amount)))
.ok_or_else(|| "Failed to subtract days from the current time".to_string()),
(OffsetDirection::Future, TimeUnit::Month) => now_local
.checked_add_months(Months::new(offset.amount))
.ok_or_else(|| "Failed to add months to the current time".to_string()),
(OffsetDirection::Past, TimeUnit::Month) => now_local
.checked_sub_months(Months::new(offset.amount))
.ok_or_else(|| "Failed to subtract months from the current time".to_string()),
(OffsetDirection::Future, TimeUnit::Year) => now_local
.checked_add_months(Months::new(years_to_months(offset.amount)?))
.ok_or_else(|| "Failed to add years to the current time".to_string()),
(OffsetDirection::Past, TimeUnit::Year) => now_local
.checked_sub_months(Months::new(years_to_months(offset.amount)?))
.ok_or_else(|| "Failed to subtract years from the current time".to_string()),
}
}
fn years_to_months(years: u32) -> Result<u32, String> {
years
.checked_mul(12)
.ok_or_else(|| "amount is too large to convert years into months".to_string())
}
fn format_human_time(value: &DateTime<chrono_tz::Tz>) -> String {
value.format("%Y-%m-%d %H:%M:%S %Z").to_string()
}
fn pluralized_unit(unit: TimeUnit, amount: u32) -> &'static str {
match (unit, amount) {
(TimeUnit::Minute, 1) => "minute",
(TimeUnit::Minute, _) => "minutes",
(TimeUnit::Hour, 1) => "hour",
(TimeUnit::Hour, _) => "hours",
(TimeUnit::Day, 1) => "day",
(TimeUnit::Day, _) => "days",
(TimeUnit::Month, 1) => "month",
(TimeUnit::Month, _) => "months",
(TimeUnit::Year, 1) => "year",
(TimeUnit::Year, _) => "years",
}
}
#[derive(Clone, Copy)]
struct OffsetRequest {
direction: OffsetDirection,
amount: u32,
unit: TimeUnit,
}
#[derive(Clone, Copy)]
enum OffsetDirection {
Past,
Future,
}
impl OffsetDirection {
fn parse(value: &str) -> Result<Self, String> {
match value.trim().to_ascii_lowercase().as_str() {
"past" | "before" | "ago" => Ok(Self::Past),
"future" | "after" | "later" => Ok(Self::Future),
other => Err(format!(
"Invalid direction: {other}. Expected one of: past, future, before, after, ago, later"
)),
}
}
fn as_str(&self) -> &'static str {
match self {
Self::Past => "past",
Self::Future => "future",
}
}
}
#[derive(Clone, Copy)]
enum TimeUnit {
Minute,
Hour,
Day,
Month,
Year,
}
impl TimeUnit {
fn parse(value: &str) -> Result<Self, String> {
match value.trim().to_ascii_lowercase().as_str() {
"minute" | "minutes" => Ok(Self::Minute),
"hour" | "hours" => Ok(Self::Hour),
"day" | "days" => Ok(Self::Day),
"month" | "months" => Ok(Self::Month),
"year" | "years" => Ok(Self::Year),
other => Err(format!(
"Invalid unit: {other}. Expected one of: minute, hour, day, month, year"
)),
}
}
fn as_str(&self) -> &'static str {
match self {
Self::Minute => "minute",
Self::Hour => "hour",
Self::Day => "day",
Self::Month => "month",
Self::Year => "year",
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::{TimeZone, Timelike};
fn fixed_now() -> DateTime<Utc> {
Utc.with_ymd_and_hms(2026, 4, 27, 4, 30, 0)
.single()
.unwrap()
}
#[test]
fn test_get_current_time_uses_default_timezone() {
let tool = TimeTool::new("Asia/Shanghai");
let result = tool.execute_at(fixed_now(), json!({}));
assert!(result.success);
let payload: Value = serde_json::from_str(&result.output).unwrap();
assert_eq!(payload["kind"], "current");
assert_eq!(payload["timezone"], "Asia/Shanghai");
assert_eq!(payload["result_time"], "2026-04-27T12:30:00+08:00");
}
#[test]
fn test_get_future_hours() {
let tool = TimeTool::new("Asia/Shanghai");
let result = tool.execute_at(
fixed_now(),
json!({"direction": "future", "amount": 2, "unit": "hours"}),
);
assert!(result.success);
let payload: Value = serde_json::from_str(&result.output).unwrap();
assert_eq!(payload["result_time"], "2026-04-27T14:30:00+08:00");
assert_eq!(payload["offset"]["direction"], "future");
assert_eq!(payload["offset"]["unit"], "hour");
}
#[test]
fn test_get_past_minutes() {
let tool = TimeTool::new("Asia/Shanghai");
let result = tool.execute_at(
fixed_now(),
json!({"direction": "ago", "amount": 15, "unit": "minute"}),
);
assert!(result.success);
let payload: Value = serde_json::from_str(&result.output).unwrap();
assert_eq!(payload["result_time"], "2026-04-27T12:15:00+08:00");
}
#[test]
fn test_get_future_days_crosses_date_boundary() {
let tool = TimeTool::new("Asia/Shanghai");
let result = tool.execute_at(
fixed_now(),
json!({"direction": "after", "amount": 3, "unit": "days"}),
);
assert!(result.success);
let payload: Value = serde_json::from_str(&result.output).unwrap();
assert_eq!(payload["result_time"], "2026-04-30T12:30:00+08:00");
}
#[test]
fn test_get_past_month_rolls_back_to_last_valid_day() {
let tool = TimeTool::new("Asia/Shanghai");
let now = Utc.with_ymd_and_hms(2026, 3, 31, 1, 0, 0).single().unwrap();
let result = tool.execute_at(
now,
json!({"direction": "past", "amount": 1, "unit": "month"}),
);
assert!(result.success);
let payload: Value = serde_json::from_str(&result.output).unwrap();
assert_eq!(payload["result_time"], "2026-02-28T09:00:00+08:00");
}
#[test]
fn test_get_future_year_handles_leap_day() {
let tool = TimeTool::new("Asia/Shanghai");
let now = Utc.with_ymd_and_hms(2024, 2, 29, 4, 0, 0).single().unwrap();
let result = tool.execute_at(
now,
json!({"direction": "future", "amount": 1, "unit": "year"}),
);
assert!(result.success);
let payload: Value = serde_json::from_str(&result.output).unwrap();
assert_eq!(payload["result_time"], "2025-02-28T12:00:00+08:00");
}
#[test]
fn test_timezone_override_is_used() {
let tool = TimeTool::new("Asia/Shanghai");
let result = tool.execute_at(fixed_now(), json!({"timezone": "Europe/Berlin"}));
assert!(result.success);
let payload: Value = serde_json::from_str(&result.output).unwrap();
assert_eq!(payload["timezone"], "Europe/Berlin");
let result_time = payload["result_time"].as_str().unwrap();
assert!(result_time.ends_with("+02:00"));
}
#[test]
fn test_missing_relative_unit_returns_error() {
let tool = TimeTool::new("Asia/Shanghai");
let result = tool.execute_at(fixed_now(), json!({"direction": "future", "amount": 1}));
assert!(!result.success);
assert_eq!(
result.error.as_deref(),
Some("Missing required parameter: unit when requesting a relative time")
);
}
#[test]
fn test_invalid_timezone_returns_error() {
let tool = TimeTool::new("Asia/Shanghai");
let result = tool.execute_at(fixed_now(), json!({"timezone": "Mars/Base"}));
assert!(!result.success);
assert!(result.error.unwrap().contains("Invalid timezone"));
}
#[test]
fn test_zero_offset_keeps_same_local_time() {
let tool = TimeTool::new("Asia/Shanghai");
let result = tool.execute_at(
fixed_now(),
json!({"direction": "future", "amount": 0, "unit": "hours"}),
);
assert!(result.success);
let payload: Value = serde_json::from_str(&result.output).unwrap();
let result_time = chrono::DateTime::parse_from_rfc3339(payload["result_time"].as_str().unwrap())
.unwrap();
assert_eq!(result_time.hour(), 12);
assert_eq!(result_time.minute(), 30);
}
}