diff --git a/.gitignore b/.gitignore index 15b9029..a195419 100644 --- a/.gitignore +++ b/.gitignore @@ -46,3 +46,6 @@ CLAUDE.md AGENTS.md .planning/ .claude/ + +# Node +node_modules/ diff --git a/benchmarks/bench.py b/benchmarks/bench.py index 9398aab..a0aa70b 100644 --- a/benchmarks/bench.py +++ b/benchmarks/bench.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 """Benchmark script for hexus. Run inside the docker container: @@ -166,13 +165,12 @@ def make_text(i: int) -> str: print(f" Cross-agent (None) sees: {len(rows_all)} rows") # Cleanup - with store._get_pool().connection() as conn: - with conn.cursor() as cur: - cur.execute( - "DELETE FROM memory_entries WHERE agent_identity IN (%s, %s)", - (agent, agent2), - ) - conn.commit() + with store._get_pool().connection() as conn, conn.cursor() as cur: + cur.execute( + "DELETE FROM memory_entries WHERE agent_identity IN (%s, %s)", + (agent, agent2), + ) + conn.commit() store.close() diff --git a/hexus/__init__.py b/hexus/__init__.py index 674d83c..2def22b 100644 --- a/hexus/__init__.py +++ b/hexus/__init__.py @@ -46,12 +46,12 @@ import os import re from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any try: from agent.memory_provider import MemoryProvider - from tools.registry import tool_error from hermes_cli.config import cfg_get + from tools.registry import tool_error except ( ImportError ): # pragma: no cover - standalone smoke tests do not install hermes-agent @@ -59,12 +59,12 @@ tool_error = None # type: ignore[assignment] cfg_get = None # type: ignore[assignment] -from .embed import embed, EmbeddingError -from .store import MemoryStore import hashlib -from .writer import AsyncWriter, _PendingWrite -from .pipeline.router import ContentRouter +from .embed import EmbeddingError, embed +from .pipeline.router import ContentRouter +from .store import MemoryStore +from .writer import AsyncWriter, _PendingWrite # Boilerplate / acknowledgement-only turns that are not worth embedding or # storing. Case-insensitive whole-string match after strip. Combined with @@ -443,7 +443,7 @@ def _load_plugin_config() -> dict: cfg_get(data, "plugins", "hexus", default={}) or {} ) return expanded if isinstance(expanded, dict) else {} - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # noqa: BLE001 # noqa: BLE001 return {} @@ -502,14 +502,14 @@ def __init__(self, config: dict | None = None): "install this package inside Hermes Agent to use the provider." ) self._config = {**DEFAULTS, **(config or {})} - self._store: Optional[MemoryStore] = None - self._writer: Optional[AsyncWriter] = None + self._store: MemoryStore | None = None + self._writer: AsyncWriter | None = None self._agent_identity: str = "default" self._session_id: str = "" self._healthy = False self._embed_warned = False - self._last_md_mtimes: Dict[str, float] = {} - self._hermes_home: Optional[str] = None + self._last_md_mtimes: dict[str, float] = {} + self._hermes_home: str | None = None self._content_router = ContentRouter() @property @@ -609,7 +609,7 @@ def initialize(self, session_id: str, **kwargs) -> None: "embed_url" ): try: - from .embedder import get_default_embedder, DEFAULT_MODEL + from .embedder import DEFAULT_MODEL, get_default_embedder get_default_embedder( model_name=self._config.get("embed_model") or DEFAULT_MODEL @@ -670,7 +670,7 @@ def _check_and_sync_markdown_files(self) -> None: from hermes_constants import get_hermes_home self._hermes_home = str(get_hermes_home()) - except Exception: + except Exception: # noqa: BLE001 # noqa: BLE001 return memories_dir = Path(self._hermes_home) / "memories" @@ -686,7 +686,7 @@ def _check_and_sync_markdown_files(self) -> None: ): changed = True self._last_md_mtimes[fname] = mtime - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.debug("hexus failed to get mtime for %s: %s", fname, exc) if changed: @@ -755,7 +755,7 @@ def _sync_skills_from_disk(self) -> None: self._last_md_mtimes[rel_path] = mtime logger.info("hexus: synced skill '%s' from disk", skill_name) - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.debug("hexus failed to sync skill %s: %s", skill_file, exc) # -- System prompt + ambient recall -------------------------------------- @@ -766,7 +766,7 @@ def system_prompt_block(self) -> str: try: count_scoped = self._store.count(agent_identity=self._agent_identity) count_all = self._store.count() - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # noqa: BLE001 # noqa: BLE001 count_scoped = count_all = 0 if count_all == 0: return ( @@ -866,9 +866,7 @@ def _is_noise(content: str, *, min_chars: int) -> bool: return True if len(stripped) < min_chars: return True - if _NOISE_RE.match(stripped): - return True - return False + return bool(_NOISE_RE.match(stripped)) # -- Built-in memory mirror (THE main integration point) ---------------- @@ -877,7 +875,7 @@ def on_memory_write( action: str, target: str, content: str, - metadata: Optional[Dict[str, Any]] = None, + metadata: dict[str, Any] | None = None, ) -> None: """Mirror built-in `memory` tool writes to Postgres (non-blocking). @@ -908,7 +906,7 @@ def on_memory_write( metadata=meta, ) - def _worker(self, item: "_PendingWrite") -> None: + def _worker(self, item: _PendingWrite) -> None: """Drain-thread worker: embed + DB write for a single queued item. Must NOT raise — the AsyncWriter logs + survives if we do, but @@ -1074,7 +1072,7 @@ def _worker(self, item: "_PendingWrite") -> None: str(exc)[:200], ) - def _generate_session_summary(self, messages_json: str) -> Optional[str]: + def _generate_session_summary(self, messages_json: str) -> str | None: import urllib.request try: @@ -1148,7 +1146,7 @@ def _generate_session_summary(self, messages_json: str) -> Optional[str]: "hexus: session summary successfully generated: %s", summary_content ) return summary_content - except Exception as e: + except Exception as e: # noqa: BLE001 logger.debug("hexus failed to generate session summary via LLM: %s", e) return None @@ -1180,7 +1178,7 @@ def on_delegation( metadata=meta, ) - def on_pre_compress(self, messages: List[Dict[str, Any]]) -> str: + def on_pre_compress(self, messages: list[dict[str, Any]]) -> str: """Called before context compaction discards older messages.""" if not self._healthy or not self._writer or not messages: return "" @@ -1196,7 +1194,7 @@ def on_pre_compress(self, messages: List[Dict[str, Any]]) -> str: ) return "" - def on_session_end(self, messages: List[Dict[str, Any]]) -> None: + def on_session_end(self, messages: list[dict[str, Any]]) -> None: """Called when a session ends.""" if not self._healthy or not self._writer: return @@ -1236,7 +1234,7 @@ def on_session_end(self, messages: List[Dict[str, Any]]) -> None: # -- Bulk sync (v0.1.1) -------------------------------------------------- - def _bulk_sync_from_disk(self, hermes_home: Optional[str]) -> None: + def _bulk_sync_from_disk(self, hermes_home: str | None) -> None: """Import MEMORY.md + USER.md entries from disk into memory_entries. Called by initialize(). Runs synchronously (not via async writer) @@ -1252,7 +1250,7 @@ def _bulk_sync_from_disk(self, hermes_home: Optional[str]) -> None: from hermes_constants import get_hermes_home hermes_home = str(get_hermes_home()) - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # noqa: BLE001 # noqa: BLE001 return memories_dir = Path(hermes_home) / "memories" @@ -1295,7 +1293,7 @@ def _fn(text: str): # -- Tool surface -------------------------------------------------------- - def get_tool_schemas(self) -> List[Dict[str, Any]]: + def get_tool_schemas(self) -> list[dict[str, Any]]: return [ RECALL_MEMORY_SCHEMA, RECALL_CONVERSATION_SCHEMA, @@ -1310,7 +1308,7 @@ def get_tool_schemas(self) -> List[Dict[str, Any]]: MEMORY_STATS_SCHEMA, ] - def handle_tool_call(self, tool_name: str, args: Dict[str, Any], **kwargs) -> str: + def handle_tool_call(self, tool_name: str, args: dict[str, Any], **kwargs) -> str: if tool_name == "recall_conversation": return self._handle_recall_conversation(args) if tool_name == "recall_delegation": @@ -1352,7 +1350,7 @@ def handle_tool_call(self, tool_name: str, args: Dict[str, Any], **kwargs) -> st args.get("scope") or self._config.get("scope_default") or "current" ).strip() if scope == "current": - agent_filter: Optional[str] = self._agent_identity + agent_filter: str | None = self._agent_identity elif scope == "all": agent_filter = None else: @@ -1360,7 +1358,7 @@ def handle_tool_call(self, tool_name: str, args: Dict[str, Any], **kwargs) -> st # Target resolution: 'memory'/'user'/'both'. target_arg = (args.get("target") or "both").strip() - target_filter: Optional[str] = None if target_arg == "both" else target_arg + target_filter: str | None = None if target_arg == "both" else target_arg if target_filter not in (None, "memory", "user"): return tool_error(f"Invalid target: {target_arg!r}") @@ -1404,7 +1402,7 @@ def handle_tool_call(self, tool_name: str, args: Dict[str, Any], **kwargs) -> st ) return json.dumps({"results": results, "count": len(results)}) - def _handle_recall_conversation(self, args: Dict[str, Any]) -> str: + def _handle_recall_conversation(self, args: dict[str, Any]) -> str: """Tool handler for recall_conversation over the conversations table.""" if not self._healthy or not self._store: return json.dumps({"results": [], "count": 0, "error": "hexus unavailable"}) @@ -1418,8 +1416,8 @@ def _handle_recall_conversation(self, args: Dict[str, Any]) -> str: limit = 5 scope = (args.get("scope") or "current").strip() - agent_filter: Optional[str] = None - session_filter: Optional[str] = None + agent_filter: str | None = None + session_filter: str | None = None if scope == "current": agent_filter = self._agent_identity elif scope == "session": @@ -1464,7 +1462,7 @@ def _handle_recall_conversation(self, args: Dict[str, Any]) -> str: ) return json.dumps({"results": results, "count": len(results)}) - def _handle_recall_delegation(self, args: Dict[str, Any]) -> str: + def _handle_recall_delegation(self, args: dict[str, Any]) -> str: """Tool handler for recall_delegation over the delegations table.""" if not self._healthy or not self._store: return json.dumps({"results": [], "count": 0, "error": "hexus unavailable"}) @@ -1478,7 +1476,7 @@ def _handle_recall_delegation(self, args: Dict[str, Any]) -> str: limit = 5 scope = (args.get("scope") or "current").strip() - agent_filter: Optional[str] = None + agent_filter: str | None = None if scope == "current": agent_filter = self._agent_identity elif scope == "all": @@ -1521,7 +1519,7 @@ def _handle_recall_delegation(self, args: Dict[str, Any]) -> str: ) return json.dumps({"results": results, "count": len(results)}) - def _handle_entity_graph(self, args: Dict[str, Any]) -> str: + def _handle_entity_graph(self, args: dict[str, Any]) -> str: """Tool handler for entity_graph.""" if not self._healthy or not self._store: return json.dumps({"error": "hexus unavailable"}) @@ -1539,7 +1537,7 @@ def _handle_entity_graph(self, args: Dict[str, Any]) -> str: limit = 5 scope = (args.get("scope") or "current").strip() - agent_filter: Optional[str] = None + agent_filter: str | None = None if scope == "current": agent_filter = self._agent_identity elif scope == "all": @@ -1555,10 +1553,10 @@ def _handle_entity_graph(self, args: Dict[str, Any]) -> str: limit=limit, ) return json.dumps(res) - except Exception as exc: + except Exception as exc: # noqa: BLE001 return json.dumps({"error": f"db: {exc}"}) - def _handle_graph_walk(self, args: Dict[str, Any]) -> str: + def _handle_graph_walk(self, args: dict[str, Any]) -> str: """Tool handler for graph_walk.""" if not self._healthy or not self._store: return json.dumps({"error": "hexus unavailable"}) @@ -1581,7 +1579,7 @@ def _handle_graph_walk(self, args: Dict[str, Any]) -> str: limit = 5 scope = (args.get("scope") or "current").strip() - agent_filter: Optional[str] = None + agent_filter: str | None = None if scope == "current": agent_filter = self._agent_identity elif scope == "all": @@ -1598,10 +1596,10 @@ def _handle_graph_walk(self, args: Dict[str, Any]) -> str: limit=limit, ) return json.dumps({"results": res}) - except Exception as exc: + except Exception as exc: # noqa: BLE001 return json.dumps({"error": f"db: {exc}"}) - def _handle_common_topics(self, args: Dict[str, Any]) -> str: + def _handle_common_topics(self, args: dict[str, Any]) -> str: """Tool handler for common_topics.""" if not self._healthy or not self._store: return json.dumps({"error": "hexus unavailable"}) @@ -1617,7 +1615,7 @@ def _handle_common_topics(self, args: Dict[str, Any]) -> str: limit = 10 scope = (args.get("scope") or "current").strip() - agent_filter: Optional[str] = None + agent_filter: str | None = None if scope == "current": agent_filter = self._agent_identity elif scope == "all": @@ -1632,10 +1630,10 @@ def _handle_common_topics(self, args: Dict[str, Any]) -> str: limit=limit, ) return json.dumps({"results": res}) - except Exception as exc: + except Exception as exc: # noqa: BLE001 return json.dumps({"error": f"db: {exc}"}) - def _handle_confirm_memory(self, args: Dict[str, Any]) -> str: + def _handle_confirm_memory(self, args: dict[str, Any]) -> str: """Tool handler for confirm_memory.""" if not self._healthy or not self._store: return json.dumps({"error": "hexus unavailable"}) @@ -1651,10 +1649,10 @@ def _handle_confirm_memory(self, args: Dict[str, Any]) -> str: try: success = self._store.confirm_entry(entry_id, self._agent_identity) return json.dumps({"id": entry_id, "success": success}) - except Exception as exc: + except Exception as exc: # noqa: BLE001 return json.dumps({"error": f"db: {exc}"}) - def _handle_reject_memory(self, args: Dict[str, Any]) -> str: + def _handle_reject_memory(self, args: dict[str, Any]) -> str: """Tool handler for reject_memory.""" if not self._healthy or not self._store: return json.dumps({"error": "hexus unavailable"}) @@ -1670,10 +1668,10 @@ def _handle_reject_memory(self, args: Dict[str, Any]) -> str: try: success = self._store.reject_entry(entry_id, self._agent_identity) return json.dumps({"id": entry_id, "success": success}) - except Exception as exc: + except Exception as exc: # noqa: BLE001 return json.dumps({"error": f"db: {exc}"}) - def _handle_summarize_session(self, args: Dict[str, Any]) -> str: + def _handle_summarize_session(self, args: dict[str, Any]) -> str: """Tool handler for summarize_session.""" if not self._healthy or not self._store: return json.dumps({"error": "hexus unavailable"}) @@ -1694,10 +1692,10 @@ def _handle_summarize_session(self, args: Dict[str, Any]) -> str: agent_identity=self._agent_identity, ) return json.dumps(res) - except Exception as exc: + except Exception as exc: # noqa: BLE001 return json.dumps({"error": f"db: {exc}"}) - def _handle_headroom_retrieve(self, args: Dict[str, Any]) -> str: + def _handle_headroom_retrieve(self, args: dict[str, Any]) -> str: """Tool handler for headroom_retrieve.""" if not self._healthy or not self._store: return json.dumps({"error": "hexus unavailable"}) @@ -1715,10 +1713,10 @@ def _handle_headroom_retrieve(self, args: Dict[str, Any]) -> str: if content is None: return json.dumps({"id": entry_id, "found": False, "content": None}) return json.dumps({"id": entry_id, "found": True, "content": content}) - except Exception as exc: + except Exception as exc: # noqa: BLE001 return json.dumps({"error": f"db: {exc}"}) - def _handle_memory_stats(self, args: Dict[str, Any]) -> str: + def _handle_memory_stats(self, args: dict[str, Any]) -> str: """Tool handler for memory_stats.""" if not self._healthy or not self._store: return json.dumps({"error": "hexus unavailable"}) @@ -1727,12 +1725,12 @@ def _handle_memory_stats(self, args: Dict[str, Any]) -> str: res = memory_stats(self._store, args) return json.dumps(res) - except Exception as exc: + except Exception as exc: # noqa: BLE001 return json.dumps({"error": f"stats check failed: {exc}"}) # -- Setup hooks --------------------------------------------------------- - def get_config_schema(self) -> List[Dict[str, Any]]: + def get_config_schema(self) -> list[dict[str, Any]]: return [ { "key": "dsn", @@ -1809,14 +1807,14 @@ def get_config_schema(self) -> List[Dict[str, Any]]: }, ] - def save_config(self, values: Dict[str, Any], hermes_home: str) -> None: + def save_config(self, values: dict[str, Any], hermes_home: str) -> None: from pathlib import Path config_path = Path(hermes_home) / "config.yaml" try: import yaml - existing: Dict[str, Any] = {} + existing: dict[str, Any] = {} if config_path.exists(): with open(config_path, encoding="utf-8-sig") as fh: existing = yaml.safe_load(fh) or {} @@ -1829,7 +1827,7 @@ def save_config(self, values: Dict[str, Any], hermes_home: str) -> None: # -- Helpers ------------------------------------------------------------- - def _maybe_embed(self, content: str) -> Optional[List[float]]: + def _maybe_embed(self, content: str) -> list[float] | None: if not self._config.get("embed_on_write", True): return None try: diff --git a/hexus/ccr/cache.py b/hexus/ccr/cache.py index 04218e0..bd8f932 100644 --- a/hexus/ccr/cache.py +++ b/hexus/ccr/cache.py @@ -1,6 +1,5 @@ # Forked from andreab67/hermes-memory-pgvector (BSD-3-Clause) import threading -from typing import Dict, Optional class CCRCache: @@ -8,10 +7,10 @@ class CCRCache: def __init__(self, maxsize: int = 1000): self._maxsize = maxsize - self._cache: Dict[int, str] = {} + self._cache: dict[int, str] = {} self._lock = threading.Lock() - def get(self, memory_id: int) -> Optional[str]: + def get(self, memory_id: int) -> str | None: with self._lock: if memory_id in self._cache: val = self._cache.pop(memory_id) diff --git a/hexus/embed.py b/hexus/embed.py index 31d1582..196fde4 100644 --- a/hexus/embed.py +++ b/hexus/embed.py @@ -21,7 +21,6 @@ import logging import urllib.error import urllib.request -from typing import List, Optional logger = logging.getLogger(__name__) @@ -49,11 +48,11 @@ class EmbeddingError(Exception): def embed( text: str, *, - base_url: Optional[str] = None, - model: Optional[str] = None, + base_url: str | None = None, + model: str | None = None, timeout: float = 10.0, expected_dim: int = EXPECTED_DIM, -) -> List[float]: +) -> list[float]: """Return an embedding for `text`. Dispatch: @@ -88,9 +87,11 @@ def embed( # path (e.g. a CI env that just runs unit tests against a mock # endpoint). from .embedder import ( + DEFAULT_MODEL, get_default_embedder, + ) + from .embedder import ( EmbedderError as _LocalErr, - DEFAULT_MODEL, ) embedder = get_default_embedder(model_name=model or DEFAULT_MODEL) @@ -129,7 +130,7 @@ def embed( def _post( url: str, body: dict, *, timeout: float, expected_dim: int, extract -) -> List[float]: +) -> list[float]: data = json.dumps(body).encode("utf-8") req = urllib.request.Request( url, @@ -159,7 +160,7 @@ def _post( return vec -def to_hexus_literal(vec: List[float]) -> str: +def to_hexus_literal(vec: list[float]) -> str: """Render a Python list of floats as a hexus input literal. psycopg can also handle this via type adapters, but the literal form diff --git a/hexus/embedder.py b/hexus/embedder.py index 72e551b..ca5df5d 100644 --- a/hexus/embedder.py +++ b/hexus/embedder.py @@ -33,7 +33,6 @@ os.environ.setdefault("USER", "agy") import threading from dataclasses import dataclass, replace -from typing import Dict, List, Optional, Tuple logger = logging.getLogger(__name__) @@ -96,7 +95,7 @@ class EmbedStats: tokens_dropped: int = 0 # approx tokens lost to truncation (Σ tc-max_seq) max_tokens_seen: int = 0 # largest single-text token count observed - def as_dict(self) -> Dict[str, int]: + def as_dict(self) -> dict[str, int]: return { "texts_embedded": self.texts_embedded, "texts_over_limit": self.texts_over_limit, @@ -108,7 +107,7 @@ def as_dict(self) -> Dict[str, int]: } -def _resolve_long_text_mode(mode: Optional[str]) -> str: +def _resolve_long_text_mode(mode: str | None) -> str: """Resolve the configured mode: explicit arg > env var > default. An unrecognised value falls back to the default with a warning rather @@ -156,9 +155,9 @@ def __init__( self, model_name: str = DEFAULT_MODEL, *, - cache_dir: Optional[str] = None, + cache_dir: str | None = None, device: str = "cpu", - long_text_mode: Optional[str] = None, + long_text_mode: str | None = None, ): self._model_name = model_name self._cache_dir = cache_dir @@ -201,8 +200,10 @@ def dim(self) -> int: # transformer config. Fall through to the constant if not. try: return int(self._model.get_sentence_embedding_dimension()) - except Exception: # noqa: BLE001 - pass + except Exception as exc: # noqa: BLE001 + logger.debug( + "hexus embedder: failed to get embedding dimension: %s", exc + ) return DEFAULT_DIM if self._model_name == DEFAULT_MODEL else 0 @property @@ -229,7 +230,7 @@ def reset_stats(self) -> None: with self._stats_lock: self._stats = EmbedStats() - def embed(self, texts: List[str]) -> List[List[float]]: + def embed(self, texts: list[str]) -> list[list[float]]: """Embed a list of texts → list of float vectors. Empty / whitespace-only inputs are silently dropped (returned @@ -266,7 +267,7 @@ def embed(self, texts: List[str]) -> List[List[float]]: convert_to_numpy=True, show_progress_bar=False, ) - except Exception as exc: # noqa: BLE001 — fail-soft, surface in logs + except Exception as exc: raise EmbedderError(f"local embed failed: {exc}") from exc # Reassemble one vector per original input (chunk windows collapse @@ -291,8 +292,8 @@ def embed(self, texts: List[str]) -> List[List[float]]: # -- Long-input handling (issue #7) ------------------------------------- def _plan_encode( - self, texts: List[str], model - ) -> Tuple[List[str], List[Tuple[int, int, Optional[List[float]]]]]: + self, texts: list[str], model + ) -> tuple[list[str], list[tuple[int, int, list[float] | None]]]: """Turn `texts` into (`pieces` to encode, `plan` to reassemble them). Each plan entry is `(start, count, weights)`: @@ -313,8 +314,8 @@ def _plan_encode( ) mode = self._long_text_mode - pieces: List[str] = [] - plan: List[Tuple[int, int, Optional[List[float]]]] = [] + pieces: list[str] = [] + plan: list[tuple[int, int, list[float] | None]] = [] for i, text in enumerate(texts): with self._stats_lock: @@ -329,8 +330,7 @@ def _plan_encode( # Over the limit — count it (in every mode) and record the peak. with self._stats_lock: self._stats.texts_over_limit += 1 - if tc > self._stats.max_tokens_seen: - self._stats.max_tokens_seen = tc + self._stats.max_tokens_seen = max(self._stats.max_tokens_seen, tc) over_count = self._stats.texts_over_limit if mode == LONG_TEXT_MODE_CHUNK: @@ -361,8 +361,8 @@ def _plan_encode( return pieces, plan def _assemble( - self, raw, plan: List[Tuple[int, int, Optional[List[float]]]] - ) -> List[List[float]]: + self, raw, plan: list[tuple[int, int, list[float] | None]] + ) -> list[list[float]]: """Collapse encoded `pieces` back to one vector per original input. Single-piece entries are returned verbatim (byte-identical to the @@ -379,7 +379,7 @@ def _assemble( import numpy as np - vectors: List[List[float]] = [] + vectors: list[list[float]] = [] for start, count, weights in plan: if count == 1: vectors.append(raw[start].tolist()) @@ -406,7 +406,7 @@ def _resolve_max_seq(model) -> int: except (TypeError, ValueError): return 0 - def _token_counts(self, texts: List[str], tokenizer) -> Optional[List[int]]: + def _token_counts(self, texts: list[str], tokenizer) -> list[int] | None: """True (untruncated) token count per text, or None if unavailable. `verbose=False` suppresses HuggingFace's "sequence longer than model @@ -427,7 +427,7 @@ def _token_counts(self, texts: List[str], tokenizer) -> Optional[List[int]]: ) return None - def _chunk_text(self, text: str, tokenizer, max_seq: int) -> List[Tuple[str, int]]: + def _chunk_text(self, text: str, tokenizer, max_seq: int) -> list[tuple[str, int]]: """Split `text` into overlapping token windows → [(chunk_text, n_tokens)]. We reserve room for the special tokens the tokenizer re-adds when each @@ -451,7 +451,7 @@ def _chunk_text(self, text: str, tokenizer, max_seq: int) -> List[Tuple[str, int if len(ids) <= window: return [(text, len(ids))] - chunks: List[Tuple[str, int]] = [] + chunks: list[tuple[str, int]] = [] for start in range(0, len(ids), stride): window_ids = ids[start : start + window] if not window_ids: @@ -550,7 +550,7 @@ def _load_model(self): self._device, ) return self._model - except Exception as exc: # noqa: BLE001 + except Exception as exc: self._load_failed = True raise EmbedderError( f"failed to load sentence-transformers model {self._model_name}: {exc}" @@ -564,17 +564,17 @@ def _load_model(self): # Caching is keyed on (model_name, cache_dir, device) so a request for a # different model returns a different embedder (mostly relevant for tests # — production uses one model). The dict is small in practice. -_singletons: dict[tuple[str, Optional[str], str, str], "LocalBertEmbedder"] = {} +_singletons: dict[tuple[str, str | None, str, str], LocalBertEmbedder] = {} _singleton_lock = threading.Lock() def get_default_embedder( model_name: str = DEFAULT_MODEL, *, - cache_dir: Optional[str] = None, - device: Optional[str] = None, - long_text_mode: Optional[str] = None, -) -> "LocalBertEmbedder": + cache_dir: str | None = None, + device: str | None = None, + long_text_mode: str | None = None, +) -> LocalBertEmbedder: """Return the process-wide default embedder for these args, constructing it on first call. Subsequent calls with the same (model_name, cache_dir, device, long_text_mode) return the same instance. @@ -587,7 +587,6 @@ def get_default_embedder( part of the cache key so a request for a different mode returns a distinct embedder rather than silently reusing another mode's instance. """ - global _singletons if device is None: device = os.environ.get("HEXUS_EMBED_DEVICE", "cpu") mode = _resolve_long_text_mode(long_text_mode) diff --git a/hexus/entity_extractor.py b/hexus/entity_extractor.py index 95aff9f..04b95d4 100644 --- a/hexus/entity_extractor.py +++ b/hexus/entity_extractor.py @@ -1,5 +1,4 @@ import re -from typing import List, Dict DEFAULT_PATTERNS = { "url": r'https?://[^\s<>"]+', @@ -14,12 +13,12 @@ class EntityExtractor: - def __init__(self, patterns: Dict[str, str] = None, enabled: bool = True): + def __init__(self, patterns: dict[str, str] | None = None, enabled: bool = True): self.enabled = enabled self.patterns = {**DEFAULT_PATTERNS, **(patterns or {})} self._compiled = {t: re.compile(p) for t, p in self.patterns.items()} - def extract_entities(self, text: str) -> List[Dict[str, str]]: + def extract_entities(self, text: str) -> list[dict[str, str]]: if not self.enabled or not text: return [] diff --git a/hexus/pipeline/router.py b/hexus/pipeline/router.py index 7027704..cf59e5d 100644 --- a/hexus/pipeline/router.py +++ b/hexus/pipeline/router.py @@ -1,6 +1,6 @@ # Forked from andreab67/hermes-memory-pgvector (BSD-3-Clause) -import re import json +import re class ContentRouter: @@ -63,7 +63,7 @@ def _compress_json(self, text: str) -> str: return f"[Compressed JSON Object] keys: {', '.join(keys)}\nSample: {json.dumps(truncated)}" elif isinstance(data, list): return f"[Compressed JSON Array] length: {len(data)}\nFirst item: {json.dumps(data[0]) if data else 'empty'}" - except Exception: + except Exception: # noqa: BLE001, S110 pass return text[: self.threshold_chars] + "\n... [Truncated JSON]" @@ -98,9 +98,9 @@ def _compress_code(self, text: str) -> str: lines = text.splitlines() compressed_lines = [] for line in lines: - if re.match(r"^\s*(def|class|import|from|async\s+def)\b", line): - compressed_lines.append(line) - elif re.match(r"^\s*#.*", line) and len(compressed_lines) < 10: + if re.match(r"^\s*(def|class|import|from|async\s+def)\b", line) or ( + re.match(r"^\s*#.*", line) and len(compressed_lines) < 10 + ): compressed_lines.append(line) if len(compressed_lines) < 3: diff --git a/hexus/store.py b/hexus/store.py index 382a184..62b5925 100644 --- a/hexus/store.py +++ b/hexus/store.py @@ -24,9 +24,9 @@ import os import threading from dataclasses import dataclass, replace -from datetime import datetime, timezone, timedelta +from datetime import UTC, datetime, timedelta from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple +from typing import Any import psycopg from psycopg.rows import dict_row @@ -62,7 +62,7 @@ _RECALL_COUNT_TABLES = frozenset({"memory_entries", "conversations", "delegations"}) -def _resolve_isolation(value: Optional[str] = None) -> str: +def _resolve_isolation(value: str | None = None) -> str: """Resolve the multi-agent isolation policy. 'shared' (default): reads/recall/search may span every agent's memory — @@ -155,7 +155,7 @@ class RerankStats: tokens_dropped: int = 0 # approx doc tokens never scored (truncate + capped) max_tokens_seen: int = 0 # largest single-doc token count observed - def as_dict(self) -> Dict[str, int]: + def as_dict(self) -> dict[str, int]: return { "docs_reranked": self.docs_reranked, "docs_over_limit": self.docs_over_limit, @@ -187,7 +187,7 @@ def reset_rerank_stats() -> None: _rerank_over_limit_total = 0 -def _resolve_rerank_mode(mode: Optional[str]) -> str: +def _resolve_rerank_mode(mode: str | None) -> str: raw = mode or os.environ.get(RERANK_MODE_ENV) or DEFAULT_RERANK_MODE candidate = raw.strip().lower() if candidate not in VALID_RERANK_MODES: @@ -224,7 +224,7 @@ def _cross_encoder_max_len(model) -> int: return RERANK_DEFAULT_MAX_LEN -def _split_doc_for_rerank(doc: str, tokenizer, budget: int) -> Tuple[List[str], int]: +def _split_doc_for_rerank(doc: str, tokenizer, budget: int) -> tuple[list[str], int]: """Split `doc` into ≤RERANK_MAX_PASSAGES overlapping token windows. Returns (passages, tokens_unscored) where tokens_unscored is the tail @@ -232,14 +232,14 @@ def _split_doc_for_rerank(doc: str, tokenizer, budget: int) -> Tuple[List[str], """ try: ids = tokenizer.encode(doc, add_special_tokens=False, verbose=False) - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # noqa: BLE001 return [doc], 0 if len(ids) <= budget: return [doc], 0 overlap = min(RERANK_PASSAGE_OVERLAP_TOKENS, budget - 1) if budget > 1 else 0 stride = max(1, budget - overlap) - passages: List[str] = [] + passages: list[str] = [] covered = 0 for start in range(0, len(ids), stride): window = ids[start : start + budget] @@ -256,8 +256,8 @@ def _split_doc_for_rerank(doc: str, tokenizer, budget: int) -> Tuple[List[str], def rerank_scores( - model, query_text: Optional[str], docs: List[str], *, mode: Optional[str] = None -) -> List[float]: + model, query_text: str | None, docs: list[str], *, mode: str | None = None +) -> list[float]: """Score each (query, doc) with the cross-encoder, one score per doc. Handles docs longer than the cross-encoder window per `mode` @@ -273,7 +273,7 @@ def rerank_scores( # Doc budget = window minus the query and the pair's special tokens. If we # can't tokenize, skip the guard and let the model truncate (old path). - budget: Optional[int] = None + budget: int | None = None if tokenizer is not None: try: max_len = _cross_encoder_max_len(model) @@ -282,15 +282,15 @@ def rerank_scores( ) try: special = tokenizer.num_special_tokens_to_add(pair=True) - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # noqa: BLE001 special = 3 budget = max(1, max_len - q_len - special) except Exception as exc: # noqa: BLE001 — best-effort guard logger.debug("rerank length guard unavailable (%s); truncating", exc) budget = None - pairs: List[List[str]] = [] - plan: List[Tuple[int, int]] = [] # (start, count) into pairs, per doc + pairs: list[list[str]] = [] + plan: list[tuple[int, int]] = [] # (start, count) into pairs, per doc for doc in docs: doc = doc or "" with _rerank_stats_lock: @@ -303,7 +303,7 @@ def rerank_scores( try: d_len = len(tokenizer.encode(doc, add_special_tokens=False, verbose=False)) - except Exception: # noqa: BLE001 + except Exception: # noqa: BLE001 # noqa: BLE001 d_len = 0 if d_len <= budget: pairs.append([query_text, doc]) @@ -313,8 +313,7 @@ def rerank_scores( # Over the doc budget — count it in every mode. with _rerank_stats_lock: _rerank_stats.docs_over_limit += 1 - if d_len > _rerank_stats.max_tokens_seen: - _rerank_stats.max_tokens_seen = d_len + _rerank_stats.max_tokens_seen = max(_rerank_stats.max_tokens_seen, d_len) _rerank_over_limit_total += 1 over_count = _rerank_over_limit_total @@ -361,7 +360,7 @@ def rerank_scores( ) raw = model.predict(pairs) - scores: List[float] = [] + scores: list[float] = [] for start, count in plan: if count == 1: scores.append(float(raw[start])) @@ -383,8 +382,8 @@ def __init__( max_idle: float = 30.0, max_lifetime: float = 300.0, entity_extractor_enabled: bool = True, - entity_extractor_patterns: Optional[Dict[str, str]] = None, - isolation: Optional[str] = None, + entity_extractor_patterns: dict[str, str] | None = None, + isolation: str | None = None, ): """Open a lazily-initialized, self-draining ConnectionPool. @@ -404,7 +403,7 @@ def __init__( """ self._dsn = dsn self._lock = threading.Lock() - self._pool: Optional[ConnectionPool] = None + self._pool: ConnectionPool | None = None self._min_size = min_size self._max_size = max_size self._timeout = timeout @@ -418,7 +417,7 @@ def __init__( if env_patterns is not None: try: entity_extractor_patterns = json.loads(env_patterns) - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.warning( "Failed to parse HEXUS_ENTITY_EXTRACTOR_PATTERNS: %s", exc ) @@ -483,7 +482,7 @@ class SchemaNotApplied(RuntimeError): def ensure_schema(self) -> None: """Verify the schema is in place. Does NOT run DDL.""" - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute("SELECT to_regclass('memory_entries')") if cur.fetchone()[0] is None: @@ -507,7 +506,7 @@ def adapt_vector_precision(self) -> None: # Determine target type target_type = "halfvec(384)" if precision == "float16" else "vector(384)" - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: # Check current type of memory_entries.embedding cur.execute(""" @@ -557,7 +556,7 @@ def adapt_vector_precision(self) -> None: logger.info( "Successfully altered database columns to %s", target_type ) - except Exception as exc: + except Exception as exc: # noqa: BLE001 conn.rollback() logger.warning( "Failed to alter database column types (insufficient permissions?): %s", @@ -651,7 +650,7 @@ def adapt_vector_precision(self) -> None: WITH (m = 16, ef_construction = 64); """) conn.commit() - except Exception as exc: + except Exception as exc: # noqa: BLE001 conn.rollback() logger.warning( "Failed to create/ensure quantization indexes: %s", exc @@ -790,7 +789,7 @@ def apply_migration_as_admin(self, *, admin_dsn: str) -> None: """One-shot admin path: run the full migrations with privileged creds.""" migrations_dir = Path(__file__).parent / "migrations" sql_files = sorted(migrations_dir.glob("*.sql")) - with psycopg.connect(admin_dsn, autocommit=True) as conn: + with psycopg.connect(admin_dsn, autocommit=True) as conn: # noqa: SIM117 with conn.cursor() as cur: for sql_file in sql_files: sql = sql_file.read_text(encoding="utf-8") @@ -806,11 +805,11 @@ def add( agent_identity: str, target: str, content: str, - embedding: Optional[List[float]] = None, - metadata: Optional[Dict[str, Any]] = None, - compressed: Optional[str] = None, - content_hash: Optional[bytes] = None, - ) -> Optional[int]: + embedding: list[float] | None = None, + metadata: dict[str, Any] | None = None, + compressed: str | None = None, + content_hash: bytes | None = None, + ) -> int | None: """Insert a memory entry. Returns row id, or None if duplicate (no-op).""" meta = dict(metadata or {}) if "entities" not in meta: @@ -825,7 +824,7 @@ def add( vec_literal = to_hexus_literal(embedding) if embedding is not None else None - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: # Deduplication check: does a row with this content_hash, target and agent_identity exist? cur.execute( @@ -930,9 +929,9 @@ def replace( target: str, old_text: str, new_content: str, - new_embedding: Optional[List[float]] = None, - compressed: Optional[str] = None, - content_hash: Optional[bytes] = None, + new_embedding: list[float] | None = None, + compressed: str | None = None, + content_hash: bytes | None = None, ) -> int: """Update entries in (agent_identity, target) where content contains old_text.""" vec_literal = ( @@ -947,7 +946,7 @@ def replace( like_pattern = f"%{_escape_like(old_text)}%" - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: # Find matching rows to update cache cur.execute( @@ -1007,7 +1006,7 @@ def remove( """ like_pattern = f"%{_escape_like(old_text)}%" - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute( r""" @@ -1027,18 +1026,18 @@ def remove( def list_entries( self, *, - agent_identity: Optional[str] = None, - target: Optional[str] = None, + agent_identity: str | None = None, + target: str | None = None, limit: int = 100, offset: int = 0, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """List entries in an agent's scope. agent_identity=None/empty → list across ALL agents (matches `search` and `count`). target=None → both stores. """ - clauses: List[str] = [] - params: List[Any] = [] + clauses: list[str] = [] + params: list[Any] = [] if agent_identity: clauses.append("agent_identity = %s") params.append(agent_identity) @@ -1049,7 +1048,7 @@ def list_entries( params.append(limit) params.append(offset) - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute( f""" @@ -1067,18 +1066,18 @@ def list_entries( def search( self, *, - query_embedding: List[float], - agent_identity: Optional[str] = None, - target: Optional[str] = None, + query_embedding: list[float], + agent_identity: str | None = None, + target: str | None = None, limit: int = 5, min_similarity: float = 0.0, decay_half_life_days: float = 0.0, recall_boost_weight: float = 0.0, - platform: Optional[str] = None, + platform: str | None = None, min_confidence: float = 0.0, rerank: bool = False, - query_text: Optional[str] = None, - ) -> List[Dict[str, Any]]: + query_text: str | None = None, + ) -> list[dict[str, Any]]: """Semantic recall via cosine distance. agent_identity=None → search across ALL agents (cross-theme recall). @@ -1086,8 +1085,8 @@ def search( Returns rows with `score` = 1 - cosine_distance ∈ [0, 1]. """ vec_literal = to_hexus_literal(query_embedding) - clauses: List[str] = [] - params: List[Any] = [] + clauses: list[str] = [] + params: list[Any] = [] if agent_identity: clauses.append("agent_identity = %s") params.append(agent_identity) @@ -1109,7 +1108,7 @@ def search( db_limit = max(limit * 10, 50) if rerank: db_limit = max(db_limit, 100) - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute( f""" @@ -1187,20 +1186,20 @@ def search( def hybrid_search( self, *, - query_embedding: List[float], + query_embedding: list[float], query_text: str, - agent_identity: Optional[str] = None, - target: Optional[str] = None, + agent_identity: str | None = None, + target: str | None = None, limit: int = 5, vector_weight: float = 0.7, text_weight: float = 0.3, min_similarity: float = 0.0, decay_half_life_days: float = 0.0, recall_boost_weight: float = 0.0, - platform: Optional[str] = None, + platform: str | None = None, min_confidence: float = 0.0, rerank: bool = False, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Blend semantic vector search and full-text search.""" if not query_text or not query_text.strip(): rows = self.search( @@ -1223,8 +1222,8 @@ def hybrid_search( vec_literal = to_hexus_literal(query_embedding) - clauses: List[str] = [] - params: List[Any] = [] + clauses: list[str] = [] + params: list[Any] = [] if agent_identity: clauses.append("agent_identity = %s") params.append(agent_identity) @@ -1274,23 +1273,20 @@ def hybrid_search( LIMIT %s """ - v_params = [vec_literal] - for p in params: - v_params.append(p) - v_params.extend([vec_literal, db_limit]) + v_params = [vec_literal, *params, vec_literal, db_limit] t_params = [query_text, query_text] + params + [db_limit] all_params = v_params + t_params + [vector_weight, text_weight, db_limit] - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute(sql, all_params) rows = list(cur.fetchall()) # Blended Score Calculation: # Combined Score = 0.6 * S_vector + 0.3 * S_BM25 + 0.1 * S_recency - now = datetime.now(timezone.utc) + now = datetime.now(UTC) for r in rows: ts_val = r.get("updated_at") or r.get("ts") or r.get("created_at") if isinstance(ts_val, str): @@ -1298,11 +1294,11 @@ def hybrid_search( from datetime import datetime as dt ts_val = dt.fromisoformat(ts_val) - except Exception: + except Exception: # noqa: BLE001, S110 pass if ts_val: if ts_val.tzinfo is None: - ts_val = ts_val.replace(tzinfo=timezone.utc) + ts_val = ts_val.replace(tzinfo=UTC) age_days = (now - ts_val).total_seconds() / 86400.0 if decay_half_life_days > 0.0: r["recency_score"] = math.exp( @@ -1348,8 +1344,8 @@ def hybrid_search( return rows def fetch_full( - self, memory_id: int, agent_identity: Optional[str] = None - ) -> Optional[str]: + self, memory_id: int, agent_identity: str | None = None + ) -> str | None: """Fetch the original full content of a memory entry, checking CCRCache first. When isolation is 'strict' and an `agent_identity` is supplied, the @@ -1365,7 +1361,7 @@ def fetch_full( if cached is not None: return cached - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: if scoped: cur.execute( @@ -1396,9 +1392,9 @@ def bulk_upsert_md( *, agent_identity: str, target: str, - file_path: "Path | str", + file_path: Path | str, embed_fn, - ) -> Dict[str, int]: + ) -> dict[str, int]: """Parse a MEMORY.md / USER.md file and upsert each entry. Idempotent + cheap on re-run: we SELECT the existing content set @@ -1426,7 +1422,7 @@ def bulk_upsert_md( # Single bulk SELECT of existing content for this scope. Beats N+1 # by a wide margin and keeps re-init nearly free. - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute( "SELECT content FROM memory_entries WHERE agent_identity = %s AND target = %s", @@ -1443,7 +1439,7 @@ def bulk_upsert_md( vec = None try: vec = embed_fn(entry) if embed_fn else None - except Exception: # noqa: BLE001 — fail-soft on bulk embed + except Exception: # noqa: BLE001 # noqa: BLE001 — fail-soft on bulk embed vec = None row_id = self.add( agent_identity=agent_identity, @@ -1468,8 +1464,8 @@ def append_turn( agent_identity: str, role: str, content: str, - embedding: Optional[List[float]] = None, - metadata: Optional[Dict[str, Any]] = None, + embedding: list[float] | None = None, + metadata: dict[str, Any] | None = None, ) -> int: """Insert one chat turn. Returns row id. @@ -1484,7 +1480,7 @@ def append_turn( meta_json = json.dumps(meta) vec_literal = to_hexus_literal(embedding) if embedding is not None else None - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute( """ @@ -1502,19 +1498,19 @@ def append_turn( def search_turns( self, *, - query_embedding: List[float], - agent_identity: Optional[str] = None, - session_id: Optional[str] = None, + query_embedding: list[float], + agent_identity: str | None = None, + session_id: str | None = None, limit: int = 5, min_similarity: float = 0.0, decay_half_life_days: float = 0.0, recall_boost_weight: float = 0.0, - platform: Optional[str] = None, - ) -> List[Dict[str, Any]]: + platform: str | None = None, + ) -> list[dict[str, Any]]: """Semantic recall over conversation turns. Same shape as `search()`.""" vec_literal = to_hexus_literal(query_embedding) - clauses: List[str] = [] - params: List[Any] = [] + clauses: list[str] = [] + params: list[Any] = [] if agent_identity: clauses.append("agent_identity = %s") params.append(agent_identity) @@ -1534,7 +1530,7 @@ def search_turns( if self._vector_precision == "binary": db_limit = max(limit * 10, 50) - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute( f""" @@ -1549,7 +1545,7 @@ def search_turns( ) rows = list(cur.fetchall()) else: - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute( f""" @@ -1586,18 +1582,18 @@ def search_turns( def hybrid_search_turns( self, *, - query_embedding: List[float], + query_embedding: list[float], query_text: str, - agent_identity: Optional[str] = None, - session_id: Optional[str] = None, + agent_identity: str | None = None, + session_id: str | None = None, limit: int = 5, vector_weight: float = 0.7, text_weight: float = 0.3, min_similarity: float = 0.0, decay_half_life_days: float = 0.0, recall_boost_weight: float = 0.0, - platform: Optional[str] = None, - ) -> List[Dict[str, Any]]: + platform: str | None = None, + ) -> list[dict[str, Any]]: """Blend semantic vector search and full-text search over conversation turns.""" if not query_text or not query_text.strip(): rows = self.search_turns( @@ -1616,8 +1612,8 @@ def hybrid_search_turns( return rows vec_literal = to_hexus_literal(query_embedding) - clauses: List[str] = [] - params: List[Any] = [] + clauses: list[str] = [] + params: list[Any] = [] if agent_identity: clauses.append("agent_identity = %s") params.append(agent_identity) @@ -1664,16 +1660,13 @@ def hybrid_search_turns( LIMIT %s """ - v_params = [vec_literal] - for p in params: - v_params.append(p) - v_params.extend([vec_literal, limit]) + v_params = [vec_literal, *params, vec_literal, limit] t_params = [query_text, query_text] + params + [limit] all_params = v_params + t_params + [vector_weight, text_weight, limit] - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute(sql, all_params) rows = list(cur.fetchall()) @@ -1701,13 +1694,13 @@ def record_delegation( agent_identity: str = "default", task: str, result: str, - embedding: Optional[List[float]] = None, - metadata: Optional[Dict[str, Any]] = None, + embedding: list[float] | None = None, + metadata: dict[str, Any] | None = None, ) -> int: """Insert a delegation entry. Returns row id.""" vec_literal = to_hexus_literal(embedding) if embedding is not None else None meta_json = json.dumps(metadata or {}) - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute( """ @@ -1733,19 +1726,19 @@ def record_delegation( def search_delegations( self, *, - query_embedding: List[float], - agent_identity: Optional[str] = None, - parent_session_id: Optional[str] = None, + query_embedding: list[float], + agent_identity: str | None = None, + parent_session_id: str | None = None, limit: int = 5, min_similarity: float = 0.0, decay_half_life_days: float = 0.0, recall_boost_weight: float = 0.0, - platform: Optional[str] = None, - ) -> List[Dict[str, Any]]: + platform: str | None = None, + ) -> list[dict[str, Any]]: """Semantic recall over delegations.""" vec_literal = to_hexus_literal(query_embedding) - clauses: List[str] = [] - params: List[Any] = [] + clauses: list[str] = [] + params: list[Any] = [] if agent_identity: clauses.append("agent_identity = %s") params.append(agent_identity) @@ -1765,7 +1758,7 @@ def search_delegations( if self._vector_precision == "binary": db_limit = max(limit * 10, 50) - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute( f""" @@ -1780,7 +1773,7 @@ def search_delegations( ) rows = list(cur.fetchall()) else: - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute( f""" @@ -1817,11 +1810,11 @@ def search_delegations( def cleanup_stale_records( self, *, - conversations_ttl_days: Optional[int] = None, - memories_ttl_days: Optional[int] = None, - delegations_ttl_days: Optional[int] = None, + conversations_ttl_days: int | None = None, + memories_ttl_days: int | None = None, + delegations_ttl_days: int | None = None, dry_run: bool = False, - ) -> Dict[str, int]: + ) -> dict[str, int]: """Delete records older than the specified TTL. Returns counts of deleted items. @@ -1837,12 +1830,12 @@ def cleanup_stale_records( ("memory_entries", "updated_at", memories_ttl_days), ("delegations", "ts", delegations_ttl_days), ] - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: for table, ts_col, ttl in targets: if ttl is None or ttl <= 0: continue - limit_date = datetime.now(timezone.utc) - timedelta(days=ttl) + limit_date = datetime.now(UTC) - timedelta(days=ttl) assert table in { "conversations", "memory_entries", @@ -1873,11 +1866,11 @@ def cleanup_stale_records( def count_turns( self, *, - agent_identity: Optional[str] = None, - session_id: Optional[str] = None, + agent_identity: str | None = None, + session_id: str | None = None, ) -> int: - clauses: List[str] = [] - params: List[Any] = [] + clauses: list[str] = [] + params: list[Any] = [] if agent_identity: clauses.append("agent_identity = %s") params.append(agent_identity) @@ -1885,7 +1878,7 @@ def count_turns( clauses.append("session_id = %s") params.append(session_id) where = ("WHERE " + " AND ".join(clauses)) if clauses else "" - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute(f"SELECT COUNT(*) FROM conversations {where}", params) return int(cur.fetchone()[0]) @@ -1893,11 +1886,11 @@ def count_turns( def count( self, *, - agent_identity: Optional[str] = None, - target: Optional[str] = None, + agent_identity: str | None = None, + target: str | None = None, ) -> int: - clauses: List[str] = [] - params: List[Any] = [] + clauses: list[str] = [] + params: list[Any] = [] if agent_identity: clauses.append("agent_identity = %s") params.append(agent_identity) @@ -1906,15 +1899,15 @@ def count( params.append(target) where = ("WHERE " + " AND ".join(clauses)) if clauses else "" - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute(f"SELECT COUNT(*) FROM memory_entries {where}", params) return int(cur.fetchone()[0]) - def health(self) -> Dict[str, Any]: + def health(self) -> dict[str, Any]: """Liveness probe — pool reachable + table exists. Never raises.""" try: - with self._get_pool().connection(timeout=3.0) as conn: + with self._get_pool().connection(timeout=3.0) as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute("SELECT to_regclass('memory_entries') IS NOT NULL") has_table = bool(cur.fetchone()[0]) @@ -1933,7 +1926,7 @@ def health(self) -> Dict[str, Any]: except Exception as exc: # noqa: BLE001 return {"ok": False, "error": str(exc)[:200], "row_count": 0} - def get_metrics_data(self) -> Dict[str, Any]: + def get_metrics_data(self) -> dict[str, Any]: """Fetch detailed metrics from the database for Prometheus output.""" data = { "memory_entries": [], @@ -1948,13 +1941,13 @@ def get_metrics_data(self) -> Dict[str, Any]: } # Helper to execute query and return list of rows - def query_safe(sql: str, params: Optional[list] = None) -> list: + def query_safe(sql: str, params: list | None = None) -> list: try: - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute(sql, params or []) return list(cur.fetchall()) - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.warning("Metrics query failed (%s): %s", sql[:50], exc) return [] @@ -2042,8 +2035,8 @@ def query_safe(sql: str, params: Optional[list] = None) -> list: return data def _apply_recall_boost( - self, rows: List[Dict[str, Any]], boost_weight: float - ) -> List[Dict[str, Any]]: + self, rows: list[dict[str, Any]], boost_weight: float + ) -> list[dict[str, Any]]: if boost_weight <= 0.0: return rows for r in rows: @@ -2057,11 +2050,11 @@ def _apply_recall_boost( return rows def _apply_temporal_decay( - self, rows: List[Dict[str, Any]], half_life_days: float - ) -> List[Dict[str, Any]]: + self, rows: list[dict[str, Any]], half_life_days: float + ) -> list[dict[str, Any]]: if half_life_days <= 0.0: return rows - now = datetime.now(timezone.utc) + now = datetime.now(UTC) for r in rows: ts_val = r.get("updated_at") or r.get("ts") or r.get("created_at") if isinstance(ts_val, str): @@ -2069,14 +2062,14 @@ def _apply_temporal_decay( from datetime import datetime as dt ts_val = dt.fromisoformat(ts_val) - except Exception: + except Exception: # noqa: BLE001, S112 continue if not ts_val: continue # Ensure ts_val has timezone info (psycopg datetimes are timezone-aware, now is utc) if ts_val.tzinfo is None: - ts_val = ts_val.replace(tzinfo=timezone.utc) + ts_val = ts_val.replace(tzinfo=UTC) age_days = (now - ts_val).total_seconds() / 86400.0 # exponential decay: score * 2^(-age/half_life) @@ -2085,8 +2078,8 @@ def _apply_temporal_decay( return rows def _apply_min_confidence( - self, rows: List[Dict[str, Any]], min_confidence: float - ) -> List[Dict[str, Any]]: + self, rows: list[dict[str, Any]], min_confidence: float + ) -> list[dict[str, Any]]: if min_confidence <= 0.0: return rows filtered = [] @@ -2109,9 +2102,7 @@ def _apply_min_confidence( filtered.append(r) return filtered - def confirm_entry( - self, entry_id: int, agent_identity: Optional[str] = None - ) -> bool: + def confirm_entry(self, entry_id: int, agent_identity: str | None = None) -> bool: """Increment confirm_count in metadata JSONB for the given entry ID. When `agent_identity` is supplied the mutation is scoped to that agent, @@ -2120,7 +2111,7 @@ def confirm_entry( """ return self._bump_entry_count(entry_id, "confirm_count", agent_identity) - def reject_entry(self, entry_id: int, agent_identity: Optional[str] = None) -> bool: + def reject_entry(self, entry_id: int, agent_identity: str | None = None) -> bool: """Increment reject_count in metadata JSONB for the given entry ID. Scoped to `agent_identity` when supplied — see `confirm_entry`. @@ -2128,7 +2119,7 @@ def reject_entry(self, entry_id: int, agent_identity: Optional[str] = None) -> b return self._bump_entry_count(entry_id, "reject_count", agent_identity) def _bump_entry_count( - self, entry_id: int, field: str, agent_identity: Optional[str] + self, entry_id: int, field: str, agent_identity: str | None ) -> bool: # `field` is one of the fixed literals passed by confirm/reject above, # never caller input, so interpolating it into the JSONB path is safe. @@ -2143,18 +2134,18 @@ def _bump_entry_count( ) WHERE id = %s """ - params: List[Any] = [entry_id] + params: list[Any] = [entry_id] if agent_identity is not None: sql += " AND agent_identity = %s\n" params.append(agent_identity) - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute(sql, tuple(params)) updated = cur.rowcount conn.commit() return updated > 0 - def increment_recall_counts(self, table: str, ids: List[int]) -> None: + def increment_recall_counts(self, table: str, ids: list[int]) -> None: if not ids: return if table not in _RECALL_COUNT_TABLES: @@ -2173,11 +2164,11 @@ def increment_recall_counts(self, table: str, ids: List[int]) -> None: WHERE id = ANY(%s) """ try: - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute(sql, (ids,)) conn.commit() - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.warning("Failed to increment recall counts for %s: %s", table, exc) def entity_graph( @@ -2185,9 +2176,9 @@ def entity_graph( *, entity_type: str, entity_value: str, - agent_identity: Optional[str] = None, + agent_identity: str | None = None, limit: int = 5, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Find other entities that co-occur with the given entity.""" query_entity_json = json.dumps([{"type": entity_type, "value": entity_value}]) @@ -2222,7 +2213,7 @@ def entity_graph( entity_value, limit, ] - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute(sql, params) related = list(cur.fetchall()) @@ -2236,10 +2227,10 @@ def graph_walk( *, entity_type: str, entity_value: str, - agent_identity: Optional[str] = None, + agent_identity: str | None = None, max_depth: int = 2, limit: int = 5, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Perform recursive CTE path traversal from a start entity.""" sql = """ WITH RECURSIVE graph_walk AS ( @@ -2297,7 +2288,7 @@ def graph_walk( entity_value, limit, ] - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute(sql, params) return list(cur.fetchall()) @@ -2305,10 +2296,10 @@ def graph_walk( def common_topics( self, *, - agent_identity: Optional[str] = None, + agent_identity: str | None = None, min_strength: int = 2, limit: int = 10, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Find clusters of heavily co-occurring entities/topics.""" sql = """ SELECT @@ -2328,7 +2319,7 @@ def common_topics( LIMIT %s """ params = [agent_identity, agent_identity, min_strength, limit] - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor(row_factory=dict_row) as cur: cur.execute(sql, params) return list(cur.fetchall()) @@ -2338,8 +2329,8 @@ def summarize_session( *, session_id: str, limit: int = 5, - agent_identity: Optional[str] = None, - ) -> Dict[str, Any]: + agent_identity: str | None = None, + ) -> dict[str, Any]: """Compute the vector centroid of a session's turns and find the K closest turns. When `agent_identity` is supplied the session lookup is scoped to that @@ -2397,13 +2388,14 @@ def summarize_session( # that runs when the agent is idle to group/summarize low-confidence or heavily co-occurring memory entries. def consolidate_low_confidence_memories( - self, agent_identity: Optional[str] = None - ) -> Dict[str, Any]: + self, agent_identity: str | None = None + ) -> dict[str, Any]: """Query low-confidence (frequently rejected) memories and send them to the LLM for pruning/merging.""" import json - import urllib.request - import urllib.error import os + import urllib.error + import urllib.request + from hexus.store import dict_row api_base = os.environ.get("LLM_API_BASE") or "http://headroom:8787/v1" @@ -2435,14 +2427,14 @@ def consolidate_low_confidence_memories( params.append(agent_identity) query += " LIMIT 20" - with self._get_pool().connection() as conn: - with ( - conn.cursor(row_factory=dict_row) - if hasattr(conn, "cursor") - else conn.cursor() as cur - ): - cur.execute(query, params) - rows = list(cur.fetchall()) + with ( + self._get_pool().connection() as conn, + conn.cursor(row_factory=dict_row) + if hasattr(conn, "cursor") + else conn.cursor() as cur, + ): + cur.execute(query, params) + rows = list(cur.fetchall()) if not rows: return {"status": "ok", "processed": 0, "deletions": 0, "replacements": 0} @@ -2493,7 +2485,7 @@ def consolidate_low_confidence_memories( with urllib.request.urlopen(req, timeout=30) as resp: resp_data = json.loads(resp.read().decode("utf-8")) llm_response = resp_data["choices"][0]["message"]["content"].strip() - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.error( "Failed to query LLM for low-confidence memory consolidation: %s", exc ) @@ -2509,7 +2501,7 @@ def consolidate_low_confidence_memories( lines = lines[:-1] llm_response = "\n".join(lines).strip() data = json.loads(llm_response) - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.error( "Failed to parse LLM consolidation response JSON: %s. Response: %r", exc, @@ -2527,7 +2519,7 @@ def consolidate_low_confidence_memories( if deletions: valid_del_ids = [r["id"] for r in rows if r["id"] in deletions] if valid_del_ids: - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute( "DELETE FROM memory_entries WHERE id = ANY(%s)", @@ -2544,7 +2536,7 @@ def consolidate_low_confidence_memories( valid_rep_ids = [r["id"] for r in rows if r["id"] in rep_ids] if valid_rep_ids and rep_content: - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute( "DELETE FROM memory_entries WHERE id = ANY(%s)", @@ -2568,13 +2560,14 @@ def consolidate_low_confidence_memories( } def consolidate_cooccurring_memories( - self, agent_identity: Optional[str] = None - ) -> Dict[str, Any]: + self, agent_identity: str | None = None + ) -> dict[str, Any]: """Query clusters of co-occurring entities and consolidate their memories using the LLM.""" import json - import urllib.request - import urllib.error import os + import urllib.error + import urllib.request + from hexus.store import dict_row api_base = os.environ.get("LLM_API_BASE") or "http://headroom:8787/v1" @@ -2615,14 +2608,14 @@ def consolidate_cooccurring_memories( meta_b = json.dumps({"entities": [{"type": type_b, "value": val_b}]}) params = [agent_identity, agent_identity, meta_a, meta_b] - with self._get_pool().connection() as conn: - with ( - conn.cursor(row_factory=dict_row) - if hasattr(conn, "cursor") - else conn.cursor() as cur - ): - cur.execute(query, params) - rows = list(cur.fetchall()) + with ( + self._get_pool().connection() as conn, + conn.cursor(row_factory=dict_row) + if hasattr(conn, "cursor") + else conn.cursor() as cur, + ): + cur.execute(query, params) + rows = list(cur.fetchall()) rows = [r for r in rows if r["id"] not in processed_ids] @@ -2675,7 +2668,7 @@ def consolidate_cooccurring_memories( with urllib.request.urlopen(req, timeout=30) as resp: resp_data = json.loads(resp.read().decode("utf-8")) llm_response = resp_data["choices"][0]["message"]["content"].strip() - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.error("Failed to query LLM for topic consolidation: %s", exc) continue @@ -2688,7 +2681,7 @@ def consolidate_cooccurring_memories( lines = lines[:-1] llm_response = "\n".join(lines).strip() data = json.loads(llm_response) - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.error( "Failed to parse LLM topic consolidation JSON: %s. Response: %r", exc, @@ -2702,7 +2695,7 @@ def consolidate_cooccurring_memories( valid_ids = [r["id"] for r in rows if r["id"] in ids_to_replace] if len(valid_ids) >= 2 and consolidated_content: - with self._get_pool().connection() as conn: + with self._get_pool().connection() as conn: # noqa: SIM117 with conn.cursor() as cur: cur.execute( "DELETE FROM memory_entries WHERE id = ANY(%s)", diff --git a/hexus/webhook/dispatcher.py b/hexus/webhook/dispatcher.py index 990aa4d..b29273e 100644 --- a/hexus/webhook/dispatcher.py +++ b/hexus/webhook/dispatcher.py @@ -2,10 +2,11 @@ import hmac import json import logging -import time import threading +import time +from typing import Any + import requests -from typing import Any, Dict, Optional logger = logging.getLogger(__name__) @@ -17,9 +18,9 @@ def sign_payload(payload_bytes: bytes, secret: str) -> str: def dispatch_webhook_sync( url: str, - secret: Optional[str], + secret: str | None, event: str, - payload: Dict[str, Any], + payload: dict[str, Any], max_retries: int = 3, initial_backoff: float = 1.0, ) -> None: @@ -32,7 +33,7 @@ def dispatch_webhook_sync( try: body = json.dumps(full_payload) - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.error("Failed to serialize webhook payload for event %s: %s", event, exc) return @@ -75,11 +76,10 @@ def dispatch_webhook_sync( attempt + 1, exc, ) - except Exception as exc: + except Exception: logger.exception( - "Unexpected error in webhook dispatch thread for event %s: %s", + "Unexpected error in webhook dispatch thread for event %s", event, - exc, ) return @@ -96,10 +96,10 @@ def dispatch_webhook_sync( def dispatch_webhook( - url: Optional[str], - secret: Optional[str], + url: str | None, + secret: str | None, event: str, - payload: Dict[str, Any], + payload: dict[str, Any], ) -> None: """Asynchronously dispatch webhook in a background daemon thread.""" if not url: diff --git a/hexus/writer.py b/hexus/writer.py index 1bb8173..cdf3906 100644 --- a/hexus/writer.py +++ b/hexus/writer.py @@ -23,8 +23,9 @@ import logging import queue import threading -from typing import Any, Callable, Dict, Optional import weakref +from collections.abc import Callable +from typing import Any logger = logging.getLogger(__name__) @@ -33,7 +34,7 @@ # Pending-write payload — kept small so the queue stays bounded. class _PendingWrite: - __slots__ = ("action", "agent_identity", "target", "content", "extra", "metadata") + __slots__ = ("action", "agent_identity", "content", "extra", "metadata", "target") def __init__( self, @@ -42,8 +43,8 @@ def __init__( agent_identity: str, target: str, content: str, - extra: Optional[Dict[str, Any]] = None, - metadata: Optional[Dict[str, Any]] = None, + extra: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, ): self.action = action self.agent_identity = agent_identity @@ -68,10 +69,8 @@ def __init__( self, worker_fn: Callable[[_PendingWrite], None], *, maxsize: int = 256 ): self._worker_fn = worker_fn - self._queue: "queue.Queue[Optional[_PendingWrite]]" = queue.Queue( - maxsize=maxsize - ) - self._thread: Optional[threading.Thread] = None + self._queue: queue.Queue[_PendingWrite | None] = queue.Queue(maxsize=maxsize) + self._thread: threading.Thread | None = None self._stop = threading.Event() self._dropped = 0 self._dropped_warned = False @@ -88,8 +87,8 @@ def enqueue( agent_identity: str, target: str, content: str, - extra: Optional[Dict[str, Any]] = None, - metadata: Optional[Dict[str, Any]] = None, + extra: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, ) -> bool: """Enqueue a write. Returns True on accept, False on drop (queue full). @@ -138,7 +137,7 @@ def shutdown(self, timeout: float = 5.0) -> None: if self._thread.is_alive(): logger.warning("hexus writer thread did not drain within %.1fs", timeout) - def stats(self) -> Dict[str, Any]: + def stats(self) -> dict[str, Any]: with self._lock: lats = list(self._latencies) p50 = float("nan") diff --git a/mcp_server/cli.py b/mcp_server/cli.py index 78becd7..a4f3891 100644 --- a/mcp_server/cli.py +++ b/mcp_server/cli.py @@ -21,7 +21,6 @@ import logging import os import sys -from typing import List, Optional logger = logging.getLogger("mcp_server") @@ -132,6 +131,7 @@ def cmd_doctor(args: argparse.Namespace) -> int: return 2 from hexus.store import MemoryStore + from . import tools _configure_logging(args.log_level) @@ -152,7 +152,7 @@ def _redact_dsn(dsn: str) -> str: return re.sub(r"(password\s*=\s*)([^\s]+)", r"\1***", dsn, flags=re.IGNORECASE) -def main(argv: Optional[List[str]] = None) -> int: +def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( prog="hexus-mcp", description=( diff --git a/mcp_server/import_cli.py b/mcp_server/import_cli.py index 12620cc..57cd896 100644 --- a/mcp_server/import_cli.py +++ b/mcp_server/import_cli.py @@ -6,9 +6,9 @@ import json import os import sys -from typing import List, Optional -from hexus.store import MemoryStore + from hexus.embed import embed +from hexus.store import MemoryStore def import_mem0( @@ -26,7 +26,7 @@ def import_mem0( try: with open(file_path, "r", encoding="utf-8") as f: data = json.load(f) - except Exception as exc: + except Exception as exc: # noqa: BLE001 print(f"ERROR: Failed to parse JSON: {exc}", file=sys.stderr) sys.exit(1) @@ -53,7 +53,7 @@ def import_mem0( try: vec = embed(content) - except Exception as exc: + except Exception as exc: # noqa: BLE001 print( f"Warning: Failed to embed '{content[:40]}...': {exc}. Inserting without embedding." ) @@ -71,7 +71,7 @@ def import_mem0( success += 1 else: skipped += 1 - except Exception as exc: + except Exception as exc: # noqa: BLE001 print(f"ERROR: Failed to insert item {idx}: {exc}", file=sys.stderr) errors += 1 @@ -93,7 +93,7 @@ def import_honcho( try: with open(file_path, "r", encoding="utf-8") as f: data = json.load(f) - except Exception as exc: + except Exception as exc: # noqa: BLE001 print(f"ERROR: Failed to parse JSON: {exc}", file=sys.stderr) sys.exit(1) @@ -121,7 +121,7 @@ def import_honcho( try: vec = embed(content) - except Exception as exc: + except Exception as exc: # noqa: BLE001 print( f"Warning: Failed to embed '{content[:40]}...': {exc}. Inserting without embedding." ) @@ -139,7 +139,7 @@ def import_honcho( success += 1 else: skipped += 1 - except Exception as exc: + except Exception as exc: # noqa: BLE001 print(f"ERROR: Failed to insert item {idx}: {exc}", file=sys.stderr) errors += 1 @@ -160,7 +160,7 @@ def import_holographic( try: with open(file_path, "r", encoding="utf-8") as f: data = json.load(f) - except Exception as exc: + except Exception as exc: # noqa: BLE001 print(f"ERROR: Failed to parse JSON: {exc}", file=sys.stderr) sys.exit(1) @@ -188,7 +188,7 @@ def import_holographic( try: vec = embed(content) - except Exception as exc: + except Exception as exc: # noqa: BLE001 print( f"Warning: Failed to embed '{content[:40]}...': {exc}. Inserting without embedding." ) @@ -206,7 +206,7 @@ def import_holographic( success += 1 else: skipped += 1 - except Exception as exc: + except Exception as exc: # noqa: BLE001 print(f"ERROR: Failed to insert item {idx}: {exc}", file=sys.stderr) errors += 1 @@ -234,12 +234,12 @@ def import_markdown( print( f"Import complete: {res['inserted']} inserted, {res['skipped']} skipped/duplicates (parsed {res['parsed']} total)." ) - except Exception as exc: + except Exception as exc: # noqa: BLE001 print(f"ERROR: Bulk import failed: {exc}", file=sys.stderr) sys.exit(1) -def main(argv: Optional[List[str]] = None) -> int: +def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( prog="hexus-import", description="Bulk import tool for Hexus memory store.", diff --git a/mcp_server/server.py b/mcp_server/server.py index e357856..d86cfa5 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -21,12 +21,11 @@ from __future__ import annotations import logging -from typing import Any, Dict, Optional, List +import os +from typing import Any from hexus.store import MemoryStore - -import os from . import tools logger = logging.getLogger(__name__) @@ -35,6 +34,7 @@ def _generate_metrics(store: MemoryStore) -> str: """Generate Prometheus metrics from DB and AsyncWriter status.""" import math + from hexus.writer import _active_writers # 1. Base liveness and totals @@ -57,7 +57,7 @@ def _generate_metrics(store: MemoryStore) -> str: # 2. Detailed Database Metrics from store.get_metrics_data() try: db_data = store.get_metrics_data() - except Exception as exc: + except Exception as exc: # noqa: BLE001 lines.append(f"# ERROR: Failed to query detailed metrics: {exc}") db_data = {} @@ -225,7 +225,7 @@ def clean_lbl(val: Any) -> str: for obj in _active_writers: queue_stats = obj.stats() break - except Exception as exc: + except Exception as exc: # noqa: BLE001 lines.append(f"# ERROR: Failed to extract writer queue stats: {exc}") if queue_stats: @@ -258,7 +258,7 @@ def clean_lbl(val: Any) -> str: lines.append(f'hexus_writer_latency_seconds{{quantile="0.95"}} {p95}') # 4. Background Cleanup Stats - cleanup_interval = int(os.environ.get("HEXUS_CLEANUP_INTERVAL_HOURS", 24)) + cleanup_interval = int(os.environ.get("HEXUS_CLEANUP_INTERVAL_HOURS", "24")) memories_ttl = os.environ.get("HEXUS_CLEANUP_MEMORIES_TTL_DAYS") memories_ttl = int(memories_ttl) if memories_ttl else None conversations_ttl = os.environ.get("HEXUS_CLEANUP_CONVERSATIONS_TTL_DAYS") @@ -303,7 +303,7 @@ def clean_lbl(val: Any) -> str: # 5. Background Consolidation Stats consolidation_interval = int( - os.environ.get("HEXUS_CONSOLIDATION_INTERVAL_HOURS", 12) + os.environ.get("HEXUS_CONSOLIDATION_INTERVAL_HOURS", "12") ) consolidation_thread_alive = any( t.name == "hexus-consolidation-thread" for t in threading.enumerate() @@ -421,7 +421,7 @@ def _build_server( store: MemoryStore, *, name: str = "hexus", - instructions: Optional[str] = None, + instructions: str | None = None, ): """Build and return a configured `mcp.server.fastmcp.FastMCP` instance. @@ -433,6 +433,7 @@ def _build_server( # Imported lazily so `pip install hexus` (no [mcp] extra) # doesn't pull mcp as a transitive runtime dep. import os + from mcp.server.fastmcp import FastMCP if instructions is None: @@ -471,7 +472,7 @@ def _build_server( } # -- Scheduled Background Cleanup -------------------------------------- - cleanup_interval = int(os.environ.get("HEXUS_CLEANUP_INTERVAL_HOURS", 24)) + cleanup_interval = int(os.environ.get("HEXUS_CLEANUP_INTERVAL_HOURS", "24")) memories_ttl = os.environ.get("HEXUS_CLEANUP_MEMORIES_TTL_DAYS") memories_ttl = int(memories_ttl) if memories_ttl else None conversations_ttl = os.environ.get("HEXUS_CLEANUP_CONVERSATIONS_TTL_DAYS") @@ -526,12 +527,13 @@ def background_cleanup_loop(): # -- Scheduled Background Consolidation -------------------------------- consolidation_interval = int( - os.environ.get("HEXUS_CONSOLIDATION_INTERVAL_HOURS", 12) + os.environ.get("HEXUS_CONSOLIDATION_INTERVAL_HOURS", "12") ) if consolidation_interval > 0: import threading import time + from hexus.writer import _active_writers def background_consolidation_loop(): @@ -606,7 +608,7 @@ def background_consolidation_loop(): res_co, ) break # Success - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.error( "Error during scheduled background consolidation (attempt %d/%d): %s", attempt, @@ -629,7 +631,7 @@ def background_consolidation_loop(): # -- tools ------------------------------------------------------------- @mcp.tool() - def memory_health() -> Dict[str, Any]: + def memory_health() -> dict[str, Any]: """Liveness + capability check. Returns DB status, embedder model/dim, row counts.""" return tools.memory_health(store, {}) @@ -638,10 +640,10 @@ def memory_retain( contents: list[str], target: str = "memory", agent_identity: str = "", - metadata: Optional[Dict[str, Any]] = None, + metadata: dict[str, Any] | None = None, doc_type: str = "memory", source_url: str = "", - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Add one or many memory entries. Each content becomes one row. Args: @@ -682,9 +684,9 @@ def memory_recall( target: str = "", min_similarity: float = 0.0, min_confidence: float = 0.0, - decay_half_life_days: Optional[float] = None, - recall_boost_weight: Optional[float] = None, - ) -> Dict[str, Any]: + decay_half_life_days: float | None = None, + recall_boost_weight: float | None = None, + ) -> dict[str, Any]: """Semantic search over memory entries. Args: @@ -725,9 +727,9 @@ def memory_hybrid_search( target: str = "", min_similarity: float = 0.0, min_confidence: float = 0.0, - decay_half_life_days: Optional[float] = None, - recall_boost_weight: Optional[float] = None, - ) -> Dict[str, Any]: + decay_half_life_days: float | None = None, + recall_boost_weight: float | None = None, + ) -> dict[str, Any]: """Hybrid search blending semantic vector search and full-text search over memory entries. Args: @@ -768,7 +770,7 @@ def memory_search( target: str = "", limit: int = 20, offset: int = 0, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Browse memory entries without semantic search (list / paginate). Returns: {"count", "limit", "offset", "rows": [...]} @@ -788,7 +790,7 @@ def memory_forget( id: int, confirm: bool = False, agent_identity: str = "", - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Delete a memory entry by id. Pass confirm=true to actually delete. Dry-run by default (returns what would happen). Restricted to the @@ -811,9 +813,9 @@ def memory_recall_turns( agent_identity: str = "", session_id: str = "", min_similarity: float = 0.0, - decay_half_life_days: Optional[float] = None, - recall_boost_weight: Optional[float] = None, - ) -> Dict[str, Any]: + decay_half_life_days: float | None = None, + recall_boost_weight: float | None = None, + ) -> dict[str, Any]: """Semantic search over past chat turns (every user/assistant exchange). Args: @@ -851,9 +853,9 @@ def memory_hybrid_recall_turns( agent_identity: str = "", session_id: str = "", min_similarity: float = 0.0, - decay_half_life_days: Optional[float] = None, - recall_boost_weight: Optional[float] = None, - ) -> Dict[str, Any]: + decay_half_life_days: float | None = None, + recall_boost_weight: float | None = None, + ) -> dict[str, Any]: """Hybrid search blending semantic vector search and full-text search over conversation turns. Args: @@ -893,8 +895,8 @@ def memory_append_turn( role: str, content: str, agent_identity: str = "", - metadata: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: + metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: """Append one chat turn. Use this to capture a (user, assistant) exchange into the conversation log for later semantic recall. @@ -925,8 +927,8 @@ def memory_record_delegation( task: str, result: str, agent_identity: str = "", - metadata: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: + metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: """Record a subagent delegation. Args: @@ -960,7 +962,7 @@ def memory_recall_delegations( min_similarity: float = 0.0, decay_half_life_days: float = 0.0, recall_boost_weight: float = 0.0, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Recall subagent delegations by semantic similarity query. Args: @@ -992,7 +994,7 @@ def memory_count( agent_identity: str = "", target: str = "", session_id: str = "", - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Return row counts for memory_entries and conversations, scoped as requested. Args: @@ -1013,11 +1015,11 @@ def memory_count( @mcp.tool() def memory_cleanup( - conversations_ttl_days: Optional[int] = None, - memories_ttl_days: Optional[int] = None, - delegations_ttl_days: Optional[int] = None, + conversations_ttl_days: int | None = None, + memories_ttl_days: int | None = None, + delegations_ttl_days: int | None = None, confirm: bool = False, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Delete stale records from conversations, memory_entries, and delegations based on TTL. This is a **fleet-wide, unscoped** destructive operation — it deletes @@ -1060,7 +1062,7 @@ def memory_cleanup( @mcp.tool() def memory_consolidate( agent_identity: str = "", - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Trigger memory consolidation for low-confidence or heavily co-occurring entries. Args: @@ -1114,7 +1116,7 @@ def memory_entity_graph( entity_value: str, agent_identity: str = "", limit: int = 5, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Find other entities that co-occur with a target entity. Args: @@ -1140,7 +1142,7 @@ def memory_graph_walk( agent_identity: str = "", max_depth: int = 2, limit: int = 5, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Traverse the co-occurrence graph up to N hops away from a start entity. Args: @@ -1166,7 +1168,7 @@ def memory_common_topics( agent_identity: str = "", min_strength: int = 2, limit: int = 10, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Retrieve clusters/cliques of heavily co-occurring entities. Args: @@ -1184,7 +1186,7 @@ def memory_common_topics( ) @mcp.tool() - def memory_confirm(id: int) -> Dict[str, Any]: + def memory_confirm(id: int) -> dict[str, Any]: """Increment confirm_count in metadata JSONB for the given entry ID. Args: @@ -1198,7 +1200,7 @@ def memory_confirm(id: int) -> Dict[str, Any]: ) @mcp.tool() - def memory_reject(id: int) -> Dict[str, Any]: + def memory_reject(id: int) -> dict[str, Any]: """Increment reject_count in metadata JSONB for the given entry ID. Args: @@ -1215,7 +1217,7 @@ def memory_reject(id: int) -> Dict[str, Any]: def memory_summarize_session( session_id: str, limit: int = 5, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Compute the vector centroid of a session's turns and find the K closest turns. Args: @@ -1231,7 +1233,7 @@ def memory_summarize_session( ) @mcp.tool() - def memory_retrieve(id: int) -> Dict[str, Any]: + def memory_retrieve(id: int) -> dict[str, Any]: """Retrieve the original full content of a memory entry by its integer ID. Args: @@ -1245,7 +1247,7 @@ def memory_retrieve(id: int) -> Dict[str, Any]: ) @mcp.tool() - def headroom_retrieve(id: int) -> Dict[str, Any]: + def headroom_retrieve(id: int) -> dict[str, Any]: """Retrieve the original full content of a memory entry by its integer ID. Args: @@ -1259,24 +1261,153 @@ def headroom_retrieve(id: int) -> Dict[str, Any]: ) @mcp.tool() - def memory_stats() -> Dict[str, Any]: + def memory_stats() -> dict[str, Any]: """Return metrics from Hexus database and background async queue stats.""" return tools.memory_stats(store, {}) - # -- patch ASGI app for Prometheus /metrics endpoint ------------------- + # -- patch ASGI app for REST API + Prometheus /metrics endpoint ------------------- _orig_get_asgi_app = mcp.streamable_http_app - def get_asgi_app_with_metrics(*args, **kwargs): + def get_asgi_app_with_rest_api(*args, **kwargs): app = _orig_get_asgi_app(*args, **kwargs) - from starlette.responses import Response - from . import tools + from starlette.requests import Request + from starlette.responses import JSONResponse, Response tools.http_transport_active = True + async def health(request): + """Health check endpoint for pi extension and other clients.""" + try: + result = tools.memory_health(store, {}) + return JSONResponse(result) + except Exception as exc: # noqa: BLE001 + return JSONResponse( + {"status": "error", "error": str(exc)}, status_code=500 + ) + + async def recall(request: Request): + """Semantic search over memory entries. + + POST /api/recall + Body: {"query": str, "top_k"?: int, "agent_identity"?: str, + "target"?: "memory"|"user", "min_similarity"?: float} + Returns: {"query", "count", "results": [...]} + """ + try: + body = await request.json() + except ValueError: + return JSONResponse({"error": "invalid JSON body"}, status_code=400) + + try: + result = tools.memory_recall(store, body) + if "error" in result: + return JSONResponse(result, status_code=400) + return JSONResponse(result) + except Exception as exc: # noqa: BLE001 + return JSONResponse({"error": str(exc)}, status_code=500) + + async def retain(request: Request): + """Store one or more memory entries. + + POST /api/retain + Body (new simplified): {"contents": [str], "target"?: str, "metadata"?: dict, "agent_identity"?: str} + Body (legacy): {"contents": [{"content": str, "target"?: str, "metadata"?: dict}], "agent_identity"?: str} + Returns: {"inserted", "duplicates", "errors"} + """ + try: + body = await request.json() + except ValueError: + return JSONResponse({"error": "invalid JSON body"}, status_code=400) + + contents = body.get("contents", []) + agent_identity = body.get("agent_identity") + target = body.get("target") + metadata = body.get("metadata") + + if not isinstance(contents, list) or not contents: + return JSONResponse( + {"error": "contents must be a non-empty list"}, status_code=400 + ) + + # Normalize: support both string[] (new) and {content}[] (legacy) + flat_contents = [] + flat_metas = [] + flat_targets = [] + for item in contents: + if isinstance(item, str): + flat_contents.append(item) + flat_metas.append(metadata) + flat_targets.append(target) + elif isinstance(item, dict): + c = item.get("content") + if not isinstance(c, str): + return JSONResponse( + { + "error": "contents items must be strings or contain a 'content' key with a string value" + }, + status_code=400, + ) + flat_contents.append(c) + # Per-item target/metadata win over the shared defaults. + # Use `is not None` (not `or`) so empty dict {} / empty + # string targets are preserved rather than silently dropped. + item_target = item.get("target") + flat_targets.append( + item_target if item_target is not None else target + ) + item_meta = item.get("metadata") + flat_metas.append(item_meta if item_meta is not None else metadata) + else: + return JSONResponse( + {"error": "contents items must be strings or objects"}, + status_code=400, + ) + + args = { + "contents": flat_contents, + "targets": flat_targets, + "metadata": flat_metas, + } + if agent_identity: + args["agent_identity"] = agent_identity + + try: + result = tools.memory_retain(store, args) + return JSONResponse(result) + except Exception as exc: # noqa: BLE001 + return JSONResponse({"error": str(exc)}, status_code=500) + + async def append_turn(request: Request): + """Store a conversation turn. + + POST /api/append-turn + Body: {"session_id": str, "role": str, "content": str, + "agent_identity"?: str, "metadata"?: dict} + Returns: {"id", "session_id", "role"} + """ + try: + body = await request.json() + except ValueError: + return JSONResponse({"error": "invalid JSON body"}, status_code=400) + + try: + result = tools.memory_append_turn(store, body) + if "error" in result: + return JSONResponse(result, status_code=400) + return JSONResponse(result) + except Exception as exc: # noqa: BLE001 + return JSONResponse({"error": str(exc)}, status_code=500) + async def metrics(request): - return Response(_generate_metrics(store), media_type="text/plain") + # Plain text, NOT JSONResponse — JSONResponse would quote/escape the + # whole payload and break Prometheus scraping. + return Response(content=_generate_metrics(store), media_type="text/plain") - app.add_route("/metrics", metrics) + app.add_route("/api/health", health, ["GET"]) + app.add_route("/api/recall", recall, ["POST"]) + app.add_route("/api/retain", retain, ["POST"]) + app.add_route("/api/append-turn", append_turn, ["POST"]) + app.add_route("/metrics", metrics, ["GET"]) # Optional bearer-token auth on the HTTP transport. When # HEXUS_API_TOKEN is set, every HTTP request (MCP calls + /metrics) @@ -1302,7 +1433,7 @@ async def metrics(request): ) return app - mcp.streamable_http_app = get_asgi_app_with_metrics + mcp.streamable_http_app = get_asgi_app_with_rest_api return mcp @@ -1311,7 +1442,7 @@ def build_server( dsn: str, *, name: str = "hexus", - instructions: Optional[str] = None, + instructions: str | None = None, ) -> Any: """Build (but don't run) an MCP server for the given DSN. diff --git a/mcp_server/tools.py b/mcp_server/tools.py index a2fe351..a88de9e 100644 --- a/mcp_server/tools.py +++ b/mcp_server/tools.py @@ -24,7 +24,7 @@ import logging import os from contextvars import ContextVar -from typing import Any, Dict, List, Optional +from typing import Any from hexus.embed import EmbeddingError, embed from hexus.store import MemoryStore @@ -57,14 +57,14 @@ def default_agent_identity() -> str: # ----------------------------------------------------------------------- -def memory_health(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_health(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Liveness + capability check. Useful for MCP client setup probes. Returns the store's status, the embedder model name + dim, and a row count. Always succeeds if Postgres is reachable, even if the embedder isn't loaded yet (lazy load). """ - from hexus.embedder import DEFAULT_MODEL, DEFAULT_DIM + from hexus.embedder import DEFAULT_DIM, DEFAULT_MODEL try: store.ensure_schema() @@ -94,7 +94,7 @@ def memory_health(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: } -def _coerce_agent_identity(args: Dict[str, Any]) -> str: +def _coerce_agent_identity(args: dict[str, Any]) -> str: """Read agent_identity from args, defaulting to env / 'default'.""" a = args.get("agent_identity") if isinstance(a, str) and a.strip(): @@ -124,7 +124,7 @@ def _coerce_agent_identity(args: Dict[str, Any]) -> str: # only writer is the server's request middleware. stdio / in-process callers # leave it None and keep the pre-#19 behavior. An explicit `_caller_identity` # in args (in-process callers, tests) still takes precedence over the var. -current_caller: ContextVar[Optional[str]] = ContextVar( +current_caller: ContextVar[str | None] = ContextVar( "hexus_current_caller", default=None ) @@ -132,7 +132,7 @@ def _coerce_agent_identity(args: Dict[str, Any]) -> str: http_transport_active = False -def _caller_identity(args: Dict[str, Any]) -> Optional[str]: +def _caller_identity(args: dict[str, Any]) -> str | None: """Server-derived transport identity, or None when none was set. Precedence: explicit `_caller_identity` in args (trusted in-process @@ -153,13 +153,13 @@ def _caller_identity(args: Dict[str, Any]) -> Optional[str]: return None -def _write_identity(args: Dict[str, Any]) -> str: +def _write_identity(args: dict[str, Any]) -> str: """Scope an agent writes into. Transport identity is authoritative; else fall back to the explicit arg, then the env default.""" return _caller_identity(args) or _coerce_agent_identity(args) -def _scope_identity(args: Dict[str, Any]) -> Optional[str]: +def _scope_identity(args: dict[str, Any]) -> str | None: """Identity to scope a by-id mutation/read to, or None to stay unscoped. Transport identity wins; otherwise an explicit non-empty `agent_identity` @@ -176,7 +176,7 @@ def _scope_identity(args: Dict[str, Any]) -> Optional[str]: return None -def _read_identity(store: MemoryStore, args: Dict[str, Any]) -> Optional[str]: +def _read_identity(store: MemoryStore, args: dict[str, Any]) -> str | None: """Effective agent_identity filter for a read/search, resolving the empty-identity asymmetry (issue #19 item C) consistently via the store's isolation policy: @@ -194,7 +194,7 @@ def _read_identity(store: MemoryStore, args: Dict[str, Any]) -> Optional[str]: return None -def _coerce_target(args: Dict[str, Any]) -> Optional[str]: +def _coerce_target(args: dict[str, Any]) -> str | None: """target ∈ {'memory', 'user', None}. Anything else is rejected. Empty string is treated as None to match MCP clients that send @@ -208,12 +208,13 @@ def _coerce_target(args: Dict[str, Any]) -> Optional[str]: return t -def memory_retain(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_retain(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Add one or many memory entries. Mirrors the plugin's `on_memory_write`. args: contents: list[str] — text to store (one row per element) target: 'memory' | 'user' | None (default = 'memory') + targets: list[str] | None — per-item target, overrides `target` metadata: dict | list[dict] | None — per-item metadata agent_identity: str | None (default = env / 'default') doc_type: 'document' | 'note' | 'memory' (default 'memory') — stored in metadata @@ -228,7 +229,23 @@ def memory_retain(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: if not isinstance(c, str) or not c.strip(): raise ValueError(f"contents[{i}] must be a non-empty string") - target = _coerce_target(args) or "memory" + # Normalize targets to one entry per content. A per-item `targets` list + # takes precedence; a scalar `target` applies to all items. + targets_in = args.get("targets") + if targets_in is None: + target_list = [_coerce_target(args) or "memory"] * len(contents) + elif isinstance(targets_in, list): + if len(targets_in) != len(contents): + raise ValueError("targets list length must match contents length") + target_list = [] + for t in targets_in: + coerced = _coerce_target({"target": t}) + if coerced is None: + coerced = "memory" + target_list.append(coerced) + else: + raise ValueError("targets must be a list or None") + agent = _write_identity(args) doc_type = args.get("doc_type", "memory") source_url = args.get("source_url") @@ -236,7 +253,7 @@ def memory_retain(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: # Normalize metadata to one entry per content if metadata_in is None: - metas: List[Optional[Dict[str, Any]]] = [None] * len(contents) + metas: list[dict[str, Any] | None] = [None] * len(contents) elif isinstance(metadata_in, list): if len(metadata_in) != len(contents): raise ValueError("metadata list length must match contents length") @@ -247,7 +264,7 @@ def memory_retain(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: raise ValueError("metadata must be a dict, a list of dicts, or None") # Stamp each item with doc_type + source_url - stamped: List[Dict[str, Any]] = [] + stamped: list[dict[str, Any]] = [] for m in metas: out = dict(m) if m else {} if doc_type and "doc_type" not in out: @@ -266,8 +283,9 @@ def memory_retain(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: inserted = 0 duplicates = 0 - errors: List[str] = [] - for content, vec, meta in zip(contents, vectors, stamped): + errors: list[str] = [] + for i, (content, vec, meta) in enumerate(zip(contents, vectors, stamped)): + target = target_list[i] try: row_id = store.add( agent_identity=agent, @@ -300,7 +318,7 @@ def memory_retain(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: return {"inserted": inserted, "duplicates": duplicates, "errors": errors} -def memory_recall(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_recall(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Semantic search over memory_entries. args: @@ -317,10 +335,8 @@ def memory_recall(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: raise ValueError("query must be a non-empty string") top_k = int(args.get("top_k", 5)) - if top_k < 1: - top_k = 1 - if top_k > 100: - top_k = 100 + top_k = max(top_k, 1) + top_k = min(top_k, 100) agent = _read_identity(store, args) @@ -336,12 +352,12 @@ def memory_recall(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: decay_val = args.get("decay_half_life_days") if decay_val is None: - decay_val = os.environ.get("HEXUS_DECAY_HALF_LIFE_DAYS", 0.0) + decay_val = os.environ.get("HEXUS_DECAY_HALF_LIFE_DAYS", "0.0") decay_half_life_days = max(0.0, float(decay_val)) boost_val = args.get("recall_boost_weight") if boost_val is None: - boost_val = os.environ.get("HEXUS_RECALL_BOOST_WEIGHT", 0.0) + boost_val = os.environ.get("HEXUS_RECALL_BOOST_WEIGHT", "0.0") recall_boost_weight = max(0.0, float(boost_val)) try: @@ -366,7 +382,7 @@ def memory_recall(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: } -def memory_search(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_search(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """List entries (no embedding) — browse / paginate / inspect. args: @@ -380,13 +396,10 @@ def memory_search(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: agent = _read_identity(store, args) target = _coerce_target(args) limit = int(args.get("limit", 20)) - if limit < 1: - limit = 1 - if limit > 200: - limit = 200 + limit = max(limit, 1) + limit = min(limit, 200) offset = int(args.get("offset", 0)) - if offset < 0: - offset = 0 + offset = max(offset, 0) rows = store.list_entries( agent_identity=agent, target=target, limit=limit, offset=offset @@ -400,7 +413,7 @@ def memory_search(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: } -def memory_forget(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_forget(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Delete a memory entry by id. Requires `confirm=true` to actually delete; without it the call is a dry-run that reports what would happen. This makes "drop everything matching a query" hard to do @@ -416,14 +429,13 @@ def memory_forget(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: "hint": "pass confirm=true to actually delete", } agent = _write_identity(args) - with store._get_pool().connection() as conn: # noqa: SLF001 — admin path - with conn.cursor() as cur: - cur.execute( - "DELETE FROM memory_entries WHERE id = %s AND agent_identity = %s RETURNING id, target, content", - (entry_id, agent), - ) - row = cur.fetchone() - conn.commit() + with store._get_pool().connection() as conn, conn.cursor() as cur: + cur.execute( + "DELETE FROM memory_entries WHERE id = %s AND agent_identity = %s RETURNING id, target, content", + (entry_id, agent), + ) + row = cur.fetchone() + conn.commit() if row: webhook_url = os.environ.get("HEXUS_WEBHOOK_URL") @@ -449,7 +461,7 @@ def memory_forget(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: } -def memory_recall_turns(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_recall_turns(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Semantic search over conversation turns. Mirrors `recall_conversation`. args: @@ -465,10 +477,8 @@ def memory_recall_turns(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, A if not isinstance(query, str) or not query.strip(): raise ValueError("query must be a non-empty string") top_k = int(args.get("top_k", 5)) - if top_k < 1: - top_k = 1 - if top_k > 100: - top_k = 100 + top_k = max(top_k, 1) + top_k = min(top_k, 100) agent = args.get("agent_identity") if isinstance(agent, str) and agent.strip() == "": agent = None @@ -479,12 +489,12 @@ def memory_recall_turns(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, A decay_val = args.get("decay_half_life_days") if decay_val is None: - decay_val = os.environ.get("HEXUS_DECAY_HALF_LIFE_DAYS", 0.0) + decay_val = os.environ.get("HEXUS_DECAY_HALF_LIFE_DAYS", "0.0") decay_half_life_days = max(0.0, float(decay_val)) boost_val = args.get("recall_boost_weight") if boost_val is None: - boost_val = os.environ.get("HEXUS_RECALL_BOOST_WEIGHT", 0.0) + boost_val = os.environ.get("HEXUS_RECALL_BOOST_WEIGHT", "0.0") recall_boost_weight = max(0.0, float(boost_val)) try: @@ -508,7 +518,7 @@ def memory_recall_turns(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, A } -def memory_append_turn(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_append_turn(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Append one chat turn. Mirrors the plugin's `sync_turn` capture. args: @@ -550,7 +560,7 @@ def memory_append_turn(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, An return {"id": row_id, "session_id": session_id, "role": role} -def memory_count(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_count(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Return row counts for entries and turns, scoped as requested. args: @@ -580,7 +590,7 @@ def memory_count(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: # ----------------------------------------------------------------------- -def _embed_batch(texts: List[str]) -> List[List[float]]: +def _embed_batch(texts: list[str]) -> list[list[float]]: """Embed a list of texts. For the local path this batches in one model.encode() call. The HTTP path embeds one at a time (limitation of the upstream embed() function) — fine for low-volume MCP traffic, @@ -589,7 +599,7 @@ def _embed_batch(texts: List[str]) -> List[List[float]]: return [embed(t) for t in texts] -def _row_to_dict(row: Any, *, include_embedding: bool = False) -> Dict[str, Any]: +def _row_to_dict(row: Any, *, include_embedding: bool = False) -> dict[str, Any]: """Coerce a DB row (psycopg dict_row or plain tuple) into a JSON-safe dict.""" if hasattr(row, "keys"): d = dict(row) @@ -606,7 +616,7 @@ def _row_to_dict(row: Any, *, include_embedding: bool = False) -> Dict[str, Any] return d -def memory_hybrid_search(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_hybrid_search(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Blend semantic vector search and full-text search over memory_entries. args: @@ -641,12 +651,12 @@ def memory_hybrid_search(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, decay_val = args.get("decay_half_life_days") if decay_val is None: - decay_val = os.environ.get("HEXUS_DECAY_HALF_LIFE_DAYS", 0.0) + decay_val = os.environ.get("HEXUS_DECAY_HALF_LIFE_DAYS", "0.0") decay_half_life_days = max(0.0, float(decay_val)) boost_val = args.get("recall_boost_weight") if boost_val is None: - boost_val = os.environ.get("HEXUS_RECALL_BOOST_WEIGHT", 0.0) + boost_val = os.environ.get("HEXUS_RECALL_BOOST_WEIGHT", "0.0") recall_boost_weight = max(0.0, float(boost_val)) try: @@ -675,8 +685,8 @@ def memory_hybrid_search(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, def memory_hybrid_recall_turns( - store: MemoryStore, args: Dict[str, Any] -) -> Dict[str, Any]: + store: MemoryStore, args: dict[str, Any] +) -> dict[str, Any]: """Blend semantic vector search and full-text search over conversation turns. args: @@ -713,12 +723,12 @@ def memory_hybrid_recall_turns( decay_val = args.get("decay_half_life_days") if decay_val is None: - decay_val = os.environ.get("HEXUS_DECAY_HALF_LIFE_DAYS", 0.0) + decay_val = os.environ.get("HEXUS_DECAY_HALF_LIFE_DAYS", "0.0") decay_half_life_days = max(0.0, float(decay_val)) boost_val = args.get("recall_boost_weight") if boost_val is None: - boost_val = os.environ.get("HEXUS_RECALL_BOOST_WEIGHT", 0.0) + boost_val = os.environ.get("HEXUS_RECALL_BOOST_WEIGHT", "0.0") recall_boost_weight = max(0.0, float(boost_val)) try: @@ -746,8 +756,8 @@ def memory_hybrid_recall_turns( def memory_record_delegation( - store: MemoryStore, args: Dict[str, Any] -) -> Dict[str, Any]: + store: MemoryStore, args: dict[str, Any] +) -> dict[str, Any]: """Record a subagent delegation. args: @@ -800,8 +810,8 @@ def memory_record_delegation( def memory_recall_delegations( - store: MemoryStore, args: Dict[str, Any] -) -> Dict[str, Any]: + store: MemoryStore, args: dict[str, Any] +) -> dict[str, Any]: """Recall subagent delegations. args: @@ -870,7 +880,7 @@ def memory_recall_delegations( } -def memory_entity_graph(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_entity_graph(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Find other entities that co-occur with a target entity.""" entity_type = args.get("entity_type") if not isinstance(entity_type, str) or not entity_type.strip(): @@ -895,7 +905,7 @@ def memory_entity_graph(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, A ) -def memory_graph_walk(store: MemoryStore, args: Dict[str, Any]) -> List[Dict[str, Any]]: +def memory_graph_walk(store: MemoryStore, args: dict[str, Any]) -> list[dict[str, Any]]: """Traverse the co-occurrence graph up to N hops away from a start entity.""" entity_type = args.get("entity_type") if not isinstance(entity_type, str) or not entity_type.strip(): @@ -925,8 +935,8 @@ def memory_graph_walk(store: MemoryStore, args: Dict[str, Any]) -> List[Dict[str def memory_common_topics( - store: MemoryStore, args: Dict[str, Any] -) -> List[Dict[str, Any]]: + store: MemoryStore, args: dict[str, Any] +) -> list[dict[str, Any]]: """Retrieve clusters/cliques of heavily co-occurring entities.""" agent = args.get("agent_identity") if isinstance(agent, str) and agent.strip() == "": @@ -945,7 +955,7 @@ def memory_common_topics( ) -def memory_confirm(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_confirm(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Increment confirm_count in metadata JSONB for the given entry ID.""" entry_id = args.get("id") if entry_id is None: @@ -959,7 +969,7 @@ def memory_confirm(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: return {"id": entry_id, "success": success} -def memory_reject(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_reject(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Increment reject_count in metadata JSONB for the given entry ID.""" entry_id = args.get("id") if entry_id is None: @@ -974,8 +984,8 @@ def memory_reject(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: def memory_summarize_session( - store: MemoryStore, args: Dict[str, Any] -) -> Dict[str, Any]: + store: MemoryStore, args: dict[str, Any] +) -> dict[str, Any]: """Compute the vector centroid of a session's turns and find the K closest turns.""" session_id = args.get("session_id") if not isinstance(session_id, str) or not session_id.strip(): @@ -991,7 +1001,7 @@ def memory_summarize_session( ) -def memory_retrieve(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_retrieve(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Retrieve the original full content of a memory entry by its integer ID. args: @@ -1011,7 +1021,7 @@ def memory_retrieve(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: return {"id": entry_id, "found": True, "content": content} -def memory_stats(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: +def memory_stats(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: """Return metrics from Hexus database and background async queue stats.""" db_stats = { "memory_entries_count": store.count(agent_identity=None, target=None), diff --git a/pi-extension/README.md b/pi-extension/README.md new file mode 100644 index 0000000..4496ddb --- /dev/null +++ b/pi-extension/README.md @@ -0,0 +1,102 @@ +# hexus-pi — Vector Memory Extension for pi Agent + +Integrates hexus (Postgres + pgvector memory) into [pi](https://github.com/earendil-works/pi) agent harness. + +## Features + +1. **Memory Recall** — Embeds user prompts and searches hexus for relevant past memories, injecting them into the system prompt +2. **Turn Capture** — Stores conversation turns in hexus for future recall +3. **Session Reflection** — Automatically extracts durable facts from long/idle sessions and saves them to memory + +## Installation + +### Option 1: Symlink (Development) + +```bash +ln -s /path/to/hexus/pi-extension ~/.pi/agent/extensions/hexus +``` + +### Option 2: Copy + +```bash +cp -r /path/to/hexus/pi-extension ~/.pi/agent/extensions/hexus +``` + +### Option 3: Publish as pi Package + +(WIP) Publish to npm and install via `pi install hexus` + +## Configuration + +Edit `~/.pi/agent/extensions/hexus/config.json`: + +```json +{ + "apiUrl": "http://localhost:8000", + "agentIdentity": "pi", + "recallLimit": 5, + "minSimilarity": 0.3, + "enabled": true, + "storeTurns": true +} +``` + +Or via environment variables: + +| Variable | Default | Description | +|----------|---------|-------------| +| `HEXUS_API_URL` | `http://localhost:8000` | hexus API base URL | +| `HEXUS_AGENT_IDENTITY` | `pi` | Memory namespace | +| `HEXUS_ENABLED` | `true` | Enable/disable extension | +| `HEXUS_STORE_TURNS` | `true` | Store conversation turns | +| `HEXUS_REFLECTION_ENABLED` | `true` | Enable session reflection | +| `HEXUS_REFLECTION_TOKEN_THRESHOLD` | `8000` | Min tokens before reflection | +| `HEXUS_REFLECTION_MIN_TURNS` | `10` | Turns between reflections | +| `HEXUS_REFLECTION_IDLE_SECONDS` | `10` | Idle time before reflection | +| `HEXUS_REFLECTION_MODEL` | `tobiTradez/minimax-m2.7-highspeed` | Model as `"provider/modelId"` | + +## hexus API Endpoints Required + +This extension requires the REST API endpoints added to hexus: + +- `GET /api/health` — Health check +- `POST /api/recall` — Vector search +- `POST /api/retain` — Store memory entries +- `POST /api/append-turn` — Store conversation turns + +These are available in hexus v0.9.2+. + +## Tools + +The extension registers two tools for explicit memory operations: + +- `hexus_recall` — Search memory for past entries +- `hexus_retain` — Save important info to memory + +## Reflection + +Session reflection automatically extracts durable facts from conversations when: + +1. Session exceeds token threshold (default: 8000 tokens) +2. 10+ turns since last reflection +3. User idle for 10+ seconds + +The reflection model analyzes the conversation and saves facts like: +- User preferences and projects +- Key decisions made +- Important context for future sessions + +## Multi-Harness Strategy + +This extension is designed to work with pi, but the same hexus backend can be used by other harnesses: + +- **Hermes** — Native hexus plugin (`pip install hexus`) +- **pi** — This extension +- **Claude Desktop / Cursor** — MCP server (`hexus-mcp serve --transport stdio`) +- **Others** — HTTP REST API + +See main [hexus README](../README.md) for full architecture. + +## License + +BSD-3-Clause diff --git a/pi-extension/config.json b/pi-extension/config.json new file mode 100644 index 0000000..a45efb7 --- /dev/null +++ b/pi-extension/config.json @@ -0,0 +1,8 @@ +{ + "apiUrl": "http://localhost:8000", + "agentIdentity": "pi", + "recallLimit": 5, + "minSimilarity": 0.3, + "enabled": true, + "storeTurns": true +} diff --git a/pi-extension/config.ts b/pi-extension/config.ts new file mode 100644 index 0000000..dc48903 --- /dev/null +++ b/pi-extension/config.ts @@ -0,0 +1,118 @@ +/** + * hexus config — loaded from ~/.pi/agent/extensions/hexus/config.json + * or environment variables. + * + * Uses dynamic import() to avoid sync fs I/O at module load time + * in ESM contexts. Falls back to defaults if config file is missing. + */ + +export interface HexusConfig { + /** hexus API base URL */ + apiUrl: string; + /** Default agent_identity for memory operations */ + agentIdentity: string; + /** Number of memories to recall per turn */ + recallLimit: number; + /** Minimum similarity score (0-1) */ + minSimilarity: number; + /** Whether to inject memory into system prompt */ + enabled: boolean; + /** Whether to store turns in hexus */ + storeTurns: boolean; +} + +const CONFIG_PATH = "hexus/config.json"; + +// Default config — used when no config file exists +const DEFAULTS: HexusConfig = { + apiUrl: "http://localhost:8000", + agentIdentity: "pi", + recallLimit: 5, + minSimilarity: 0.3, + enabled: true, + storeTurns: true, +}; + +function applyEnvOverrides(cfg: HexusConfig): HexusConfig { + const envUrl = process.env["HEXUS_API_URL"]; + const envIdentity = process.env["HEXUS_AGENT_IDENTITY"]; + const envEnabled = process.env["HEXUS_ENABLED"]; + const envStoreTurns = process.env["HEXUS_STORE_TURNS"]; + + return { + apiUrl: envUrl ?? cfg.apiUrl, + agentIdentity: envIdentity ?? cfg.agentIdentity, + recallLimit: cfg.recallLimit, + minSimilarity: cfg.minSimilarity, + enabled: envEnabled !== undefined ? envEnabled !== "false" : cfg.enabled, + storeTurns: envStoreTurns !== undefined ? envStoreTurns !== "false" : cfg.storeTurns, + }; +} + +async function loadConfigAsync(): Promise { + try { + // Dynamically import to avoid sync I/O at module load + const { readFileSync, existsSync } = await import("node:fs"); + const { join } = await import("node:path"); + + // Get CONFIG_DIR_NAME from pi-coding-agent at runtime + let configDir = ".pi"; + try { + const piModule = await import("@earendil-works/pi-coding-agent"); + configDir = piModule.CONFIG_DIR_NAME ?? ".pi"; + } catch { + // Module import failed, use default + } + + const home = process.env["HOME"] ?? "/home/codenamekt"; + + // Try project-local first, then global + const projectConfig = join(process.cwd(), configDir, CONFIG_PATH); + const globalConfig = join(home, configDir, "agent", "extensions", CONFIG_PATH); + + let configPath: string | undefined; + if (existsSync(projectConfig)) { + configPath = projectConfig; + } else if (existsSync(globalConfig)) { + configPath = globalConfig; + } + + if (configPath) { + const content = readFileSync(configPath, "utf-8"); + const fileConfig = JSON.parse(content); + return applyEnvOverrides({ + apiUrl: fileConfig.apiUrl ?? DEFAULTS.apiUrl, + agentIdentity: fileConfig.agentIdentity ?? DEFAULTS.agentIdentity, + recallLimit: fileConfig.recallLimit ?? DEFAULTS.recallLimit, + minSimilarity: fileConfig.minSimilarity ?? DEFAULTS.minSimilarity, + enabled: fileConfig.enabled ?? DEFAULTS.enabled, + storeTurns: fileConfig.storeTurns ?? DEFAULTS.storeTurns, + }); + } + } catch { + // Config file not found or parse error, fall through to defaults + } + + return applyEnvOverrides(DEFAULTS); +} + +let _config: HexusConfig | undefined; +let _loadPromise: Promise | undefined; + +export function getConfig(): HexusConfig { + // After initial async load, this returns the cached config. + // Before async load completes, it returns defaults (avoids blocking). + return _config ?? DEFAULTS; +} + +/** Kick off async config load. Call once at extension init. */ +export function initConfig(): Promise { + if (_config) return Promise.resolve(_config); + if (!_loadPromise) { + _loadPromise = loadConfigAsync().then((cfg) => { + _config = cfg; + return cfg; + }); + } + return _loadPromise; +} diff --git a/pi-extension/http-client.ts b/pi-extension/http-client.ts new file mode 100644 index 0000000..8880049 --- /dev/null +++ b/pi-extension/http-client.ts @@ -0,0 +1,181 @@ +/** + * hexus HTTP client — calls the hexus REST API + */ + +import { getConfig } from "./config"; + +export interface MemoryResult { + id: number; + agent_identity: string; + target: string; + content: string; + score: number; + created_at: string; + metadata?: Record; +} + +export interface RecallResponse { + query: string; + count: number; + results: MemoryResult[]; +} + +export interface HealthResponse { + status: string; + schema_ok: boolean; + embedder: { + model: string; + dim: number; + eager_loaded: boolean; + }; + row_counts: { + memory_entries: number; + }; +} + +export interface RetainResponse { + inserted: number; + duplicates: number; + errors: string[]; +} + +export interface AppendTurnResponse { + id: number; + session_id: string; + role: string; +} + +class HexusClient { + readonly baseUrl: string; + private healthCache: { data: HealthResponse | null; expiry: number } = { data: null, expiry: 0 }; + private offlineUntil = 0; // Unix ms; skip requests when offline + + constructor(baseUrl: string) { + this.baseUrl = baseUrl.replace(/\/$/, ""); // Remove trailing slash + } + + /** Returns cached health if fresh; updates cache on miss or expiry. */ + async health(forceRefresh = false): Promise { + const now = Date.now(); + if (!forceRefresh && this.healthCache.data && now < this.healthCache.expiry) { + return this.healthCache.data; + } + try { + const data = await this.request("/api/health", {}, 3000); + this.healthCache = { data, expiry: now + 30_000 }; + this.offlineUntil = 0; + return data; + } catch { + // Mark offline for 5s to avoid hammering a dead server + this.offlineUntil = now + 5_000; + return this.healthCache.data; // return stale cache if available + } + } + + private async request( + path: string, + options: RequestInit = {}, + timeoutMs = 5000 + ): Promise { + const now = Date.now(); + if (this.offlineUntil && now < this.offlineUntil) { + throw new Error("hexus offline"); + } + + const url = `${this.baseUrl}${path}`; + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), timeoutMs); + + try { + const response = await fetch(url, { + ...options, + signal: controller.signal, + headers: { + "Content-Type": "application/json", + ...options.headers, + }, + }); + clearTimeout(timeout); + + if (!response.ok) { + const text = await response.text().catch(() => "Unknown error"); + throw new Error( + `hexus API error: ${response.status} ${response.statusText} - ${text}` + ); + } + + return response.json() as Promise; + } catch (err) { + clearTimeout(timeout); + if (err instanceof Error && err.name === "AbortError") { + throw new Error(`hexus request timed out after ${timeoutMs}ms`); + } + throw err; + } + } + + // Deprecated: use health() directly — kept for callers that ignore null + async healthSync(): Promise { + return (await this.health(true)) ?? { status: "offline", schema_ok: false, embedder: { model: "", dim: 0, eager_loaded: false }, row_counts: { memory_entries: 0 } }; + } + + async recall(params: { + query: string; + top_k?: number; + agent_identity?: string; + target?: string; + min_similarity?: number; + }): Promise { + return this.request("/api/recall", { + method: "POST", + body: JSON.stringify(params), + }); + } + + async retain(params: { + /** List of memory content strings (not objects). */ + contents: string[]; + /** Target applied to all items. */ + target?: string; + /** Per-item or shared metadata. */ + metadata?: Record | Record[]; + agent_identity?: string; + }): Promise { + return this.request("/api/retain", { + method: "POST", + body: JSON.stringify(params), + }); + } + + async appendTurn(params: { + session_id: string; + role: string; + content: string; + agent_identity?: string; + metadata?: Record; + }): Promise { + return this.request("/api/append-turn", { + method: "POST", + body: JSON.stringify(params), + }); + } +} + +let _client: HexusClient | undefined; + +export function getClient(): HexusClient { + const config = getConfig(); + if (!_client) { + _client = new HexusClient(config.apiUrl); + } else if (_client.baseUrl !== config.apiUrl.replace(/\/$/, "")) { + // Config loaded asynchronously after the client was first created + // (getConfig() returns defaults until initConfig() resolves), so the + // singleton may hold a stale URL — rebuild it with the real one. + _client = new HexusClient(config.apiUrl); + } + return _client; +} + +export function resetClient(): void { + _client = undefined; +} diff --git a/pi-extension/index.ts b/pi-extension/index.ts new file mode 100644 index 0000000..d636255 --- /dev/null +++ b/pi-extension/index.ts @@ -0,0 +1,433 @@ +/** + * hexus — Vector memory extension for pi + * + * Integrates hexus (Postgres + pgvector memory) into pi agent. + * Features: + * 1. Embeds user's prompt → calls hexus /api/recall → injects relevant memories + * 2. Stores conversation turns in hexus for future recall + * 3. Session reflection when: + * - Session exceeds token threshold (default 8000) + * - 10+ turns since last reflection + * - User idle for 10+ seconds + */ + +import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; +import { Type } from "typebox"; +import { complete } from "@earendil-works/pi-ai/compat"; +import { getConfig, initConfig } from "./config"; +import { getClient, type MemoryResult } from "./http-client"; + +// --------------------------------------------------------------------------- +// Reflection Config +// --------------------------------------------------------------------------- + +interface ReflectionConfig { + enabled: boolean; + tokenThreshold: number; + minTurnsBetweenReflections: number; + idleSeconds: number; + /** Model in "provider/modelId" format (e.g. "tobiTradez/minimax-m2.7-highspeed") */ + model: string; +} + +function getReflectionConfig(): ReflectionConfig { + const enabled = process.env["HEXUS_REFLECTION_ENABLED"] !== "false"; + const tokenThreshold = parseInt(process.env["HEXUS_REFLECTION_TOKEN_THRESHOLD"] ?? "", 10); + const minTurns = parseInt(process.env["HEXUS_REFLECTION_MIN_TURNS"] ?? "", 10); + const idleSeconds = parseInt(process.env["HEXUS_REFLECTION_IDLE_SECONDS"] ?? "", 10); + + if (isNaN(tokenThreshold)) console.warn("hexus: HEXUS_REFLECTION_TOKEN_THRESHOLD not numeric, using 8000"); + if (isNaN(minTurns)) console.warn("hexus: HEXUS_REFLECTION_MIN_TURNS not numeric, using 10"); + if (isNaN(idleSeconds)) console.warn("hexus: HEXUS_REFLECTION_IDLE_SECONDS not numeric, using 10"); + + return { + enabled, + tokenThreshold: isNaN(tokenThreshold) ? 8000 : tokenThreshold, + minTurnsBetweenReflections: isNaN(minTurns) ? 10 : minTurns, + idleSeconds: isNaN(idleSeconds) ? 10 : idleSeconds, + // HEXUS_REFLECTION_MODEL may already be a full "provider/modelId" string + // (e.g. "tobiTradez/minimax-m2.7-highspeed"). Only prefix the provider when + // it's a bare model ID — otherwise "openai/gpt-4o" would become + // "headroom/openai/gpt-4o" (provider=headroom, modelId=openai/gpt-4o). + model: (() => { + const model = process.env["HEXUS_REFLECTION_MODEL"] ?? "tobiTradez/minimax-m2.7-highspeed"; + return model.includes("/") + ? model + : `${process.env["HEXUS_REFLECTION_MODEL_PROVIDER"] ?? "headroom"}/${model}`; + })(), + }; +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/** Simple hash of a string — used for session IDs to avoid leaking filepath. */ +function hashString(str: string): string { + let hash = 5381; + for (let i = 0; i < str.length; i++) { + hash = ((hash << 5) + hash) ^ str.charCodeAt(i); + } + // Return hex string (no negative values) + return Math.abs(hash >>> 0).toString(16).padStart(8, "0"); +} + +/** Truncate text at word boundary, not mid-word. */ +function truncateAtWord(text: string, maxLen: number): string { + if (text.length <= maxLen) return text; + const truncated = text.slice(0, maxLen); + const lastSpace = truncated.lastIndexOf(" "); + return (lastSpace > maxLen * 0.7 ? truncated.slice(0, lastSpace) : truncated) + "…"; +} + +function formatMemoryContext(results: MemoryResult[], agentIdentity: string): string { + if (!results.length) return ""; + + const lines = [`## Relevant Memory (hexus, ${agentIdentity})`, ""]; + for (const r of results) { + const score = ((r.score || 0) * 100).toFixed(0); + const content = truncateAtWord((r.content || "").replace(/\n+/g, " ").trim(), 400); + lines.push(`- **[${score}%]** (${r.target || "memory"}) ${content}`); + } + lines.push("", "*Use `hexus_recall` to search, or `hexus_retain` to save.*"); + return lines.join("\n"); +} + +function extractText(content: unknown): string[] { + if (typeof content === "string") return [content]; + if (!Array.isArray(content)) return []; + return content + .filter((b): b is { type: "text"; text: string } => (b as any)?.type === "text") + .map((b) => b.text) + .filter((t): t is string => typeof t === "string"); +} + +function buildConversationText(entries: any[]): string { + // getBranch() may return undefined/null before session state is populated — + // guard so `for...of` doesn't throw TypeError: undefined is not iterable. + if (!Array.isArray(entries)) return ""; + const sections: string[] = []; + for (const e of entries) { + if (e.type !== "message" || !["user", "assistant"].includes(e.message?.role)) continue; + const texts = extractText(e.message?.content ?? []); + if (texts.length) { + sections.push(`${e.message.role === "user" ? "User" : "Assistant"}: ${texts.join("\n")}`); + } + } + return sections.join("\n\n"); +} + +const REFLECTION_PROMPT = `Extract 1-5 durable facts from this conversation that should be remembered. + +Rules: +- New info about the user (preferences, projects, people, events) +- Key decisions or agreements made +- Important context for future sessions +- Things the user explicitly asked to remember + +Format as JSON array: +[{"content": "fact here", "target": "memory|user"}, ...] + +Empty array if nothing worth remembering. Only output JSON.`; + +// --------------------------------------------------------------------------- +// Extension +// --------------------------------------------------------------------------- + +export default function hexus(pi: ExtensionAPI) { + // Bootstrap config asynchronously — initConfig() runs once and caches. + // getConfig() returns defaults until the promise resolves, then returns real config. + initConfig().catch((err) => console.warn("hexus: config load failed:", err)); + const config = getConfig(); + const reflConfig = getReflectionConfig(); + + // State + let sessionId: string | undefined; + let turnsSinceReflection = 0; + let lastReflectionTurn = 0; + let idleTimer: ReturnType | null = null; + let isReflecting = false; + let recallInFlight = false; // Dedupe concurrent recall calls + + // ------------------------------------------------------------------------- + // Reflection + // ------------------------------------------------------------------------- + + async function runReflection(ctx: ExtensionAPI): Promise { + if (isReflecting) return; + isReflecting = true; + + const client = getClient(); + const branch = ctx.sessionManager.getBranch(); + const convText = buildConversationText(branch); + + if (!convText.trim()) { isReflecting = false; return; } + + ctx.ui.setStatus("hexus", "hexus: reflecting..."); + ctx.ui.notify("Running session reflection...", "info"); + + try { + // Split on the FIRST "/" only — model IDs can contain their own "/" + // (e.g. "headroom/tobiTradez/minimax-m2.7-highspeed") and a missing + // separator must not yield an undefined modelId. + const slashIdx = reflConfig.model.indexOf("/"); + const provider = slashIdx > 0 ? reflConfig.model.slice(0, slashIdx) : "headroom"; + const modelId = slashIdx > 0 ? reflConfig.model.slice(slashIdx + 1) : reflConfig.model; + let model = ctx.modelRegistry.find(provider, modelId); + if (!model) { + // Model not found — refresh registry in case litellm is still loading + await ctx.modelRegistry.refresh(); + model = ctx.modelRegistry.find(provider, modelId); + } + if (!model) { + // Try current model as fallback + model = ctx.model; + } + if (!model) { + console.warn(`hexus: model ${reflConfig.model} not found in registry`); + isReflecting = false; + return; + } + + const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model); + if (!auth.ok || !auth.apiKey) { isReflecting = false; return; } + + const resp = await complete(model, + { messages: [{ role: "user" as const, content: [{ type: "text" as const, text: `Extract facts:\n\n${convText.slice(-4000)}` }], timestamp: Date.now() }], systemPrompt: REFLECTION_PROMPT }, + { apiKey: auth.apiKey, headers: auth.headers, env: auth.env } + ); + + const respText = resp.content.filter((c): c is { type: "text"; text: string } => c.type === "text").map(c => c.text).join("\n"); + + let facts: Array<{ content: string; target: string }> = []; + try { + const m = respText.match(/\[[\s\S]*\]/); + if (m) facts = JSON.parse(m[0]); + } catch {} + + if (facts.length > 0) { + const result = await client.retain({ + contents: facts.map(f => f.content), + target: "memory", + metadata: facts.map(f => ({ source: "reflection", target: f.target, reflection_turns: turnsSinceReflection })), + agent_identity: config.agentIdentity, + }); + ctx.ui.notify(`Reflection: saved ${result.inserted} fact${result.inserted !== 1 ? "s" : ""}`, "info"); + } else { + ctx.ui.notify("Reflection: nothing new to remember", "info"); + } + + lastReflectionTurn = turnsSinceReflection; + turnsSinceReflection = 0; + } catch (err) { + console.error("hexus reflection error:", err); + } finally { + isReflecting = false; + const health = await client.health(true).catch(() => null); + ctx.ui.setStatus("hexus", `hexus: ${health?.row_counts.memory_entries ?? "?"} memories`); + if (reflConfig.enabled) scheduleReflection(ctx); + } + } + + function scheduleReflection(ctx: ExtensionAPI): void { + if (idleTimer) clearTimeout(idleTimer); + idleTimer = setTimeout(async () => { + if (!reflConfig.enabled || isReflecting) return; + const usage = ctx.getContextUsage(); + const tokens = usage?.tokens ?? 0; + const sinceLast = turnsSinceReflection - lastReflectionTurn; + if (tokens >= reflConfig.tokenThreshold && sinceLast >= reflConfig.minTurnsBetweenReflections) { + await runReflection(ctx); + } else { + scheduleReflection(ctx); + } + }, reflConfig.idleSeconds * 1000); + } + + // ------------------------------------------------------------------------- + // Events + // ------------------------------------------------------------------------- + + pi.on("session_start", async (_event, ctx) => { + turnsSinceReflection = 0; + lastReflectionTurn = 0; + const client = getClient(); + + const health = await client.health(true).catch(() => null); + if (health?.status === "ok") { + ctx.ui.notify(`hexus connected (${health.row_counts.memory_entries} memories)`, "info"); + ctx.ui.setStatus("hexus", `hexus: ${health.row_counts.memory_entries} memories`); + if (reflConfig.enabled) scheduleReflection(ctx); + } else { + ctx.ui.notify("hexus: connection failed", "warning"); + ctx.ui.setStatus("hexus", "hexus: offline"); + } + }); + + // FIX: Turn counting moved here from input event. + // input fires per keystroke, turn_end fires once per completed turn. + // We also store the assistant message in the same handler. + pi.on("turn_end", async (event, ctx) => { + turnsSinceReflection++; + + // Reset idle timer after a completed turn + if (idleTimer) { clearTimeout(idleTimer); idleTimer = null; } + if (reflConfig.enabled) scheduleReflection(ctx); + + // Store assistant message in hexus + if (config.storeTurns && sessionId) { + const client = getClient(); + const text = (event.message?.content || []) + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map(c => c.text).join("\n"); + + if (text.length >= 20) { + client.appendTurn({ + session_id: sessionId, + role: "assistant", + content: text.slice(0, 4000), + agent_identity: config.agentIdentity, + }).catch(() => {}); + } + } + }); + + pi.on("input", async (event, ctx) => { + // Only reset idle timer on new user input — don't schedule reflection here + // (turn_end handles that to avoid double-scheduling) + if (idleTimer) { clearTimeout(idleTimer); idleTimer = null; } + if (reflConfig.enabled && event.text?.trim()) scheduleReflection(ctx); + }); + + pi.on("before_agent_start", async (event, ctx) => { + const client = getClient(); + + if (!sessionId) { + const sf = ctx.sessionManager.getSessionFile(); + // FIX: Hash the filepath rather than base64-encoding it (avoids leaking paths) + sessionId = sf ? `pi:${hashString(sf)}` : `pi:${hashString(String(Date.now()))}`; + } + + const prompt = event.prompt?.trim(); + if (!prompt) return; + + // FIX: Dedupe concurrent recall calls within the same turn. + if (recallInFlight) return; + recallInFlight = true; + + try { + const health = await client.health().catch(() => null); + if (health?.status !== "ok") { ctx.ui.setStatus("hexus", "hexus: offline"); return; } + ctx.ui.setStatus("hexus", `hexus: ${health.row_counts.memory_entries} memories`); + + const recall = await client.recall({ + query: prompt, + top_k: config.recallLimit, + agent_identity: config.agentIdentity, + min_similarity: config.minSimilarity, + }); + + if (recall.results.length > 0) { + return { message: { customType: "hexus-memory", content: formatMemoryContext(recall.results, config.agentIdentity), display: false } }; + } + } catch (err) { + ctx.ui.setStatus("hexus", `hexus: error (${err instanceof Error ? err.message : "?"})`); + } finally { + recallInFlight = false; + } + }); + // ------------------------------------------------------------------------- + + pi.registerTool({ + name: "hexus_recall", + label: "Recall Memory", + description: "Search hexus vector memory for relevant past entries.", + parameters: Type.Object({ + query: Type.String({ description: "Search query" }), + limit: Type.Optional(Type.Number({ description: "Max results (default 5)" })), + scope: Type.Optional(Type.String({ description: "'current' or 'all'" })), + }), + async execute(_toolCallId, params, _signal, onUpdate) { + const client = getClient(); + onUpdate?.({ content: [{ type: "text", text: "Searching memory..." }] }); + + const [recall, health] = await Promise.all([ + client.recall({ + query: params.query, + top_k: Math.min(params.limit ?? 5, 20), + agent_identity: params.scope === "all" ? undefined : config.agentIdentity, + }), + client.health().catch(() => null), + ]); + + if (!recall.results.length) { + return { + content: [{ type: "text", text: `No memories found for: \"${params.query}\"` }], + details: { count: 0, memoryCount: health?.row_counts.memory_entries ?? null }, + }; + } + + const lines = [`Found ${recall.count} memories:\n`]; + for (const r of recall.results) { + lines.push(`- [${((r.score || 0) * 100).toFixed(0)}%] ${truncateAtWord(r.content, 200)}`); + } + if (health?.status === "ok") { + lines.push(`\n_hexus: ${health.row_counts.memory_entries} total memories_`); + } + return { content: [{ type: "text", text: lines.join("\n") }], details: { count: recall.count, memoryCount: health?.row_counts.memory_entries ?? null } }; + }, + }); + + pi.registerTool({ + name: "hexus_retain", + label: "Retain Memory", + description: "Save important info to hexus memory.", + parameters: Type.Object({ + content: Type.String({ description: "What to remember" }), + target: Type.Optional(Type.String({ description: "'memory' or 'user'" })), + }), + async execute(_toolCallId, params, _signal, onUpdate) { + const client = getClient(); + onUpdate?.({ content: [{ type: "text", text: "Saving..." }] }); + + const [result, health] = await Promise.all([ + client.retain({ + contents: [params.content], + target: params.target ?? "memory", + agent_identity: config.agentIdentity, + }), + client.health(true).catch(() => null), // Refresh health after write + ]); + + const status = result.inserted > 0 + ? `Saved (${result.inserted} new, ${result.duplicates} dupes)` + : "Already exists"; + const healthLine = health?.status === "ok" ? ` | hexus: ${health.row_counts.memory_entries} memories` : ""; + + return { + content: [{ type: "text", text: `${status}${healthLine}` }], + details: result, + }; + }, + }); + + // Manual reflection trigger for testing + pi.registerCommand("reflect", { + description: "Manually trigger session reflection to extract and store facts", + handler: async (_args, ctx) => { + if (!reflConfig.enabled) { + ctx.ui.notify("Reflection is disabled (HEXUS_REFLECTION_ENABLED=false)", "warning"); + return; + } + ctx.ui.notify("Running reflection...", "info"); + await runReflection(ctx); + }, + }); + + pi.on("session_shutdown", () => { + if (idleTimer) { clearTimeout(idleTimer); idleTimer = null; } + // Reset in-flight state so a new session starts clean + isReflecting = false; + recallInFlight = false; + }); +} diff --git a/pi-extension/package-lock.json b/pi-extension/package-lock.json new file mode 100644 index 0000000..e274d9a --- /dev/null +++ b/pi-extension/package-lock.json @@ -0,0 +1,21 @@ +{ + "name": "hexus", + "version": "0.1.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "hexus", + "version": "0.1.0", + "dependencies": { + "typebox": "^1.0.0" + } + }, + "node_modules/typebox": { + "version": "1.3.8", + "resolved": "https://registry.npmjs.org/typebox/-/typebox-1.3.8.tgz", + "integrity": "sha512-xYaJgF0KMvBViKWRaKTAtfR6sDt/yH6xAjGAHXYJxKxUF4pVyPXZZhIlwOg6cDIE81N7L0pyIHs136RHBm6rLQ==", + "license": "MIT" + } + } +} diff --git a/pi-extension/package.json b/pi-extension/package.json new file mode 100644 index 0000000..9b337d8 --- /dev/null +++ b/pi-extension/package.json @@ -0,0 +1,11 @@ +{ + "name": "hexus", + "version": "0.1.0", + "description": "hexus — Vector memory extension for pi", + "dependencies": { + "typebox": "^1.0.0" + }, + "pi": { + "extensions": ["./index.ts"] + } +} diff --git a/tests/test_embedder.py b/tests/test_embedder.py index a96ccb3..49bc899 100644 --- a/tests/test_embedder.py +++ b/tests/test_embedder.py @@ -15,12 +15,10 @@ import os import threading -from typing import List from unittest.mock import patch import pytest - # --------------------------------------------------------------------------- # Structural / fast tests — no model load # --------------------------------------------------------------------------- @@ -29,7 +27,7 @@ def test_constants(): """The public constants are pinned to the values the rest of the code (and the schema migration) assume.""" - from hexus.embedder import DEFAULT_MODEL, DEFAULT_DIM + from hexus.embedder import DEFAULT_DIM, DEFAULT_MODEL assert DEFAULT_MODEL == "sentence-transformers/all-MiniLM-L6-v2" assert DEFAULT_DIM == 384 @@ -134,7 +132,7 @@ def test_singleton_is_thread_safe(): from hexus.embedder import get_default_embedder, reset_default_embedder reset_default_embedder() - instances: List = [] + instances: list = [] barrier = threading.Barrier(8) def grab(): @@ -226,7 +224,7 @@ def _embedder_with_model(mode, max_seq=8): def test_long_text_mode_default_is_warn(): - from hexus.embedder import LocalBertEmbedder, DEFAULT_LONG_TEXT_MODE + from hexus.embedder import DEFAULT_LONG_TEXT_MODE, LocalBertEmbedder assert LocalBertEmbedder().long_text_mode == DEFAULT_LONG_TEXT_MODE == "warn" @@ -244,6 +242,7 @@ def test_long_text_mode_env_override(monkeypatch): def test_long_text_mode_invalid_falls_back(caplog): import logging + from hexus.embedder import LocalBertEmbedder with caplog.at_level(logging.WARNING, logger="hexus.embedder"): @@ -586,9 +585,10 @@ def embed(self, texts): def test_embed_http_404_raises_embedding_error(): """The HTTP path raises EmbeddingError on a non-2xx response, not a urllib.error.HTTPError leaking out.""" + import urllib.error + from hexus import embed as embed_fn from hexus.embed import EmbeddingError - import urllib.error with patch("urllib.request.urlopen") as urlopen: urlopen.side_effect = urllib.error.HTTPError( diff --git a/tests/test_http_auth.py b/tests/test_http_auth.py index 6cc2cd7..d8eaf2b 100644 --- a/tests/test_http_auth.py +++ b/tests/test_http_auth.py @@ -16,7 +16,7 @@ # DB connection just to reach the auth wrapper. pytest.importorskip("psycopg") -from mcp_server.server import _wrap_with_bearer_auth # noqa: E402 +from mcp_server.server import _wrap_with_bearer_auth def _make_downstream(): @@ -111,8 +111,8 @@ async def send(msg): # Server-derived caller identity (issue #19 item A) # ----------------------------------------------------------------------- -from mcp_server.server import _wrap_with_identity # noqa: E402 -from mcp_server import tools # noqa: E402 +from mcp_server import tools +from mcp_server.server import _wrap_with_identity def _identity_scope(session_key=None): @@ -206,7 +206,7 @@ def test_read_identity_strict_confines_to_caller(monkeypatch): # SQL-safety helpers + isolation policy (issue #19), no DB required # ----------------------------------------------------------------------- -from hexus.store import _escape_like, _resolve_isolation # noqa: E402 +from hexus.store import _escape_like, _resolve_isolation def test_escape_like_neutralizes_wildcards(): diff --git a/tests/test_import_cli.py b/tests/test_import_cli.py index d06ef69..51fad81 100644 --- a/tests/test_import_cli.py +++ b/tests/test_import_cli.py @@ -1,6 +1,8 @@ import json -import pytest from unittest.mock import MagicMock + +import pytest + from hexus.store import MemoryStore from mcp_server.import_cli import main diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index cd2630e..d795986 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -66,7 +66,7 @@ def store(): yield s # Best-effort cleanup of any rows this test wrote. try: - with s._get_pool().connection() as conn: # noqa: SLF001 + with s._get_pool().connection() as conn: with conn.cursor() as cur: cur.execute( "DELETE FROM memory_entries WHERE agent_identity = %s", @@ -119,7 +119,7 @@ def test_memory_retain_inserts_rows(store): { "contents": ["alpha bravo", "charlie delta"], "target": "memory", - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), "metadata": {"doc_type": "document", "source_url": "https://example.com/a"}, }, ) @@ -134,7 +134,7 @@ def test_memory_retain_dedupes_on_exact_repeat(store): args = { "contents": ["unique content one"], "target": "memory", - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), } first = tools.memory_retain(store, args) second = tools.memory_retain(store, args) @@ -161,7 +161,7 @@ def test_memory_retain_rejects_bad_target(store): { "contents": ["x"], "target": "bogus", - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) @@ -176,7 +176,7 @@ def test_empty_string_target_defaults_to_both_stores(store): { "contents": ["Postgres + hexus is great for semantic search."], "target": "memory", - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) @@ -185,7 +185,7 @@ def test_empty_string_target_defaults_to_both_stores(store): { "query": "hexus semantic search", "top_k": 3, - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), "target": "", }, ) @@ -195,7 +195,7 @@ def test_empty_string_target_defaults_to_both_stores(store): count = tools.memory_count( store, { - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), "target": "", }, ) @@ -216,7 +216,7 @@ def test_memory_recall_round_trip(store): store, { "contents": docs, - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) @@ -225,7 +225,7 @@ def test_memory_recall_round_trip(store): { "query": "what does hexus do", "top_k": 3, - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) assert out["count"] == 3 @@ -250,7 +250,7 @@ def test_memory_hybrid_search_round_trip(store): store, { "contents": docs, - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) @@ -261,7 +261,7 @@ def test_memory_hybrid_search_round_trip(store): "top_k": 3, "vector_weight": 0.5, "text_weight": 0.5, - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) assert out["count"] >= 1 @@ -280,7 +280,7 @@ def test_memory_hybrid_recall_turns_round_trip(store): "session_id": "session-123", "role": "user", "content": "My favorite database is Postgres.", - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) tools.memory_append_turn( @@ -289,7 +289,7 @@ def test_memory_hybrid_recall_turns_round_trip(store): "session_id": "session-123", "role": "assistant", "content": "I prefer local BERT embeddings.", - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) @@ -300,7 +300,7 @@ def test_memory_hybrid_recall_turns_round_trip(store): "top_k": 2, "vector_weight": 0.5, "text_weight": 0.5, - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) assert out["count"] >= 1 @@ -313,7 +313,7 @@ def test_memory_hybrid_recall_turns_round_trip(store): def test_memory_delegation_round_trip(store): from mcp_server import tools - identity = agent_of(store) # noqa: SLF001 + identity = agent_of(store) rec = tools.memory_record_delegation( store, { @@ -350,7 +350,7 @@ def test_memory_recall_respects_min_similarity(store): store, { "contents": ["Postgres is a relational database."], - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) out = tools.memory_recall( @@ -358,7 +358,7 @@ def test_memory_recall_respects_min_similarity(store): { "query": "Postgres relational database", "top_k": 5, - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), "min_similarity": 0.5, }, ) @@ -373,7 +373,7 @@ def test_memory_recall_caps_top_k(store): store, { "contents": [f"entry {i}" for i in range(10)], - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) out = tools.memory_recall( @@ -381,7 +381,7 @@ def test_memory_recall_caps_top_k(store): { "query": "entry", "top_k": 10000, # way over the cap - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) # Capped at 100, but we only have 10 rows so count is 10. @@ -405,7 +405,7 @@ def test_memory_recall_cross_agent_returns_other_agents_rows(store): store, { "contents": ["the meaning of life is forty two"], - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) # Recall with agent_identity="" (None) — should find it. @@ -424,13 +424,13 @@ def test_memory_search_browse(store): store, { "contents": [f"row {i}" for i in range(5)], - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) out = tools.memory_search( store, { - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), "limit": 3, }, ) @@ -445,7 +445,7 @@ def test_memory_search_browse(store): out_page2 = tools.memory_search( store, { - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), "limit": 2, "offset": 3, }, @@ -464,22 +464,22 @@ def test_memory_forget_dry_run_by_default(store): store, { "contents": ["to be deleted"], - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) assert out["inserted"] == 1 # Look up the row id. - listed = store.list_entries(agent_identity=agent_of(store), limit=1) # noqa: SLF001 + listed = store.list_entries(agent_identity=agent_of(store), limit=1) row_id = listed[0]["id"] dry = tools.memory_forget( store, - {"id": row_id, "agent_identity": agent_of(store), "confirm": False}, # noqa: SLF001 + {"id": row_id, "agent_identity": agent_of(store), "confirm": False}, ) assert dry["dry_run"] is True assert dry["deleted"] == 0 # Row is still there. - after = store.list_entries(agent_identity=agent_of(store), limit=10) # noqa: SLF001 + after = store.list_entries(agent_identity=agent_of(store), limit=10) assert any(r["id"] == row_id for r in after) @@ -490,20 +490,20 @@ def test_memory_forget_actually_deletes_with_confirm(store): store, { "contents": ["goodbye cruel world"], - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) assert out["inserted"] == 1 - row_id = store.list_entries(agent_identity=agent_of(store), limit=1)[0]["id"] # noqa: SLF001 + row_id = store.list_entries(agent_identity=agent_of(store), limit=1)[0]["id"] real = tools.memory_forget( store, - {"id": row_id, "agent_identity": agent_of(store), "confirm": True}, # noqa: SLF001 + {"id": row_id, "agent_identity": agent_of(store), "confirm": True}, ) assert real["dry_run"] is False assert real["deleted"] == 1 # And the row is gone. - after = store.list_entries(agent_identity=agent_of(store), limit=10) # noqa: SLF001 + after = store.list_entries(agent_identity=agent_of(store), limit=10) assert not any(r["id"] == row_id for r in after) @@ -529,7 +529,7 @@ def test_memory_append_turn_and_recall_turns(store): "session_id": "sess-abc", "role": "user", "content": "I love using Postgres for memory", - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) assert isinstance(out["id"], int) and out["id"] > 0 @@ -539,7 +539,7 @@ def test_memory_append_turn_and_recall_turns(store): "session_id": "sess-abc", "role": "assistant", "content": "Glad to hear! hexus is a great fit.", - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) assert out2["id"] > out["id"] @@ -549,7 +549,7 @@ def test_memory_append_turn_and_recall_turns(store): { "query": "Postgres memory", "top_k": 5, - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) assert recall["count"] >= 1 @@ -567,7 +567,7 @@ def test_memory_append_turn_validates_role(store): "session_id": "s1", "role": "pirate", "content": "arrr", - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) @@ -579,7 +579,7 @@ def test_memory_count_scopes_correctly(store): store, { "contents": [f"row {i}" for i in range(3)], - "agent_identity": agent_of(store), # noqa: SLF001 + "agent_identity": agent_of(store), }, ) out = tools.memory_count(store, {"agent_identity": agent_of(store)}) @@ -638,7 +638,7 @@ def test_multi_agent_isolation(store): # Cleanup. for a in (agent_a, agent_b): - with store._get_pool().connection() as conn: # noqa: SLF001 + with store._get_pool().connection() as conn: with conn.cursor() as cur: cur.execute( "DELETE FROM memory_entries WHERE agent_identity = %s", (a,) @@ -769,7 +769,7 @@ def test_mcp_graph_tools_round_trip(store): """Verify that we can call entity_graph, graph_walk, and common_topics via FastMCP handlers.""" from mcp_server import tools - agent = agent_of(store) # noqa: SLF001 + agent = agent_of(store) # Insert docs containing entities tools.memory_retain( @@ -825,7 +825,7 @@ def test_mcp_confirm_reject_summarize_round_trip(store): """Verify that we can call memory_confirm, memory_reject, and memory_summarize_session via FastMCP handlers.""" from mcp_server import tools - agent = agent_of(store) # noqa: SLF001 + agent = agent_of(store) # Retain entry tools.memory_retain( @@ -946,8 +946,9 @@ def test_mcp_memory_stats_round_trip(store): def test_mcp_server_cleanup_thread_startup(monkeypatch): """Verify that build_server spawns the background cleanup thread if configured.""" - import threading import os + import threading + from mcp_server.server import build_server # Set environment variables to enable cleanup @@ -969,8 +970,8 @@ def test_mcp_server_cleanup_thread_startup(monkeypatch): def test_mcp_recall_dynamic_decay(store): """Verify that memory_recall tool respects decay_half_life_days dynamic parameter.""" - from mcp_server import tools from hexus.store import dict_row + from mcp_server import tools agent = agent_of(store) @@ -993,7 +994,7 @@ def test_mcp_recall_dynamic_decay(store): ) # Retrieve their IDs from database to modify updated_at - with store._get_pool().connection() as conn: + with store._get_pool().connection() as conn: # noqa: SIM117 with ( conn.cursor(row_factory=dict_row) if hasattr(conn, "cursor") @@ -1054,8 +1055,9 @@ def test_mcp_recall_dynamic_decay(store): def test_mcp_server_consolidation_thread_startup(monkeypatch): """Verify that build_server spawns the background consolidation thread if configured.""" - import threading import os + import threading + from mcp_server.server import build_server monkeypatch.setenv("HEXUS_CONSOLIDATION_INTERVAL_HOURS", "1") @@ -1072,8 +1074,9 @@ def test_mcp_server_consolidation_thread_startup(monkeypatch): def test_mcp_consolidation_round_trip(store, monkeypatch): """Verify that consolidation works and updates DB and metrics correctly.""" - import urllib.request import json + import urllib.request + from hexus.store import dict_row # 1. Setup mock environment @@ -1093,12 +1096,9 @@ def test_mcp_consolidation_round_trip(store, monkeypatch): # 2. Insert some low confidence memories (rejected count > 0, confirm count = 0) agent = "test-agent-consolidation" - with store._get_pool().connection() as conn: - with conn.cursor() as cur: - cur.execute( - "DELETE FROM memory_entries WHERE agent_identity = %s", (agent,) - ) - conn.commit() + with store._get_pool().connection() as conn, conn.cursor() as cur: + cur.execute("DELETE FROM memory_entries WHERE agent_identity = %s", (agent,)) + conn.commit() # Insert entry to delete store.add( @@ -1115,17 +1115,17 @@ def test_mcp_consolidation_round_trip(store, monkeypatch): ) # Get IDs of inserted entries - with store._get_pool().connection() as conn: - with ( - conn.cursor(row_factory=dict_row) - if hasattr(conn, "cursor") - else conn.cursor() as cur - ): - cur.execute( - "SELECT id, content FROM memory_entries WHERE agent_identity = %s", - (agent,), - ) - rows = list(cur.fetchall()) + with ( + store._get_pool().connection() as conn, + conn.cursor(row_factory=dict_row) + if hasattr(conn, "cursor") + else conn.cursor() as cur, + ): + cur.execute( + "SELECT id, content FROM memory_entries WHERE agent_identity = %s", + (agent,), + ) + rows = list(cur.fetchall()) assert len(rows) == 3 @@ -1134,16 +1134,15 @@ def test_mcp_consolidation_round_trip(store, monkeypatch): id_rep2 = next(r["id"] for r in rows if "Second" in r["content"]) # Update metadata to make them low confidence (reject_count = 1, confirm_count = 0) - with store._get_pool().connection() as conn: - with conn.cursor() as cur: - cur.execute( - "UPDATE memory_entries SET metadata = %s WHERE id = ANY(%s)", - ( - json.dumps({"reject_count": 1, "confirm_count": 0}), - [id_delete, id_rep1, id_rep2], - ), - ) - conn.commit() + with store._get_pool().connection() as conn, conn.cursor() as cur: + cur.execute( + "UPDATE memory_entries SET metadata = %s WHERE id = ANY(%s)", + ( + json.dumps({"reject_count": 1, "confirm_count": 0}), + [id_delete, id_rep1, id_rep2], + ), + ) + conn.commit() # Define the mock LLM response data llm_resp_data = { @@ -1192,17 +1191,17 @@ def mock_urlopen(req, *args, **kwargs): assert res["replacements"] == 2 # Assert database is updated: id_delete, id_rep1, id_rep2 should be gone - with store._get_pool().connection() as conn: - with ( - conn.cursor(row_factory=dict_row) - if hasattr(conn, "cursor") - else conn.cursor() as cur - ): - cur.execute( - "SELECT id, content FROM memory_entries WHERE agent_identity = %s", - (agent,), - ) - new_rows = list(cur.fetchall()) + with ( + store._get_pool().connection() as conn, + conn.cursor(row_factory=dict_row) + if hasattr(conn, "cursor") + else conn.cursor() as cur, + ): + cur.execute( + "SELECT id, content FROM memory_entries WHERE agent_identity = %s", + (agent,), + ) + new_rows = list(cur.fetchall()) # Should have exactly 1 entry: the consolidated python text assert len(new_rows) == 1 @@ -1278,7 +1277,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): assert res_co["replacements"] == 3 # Assert database is updated: the 3 Docker facts are gone and the consolidated one is added - with store._get_pool().connection() as conn: + with store._get_pool().connection() as conn: # noqa: SIM117 with ( conn.cursor(row_factory=dict_row) if hasattr(conn, "cursor") @@ -1291,7 +1290,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): docker_rows = list(cur.fetchall()) assert len(docker_rows) == 0 - with store._get_pool().connection() as conn: + with store._get_pool().connection() as conn: # noqa: SIM117 with ( conn.cursor(row_factory=dict_row) if hasattr(conn, "cursor") @@ -1307,8 +1306,9 @@ def __exit__(self, exc_type, exc_val, exc_tb): def test_memory_consolidate_tool(store, monkeypatch): """Verify that the memory_consolidate tool can be called and executes consolidation.""" - import urllib.request import json + import urllib.request + from mcp_server.server import _build_server monkeypatch.setenv("HEXUS_SUMMARY_MODEL", "mock-summary-model") diff --git a/tests/test_migration.py b/tests/test_migration.py index 7656f90..ba2c8fc 100644 --- a/tests/test_migration.py +++ b/tests/test_migration.py @@ -19,7 +19,6 @@ import pytest - MIGRATION_PATH = ( Path(__file__).resolve().parent.parent / "hexus" / "migrations" / "001_schema.sql" ) @@ -247,11 +246,14 @@ def test_insert_768_dim_vector_rejected(pg): "psql", pg, "-tAc", - f"INSERT INTO memory_entries (agent_identity, target, content, embedding) " - f"VALUES ('{agent}', 'memory', 'dim mismatch test', '{vec}'::vector)", + ( + f"INSERT INTO memory_entries (agent_identity, target, content, embedding) " + f"VALUES ('{agent}', 'memory', 'dim mismatch test', '{vec}'::vector)" + ), ], capture_output=True, text=True, + check=False, env=_psql_env(pg), ) assert result.returncode != 0, "DB accepted a 768-dim vector; expected rejection" diff --git a/tests/test_quantization.py b/tests/test_quantization.py index 1d1c3ca..55ddd7e 100644 --- a/tests/test_quantization.py +++ b/tests/test_quantization.py @@ -1,6 +1,8 @@ import os -import pytest + import psycopg +import pytest + from hexus.store import MemoryStore # Skip the whole module if there's no DSN to talk to. @@ -16,11 +18,10 @@ def clean_db(): if not dsn: pytest.skip("PG_TEST_DSN not set") - with psycopg.connect(dsn, autocommit=True) as conn: - with conn.cursor() as cur: - cur.execute("DROP TABLE IF EXISTS delegations CASCADE;") - cur.execute("DROP TABLE IF EXISTS conversations CASCADE;") - cur.execute("DROP TABLE IF EXISTS memory_entries CASCADE;") + with psycopg.connect(dsn, autocommit=True) as conn, conn.cursor() as cur: + cur.execute("DROP TABLE IF EXISTS delegations CASCADE;") + cur.execute("DROP TABLE IF EXISTS conversations CASCADE;") + cur.execute("DROP TABLE IF EXISTS memory_entries CASCADE;") # Re-apply migrations using apply_migration_as_admin s = MemoryStore(dsn) @@ -39,25 +40,24 @@ def test_quantization_float16_adaptation(clean_db, monkeypatch): store.ensure_schema() # This calls adapt_vector_precision() # 1. Verify column type is halfvec(384) - with store._get_pool().connection() as conn: - with conn.cursor() as cur: - cur.execute(""" + with store._get_pool().connection() as conn, conn.cursor() as cur: + cur.execute(""" SELECT pg_catalog.format_type(atttypid, atttypmod) FROM pg_catalog.pg_attribute WHERE attrelid = 'memory_entries'::regclass AND attname = 'embedding'; """) - col_type = cur.fetchone()[0] - assert col_type == "halfvec(384)" + col_type = cur.fetchone()[0] + assert col_type == "halfvec(384)" - cur.execute(""" + cur.execute(""" SELECT pg_catalog.format_type(atttypid, atttypmod) FROM pg_catalog.pg_attribute WHERE attrelid = 'conversations'::regclass AND attname = 'embedding'; """) - conv_col_type = cur.fetchone()[0] - assert conv_col_type == "halfvec(384)" + conv_col_type = cur.fetchone()[0] + assert conv_col_type == "halfvec(384)" # 2. Verify we can insert and search successfully agent = "test-agent-float16" @@ -80,22 +80,21 @@ def test_quantization_binary_adaptation(clean_db, monkeypatch): store.ensure_schema() # This calls adapt_vector_precision() # 1. Verify binary index exists - with store._get_pool().connection() as conn: - with conn.cursor() as cur: - cur.execute(""" + with store._get_pool().connection() as conn, conn.cursor() as cur: + cur.execute(""" SELECT indexname FROM pg_indexes WHERE tablename = 'memory_entries' AND indexname = 'ix_memory_entries_embedding_binary_hnsw'; """) - assert cur.fetchone() is not None + assert cur.fetchone() is not None - # Verify standard cosine index does NOT exist - cur.execute(""" + # Verify standard cosine index does NOT exist + cur.execute(""" SELECT indexname FROM pg_indexes WHERE tablename = 'memory_entries' AND indexname = 'ix_memory_entries_embedding_hnsw'; """) - assert cur.fetchone() is None + assert cur.fetchone() is None # 2. Verify we can insert and search (using two-stage search) successfully agent = "test-agent-binary" diff --git a/tests/test_rerank.py b/tests/test_rerank.py index 06fa408..5551b05 100644 --- a/tests/test_rerank.py +++ b/tests/test_rerank.py @@ -18,16 +18,15 @@ # store-backed tests in this suite. psycopg = pytest.importorskip("psycopg") -from hexus.store import ( # noqa: E402 +from hexus.store import ( RERANK_MAX_PASSAGES, - rerank_scores, + _cross_encoder_max_len, + _resolve_rerank_mode, get_rerank_stats, + rerank_scores, reset_rerank_stats, - _resolve_rerank_mode, - _cross_encoder_max_len, ) - QUERY = "hi there" # 2 tokens → budget = 32 - 2 - 3 = 27 diff --git a/tests/test_smoke.py b/tests/test_smoke.py index 2615969..00d7f25 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -15,6 +15,7 @@ from __future__ import annotations import os +from datetime import UTC from typing import cast import pytest @@ -60,7 +61,7 @@ def test_hexus_literal_roundtrip(): def test_embed_empty_input_raises(): - from hexus.embed import embed, EmbeddingError + from hexus.embed import EmbeddingError, embed with pytest.raises(EmbeddingError): embed("", base_url="http://localhost:11434") @@ -118,12 +119,9 @@ def store(): # Cleanup import psycopg - with psycopg.connect(dsn) as conn: - with conn.cursor() as cur: - cur.execute( - "DELETE FROM memory_entries WHERE agent_identity = %s", (agent,) - ) - conn.commit() + with psycopg.connect(dsn) as conn, conn.cursor() as cur: + cur.execute("DELETE FROM memory_entries WHERE agent_identity = %s", (agent,)) + conn.commit() def test_health_reports_ok(store): @@ -218,13 +216,12 @@ def test_search_cross_agent_scope(store): # Cleanup the second agent — primary fixture only knows about `agent` import psycopg - with psycopg.connect(s._dsn) as conn: - with conn.cursor() as cur: - cur.execute( - "DELETE FROM memory_entries WHERE agent_identity = %s", - (other_agent,), - ) - conn.commit() + with psycopg.connect(s._dsn) as conn, conn.cursor() as cur: + cur.execute( + "DELETE FROM memory_entries WHERE agent_identity = %s", + (other_agent,), + ) + conn.commit() def test_count_filters(store): @@ -290,10 +287,9 @@ def store_with_turn_cleanup(store): yield s, agent import psycopg - with psycopg.connect(s._dsn) as conn: - with conn.cursor() as cur: - cur.execute("DELETE FROM conversations WHERE agent_identity = %s", (agent,)) - conn.commit() + with psycopg.connect(s._dsn) as conn, conn.cursor() as cur: + cur.execute("DELETE FROM conversations WHERE agent_identity = %s", (agent,)) + conn.commit() def test_append_turn_and_count(store_with_turn_cleanup): @@ -443,17 +439,17 @@ def test_search_with_decay_and_boost(store): ) # Backdate old entry by 10 days + from datetime import datetime, timedelta + import psycopg - from datetime import datetime, timedelta, timezone - old_time = datetime.now(timezone.utc) - timedelta(days=10) - with psycopg.connect(s._dsn) as conn: - with conn.cursor() as cur: - cur.execute( - "UPDATE memory_entries SET created_at = %s, updated_at = %s WHERE id = %s", - (old_time, old_time, id_old), - ) - conn.commit() + old_time = datetime.now(UTC) - timedelta(days=10) + with psycopg.connect(s._dsn) as conn, conn.cursor() as cur: + cur.execute( + "UPDATE memory_entries SET created_at = %s, updated_at = %s WHERE id = %s", + (old_time, old_time, id_old), + ) + conn.commit() # Search with decay. The new entry should score higher. rows = s.search(query_embedding=vec, agent_identity=agent, decay_half_life_days=5.0) @@ -462,11 +458,10 @@ def test_search_with_decay_and_boost(store): assert rows[0]["score"] > rows[1]["score"] # Verify that searching incremented the recall count - with psycopg.connect(s._dsn) as conn: - with conn.cursor() as cur: - cur.execute("SELECT metadata FROM memory_entries WHERE id = %s", (id_new,)) - meta = cur.fetchone()[0] or {} - assert int(meta.get("recall_count", 0)) == 1 + with psycopg.connect(s._dsn) as conn, conn.cursor() as cur: + cur.execute("SELECT metadata FROM memory_entries WHERE id = %s", (id_new,)) + meta = cur.fetchone()[0] or {} + assert int(meta.get("recall_count", 0)) == 1 def test_search_turns_with_decay_and_boost(store_with_turn_cleanup): @@ -488,17 +483,17 @@ def test_search_turns_with_decay_and_boost(store_with_turn_cleanup): ) # Backdate old turn by 10 days + from datetime import datetime, timedelta + import psycopg - from datetime import datetime, timedelta, timezone - old_time = datetime.now(timezone.utc) - timedelta(days=10) - with psycopg.connect(s._dsn) as conn: - with conn.cursor() as cur: - cur.execute( - "UPDATE conversations SET ts = %s WHERE id = %s", - (old_time, id_old), - ) - conn.commit() + old_time = datetime.now(UTC) - timedelta(days=10) + with psycopg.connect(s._dsn) as conn, conn.cursor() as cur: + cur.execute( + "UPDATE conversations SET ts = %s WHERE id = %s", + (old_time, id_old), + ) + conn.commit() # Search turns with decay. The new turn should score higher. rows = s.search_turns( @@ -509,11 +504,10 @@ def test_search_turns_with_decay_and_boost(store_with_turn_cleanup): assert rows[0]["score"] > rows[1]["score"] # Verify that search_turns incremented the recall count - with psycopg.connect(s._dsn) as conn: - with conn.cursor() as cur: - cur.execute("SELECT metadata FROM conversations WHERE id = %s", (id_new,)) - meta = cur.fetchone()[0] or {} - assert int(meta.get("recall_count", 0)) == 1 + with psycopg.connect(s._dsn) as conn, conn.cursor() as cur: + cur.execute("SELECT metadata FROM conversations WHERE id = %s", (id_new,)) + meta = cur.fetchone()[0] or {} + assert int(meta.get("recall_count", 0)) == 1 def test_cleanup_stale_records(store_with_turn_cleanup): @@ -535,30 +529,27 @@ def test_cleanup_stale_records(store_with_turn_cleanup): ) # Backdate old turn by 10 days + from datetime import datetime, timedelta + import psycopg - from datetime import datetime, timedelta, timezone - old_time = datetime.now(timezone.utc) - timedelta(days=10) - with psycopg.connect(s._dsn) as conn: - with conn.cursor() as cur: - cur.execute( - "UPDATE conversations SET ts = %s WHERE id = %s", - (old_time, id_old_conv), - ) - conn.commit() + old_time = datetime.now(UTC) - timedelta(days=10) + with psycopg.connect(s._dsn) as conn, conn.cursor() as cur: + cur.execute( + "UPDATE conversations SET ts = %s WHERE id = %s", + (old_time, id_old_conv), + ) + conn.commit() # Clean up records older than 5 days res = s.cleanup_stale_records(conversations_ttl_days=5) assert res["conversations"] == 1 - with psycopg.connect(s._dsn) as conn: - with conn.cursor() as cur: - cur.execute( - "SELECT id FROM conversations WHERE agent_identity = %s", (agent,) - ) - remaining = [r[0] for r in cur.fetchall()] - assert id_old_conv not in remaining - assert id_new_conv in remaining + with psycopg.connect(s._dsn) as conn, conn.cursor() as cur: + cur.execute("SELECT id FROM conversations WHERE agent_identity = %s", (agent,)) + remaining = [r[0] for r in cur.fetchall()] + assert id_old_conv not in remaining + assert id_new_conv in remaining def test_search_with_platform_filter(store): @@ -614,22 +605,21 @@ def test_entity_tagging_and_graph(store): # 1. Verify entity tagging extracted entities automatically import psycopg - with psycopg.connect(s._dsn) as conn: - with conn.cursor() as cur: - cur.execute( - "SELECT metadata FROM memory_entries WHERE agent_identity = %s", - (agent,), - ) - rows = cur.fetchall() - all_entities = [] - for r in rows: - meta = r[0] or {} - entities = meta.get("entities", []) - all_entities.extend(entities) + with psycopg.connect(s._dsn) as conn, conn.cursor() as cur: + cur.execute( + "SELECT metadata FROM memory_entries WHERE agent_identity = %s", + (agent,), + ) + rows = cur.fetchall() + all_entities = [] + for r in rows: + meta = r[0] or {} + entities = meta.get("entities", []) + all_entities.extend(entities) - entity_types = [e["type"] for e in all_entities] - assert "url" in entity_types - assert "file_path" in entity_types + entity_types = [e["type"] for e in all_entities] + assert "url" in entity_types + assert "file_path" in entity_types # 2. Verify entity_graph co-occurrence res = s.entity_graph( @@ -729,19 +719,18 @@ def test_append_turn_entity_extraction(store): # Verify entities are extracted in metadata import psycopg - with psycopg.connect(s._dsn) as conn: - with conn.cursor() as cur: - cur.execute( - "SELECT metadata FROM conversations WHERE session_id = %s", - (session_id,), - ) - row = cur.fetchone() - assert row is not None - meta = row[0] or {} - entities = meta.get("entities", []) - types = [e["type"] for e in entities] - assert "url" in types - assert "file_path" in types + with psycopg.connect(s._dsn) as conn, conn.cursor() as cur: + cur.execute( + "SELECT metadata FROM conversations WHERE session_id = %s", + (session_id,), + ) + row = cur.fetchone() + assert row is not None + meta = row[0] or {} + entities = meta.get("entities", []) + types = [e["type"] for e in entities] + assert "url" in types + assert "file_path" in types def test_headroom_compression_and_retrieve(store): @@ -823,10 +812,11 @@ def test_score_blending(store): ) # Old entry (decayed) + from datetime import datetime, timedelta + import psycopg - from datetime import datetime, timezone, timedelta - old_time = datetime.now(timezone.utc) - timedelta(days=10) + old_time = datetime.now(UTC) - timedelta(days=10) row_id_2 = s.add( agent_identity=agent, target="memory", @@ -834,13 +824,12 @@ def test_score_blending(store): embedding=vec, ) # Manually backdate created_at and updated_at - with psycopg.connect(s._dsn) as conn: - with conn.cursor() as cur: - cur.execute( - "UPDATE memory_entries SET created_at = %s, updated_at = %s WHERE id = %s", - (old_time, old_time, row_id_2), - ) - conn.commit() + with psycopg.connect(s._dsn) as conn, conn.cursor() as cur: + cur.execute( + "UPDATE memory_entries SET created_at = %s, updated_at = %s WHERE id = %s", + (old_time, old_time, row_id_2), + ) + conn.commit() # Query with hybrid search and decay_half_life_days active results = s.hybrid_search( @@ -898,7 +887,7 @@ def test_plugin_memory_stats(store): with patch("hexus.MemoryProvider", new=object): from hexus import HexusMemoryProvider - s, agent = store + s, _agent = store # Mock/fake config for HexusMemoryProvider provider = HexusMemoryProvider( diff --git a/tests/test_webhooks.py b/tests/test_webhooks.py index 294eb2e..c6d7e8e 100644 --- a/tests/test_webhooks.py +++ b/tests/test_webhooks.py @@ -2,18 +2,23 @@ import hmac import json import threading -from http.server import HTTPServer, BaseHTTPRequestHandler -from typing import List, Dict, Any -from unittest.mock import patch, MagicMock +from http.server import BaseHTTPRequestHandler, HTTPServer +from typing import Any +from unittest.mock import MagicMock, patch import pytest -from hexus.webhook.dispatcher import sign_payload, dispatch_webhook_sync + +from hexus.webhook.dispatcher import dispatch_webhook_sync, sign_payload class MockWebhookHandler(BaseHTTPRequestHandler): - requests_received: List[Dict[str, Any]] = [] + requests_received: list[dict[str, Any]] | None = None response_status = 200 + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.requests_received = [] + def do_POST(self): content_length = int(self.headers.get("Content-Length", 0)) body = self.rfile.read(content_length) @@ -21,7 +26,7 @@ def do_POST(self): # Parse payload try: payload = json.loads(body.decode("utf-8")) - except Exception: + except Exception: # noqa: BLE001 payload = None MockWebhookHandler.requests_received.append( @@ -131,8 +136,8 @@ def test_provider_worker_dispatches_webhooks(mock_sleep, mock_webhook_server): url, handler = mock_webhook_server from hexus import HexusMemoryProvider - from hexus.writer import _PendingWrite from hexus.webhook.dispatcher import dispatch_webhook_sync + from hexus.writer import _PendingWrite provider = HexusMemoryProvider( config={ @@ -202,9 +207,10 @@ def test_provider_worker_dispatches_webhooks(mock_sleep, mock_webhook_server): def test_mcp_tools_dispatch_webhooks(mock_sleep, mock_webhook_server): url, handler = mock_webhook_server + import os + from hexus.store import MemoryStore from mcp_server import tools - import os # Setup test environment DSN dsn = os.environ.get("PG_TEST_DSN") @@ -216,57 +222,59 @@ def test_mcp_tools_dispatch_webhooks(mock_sleep, mock_webhook_server): agent = "pytest-mcp-webhook-" + os.urandom(4).hex() # Configure webhooks via environment variables - with patch.dict( - "os.environ", - { - "HEXUS_WEBHOOK_URL": url, - "HEXUS_WEBHOOK_SECRET": "mcp_secret", - "HEXUS_AGENT_IDENTITY": agent, - }, - ): - with patch( + with ( # noqa: SIM117 + patch.dict( + "os.environ", + { + "HEXUS_WEBHOOK_URL": url, + "HEXUS_WEBHOOK_SECRET": "mcp_secret", + "HEXUS_AGENT_IDENTITY": agent, + }, + ), + patch( "hexus.webhook.dispatcher.dispatch_webhook", side_effect=dispatch_webhook_sync, - ): - # 1. Test memory_retain tool triggers webhook - with patch("mcp_server.tools._embed_batch", return_value=[[0.1] * 384]): - res = tools.memory_retain( - store, - { - "contents": ["Hello from MCP server webhook test"], - "target": "memory", - "agent_identity": agent, - }, - ) - assert res["inserted"] == 1 - - assert len(handler.requests_received) == 1 - req = handler.requests_received[-1] - assert req["headers"]["X-Hexus-Event"] == "memory_retain" - assert ( - req["payload"]["data"]["content"] - == "Hello from MCP server webhook test" - ) - - # Find the row ID to delete it - with store._get_pool().connection() as conn: - with conn.cursor() as cur: - cur.execute( - "SELECT id FROM memory_entries WHERE agent_identity = %s", - (agent,), - ) - row_id = cur.fetchone()[0] - - # 2. Test memory_forget tool triggers webhook - forget_res = tools.memory_forget( - store, {"id": row_id, "confirm": True, "agent_identity": agent} - ) - assert forget_res["deleted"] == 1 - - assert len(handler.requests_received) == 2 - req = handler.requests_received[-1] - assert req["headers"]["X-Hexus-Event"] == "memory_forget" - assert ( - req["payload"]["data"]["content"] - == "Hello from MCP server webhook test" - ) + ), + ): + # 1. Test memory_retain tool triggers webhook + with patch("mcp_server.tools._embed_batch", return_value=[[0.1] * 384]): + res = tools.memory_retain( + store, + { + "contents": ["Hello from MCP server webhook test"], + "target": "memory", + "agent_identity": agent, + }, + ) + assert res["inserted"] == 1 + + assert len(handler.requests_received) == 1 + req = handler.requests_received[-1] + assert req["headers"]["X-Hexus-Event"] == "memory_retain" + assert ( + req["payload"]["data"]["content"] + == "Hello from MCP server webhook test" + ) + + # Find the row ID to delete it + with store._get_pool().connection() as conn: # noqa: SIM117 + with conn.cursor() as cur: + cur.execute( + "SELECT id FROM memory_entries WHERE agent_identity = %s", + (agent,), + ) + row_id = cur.fetchone()[0] + + # 2. Test memory_forget tool triggers webhook + forget_res = tools.memory_forget( + store, {"id": row_id, "confirm": True, "agent_identity": agent} + ) + assert forget_res["deleted"] == 1 + + assert len(handler.requests_received) == 2 + req = handler.requests_received[-1] + assert req["headers"]["X-Hexus-Event"] == "memory_forget" + assert ( + req["payload"]["data"]["content"] + == "Hello from MCP server webhook test" + ) diff --git a/tools/graph_eye_candy.py b/tools/graph_eye_candy.py new file mode 100644 index 0000000..17c83f1 --- /dev/null +++ b/tools/graph_eye_candy.py @@ -0,0 +1,335 @@ +""" +graph_eye_candy.py — Animated force-directed graph of hexus entity co-occurrences. + +Talks directly to your Postgres via MemoryStore. No new pip deps. +Emits a single self-contained HTML file with a D3.js force layout +loaded from CDN — just open it in a browser. + +Usage: + # Overview: every co-occurring pair (constellation mode) + python tools/graph_eye_candy.py + + # Walk mode: recursive BFS from a seed entity, colored by hop depth + python tools/graph_eye_candy.py --seed domain example.com --max-depth 3 + + # Filter to one agent's memory + python tools/graph_eye_candy.py --agent my-agent --min-strength 3 + +DSN is read from $HEXUS_DSN (same env var the MCP server uses). +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path + +# Make `hexus` importable when run from repo root. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from hexus.store import MemoryStore + +# ---------- data layer ---------- + + +def fetch_overview( + store: MemoryStore, agent: str | None, min_strength: int, limit: int +) -> tuple[list[dict], list[dict]]: + """Constellation: all heavy co-occurrences, no seed required.""" + pairs = store.common_topics( + agent_identity=agent, min_strength=min_strength, limit=limit + ) + nodes: dict[str, dict] = {} + edges: list[dict] = [] + + for p in pairs: + a_id = f"{p['type_a']}:{p['value_a']}" + b_id = f"{p['type_b']}:{p['value_b']}" + nodes.setdefault( + a_id, {"id": a_id, "type": p["type_a"], "value": p["value_a"], "weight": 0} + ) + nodes.setdefault( + b_id, {"id": b_id, "type": p["type_b"], "value": p["value_b"], "weight": 0} + ) + nodes[a_id]["weight"] += p["strength"] + nodes[b_id]["weight"] += p["strength"] + edges.append( + { + "source": a_id, + "target": b_id, + "strength": p["strength"], + } + ) + + return list(nodes.values()), edges + + +def fetch_walk( + store: MemoryStore, + seed_type: str, + seed_value: str, + agent: str | None, + max_depth: int, + limit: int, +) -> tuple[list[dict], list[dict]]: + """Recursive walk: edges reconstructed from per-hop results.""" + hops = store.graph_walk( + entity_type=seed_type, + entity_value=seed_value, + agent_identity=agent, + max_depth=max_depth, + limit=limit, + ) + + nodes: dict[str, dict] = {} + edges: list[dict] = [] + seed_id = f"{seed_type}:{seed_value}" + nodes[seed_id] = {"id": seed_id, "type": seed_type, "value": seed_value, "depth": 0} + + for h in hops: + nid = f"{h['type']}:{h['value']}" + nodes.setdefault( + nid, + { + "id": nid, + "type": h["type"], + "value": h["value"], + "depth": h["min_depth"], + }, + ) + nodes[nid]["depth"] = min(nodes[nid].get("depth", 99), h["min_depth"]) + edges.append( + { + "source": seed_id, + "target": nid, + "depth": h["min_depth"], + "occurrences": h["occurrences"], + } + ) + + return list(nodes.values()), edges + + +# ---------- HTML template ---------- + +HTML_TEMPLATE = """ + + + +hexus — entity graph + + + + +
+

hexus ◦ entity graph

+
{subtitle}
+
    +
    +
    + + + +
    +
    +
    drag nodes · scroll to zoom · click node to pin its neighborhood
    + + + + + +""" + + +def render_html(nodes: list[dict], edges: list[dict], subtitle: str) -> str: + payload = {"nodes": nodes, "edges": edges} + return HTML_TEMPLATE.format(subtitle=subtitle, data_json=json.dumps(payload)) + + +# ---------- CLI ---------- + + +def main() -> int: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument( + "--dsn", + default=os.environ.get("HEXUS_DSN"), + help="Postgres DSN (default: $HEXUS_DSN)", + ) + ap.add_argument("--agent", default=None, help="Filter to one agent_identity") + ap.add_argument( + "--seed", + nargs=2, + metavar=("TYPE", "VALUE"), + help="Walk mode: graph_walk from :", + ) + ap.add_argument("--max-depth", type=int, default=2, help="Walk depth (1-5)") + ap.add_argument( + "--min-strength", + type=int, + default=2, + help="Min co-occurrence count (overview mode)", + ) + ap.add_argument("--limit", type=int, default=80, help="Max edges/nodes") + ap.add_argument("-o", "--out", default="hexus-graph.html", help="Output HTML path") + args = ap.parse_args() + + if not args.dsn: + print("error: set HEXUS_DSN or pass --dsn", file=sys.stderr) + return 2 + + store = MemoryStore(dsn=args.dsn) + + if args.seed: + seed_type, seed_value = args.seed + nodes, edges = fetch_walk( + store, seed_type, seed_value, args.agent, args.max_depth, args.limit + ) + subtitle = f"walk from {seed_type}:{seed_value} · depth ≤ {args.max_depth} · {len(nodes)} nodes / {len(edges)} edges" + if args.agent: + subtitle += f" · agent={args.agent}" + else: + nodes, edges = fetch_overview(store, args.agent, args.min_strength, args.limit) + subtitle = f"overview · min strength {args.min_strength} · {len(nodes)} nodes / {len(edges)} edges" + if args.agent: + subtitle += f" · agent={args.agent}" + + if not nodes: + print("no data — check DSN / agent filter / seed", file=sys.stderr) + return 1 + + html = render_html(nodes, edges, subtitle) + out_path = Path(args.out).resolve() + out_path.write_text(html, encoding="utf-8") + print(f"wrote {out_path} ({len(nodes)} nodes, {len(edges)} edges)") + print(f"open with: xdg-open {out_path} # or just drag into a browser") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())