tiphys.tools.scheduler

  1"""
  2Scheduler tools for the agent.
  3
  4Allows the agent to schedule, list, and manage background tasks.
  5"""
  6
  7from typing import TYPE_CHECKING, Literal
  8
  9from pydantic import BaseModel, Field
 10
 11from tiphys.context import get_context
 12from tiphys.scheduler.models import ScheduleType
 13from tiphys.tools.base import tool
 14
 15if TYPE_CHECKING:
 16    from tiphys.scheduler.service import TaskScheduler
 17
 18
 19def _get_scheduler() -> "TaskScheduler":
 20    """Get the task scheduler from the application context."""
 21    context = get_context()
 22    if not context.scheduler:
 23        raise RuntimeError("Task scheduler is not available in the current context")
 24    return context.scheduler
 25
 26
 27class ScheduleTaskArgs(BaseModel):
 28    """Arguments for scheduling a task."""
 29
 30    name: str = Field(..., description="A short, descriptive name for the task")
 31    prompt: str = Field(..., description="The instruction or prompt for the agent to execute")
 32    schedule_type: Literal["cron", "interval", "once"] = Field(
 33        ..., description="Type of schedule: 'cron', 'interval', or 'once'"
 34    )
 35    schedule_value: str = Field(
 36        ...,
 37        description="The schedule value (e.g., '*/5 * * * *' for cron, '300' for interval seconds, or ISO timestamp for once)",
 38    )
 39
 40
 41@tool(name="schedule_task", description="Schedule a new background task for the agent")
 42async def schedule_task(
 43    name: str,
 44    prompt: str,
 45    schedule_type: Literal["cron", "interval", "once"],
 46    schedule_value: str,
 47) -> str:
 48    """
 49    Schedule a task to run automatically.
 50
 51    Examples:
 52        - cron: "0 9 * * *" (Every day at 9 AM)
 53        - interval: "3600" (Every hour)
 54        - once: "2023-12-31T23:59:59" (Specific time)
 55    """
 56    scheduler = _get_scheduler()
 57
 58    # Map string type to Enum
 59    type_map = {
 60        "cron": ScheduleType.CRON,
 61        "interval": ScheduleType.INTERVAL,
 62        "once": ScheduleType.ONCE,
 63    }
 64
 65    try:
 66        task_id = await scheduler.create_task(
 67            name=name,
 68            group="main",  # TODO: In multi-tenant, use current group/user
 69            prompt=prompt,
 70            schedule_type=type_map[schedule_type],
 71            schedule_value=schedule_value,
 72        )
 73        return f"Task '{name}' scheduled successfully with ID: {task_id}"
 74    except Exception as e:
 75        return f"Failed to schedule task: {str(e)}"
 76
 77
 78class ListTasksArgs(BaseModel):
 79    """Arguments for listing tasks."""
 80
 81    status: Literal["active", "paused", "completed", "all"] = Field(
 82        "active", description="Filter tasks by status"
 83    )
 84
 85
 86@tool(name="list_tasks", description="List scheduled tasks")
 87async def list_tasks(status: str = "active") -> str:
 88    """List tasks currently managed by the scheduler."""
 89    scheduler = _get_scheduler()
 90
 91    tasks = scheduler.store.list_all()
 92
 93    if status != "all":
 94        tasks = [t for t in tasks if t.status == status]
 95
 96    if not tasks:
 97        return f"No {status} tasks found."
 98
 99    lines = [f"Found {len(tasks)} {status} tasks:"]
100    for t in tasks:
101        next_run = t.next_run.isoformat() if t.next_run else "N/A"
102        lines.append(f"- [{t.status.value}] {t.name} (ID: {t.id})")
103        lines.append(f"  Type: {t.schedule_type.value}, Value: {t.schedule_value}")
104        lines.append(f"  Next Run: {next_run}")
105        lines.append("")
106
107    return "\n".join(lines)
108
109
110class ManageTaskArgs(BaseModel):
111    """Arguments for managing a task."""
112
113    task_id: str = Field(..., description="The ID of the task to manage")
114
115
116@tool(name="delete_task", description="Delete a scheduled task")
117async def delete_task(task_id: str) -> str:
118    """Delete a task permanently."""
119    scheduler = _get_scheduler()
120
121    if scheduler.store.delete(task_id):
122        return f"Task {task_id} deleted successfully."
123    return f"Task {task_id} not found."
124
125
126@tool(name="pause_task", description="Pause a scheduled task")
127async def pause_task(task_id: str) -> str:
128    """Pause a running task."""
129    scheduler = _get_scheduler()
130
131    if await scheduler.pause_task(task_id):
132        return f"Task {task_id} paused."
133    return f"Task {task_id} not found."
134
135
136@tool(name="resume_task", description="Resume a paused task")
137async def resume_task(task_id: str) -> str:
138    """Resume a paused task."""
139    scheduler = _get_scheduler()
140
141    if await scheduler.resume_task(task_id):
142        return f"Task {task_id} resumed."
143    return f"Task {task_id} not found."