309 lines
10 KiB
Python
309 lines
10 KiB
Python
"""Agent 工具集 - 日程相关"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from datetime import date, datetime
|
||
from zoneinfo import ZoneInfo
|
||
|
||
from langchain_core.tools import tool
|
||
from sqlalchemy import select
|
||
|
||
from app.agents.context import get_current_user
|
||
from app.database import async_session
|
||
from app.models.goal import Goal, GoalStatus
|
||
from app.models.reminder import Reminder
|
||
from app.models.task import Task, TaskPriority, TaskStatus
|
||
from app.models.todo import DailyTodo, TodoSource
|
||
|
||
_executor = ThreadPoolExecutor(max_workers=4)
|
||
|
||
|
||
def _run_async(coro, timeout: int = 30):
|
||
try:
|
||
asyncio.get_running_loop()
|
||
except RuntimeError:
|
||
return asyncio.run(coro)
|
||
return _executor.submit(asyncio.run, coro).result(timeout=timeout)
|
||
|
||
|
||
def _parse_date(value: str | None) -> date:
|
||
if not value:
|
||
return date.today()
|
||
return date.fromisoformat(value)
|
||
|
||
|
||
def _parse_datetime(value: str) -> datetime:
|
||
normalized = value.strip().replace("Z", "+00:00")
|
||
return datetime.fromisoformat(normalized)
|
||
|
||
|
||
def _parse_datetime_with_timezone(value: str, time_zone: str | None) -> datetime:
|
||
"""Parse an ISO datetime and return a tz-naive datetime in the intended local time.
|
||
|
||
- If value includes an offset/Z, it will be converted to `time_zone` when provided.
|
||
- If value is naive and `time_zone` is provided, it is interpreted in that zone.
|
||
"""
|
||
parsed = _parse_datetime(value)
|
||
tz = (time_zone or "").strip()
|
||
if parsed.tzinfo is None:
|
||
if tz:
|
||
parsed = parsed.replace(tzinfo=ZoneInfo(tz))
|
||
return parsed.replace(tzinfo=None)
|
||
|
||
if tz:
|
||
parsed = parsed.astimezone(ZoneInfo(tz))
|
||
return parsed.replace(tzinfo=None)
|
||
|
||
|
||
def _normalize_title(title: str | None, content: str | None) -> str:
|
||
resolved = (title or content or "").strip()
|
||
if not resolved:
|
||
raise ValueError("title 不能为空")
|
||
return resolved
|
||
|
||
|
||
def _normalize_schedule_due_date(due_date: str | None, date_value: str | None) -> str | None:
|
||
resolved = (due_date or date_value or "").strip()
|
||
if not resolved:
|
||
return None
|
||
if "T" in resolved:
|
||
return resolved
|
||
return f"{resolved}T09:00:00"
|
||
|
||
|
||
def _format_summary(target_date: date, todos: list[DailyTodo], tasks: list[Task], reminders: list[Reminder], goals: list[Goal]) -> str:
|
||
lines = [f"日期: {target_date.isoformat()}"]
|
||
|
||
if todos:
|
||
lines.append("待办:")
|
||
lines.extend(f"- {item.title} | 完成:{'是' if item.is_completed else '否'}" for item in todos)
|
||
else:
|
||
lines.append("待办: 无")
|
||
|
||
if tasks:
|
||
lines.append("任务:")
|
||
lines.extend(
|
||
f"- {item.title} | 状态:{item.status.value if hasattr(item.status, 'value') else item.status} | 优先级:{item.priority.value if hasattr(item.priority, 'value') else item.priority} | 截止:{item.due_date.isoformat() if item.due_date else '无'}"
|
||
for item in tasks
|
||
)
|
||
else:
|
||
lines.append("任务: 无")
|
||
|
||
if reminders:
|
||
lines.append("提醒:")
|
||
lines.extend(f"- {item.title} | 时间:{item.reminder_at.isoformat()}" for item in reminders)
|
||
else:
|
||
lines.append("提醒: 无")
|
||
|
||
if goals:
|
||
lines.append("目标:")
|
||
lines.extend(
|
||
f"- {item.title} | 状态:{item.status.value if hasattr(item.status, 'value') else item.status}"
|
||
for item in goals
|
||
)
|
||
else:
|
||
lines.append("目标: 无")
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
@tool
|
||
def get_schedule_day(target_date: str | None = None) -> str:
|
||
"""获取指定日期的 todo/task/reminder/goal 聚合信息。target_date 格式 YYYY-MM-DD，默认今天。"""
|
||
uid = get_current_user()
|
||
parsed_date = _parse_date(target_date)
|
||
date_key = parsed_date.isoformat()
|
||
start_dt = datetime.combine(parsed_date, datetime.min.time())
|
||
end_dt = datetime.combine(parsed_date, datetime.max.time())
|
||
|
||
async def _get():
|
||
async with async_session() as db:
|
||
todos = (
|
||
await db.execute(
|
||
select(DailyTodo)
|
||
.where(DailyTodo.user_id == uid, DailyTodo.todo_date == date_key)
|
||
.order_by(DailyTodo.created_at.desc())
|
||
)
|
||
).scalars().all()
|
||
tasks = (
|
||
await db.execute(
|
||
select(Task)
|
||
.where(
|
||
Task.user_id == uid,
|
||
Task.due_date.is_not(None),
|
||
Task.due_date >= start_dt,
|
||
Task.due_date <= end_dt,
|
||
)
|
||
.order_by(Task.created_at.desc())
|
||
)
|
||
).scalars().all()
|
||
reminders = (
|
||
await db.execute(
|
||
select(Reminder)
|
||
.where(
|
||
Reminder.user_id == uid,
|
||
Reminder.reminder_at >= start_dt,
|
||
Reminder.reminder_at <= end_dt,
|
||
)
|
||
.order_by(Reminder.reminder_at.asc(), Reminder.created_at.asc())
|
||
)
|
||
).scalars().all()
|
||
goals = (
|
||
await db.execute(
|
||
select(Goal)
|
||
.where(Goal.user_id == uid, Goal.goal_date == date_key)
|
||
.order_by(Goal.created_at.desc())
|
||
)
|
||
).scalars().all()
|
||
return _format_summary(parsed_date, todos, tasks, reminders, goals)
|
||
|
||
try:
|
||
return _run_async(_get())
|
||
except Exception as exc:
|
||
return f"获取日程失败: {exc}"
|
||
|
||
|
||
@tool
|
||
def create_todo(title: str, todo_date: str | None = None) -> str:
|
||
"""创建指定日期的待办。todo_date 格式 YYYY-MM-DD，默认今天。"""
|
||
uid = get_current_user()
|
||
parsed_date = _parse_date(todo_date)
|
||
|
||
async def _create():
|
||
async with async_session() as db:
|
||
todo = DailyTodo(
|
||
user_id=uid,
|
||
title=title,
|
||
source=TodoSource.AI_CHAT,
|
||
todo_date=parsed_date.isoformat(),
|
||
)
|
||
db.add(todo)
|
||
await db.commit()
|
||
await db.refresh(todo)
|
||
return f"TODO创建成功: [{todo.id[:8]}] {todo.title} @ {todo.todo_date}"
|
||
|
||
try:
|
||
return _run_async(_create())
|
||
except Exception as exc:
|
||
return f"创建TODO失败: {exc}"
|
||
|
||
|
||
@tool
|
||
def create_schedule_task(
|
||
title: str = "",
|
||
description: str = "",
|
||
priority: str = "medium",
|
||
due_date: str | None = None,
|
||
content: str = "",
|
||
date: str | None = None,
|
||
) -> str:
|
||
"""创建任务。priority 支持 low/medium/high/urgent；due_date 使用 ISO datetime。兼容 content/date 别名。"""
|
||
uid = get_current_user()
|
||
resolved_title = _normalize_title(title, content)
|
||
resolved_due_date = _normalize_schedule_due_date(due_date, date)
|
||
|
||
async def _create():
|
||
async with async_session() as db:
|
||
task = Task(
|
||
user_id=uid,
|
||
title=resolved_title,
|
||
description=description or content or None,
|
||
priority=TaskPriority(priority),
|
||
due_date=_parse_datetime(resolved_due_date) if resolved_due_date else None,
|
||
status=TaskStatus.TODO,
|
||
)
|
||
db.add(task)
|
||
await db.commit()
|
||
await db.refresh(task)
|
||
due_label = task.due_date.isoformat() if task.due_date else "无截止时间"
|
||
return f"任务创建成功: [{task.id[:8]}] {task.title} | 优先级:{task.priority.value} | 截止:{due_label}"
|
||
|
||
try:
|
||
return _run_async(_create())
|
||
except Exception as exc:
|
||
return f"创建任务失败: {exc}"
|
||
|
||
|
||
@tool
|
||
def create_reminder(
|
||
title: str = "",
|
||
reminder_at: str | None = None,
|
||
note: str = "",
|
||
description: str = "",
|
||
datetime: str = "",
|
||
at: str = "",
|
||
remind_at: str = "",
|
||
content: str = "",
|
||
time_zone: str = "",
|
||
timezone: str = "",
|
||
time: str = "",
|
||
) -> str:
|
||
"""创建提醒。reminder_at 使用 ISO datetime。兼容 description/datetime/at/remind_at/time_zone 别名。"""
|
||
uid = get_current_user()
|
||
|
||
try:
|
||
resolved_title = (title or content or "").strip()
|
||
if not resolved_title:
|
||
raise ValueError("title 不能为空")
|
||
|
||
resolved_at = ((reminder_at or datetime or at or remind_at or time or "").strip())
|
||
if not resolved_at:
|
||
raise ValueError("reminder_at 不能为空")
|
||
|
||
resolved_note = (note or description or "").strip()
|
||
|
||
async def _create():
|
||
async with async_session() as db:
|
||
tz = (time_zone or timezone or "").strip()
|
||
reminder = Reminder(
|
||
user_id=uid,
|
||
title=resolved_title,
|
||
note=resolved_note or None,
|
||
reminder_at=_parse_datetime_with_timezone(resolved_at, tz),
|
||
)
|
||
db.add(reminder)
|
||
await db.commit()
|
||
await db.refresh(reminder)
|
||
return f"提醒创建成功: [{reminder.id[:8]}] {reminder.title} @ {reminder.reminder_at.isoformat()}"
|
||
|
||
return _run_async(_create())
|
||
except Exception as exc:
|
||
return f"创建提醒失败: {exc}"
|
||
|
||
|
||
@tool
|
||
def create_goal(title: str, goal_date: str | None = None, note: str = "", status: str = "active") -> str:
|
||
"""创建指定日期目标。goal_date 格式 YYYY-MM-DD，默认今天；status 支持 active/done/archived。"""
|
||
uid = get_current_user()
|
||
parsed_date = _parse_date(goal_date)
|
||
|
||
async def _create():
|
||
async with async_session() as db:
|
||
goal = Goal(
|
||
user_id=uid,
|
||
title=title,
|
||
note=note or None,
|
||
goal_date=parsed_date.isoformat(),
|
||
status=GoalStatus(status),
|
||
)
|
||
db.add(goal)
|
||
await db.commit()
|
||
await db.refresh(goal)
|
||
return f"目标创建成功: [{goal.id[:8]}] {goal.title} @ {goal.goal_date}"
|
||
|
||
try:
|
||
return _run_async(_create())
|
||
except Exception as exc:
|
||
return f"创建目标失败: {exc}"
|
||
|
||
|
||
__all__ = [
|
||
"get_schedule_day",
|
||
"create_todo",
|
||
"create_schedule_task",
|
||
"create_reminder",
|
||
"create_goal",
|
||
]
|