"""Schedule CRUD helpers for AWM workflows."""
from collections.abc import Generator
from datetime import datetime
from typing import Any
import chariot_api._openapi.awm.models as api_models
from chariot import _apis, mcp_setting
from chariot.awm import _utils, models
__all__ = [
"create_schedule",
"update_schedule",
"get_schedule",
"get_schedules",
"delete_schedule",
]
def _to_schedule(data: dict[str, Any]) -> models.Schedule:
raw = data.model_dump() if hasattr(data, "model_dump") else data
def _required(*keys: str) -> Any:
for key in keys:
value = raw.get(key)
if value is not None:
return value
raise RuntimeError(f"Received malformed schedule payload (missing {'/'.join(keys)})")
return models.Schedule(
id=_required("id"),
workflow_id=_required("workflow_id", "workflowID"),
name=raw.get("name") or "",
to_agent=_required("to_agent", "toAgent"),
prompt=_required("prompt"),
active=bool(raw.get("active", False)),
schedule=_required("schedule"),
created_by=_required("created_by", "createdBy"),
created_at=_required("created_at", "createdAt"),
expires_at=raw.get("expires_at", raw.get("expiresAt")),
expires_in=raw.get("expires_in", raw.get("expiresIn")),
)
def _upsert_payload(
to_agent: str,
prompt: str,
schedule: str,
active: bool,
name: str | None = None,
expires_at: datetime | str | None = None,
expires_in: str | None = None,
) -> api_models.InputUpsertScheduleRequest:
return api_models.InputUpsertScheduleRequest(
name=name,
to_agent=to_agent,
prompt=prompt,
active=active,
schedule=schedule,
expires_in=expires_in,
expires_at=expires_at.isoformat() if isinstance(expires_at, datetime) else expires_at,
)
def _schedules_api():
return _apis.awm.schedules_api
[docs]
@mcp_setting(mutating=True)
@_apis.login_required
def create_schedule(
workflow_id: str,
to_agent: str,
prompt: str,
schedule: str,
active: bool,
name: str | None = None,
expires_at: datetime | str | None = None,
expires_in: str | None = None,
) -> models.Schedule:
"""Create a schedule for a workflow."""
payload = _upsert_payload(
name=name,
to_agent=to_agent,
prompt=prompt,
active=active,
schedule=schedule,
expires_at=expires_at,
expires_in=expires_in,
)
response = _schedules_api().create_schedule(workflow_id, payload)
if not response.data:
raise RuntimeError("Received malformed response (missing `data`) from create_schedule")
return _to_schedule(response.data)
[docs]
@mcp_setting(mutating=True)
@_apis.login_required
def update_schedule(
schedule_id: str,
to_agent: str,
prompt: str,
schedule: str,
active: bool,
name: str | None = None,
expires_at: datetime | str | None = None,
expires_in: str | None = None,
) -> models.Schedule:
"""Update a schedule by id."""
payload = _upsert_payload(
name=name,
to_agent=to_agent,
prompt=prompt,
active=active,
schedule=schedule,
expires_at=expires_at,
expires_in=expires_in,
)
response = _schedules_api().update_schedule(schedule_id, payload)
if not response.data:
raise RuntimeError("Received malformed response (missing `data`) from update_schedule")
return _to_schedule(response.data)
[docs]
@_apis.login_required
def get_schedule(schedule_id: str) -> models.Schedule:
"""Get a schedule by id."""
response = _schedules_api().get_schedule(schedule_id)
if not response.data:
raise RuntimeError("Received malformed response (missing `data`) from get_schedule")
return _to_schedule(response.data)
[docs]
@mcp_setting(mutating=True)
@_apis.login_required
def delete_schedule(schedule_id: str) -> None:
"""Delete a schedule by id."""
_schedules_api().delete_schedule(schedule_id)
[docs]
@_apis.login_required
def get_schedules(
workflow_id: str,
name: str | None = None,
active: bool | None = None,
agents: list[str] | None = None,
limit: int | None = None,
after_key: str | None = None,
) -> Generator[models.Schedule, None, None]:
"""List schedules matching the specified filters.
:param workflow_id: Workflow id
:param name: Optional schedule name filter
:param active: Optional active/inactive filter
:param agents: Optional list of target agents
:param limit: Maximum number of results to return
:param after_key: The pagination key to fetch more items
:rtype: Generator[models.Schedule, None, None]
"""
params = {
k: v
for k, v in {
"name": name,
"active": active,
"agent": agents,
"limit": limit,
"after_key": after_key,
}.items()
if v is not None
}
def _paginate_fn(**params: Any) -> Any:
return _schedules_api().list_schedules(workflow_id=workflow_id, **params)
return _utils.paginate(
func=_paginate_fn,
params=params,
transform_item=_to_schedule,
)