155 lines
4.4 KiB
Python
155 lines
4.4 KiB
Python
import json
|
|
import os
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from redis.asyncio import Redis
|
|
|
|
REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")
|
|
REMINDER_PREFIX = "daarion:reminders"
|
|
REMINDER_BY_ID = f"{REMINDER_PREFIX}:by_id"
|
|
REMINDER_SCHEDULE = f"{REMINDER_PREFIX}:schedule"
|
|
REMINDER_TTL_SECONDS = int(os.getenv("DAARION_REMINDER_TTL_SECONDS", str(30 * 24 * 3600)))
|
|
|
|
_redis: Optional[Redis] = None
|
|
|
|
|
|
@dataclass
|
|
class Reminder:
|
|
reminder_id: str
|
|
agent_id: str
|
|
chat_id: str
|
|
user_id: str
|
|
text: str
|
|
due_ts: int
|
|
created_at: str
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
"reminder_id": self.reminder_id,
|
|
"agent_id": self.agent_id,
|
|
"chat_id": self.chat_id,
|
|
"user_id": self.user_id,
|
|
"text": self.text,
|
|
"due_ts": self.due_ts,
|
|
"created_at": self.created_at,
|
|
}
|
|
|
|
|
|
async def redis_client() -> Redis:
|
|
global _redis
|
|
if _redis is None:
|
|
_redis = Redis.from_url(REDIS_URL, decode_responses=True)
|
|
return _redis
|
|
|
|
|
|
async def close_redis() -> None:
|
|
global _redis
|
|
if _redis is not None:
|
|
await _redis.close()
|
|
_redis = None
|
|
|
|
|
|
def _iso_now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
def _iso_from_ts(ts: int) -> str:
|
|
return datetime.fromtimestamp(ts, tz=timezone.utc).isoformat()
|
|
|
|
|
|
async def create_reminder(agent_id: str, chat_id: str, user_id: str, text: str, due_ts: int) -> Dict[str, Any]:
|
|
reminder = Reminder(
|
|
reminder_id=f"rem_{uuid.uuid4().hex[:16]}",
|
|
agent_id=agent_id,
|
|
chat_id=str(chat_id),
|
|
user_id=str(user_id),
|
|
text=text.strip(),
|
|
due_ts=int(due_ts),
|
|
created_at=_iso_now(),
|
|
)
|
|
|
|
r = await redis_client()
|
|
key = f"{REMINDER_BY_ID}:{reminder.reminder_id}"
|
|
payload = json.dumps(reminder.to_dict(), ensure_ascii=False)
|
|
|
|
await r.set(key, payload, ex=REMINDER_TTL_SECONDS)
|
|
await r.zadd(REMINDER_SCHEDULE, {reminder.reminder_id: float(reminder.due_ts)})
|
|
|
|
result = reminder.to_dict()
|
|
result["due_at"] = _iso_from_ts(reminder.due_ts)
|
|
return result
|
|
|
|
|
|
async def list_reminders(agent_id: str, chat_id: str, user_id: str, limit: int = 10) -> List[Dict[str, Any]]:
|
|
r = await redis_client()
|
|
now_ts = int(time.time())
|
|
ids = await r.zrangebyscore(REMINDER_SCHEDULE, min=now_ts - 365 * 24 * 3600, max="+inf", start=0, num=max(1, limit * 5))
|
|
|
|
out: List[Dict[str, Any]] = []
|
|
for reminder_id in ids:
|
|
raw = await r.get(f"{REMINDER_BY_ID}:{reminder_id}")
|
|
if not raw:
|
|
continue
|
|
try:
|
|
item = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
if item.get("agent_id") != agent_id:
|
|
continue
|
|
if str(item.get("chat_id")) != str(chat_id):
|
|
continue
|
|
if str(item.get("user_id")) != str(user_id):
|
|
continue
|
|
item["due_at"] = _iso_from_ts(int(item.get("due_ts", 0)))
|
|
out.append(item)
|
|
if len(out) >= limit:
|
|
break
|
|
return out
|
|
|
|
|
|
async def cancel_reminder(reminder_id: str, agent_id: str, chat_id: str, user_id: str) -> bool:
|
|
r = await redis_client()
|
|
key = f"{REMINDER_BY_ID}:{reminder_id}"
|
|
raw = await r.get(key)
|
|
if not raw:
|
|
return False
|
|
try:
|
|
item = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
return False
|
|
|
|
if item.get("agent_id") != agent_id or str(item.get("chat_id")) != str(chat_id) or str(item.get("user_id")) != str(user_id):
|
|
return False
|
|
|
|
await r.delete(key)
|
|
await r.zrem(REMINDER_SCHEDULE, reminder_id)
|
|
return True
|
|
|
|
|
|
async def pop_due_reminders(limit: int = 20) -> List[Dict[str, Any]]:
|
|
r = await redis_client()
|
|
now_ts = int(time.time())
|
|
ids = await r.zrangebyscore(REMINDER_SCHEDULE, min="-inf", max=now_ts, start=0, num=max(1, limit))
|
|
out: List[Dict[str, Any]] = []
|
|
|
|
for reminder_id in ids:
|
|
removed = await r.zrem(REMINDER_SCHEDULE, reminder_id)
|
|
if removed == 0:
|
|
continue
|
|
raw = await r.get(f"{REMINDER_BY_ID}:{reminder_id}")
|
|
if not raw:
|
|
continue
|
|
await r.delete(f"{REMINDER_BY_ID}:{reminder_id}")
|
|
try:
|
|
item = json.loads(raw)
|
|
item["due_at"] = _iso_from_ts(int(item.get("due_ts", now_ts)))
|
|
out.append(item)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
return out
|