diff --git a/pyproject.toml b/pyproject.toml index b4b77e6..880f208 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ "openmemory-py~=1.2.3", "slack-sdk~=3.39.0", "workspace-mcp>=1.7.1", + "APScheduler~=3.10", ] [project.scripts] diff --git a/src/bro/cli.py b/src/bro/cli.py index e4ab442..eca0e14 100644 --- a/src/bro/cli.py +++ b/src/bro/cli.py @@ -27,6 +27,7 @@ from bro.memory import Memory from bro.knowledgebase.wiki import WikiClient from bro.mcp import GoogleWorkspaceClient, ShopifyClient +from bro.scheduler import TaskScheduler _logger = logging.getLogger(__name__) @@ -143,6 +144,9 @@ def main() -> None: app_token=os.environ["BRO_SLACK_APP_TOKEN"], bro_user_id=os.environ["BRO_SLACK_USER_ID"], ) + + scheduler = TaskScheduler(memory=memory, reasoner=rsn) + conversation = ConversationHandler( connector, user_system_prompt, @@ -150,6 +154,7 @@ def main() -> None: reasoner=rsn, memory=memory, wiki=wiki, + scheduler=scheduler, ) try: diff --git a/src/bro/conversation.py b/src/bro/conversation.py index 21ba406..df25961 100644 --- a/src/bro/conversation.py +++ b/src/bro/conversation.py @@ -16,6 +16,7 @@ from bro import util from bro.memory import Memory, tools as memory_tools from bro.knowledgebase.wiki import WikiClient, tools as wiki_tools +from bro.scheduler import TaskScheduler from bro.connector import Message, Connector, Channel, ReceivedMessage, User from bro.reasoner import Context, Reasoner @@ -106,6 +107,13 @@ consider notifying the user by sending an appropriately formatted response with the user name and `via` specified as necessary. +SCHEDULED TASKS: +Messages from `Bro Reasoner` that start with "SCHEDULED TASK:" are from automated background tasks. +For these messages: +- Process any actionable results (e.g., EMAIL CHECK RESULTS should be posted to the designated channel) +- Do NOT send acknowledgment or completion messages to users +- Only respond if there's critical information that requires immediate human attention + Important: - When writing a prompt for the reasoner, provide only the end goal, not step-by-step instructions. - Do NOT call get_reasoner_status immediately after calling task_reasoner. Wait for the reasoner to complete and report back. @@ -161,6 +169,42 @@ "parameters": {"type": "object", "properties": {}, "additionalProperties": False}, "strict": True, }, + { + "type": "function", + "name": "schedule_task", + "description": "Schedule a task to run automatically at specified times using cron syntax. " + "Example: '0 9 * * *' runs daily at 9am", + "parameters": { + "type": "object", + "properties": { + "task_prompt": {"type": "string", "description": "The task instruction for the reasoner"}, + "cron": {"type": "string", "description": "Cron expression (minute hour day month day_of_week)"}, + "task_id": {"type": "string", "description": "Unique identifier for this scheduled task"}, + }, + "required": ["task_prompt", "cron", "task_id"], + "additionalProperties": False, + }, + "strict": True, + }, + { + "type": "function", + "name": "cancel_scheduled_task", + "description": "Cancel a previously scheduled task. Use recall to find the task_id if needed (search for 'scheduled' tasks).", + "parameters": { + "type": "object", + "properties": {"task_id": {"type": "string", "description": "The ID of the task to cancel"}}, + "required": ["task_id"], + "additionalProperties": False, + }, + "strict": True, + }, + { + "type": "function", + "name": "list_scheduled_tasks", + "description": "List all currently scheduled tasks", + "parameters": {"type": "object", "properties": {}, "additionalProperties": False}, + "strict": True, + }, ] @@ -201,6 +245,7 @@ def __init__( reasoner: Reasoner, memory: Memory, wiki: WikiClient | None = None, + scheduler: TaskScheduler | None = None, ) -> None: self._msgs: list[ReceivedMessage] = [] self._current_task: Task | None = None @@ -213,6 +258,7 @@ def __init__( self._reasoner.on_task_completed_cb = self._on_task_completed_cb self._memory = memory self._wiki = wiki + self._scheduler = scheduler def _build_system_prompt(self) -> list[dict[str, Any]]: ctx: list[dict[str, Any]] = [ @@ -272,8 +318,12 @@ def _process_response_output(self, output: Any) -> None: if follow_up_output: self._process_response_output(follow_up_output) - def _on_task_completed_cb(self, message: str) -> None: + def _on_task_completed_cb(self, message: str, scheduled: bool = False) -> None: _logger.warning("🏁 " * 40 + "\n" + message) + + if scheduled: + message = f"SCHEDULED TASK: {message}" + input_data = textwrap.dedent( f"""\ via: @@ -291,8 +341,11 @@ def _on_task_completed_cb(self, message: str) -> None: } ] - self._current_task = None - _logger.info("Requesting conversation response after receiving reasoner response...") + # Clear current task for user-initiated tasks + if not scheduled: + self._current_task = None + + _logger.info(f"Requesting conversation response...") conversation_response = self._request_inference(self._context) output = conversation_response["output"] if not output: @@ -350,6 +403,24 @@ def _process(self, item: dict[str, Any]) -> str | None: case ("remember", {"text": text, "tags": tags}): result = self._memory.remember(text, tags) + case ("schedule_task", {"task_prompt": task_prompt, "cron": cron, "task_id": task_id}): + if self._scheduler: + result = self._scheduler.schedule(task_prompt, cron, task_id) + else: + result = "Scheduler not available" + + case ("cancel_scheduled_task", {"task_id": task_id}): + if self._scheduler: + result = self._scheduler.cancel(task_id) + else: + result = "Scheduler not available" + + case ("list_scheduled_tasks", {}): + if self._scheduler: + result = self._scheduler.list_tasks() + else: + result = "Scheduler not available" + case ("wiki_search", {"query": query}): if self._wiki: result = self._wiki.search(query) diff --git a/src/bro/memory.py b/src/bro/memory.py index ed69e7f..a4abec6 100644 --- a/src/bro/memory.py +++ b/src/bro/memory.py @@ -118,7 +118,18 @@ def remember(self, text: str, tags: list[str]) -> str: _logger.info(f"Adding memory with the following tags {tags}...") try: mem = self._memory.add(text, tags=tags) - _logger.debug(f"Memory stored. Memory id {mem['id']}") - return "Memory is added." + memory_id = str(mem["id"]) # Explicitly cast to str + _logger.debug(f"Memory stored. Memory id {memory_id}") + return memory_id except Exception as e: return f"Memory can't be added. Error: {e}" + + def forget(self, memory_id: str) -> bool: + """Delete a memory by ID.""" + try: + self._memory.delete(memory_id) + _logger.info(f"Deleted memory {memory_id}") + return True + except Exception as e: + _logger.error(f"Failed to delete memory {memory_id}: {e}") + return False diff --git a/src/bro/reasoner/__init__.py b/src/bro/reasoner/__init__.py index d299629..daf1c17 100644 --- a/src/bro/reasoner/__init__.py +++ b/src/bro/reasoner/__init__.py @@ -11,7 +11,7 @@ class Context: files: list[Path] -OnTaskCompleted = Callable[[str], None] +OnTaskCompleted = Callable[[str, bool], None] # (message, scheduled) class Reasoner(ABC): @@ -21,11 +21,14 @@ class Reasoner(ABC): """ @abstractmethod - def task(self, ctx: Context, /) -> bool: + def task(self, ctx: Context, /, *, scheduled: bool = False) -> bool: """ Commence a new task with the given context. The callback set via on_task_completed_cb is invoked from a worker thread with the final response once the task is finished. TODO: allow the reasoner to return files and images. + Args: + ctx: The task context + scheduled: If True, this is a scheduled task (won't trigger user notification callback) Returns True if the task is accepted, False if another task is still running. """ raise NotImplementedError diff --git a/src/bro/reasoner/openai_generic.py b/src/bro/reasoner/openai_generic.py index 1fe56de..b4f1f9b 100644 --- a/src/bro/reasoner/openai_generic.py +++ b/src/bro/reasoner/openai_generic.py @@ -388,7 +388,7 @@ """ -def _dummy_cb(_: Any) -> None: +def _dummy_cb(_: str, __: bool = False) -> None: _logger.error("The dummy callback is not supposed to be invoked") @@ -449,6 +449,7 @@ def __init__( self._context = self._build_system_prompt() self._step_number = 0 self._on_task_completed_cb: OnTaskCompleted = _dummy_cb + self._is_scheduled_task = False self._thread = threading.Thread(target=self._run_thread, daemon=True) self._thread_stop = False self._thread.start() @@ -471,12 +472,13 @@ def _build_system_prompt(self) -> list[dict[str, Any]]: ctx[0]["content"].append({"type": "input_text", "text": self._user_system_prompt}) return ctx - def task(self, ctx: Context, /) -> bool: + def task(self, ctx: Context, /, *, scheduled: bool = False) -> bool: if self._busy: return False if self._on_task_completed_cb is _dummy_cb: raise RuntimeError("Please configure the callback first") self._strategy = None + self._is_scheduled_task = scheduled self._context += [{"role": "user", "content": [{"type": "input_text", "text": ctx.prompt}]}] if ctx.files: # Ensure the files are uploaded so we can reference them in the prompt @@ -511,9 +513,10 @@ def _run_thread(self) -> None: time.sleep(1) _logger.debug("Calling the callback...") try: - self._on_task_completed_cb(final) + self._on_task_completed_cb(final, self._is_scheduled_task) except Exception as ex: _logger.exception("Unhandled exception in the callback: %s", ex) + self._is_scheduled_task = False # Reset flag else: time.sleep(1) except Exception as ex: diff --git a/src/bro/scheduler.py b/src/bro/scheduler.py new file mode 100644 index 0000000..a384277 --- /dev/null +++ b/src/bro/scheduler.py @@ -0,0 +1,119 @@ +"""Task scheduler using APScheduler + OpenMemory.""" + +import logging +from apscheduler.schedulers.background import BackgroundScheduler # type: ignore[import-untyped] +from bro.memory import Memory +from bro.reasoner import Reasoner, Context + +_logger = logging.getLogger(__name__) + + +class TaskScheduler: + def __init__(self, memory: Memory, reasoner: Reasoner) -> None: + self._memory = memory + self._reasoner = reasoner + self._scheduler = BackgroundScheduler() + self._task_memory_ids: dict[str, str] = {} # task_id -> memory_id mapping + self._scheduler.start() + self._load_scheduled_tasks() + _logger.info("Scheduler started") + + def _load_scheduled_tasks(self) -> None: + """Load scheduled tasks from memory on startup.""" + # Query returns all matching memories with their IDs + results = self._memory._memory.query("scheduled tasks", filters={"tags": ["scheduled"]}) + _logger.info(f"Found {len(results)} scheduled task memories") + + for memory_entry in results: + content = memory_entry.get("content", "") + memory_id = memory_entry.get("id", "") + + if "task_id:" in content and "prompt:" in content and "cron:" in content: + try: + parts = content.split("|") + task_id = parts[0].split("task_id:")[1].strip() + task_prompt = parts[1].split("prompt:")[1].strip() + cron = parts[2].split("cron:")[1].strip() + + # Store the memory ID for later deletion + self._task_memory_ids[task_id] = memory_id + + minute, hour, day, month, day_of_week = cron.split() + self._scheduler.add_job( + lambda p=task_prompt, tid=task_id: self._run_scheduled_task(p, tid), + "cron", + minute=minute, + hour=hour, + day=day, + month=month, + day_of_week=day_of_week, + id=task_id, + replace_existing=True, + ) + _logger.info(f"Restored scheduled task: {task_id} (memory: {memory_id})") + except Exception as e: + _logger.error(f"Failed to restore task from memory '{content}': {e}") + + def _run_scheduled_task(self, task_prompt: str, task_id: str) -> None: + """Run a scheduled task silently (no user notification).""" + _logger.info(f"Running scheduled task '{task_id}': {task_prompt}") + self._reasoner.task(Context(prompt=task_prompt, files=[]), scheduled=True) + + def schedule(self, task_prompt: str, cron: str, task_id: str) -> str: + """Schedule a task with cron syntax (e.g., '0 9 * * *' for 9am daily).""" + try: + # Store in memory and save the memory ID + memory_id = self._memory.remember( + f"task_id: {task_id} | prompt: {task_prompt} | cron: {cron}", ["procedural", "scheduled", task_id] + ) + self._task_memory_ids[task_id] = memory_id + + # Add to APScheduler + minute, hour, day, month, day_of_week = cron.split() + self._scheduler.add_job( + lambda: self._run_scheduled_task(task_prompt, task_id), + "cron", + minute=minute, + hour=hour, + day=day, + month=month, + day_of_week=day_of_week, + id=task_id, + replace_existing=True, + ) + _logger.info(f"Scheduled: {task_id} - {task_prompt} ({cron}) [memory: {memory_id}]") + return f"Successfully scheduled task '{task_id}'" + except Exception as e: + _logger.error(f"Failed to schedule {task_id}: {e}") + return f"Failed to schedule task: {e}" + + def cancel(self, task_id: str) -> str: + """Cancel a scheduled task.""" + try: + # Remove from APScheduler + self._scheduler.remove_job(task_id) + + # Delete from memory + memory_id = self._task_memory_ids.get(task_id) + if memory_id: + self._memory.forget(memory_id) + del self._task_memory_ids[task_id] + _logger.info(f"Cancelled: {task_id} (deleted memory {memory_id})") + else: + _logger.warning(f"Cancelled {task_id} but no memory ID found") + + return f"Successfully cancelled task '{task_id}'" + except Exception as e: + _logger.error(f"Failed to cancel {task_id}: {e}") + return f"Failed to cancel task: {e}" + + def list_tasks(self) -> str: + """List all scheduled tasks.""" + jobs = self._scheduler.get_jobs() + if not jobs: + return "No scheduled tasks" + + result = "Scheduled tasks:\n" + for job in jobs: + result += f"- {job.id}: next run at {job.next_run_time}\n" + return result