Rebase onto upstream (a4d95fd) #12

Closed
wylab wants to merge 144 commits from rebase-onto-upstream into main
2 changed files with 192 additions and 107 deletions
Showing only changes of commit 236d67f690 - Show all commits
+87 -97
View File
@@ -1,72 +1,65 @@
"""Heartbeat service - periodic agent wake-up to check for tasks."""
from __future__ import annotations
import asyncio
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Coroutine
from typing import Any, Callable, Coroutine
from loguru import logger
if TYPE_CHECKING:
from nanobot.providers.base import LLMProvider
# Default interval: 30 minutes
DEFAULT_HEARTBEAT_INTERVAL_S = 30 * 60
_HEARTBEAT_TOOL = [
{
"type": "function",
"function": {
"name": "heartbeat",
"description": "Report heartbeat decision after reviewing tasks.",
"parameters": {
"type": "object",
"properties": {
"action": {
"type": "string",
"enum": ["skip", "run"],
"description": "skip = nothing to do, run = has active tasks",
},
"tasks": {
"type": "string",
"description": "Natural-language summary of active tasks (required for run)",
},
},
"required": ["action"],
},
},
}
]
# The prompt sent to agent during heartbeat
HEARTBEAT_PROMPT = """Read HEARTBEAT.md in your workspace (if it exists).
Follow any instructions or tasks listed there.
If nothing needs attention, reply with just: HEARTBEAT_OK"""
# Token that indicates "nothing to do"
HEARTBEAT_OK_TOKEN = "HEARTBEAT_OK"
def _is_heartbeat_empty(content: str | None) -> bool:
"""Check if HEARTBEAT.md has no actionable content."""
if not content:
return True
# Lines to skip: empty, headers, HTML comments, empty checkboxes
skip_patterns = {"- [ ]", "* [ ]", "- [x]", "* [x]"}
for line in content.split("\n"):
line = line.strip()
if not line or line.startswith("#") or line.startswith("<!--") or line in skip_patterns:
continue
return False # Found actionable content
return True
class HeartbeatService:
"""
Periodic heartbeat service that wakes the agent to check for tasks.
Phase 1 (decision): reads HEARTBEAT.md and asks the LLM — via a virtual
tool call — whether there are active tasks. This avoids free-text parsing
and the unreliable HEARTBEAT_OK token.
Phase 2 (execution): only triggered when Phase 1 returns ``run``. The
``on_execute`` callback runs the task through the full agent loop and
returns the result to deliver.
The agent reads HEARTBEAT.md from the workspace and executes any
tasks listed there. If nothing needs attention, it replies HEARTBEAT_OK.
"""
def __init__(
self,
workspace: Path,
provider: LLMProvider,
model: str,
on_execute: Callable[[str], Coroutine[Any, Any, str]] | None = None,
on_notify: Callable[[str], Coroutine[Any, Any, None]] | None = None,
interval_s: int = 30 * 60,
on_heartbeat: Callable[[str, dict[str, Any] | None], Coroutine[Any, Any, str]] | None = None,
interval_s: int = DEFAULT_HEARTBEAT_INTERVAL_S,
enabled: bool = True,
session_manager: "SessionManager | None" = None,
target_session_key: str = "telegram:239824268",
idle_threshold_s: int = 30 * 60, # 30 minutes
):
self.workspace = workspace
self.provider = provider
self.model = model
self.on_execute = on_execute
self.on_notify = on_notify
self.on_heartbeat = on_heartbeat
self.interval_s = interval_s
self.enabled = enabled
self.session_manager = session_manager
self.target_session_key = target_session_key
self.idle_threshold_s = idle_threshold_s
self._running = False
self._task: asyncio.Task | None = None
@@ -75,48 +68,23 @@ class HeartbeatService:
return self.workspace / "HEARTBEAT.md"
def _read_heartbeat_file(self) -> str | None:
"""Read HEARTBEAT.md content."""
if self.heartbeat_file.exists():
try:
return self.heartbeat_file.read_text(encoding="utf-8")
return self.heartbeat_file.read_text()
except Exception:
return None
return None
async def _decide(self, content: str) -> tuple[str, str]:
"""Phase 1: ask LLM to decide skip/run via virtual tool call.
Returns (action, tasks) where action is 'skip' or 'run'.
"""
response = await self.provider.chat(
messages=[
{"role": "system", "content": "You are a heartbeat agent. Call the heartbeat tool to report your decision."},
{"role": "user", "content": (
"Review the following HEARTBEAT.md and decide whether there are active tasks.\n\n"
f"{content}"
)},
],
tools=_HEARTBEAT_TOOL,
model=self.model,
)
if not response.has_tool_calls:
return "skip", ""
args = response.tool_calls[0].arguments
return args.get("action", "skip"), args.get("tasks", "")
async def start(self) -> None:
"""Start the heartbeat service."""
if not self.enabled:
logger.info("Heartbeat disabled")
return
if self._running:
logger.warning("Heartbeat already running")
return
self._running = True
self._task = asyncio.create_task(self._run_loop())
logger.info("Heartbeat started (every {}s)", self.interval_s)
logger.info(f"Heartbeat started (every {self.interval_s}s)")
def stop(self) -> None:
"""Stop the heartbeat service."""
@@ -135,39 +103,61 @@ class HeartbeatService:
except asyncio.CancelledError:
break
except Exception as e:
logger.error("Heartbeat error: {}", e)
logger.error(f"Heartbeat error: {e}")
async def _tick(self) -> None:
"""Execute a single heartbeat tick."""
content = self._read_heartbeat_file()
if not content:
logger.debug("Heartbeat: HEARTBEAT.md missing or empty")
return
logger.info("Heartbeat: checking for tasks...")
# Check if user is idle (if session manager provided)
if self.session_manager and self.target_session_key:
try:
action, tasks = await self._decide(content)
session = self.session_manager.get_or_create(self.target_session_key)
if action != "run":
logger.info("Heartbeat: OK (nothing to report)")
# Find last user message timestamp
last_user_timestamp = None
for msg in reversed(session.messages):
if msg.get("role") == "user":
last_user_timestamp = msg.get("timestamp")
break
if last_user_timestamp:
from datetime import datetime
last_dt = datetime.fromisoformat(last_user_timestamp)
elapsed = (datetime.now() - last_dt).total_seconds()
if elapsed < self.idle_threshold_s:
logger.debug(f"Heartbeat: user active {int(elapsed)}s ago, skipping")
return # User is active, don't trigger heartbeat
except Exception as e:
logger.warning(f"Heartbeat: error checking idle state: {e}")
# Continue with heartbeat on error (fail open)
# Original heartbeat logic
content = self._read_heartbeat_file()
# Skip if HEARTBEAT.md is empty or doesn't exist
if _is_heartbeat_empty(content):
logger.debug("Heartbeat: no tasks (HEARTBEAT.md empty)")
return
logger.info("Heartbeat: tasks found, executing...")
if self.on_execute:
response = await self.on_execute(tasks)
if response and self.on_notify:
logger.info("Heartbeat: completed, delivering response")
await self.on_notify(response)
except Exception:
logger.exception("Heartbeat execution failed")
logger.info("Heartbeat: user idle, checking for tasks...")
if self.on_heartbeat:
try:
# Call with suppress_output metadata
response = await self.on_heartbeat(
HEARTBEAT_PROMPT,
metadata={"suppress_output": True}
)
# Note: HEARTBEAT_OK check removed - suppress mode makes it unnecessary
logger.info("Heartbeat: completed")
except Exception as e:
logger.error(f"Heartbeat execution failed: {e}")
async def trigger_now(self) -> str | None:
"""Manually trigger a heartbeat."""
content = self._read_heartbeat_file()
if not content:
if self.on_heartbeat:
return await self.on_heartbeat(HEARTBEAT_PROMPT, metadata={"suppress_output": True})
return None
action, tasks = await self._decide(content)
if action != "run" or not self.on_execute:
return None
return await self.on_execute(tasks)
+95
View File
@@ -0,0 +1,95 @@
# tests/test_heartbeat_idle.py
import pytest
import asyncio
from pathlib import Path
from datetime import datetime, timedelta
from nanobot.heartbeat.service import HeartbeatService
from nanobot.session.manager import SessionManager
from unittest.mock import AsyncMock, MagicMock
@pytest.mark.asyncio
async def test_heartbeat_skips_when_user_active():
"""Test that heartbeat doesn't trigger if user messaged recently."""
workspace = Path("/tmp/test-heartbeat")
workspace.mkdir(exist_ok=True)
# Create session with recent user message
sessions = SessionManager(workspace)
session = sessions.get_or_create("telegram:239824268")
session.add_message("user", "Recent message")
sessions.save(session)
# Create heartbeat callback
callback_called = False
async def on_heartbeat(prompt, metadata=None):
nonlocal callback_called
callback_called = True
return "response"
# Create heartbeat service
service = HeartbeatService(
workspace=workspace,
on_heartbeat=on_heartbeat,
interval_s=1, # Short interval for testing
enabled=True,
session_manager=sessions,
target_session_key="telegram:239824268"
)
# Trigger heartbeat
await service._tick()
# Callback should NOT have been called (user was active recently)
assert not callback_called
@pytest.mark.asyncio
async def test_heartbeat_triggers_when_user_idle():
"""Test that heartbeat triggers after 30min of inactivity."""
workspace = Path("/tmp/test-heartbeat")
workspace.mkdir(exist_ok=True)
# Create HEARTBEAT.md with content
heartbeat_file = workspace / "HEARTBEAT.md"
heartbeat_file.write_text("# Tasks\n- Check something\n")
# Create session with old user message (>30min ago)
sessions = SessionManager(workspace)
session = sessions.get_or_create("telegram:239824268")
# Manually set old timestamp
old_timestamp = (datetime.now() - timedelta(minutes=31)).isoformat()
session.messages.append({
"role": "user",
"content": "Old message",
"timestamp": old_timestamp
})
sessions.save(session)
# Create heartbeat callback
callback_called = False
callback_metadata = None
async def on_heartbeat(prompt, metadata=None):
nonlocal callback_called, callback_metadata
callback_called = True
callback_metadata = metadata
return "response"
# Create heartbeat service
service = HeartbeatService(
workspace=workspace,
on_heartbeat=on_heartbeat,
interval_s=1,
enabled=True,
session_manager=sessions,
target_session_key="telegram:239824268"
)
# Trigger heartbeat
await service._tick()
# Callback SHOULD have been called (user idle for >30min)
assert callback_called
assert callback_metadata == {"suppress_output": True}