161 lines
6.4 KiB
Python
161 lines
6.4 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
from apscheduler.triggers.date import DateTrigger
|
|
from apscheduler.triggers.interval import IntervalTrigger
|
|
from tzlocal import get_localzone
|
|
|
|
from traderai.memory import MemoryStore, iso_now, time_since
|
|
|
|
|
|
UEX_NOTIFICATION_JOB_ID = "uex-notification-poll"
|
|
|
|
|
|
class WakeScheduler:
|
|
def __init__(self, memory: MemoryStore) -> None:
|
|
self.memory = memory
|
|
self.scheduler = AsyncIOScheduler(timezone=get_localzone())
|
|
self.agent = None
|
|
self.uex = None
|
|
self.notification_poll_seconds = 60
|
|
|
|
def bind_agent(self, agent: Any) -> None:
|
|
self.agent = agent
|
|
|
|
def bind_uex_notifications(self, uex: Any, poll_seconds: int = 60) -> None:
|
|
self.uex = uex
|
|
self.notification_poll_seconds = max(15, poll_seconds)
|
|
|
|
def start(self) -> None:
|
|
if not self.scheduler.running:
|
|
self.scheduler.start()
|
|
self._schedule_notification_poll()
|
|
for job in self.memory.list_jobs():
|
|
self._schedule_existing(job)
|
|
|
|
def shutdown(self) -> None:
|
|
if self.scheduler.running:
|
|
self.scheduler.shutdown(wait=False)
|
|
|
|
def schedule_date(self, run_at: str, prompt: str, job_id: str | None = None) -> dict[str, Any]:
|
|
parsed = datetime.fromisoformat(run_at)
|
|
job_id = job_id or f"wake-{uuid4()}"
|
|
trigger = DateTrigger(run_date=parsed)
|
|
self.scheduler.add_job(self._run_job, trigger=trigger, id=job_id, args=[job_id, prompt], replace_existing=True)
|
|
return self.memory.add_job(job_id, prompt, "date", run_at, parsed.isoformat())
|
|
|
|
def schedule_cron(self, cron: str, prompt: str, job_id: str | None = None) -> dict[str, Any]:
|
|
job_id = job_id or f"wake-{uuid4()}"
|
|
trigger = CronTrigger.from_crontab(cron)
|
|
self.scheduler.add_job(self._run_job, trigger=trigger, id=job_id, args=[job_id, prompt], replace_existing=True)
|
|
next_run = self.scheduler.get_job(job_id).next_run_time
|
|
return self.memory.add_job(job_id, prompt, "cron", cron, next_run.isoformat() if next_run else None)
|
|
|
|
def list_jobs(self) -> list[dict[str, Any]]:
|
|
return self.memory.list_jobs()
|
|
|
|
def _schedule_existing(self, job: dict[str, Any]) -> None:
|
|
if job["trigger_type"] == "cron":
|
|
trigger = CronTrigger.from_crontab(job["trigger_value"])
|
|
elif job["trigger_type"] == "date":
|
|
trigger = DateTrigger(run_date=datetime.fromisoformat(job["trigger_value"]))
|
|
else:
|
|
return
|
|
self.scheduler.add_job(
|
|
self._run_job,
|
|
trigger=trigger,
|
|
id=job["id"],
|
|
args=[job["id"], job["prompt"]],
|
|
replace_existing=True,
|
|
)
|
|
|
|
async def _run_job(self, job_id: str, prompt: str) -> None:
|
|
last = self.memory.last_interaction()
|
|
last_text = f"{last['created_at']} ({time_since(last['created_at'])})" if last else "never"
|
|
wake_message = (
|
|
f"Scheduled wake job fired. Current time is {iso_now()}. "
|
|
f"The last chat interaction was {last_text}. Job instruction: {prompt}"
|
|
)
|
|
if self.agent is None:
|
|
self.memory.add_outbox(wake_message)
|
|
self._mark_job_finished(job_id)
|
|
return
|
|
|
|
try:
|
|
text = await self.agent.generate_wake_response(wake_message)
|
|
except Exception as exc:
|
|
text = f"Wake job failed: {exc}. Job instruction: {prompt}"
|
|
self.memory.add_outbox(text)
|
|
self._mark_job_finished(job_id)
|
|
|
|
def _mark_job_finished(self, job_id: str) -> None:
|
|
job = self.scheduler.get_job(job_id)
|
|
next_run = job.next_run_time if job else None
|
|
self.memory.mark_job_run(job_id, next_run.isoformat() if next_run else None, enabled=bool(next_run))
|
|
|
|
def _schedule_notification_poll(self) -> None:
|
|
if self.uex is None:
|
|
return
|
|
self.scheduler.add_job(
|
|
self.poll_uex_notifications,
|
|
trigger=IntervalTrigger(seconds=self.notification_poll_seconds),
|
|
id=UEX_NOTIFICATION_JOB_ID,
|
|
replace_existing=True,
|
|
next_run_time=datetime.now(),
|
|
)
|
|
|
|
async def poll_uex_notifications(self) -> list[dict[str, Any]]:
|
|
if self.uex is None:
|
|
return []
|
|
|
|
try:
|
|
response = await self.uex.get_user_notifications()
|
|
except Exception as exc:
|
|
self.memory.add_outbox(f"UEX notification poll failed: {exc}")
|
|
self.memory.set_profile("uex_last_notification_error", str(exc))
|
|
return []
|
|
notifications = response.get("notifications") or []
|
|
pending = [item for item in notifications if not item.get("date_read")]
|
|
profile = self.memory.get_profile()
|
|
seen = set(profile.get("uex_seen_notification_keys") or [])
|
|
new_pending = [item for item in pending if self._notification_key(item) not in seen]
|
|
|
|
if new_pending:
|
|
for item in new_pending:
|
|
self.memory.add_outbox(self._notification_text(item))
|
|
seen.update(self._notification_key(item) for item in new_pending)
|
|
self.memory.set_profile("uex_seen_notification_keys", sorted(seen))
|
|
self.memory.set_profile("uex_last_notification_check", iso_now())
|
|
elif notifications:
|
|
seen.update(self._notification_key(item) for item in pending)
|
|
self.memory.set_profile("uex_seen_notification_keys", sorted(seen))
|
|
self.memory.set_profile("uex_last_notification_check", iso_now())
|
|
|
|
return new_pending
|
|
|
|
@staticmethod
|
|
def _notification_key(item: dict[str, Any]) -> str:
|
|
for key in ("code", "id"):
|
|
value = item.get(key)
|
|
if value not in (None, ""):
|
|
return f"{key}:{value}"
|
|
return f"notification:{item.get('date_added')}:{item.get('message')}"
|
|
|
|
@staticmethod
|
|
def _notification_text(item: dict[str, Any]) -> str:
|
|
message = item.get("message") or "You have a pending UEX notification."
|
|
redir = item.get("redir")
|
|
code = item.get("code")
|
|
details = []
|
|
if code:
|
|
details.append(f"code `{code}`")
|
|
if redir:
|
|
details.append(f"path `{redir}`")
|
|
suffix = f" ({', '.join(details)})" if details else ""
|
|
return f"UEX notification: {message}{suffix}"
|