@atomicmemory/hermes-plugin 0.1.10

package/tools.py

@@ -0,0 +1,193 @@
+"""Tool schemas + handler dispatch for the Hermes memory provider.
+
+Lives outside `__init__.py` so the provider class stays under the workspace
+LOC ceiling. The handlers depend only on the provider's public-ish helpers
+(`_require_client`, `_user_id`, `_session_id`, `_config`, `_read_kwargs`,
+`_ingest_provenance`).
+"""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any, Callable
+
+from .config import ingest_scope_dict, read_scope_kwargs
+
+
+if TYPE_CHECKING:
+    from . import AtomicMemoryMemoryProvider
+
+
+try:
+    from tools.registry import tool_error  # type: ignore[import-not-found]
+except Exception:  # pragma: no cover — used only outside Hermes test runs.
+    def tool_error(message: str) -> str:
+        return json.dumps({"error": message})
+
+
+SEARCH_SCHEMA = {
+    "name": "atomicmemory_search",
+    "description": "Search AtomicMemory by meaning for prior user preferences, project context, decisions, and facts.",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "query": {"type": "string", "description": "What to search for."},
+            "top_k": {
+                "type": "integer",
+                "description": "Maximum results to return. Defaults to provider config, max 50.",
+            },
+        },
+        "required": ["query"],
+    },
+}
+
+CONTEXT_SCHEMA = {
+    "name": "atomicmemory_context",
+    "description": "Build an injection-ready AtomicMemory context package for a broad query.",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "query": {"type": "string", "description": "The topic to assemble context for."},
+            "token_budget": {
+                "type": "integer",
+                "description": "Approximate context budget. Defaults to provider config.",
+            },
+        },
+        "required": ["query"],
+    },
+}
+
+CONCLUDE_SCHEMA = {
+    "name": "atomicmemory_conclude",
+    "description": "Store one explicit durable fact, preference, correction, or decision in AtomicMemory.",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "conclusion": {"type": "string", "description": "The fact to store verbatim."},
+        },
+        "required": ["conclusion"],
+    },
+}
+
+
+def profile_schema(memory_scope: str) -> dict[str, Any]:
+    description = (
+        "List recent records for the current Hermes user from any AtomicMemory tool."
+        if memory_scope == "shared"
+        else "List recent Hermes-source AtomicMemory records for the current user."
+    )
+    return {
+        "name": "atomicmemory_profile",
+        "description": description,
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "limit": {
+                    "type": "integer",
+                    "description": "Maximum records to return. Defaults to provider config, max 50.",
+                },
+            },
+            "required": [],
+        },
+    }
+
+
+def _handle_search(provider: AtomicMemoryMemoryProvider, args: dict[str, Any]) -> str:
+    query = _clean_str(args.get("query"))
+    if not query:
+        return tool_error("Missing required parameter: query")
+    limit = _bounded_int(args.get("top_k"), provider._config.search_limit, 50)
+    kwargs = provider._read_kwargs(limit_override=limit)
+    page = provider._require_client().search(query=query, **kwargs)
+    if not page.memories:
+        return json.dumps({"result": "No relevant memories found.", "count": 0})
+    return json.dumps({"results": _serialize_memories(page), "count": page.count or len(page.memories)})
+
+
+def _handle_context(provider: AtomicMemoryMemoryProvider, args: dict[str, Any]) -> str:
+    query = _clean_str(args.get("query"))
+    if not query:
+        return tool_error("Missing required parameter: query")
+    token_budget = _bounded_int(args.get("token_budget"), provider._config.token_budget, 50000)
+    kwargs = read_scope_kwargs(memory_scope=provider._config.memory_scope, user_id=provider._user_id)
+    package = provider._require_client().package(
+        query=query,
+        scope=kwargs["scope"],
+        token_budget=token_budget,
+        source_site=kwargs.get("source_site"),
+    )
+    return json.dumps(
+        {
+            "result": package.injection_text,
+            "results": _serialize_memories(package),
+            "count": package.count or len(package.memories),
+            "estimated_context_tokens": package.estimated_context_tokens,
+            "citations": package.citations,
+        },
+    )
+
+
+def _handle_conclude(provider: AtomicMemoryMemoryProvider, args: dict[str, Any]) -> str:
+    conclusion = _clean_str(args.get("conclusion"))
+    if not conclusion:
+        return tool_error("Missing required parameter: conclusion")
+    provider._require_client().ingest_verbatim(
+        content=conclusion,
+        scope=ingest_scope_dict(user_id=provider._user_id),
+        provenance=provider._ingest_provenance(session_id=provider._session_id),
+        metadata={"kind": "fact"},
+    )
+    return json.dumps({"result": "Fact stored."})
+
+
+def _handle_profile(provider: AtomicMemoryMemoryProvider, args: dict[str, Any]) -> str:
+    limit = _bounded_int(args.get("limit"), provider._config.search_limit, 50)
+    kwargs = provider._read_kwargs(limit_override=limit)
+    page = provider._require_client().list_recent(**kwargs)
+    if not page.memories:
+        return json.dumps({"result": "No memories stored yet.", "count": 0})
+    return json.dumps({"results": _serialize_memories(page), "count": page.count or len(page.memories)})
+
+
+TOOL_HANDLERS: dict[str, Callable[[AtomicMemoryMemoryProvider, dict[str, Any]], str]] = {
+    "atomicmemory_search": _handle_search,
+    "atomicmemory_context": _handle_context,
+    "atomicmemory_conclude": _handle_conclude,
+    "atomicmemory_profile": _handle_profile,
+}
+
+
+def serialize_memories(page: Any) -> list[dict[str, Any]]:
+    return _serialize_memories(page)
+
+
+def format_memory_bullets(page: Any) -> str:
+    return "\n".join(f"- {item.content}" for item in page.memories if item.content)
+
+
+def _serialize_memories(page: Any) -> list[dict[str, Any]]:
+    return [
+        {
+            "id": item.id,
+            "memory": item.content,
+            "score": item.score,
+            "source_site": item.source_site,
+            "created_at": item.created_at,
+        }
+        for item in page.memories
+    ]
+
+
+def _bounded_int(value: Any, default: int, maximum: int) -> int:
+    try:
+        result = int(value)
+    except (TypeError, ValueError):
+        result = default
+    return max(1, min(result, maximum))
+
+
+def _clean_str(value: Any) -> str | None:
+    if value is None:
+        return None
+    cleaned = str(value).strip()
+    return cleaned or None

package/worker.py

@@ -0,0 +1,113 @@
+"""IngestWorker — single-writer queue + daemon thread for non-blocking sync_turn.
+
+Replaces the v1 per-call thread model that leaked orphans on join timeouts.
+Hermes' chat loop calls `submit()` (non-blocking); the daemon worker drains
+the queue serially, calling the supplied ingest function on each entry.
+
+`shutdown()` posts a sentinel and joins the worker (bounded). Drops oldest
+on a full queue with a single warning log per drop.
+"""
+
+from __future__ import annotations
+
+import logging
+import queue
+import threading
+from dataclasses import dataclass
+from typing import Any, Callable
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class IngestJob:
+    """Payload passed from sync_turn into the worker queue."""
+
+    user_content: str
+    assistant_content: str
+    session_id: str
+
+
+_SHUTDOWN_SENTINEL = object()
+
+
+class IngestWorker:
+    """Single-writer worker that drains queued ingest jobs in order."""
+
+    def __init__(
+        self,
+        *,
+        run_job: Callable[[IngestJob], None],
+        max_queue_size: int = 32,
+        on_failure: Callable[[BaseException], None] | None = None,
+    ) -> None:
+        self._run_job = run_job
+        self._on_failure = on_failure
+        self._queue: queue.Queue[Any] = queue.Queue(maxsize=max_queue_size)
+        self._thread: threading.Thread | None = None
+        self._started = threading.Event()
+        self._stopped = threading.Event()
+
+    def start(self) -> None:
+        if self._thread is not None:
+            return
+        self._thread = threading.Thread(
+            target=self._run,
+            daemon=True,
+            name="atomicmemory-ingest-worker",
+        )
+        self._thread.start()
+        self._started.set()
+
+    def submit(self, job: IngestJob) -> bool:
+        """Non-blocking enqueue. Returns False if the queue is full (drop)."""
+        if self._stopped.is_set():
+            return False
+        try:
+            self._queue.put_nowait(job)
+            return True
+        except queue.Full:
+            logger.warning("AtomicMemory ingest queue full; dropping turn for session %s", job.session_id)
+            return False
+
+    def shutdown(self, *, timeout: float = 10.0) -> None:
+        if self._thread is None:
+            return
+        self._stopped.set()
+        try:
+            self._queue.put_nowait(_SHUTDOWN_SENTINEL)
+        except queue.Full:
+            # Drain one slot to make space for the sentinel; ensures shutdown
+            # doesn't hang behind a saturated queue.
+            try:
+                self._queue.get_nowait()
+                self._queue.task_done()
+            except queue.Empty:
+                pass
+            self._queue.put_nowait(_SHUTDOWN_SENTINEL)
+        thread = self._thread
+        self._thread = None
+        thread.join(timeout=timeout)
+
+    def queue_size(self) -> int:
+        return self._queue.qsize()
+
+    def _run(self) -> None:
+        while True:
+            item = self._queue.get()
+            try:
+                if item is _SHUTDOWN_SENTINEL:
+                    return
+                try:
+                    self._run_job(item)
+                except BaseException as exc:  # noqa: BLE001 — caller decides retry policy
+                    if self._on_failure is not None:
+                        try:
+                            self._on_failure(exc)
+                        except Exception:  # noqa: BLE001
+                            logger.exception("ingest on_failure callback raised")
+                    else:
+                        logger.warning("AtomicMemory ingest job failed: %s", exc)
+            finally:
+                self._queue.task_done()