ajarbot/heartbeat.py

"""Simple Heartbeat System - Periodic agent awareness checks."""

import threading
import time
from datetime import datetime
from typing import Callable, Optional

from llm_interface import LLMInterface
from memory_system import MemorySystem

# Default heartbeat checklist template
_HEARTBEAT_TEMPLATE = """\
# Heartbeat Checklist

Run these checks every heartbeat cycle:

## Memory Checks
- Review pending tasks (status = pending)
- Check if any tasks have been pending > 24 hours

## System Checks
- Verify memory system is synced
- Log heartbeat ran successfully

## Notes
- Return HEARTBEAT_OK if nothing needs attention
- Only alert if something requires user action
"""

# Maximum number of pending tasks to include in context
MAX_PENDING_TASKS_IN_CONTEXT = 5

# Maximum characters of soul content to include in context
SOUL_PREVIEW_LENGTH = 200


class Heartbeat:
    """Periodic background checks with LLM awareness."""

    def __init__(
        self,
        memory: MemorySystem,
        llm: LLMInterface,
        interval_minutes: int = 30,
        active_hours: tuple = (8, 22),
    ) -> None:
        self.memory = memory
        self.llm = llm
        self.interval = interval_minutes * 60
        self.active_hours = active_hours
        self.running = False
        self.thread: Optional[threading.Thread] = None
        self.on_alert: Optional[Callable[[str], None]] = None

        self.heartbeat_file = memory.workspace_dir / "HEARTBEAT.md"
        if not self.heartbeat_file.exists():
            self.heartbeat_file.write_text(
                _HEARTBEAT_TEMPLATE, encoding="utf-8"
            )

    def start(self) -> None:
        """Start heartbeat in background thread."""
        if self.running:
            return

        self.running = True
        self.thread = threading.Thread(
            target=self._heartbeat_loop, daemon=True
        )
        self.thread.start()
        print(f"Heartbeat started (every {self.interval // 60}min)")

    def stop(self) -> None:
        """Stop heartbeat."""
        self.running = False
        if self.thread:
            self.thread.join()
        print("Heartbeat stopped")

    def _is_active_hours(self) -> bool:
        """Check if current time is within active hours."""
        current_hour = datetime.now().hour
        start, end = self.active_hours
        return start <= current_hour < end

    def _heartbeat_loop(self) -> None:
        """Main heartbeat loop."""
        while self.running:
            try:
                if self._is_active_hours():
                    self._run_heartbeat()
                else:
                    start, end = self.active_hours
                    print(
                        f"Heartbeat skipped "
                        f"(outside active hours {start}-{end})"
                    )
            except Exception as e:
                print(f"Heartbeat error: {e}")

            time.sleep(self.interval)

    def _build_context(self) -> str:
        """Build system context for heartbeat check."""
        soul = self.memory.get_soul()
        pending_tasks = self.memory.get_tasks(status="pending")

        context_parts = [
            "# HEARTBEAT CHECK",
            f"Current time: {datetime.now().isoformat()}",
            f"\nSOUL:\n{soul[:SOUL_PREVIEW_LENGTH]}...",
            f"\nPending tasks: {len(pending_tasks)}",
        ]

        if pending_tasks:
            context_parts.append("\nPending Tasks:")
            for task in pending_tasks[:MAX_PENDING_TASKS_IN_CONTEXT]:
                context_parts.append(f"- [{task['id']}] {task['title']}")

        return "\n".join(context_parts)

    def _run_heartbeat(self) -> None:
        """Execute one heartbeat cycle."""
        timestamp = datetime.now().strftime("%H:%M:%S")
        print(f"Heartbeat running ({timestamp})")

        checklist = self.heartbeat_file.read_text(encoding="utf-8")
        system = self._build_context()

        messages = [{
            "role": "user",
            "content": (
                f"{checklist}\n\n"
                "Process this checklist. If nothing needs attention, "
                "respond with EXACTLY 'HEARTBEAT_OK'. If something "
                "needs attention, describe it briefly."
            ),
        }]

        response = self.llm.chat(messages, system=system, max_tokens=500)

        if response.strip() != "HEARTBEAT_OK":
            print(f"Heartbeat alert: {response[:100]}...")
            if self.on_alert:
                self.on_alert(response)
            self.memory.write_memory(
                f"## Heartbeat Alert\n{response}", daily=True
            )
        else:
            print("Heartbeat OK")

    def check_now(self) -> str:
        """Run heartbeat check immediately (for testing)."""
        print("Running immediate heartbeat check...")

        checklist = self.heartbeat_file.read_text(encoding="utf-8")
        pending_tasks = self.memory.get_tasks(status="pending")
        soul = self.memory.get_soul()

        system = (
            f"Time: {datetime.now().isoformat()}\n"
            f"SOUL: {soul[:SOUL_PREVIEW_LENGTH]}...\n"
            f"Pending tasks: {len(pending_tasks)}"
        )

        messages = [{
            "role": "user",
            "content": (
                f"{checklist}\n\n"
                "Process this checklist. "
                "Return HEARTBEAT_OK if nothing needs attention."
            ),
        }]

        return self.llm.chat(messages, system=system, max_tokens=500)


if __name__ == "__main__":
    memory = MemorySystem()
    llm = LLMInterface(provider="claude")

    heartbeat = Heartbeat(
        memory, llm, interval_minutes=30, active_hours=(8, 22)
    )

    def on_alert(message: str) -> None:
        print(f"\nALERT: {message}\n")

    heartbeat.on_alert = on_alert

    result = heartbeat.check_now()
    print(f"\nResult: {result}")