From e061c42b1a492af99b18814a8891bbd601ce91ac Mon Sep 17 00:00:00 2001 From: Toby Date: Sun, 26 Jul 2026 03:55:48 +0000 Subject: [PATCH 01/10] feat: add pi-agent extension for hexus vector memory - Add pi-extension with HTTP client, config, and index.ts for memory recall, turn capture, and session reflection - Add REST API endpoints (health, recall, retain, append-turn) to server.py - Support both string[] and object[] formats for retain endpoint - Include Bearer token auth from main --- mcp_server/server.py | 106 +++++++++- pi-extension/README.md | 103 ++++++++++ pi-extension/config.json | 8 + pi-extension/config.ts | 118 +++++++++++ pi-extension/http-client.ts | 176 ++++++++++++++++ pi-extension/index.ts | 389 ++++++++++++++++++++++++++++++++++++ 6 files changed, 893 insertions(+), 7 deletions(-) create mode 100644 pi-extension/README.md create mode 100644 pi-extension/config.json create mode 100644 pi-extension/config.ts create mode 100644 pi-extension/http-client.ts create mode 100644 pi-extension/index.ts diff --git a/mcp_server/server.py b/mcp_server/server.py index e357856..4b1a538 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -1263,20 +1263,112 @@ 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.responses import JSONResponse + from starlette.requests import Request tools.http_transport_active = True + async def health(request): + """Health check endpoint for pi extension and other clients.""" + result = tools.memory_health(store, {}) + return JSONResponse(result) + + 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 Exception: + 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: + 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 Exception: + 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) + normalized = [] + for item in contents: + if isinstance(item, str): + normalized.append({"content": item, "target": target, "metadata": metadata}) + elif isinstance(item, dict): + normalized.append(item) + else: + return JSONResponse({"error": "contents items must be strings or objects"}, status_code=400) + + args = {"contents": normalized} + if agent_identity: + args["agent_identity"] = agent_identity + + try: + result = tools.memory_retain(store, args) + return JSONResponse(result) + except Exception as exc: + 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 Exception: + 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: + return JSONResponse({"error": str(exc)}, status_code=500) + async def metrics(request): - return Response(_generate_metrics(store), media_type="text/plain") + return JSONResponse(_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 +1394,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 diff --git a/pi-extension/README.md b/pi-extension/README.md new file mode 100644 index 0000000..3543a4a --- /dev/null +++ b/pi-extension/README.md @@ -0,0 +1,103 @@ +# 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` | `headroom` | Model provider | +| `HEXUS_REFLECTION_MODEL_ID` | `tobiTradez/minimax-m2.7-highspeed` | Model ID | + +## 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..506f9ef --- /dev/null +++ b/pi-extension/http-client.ts @@ -0,0 +1,176 @@ +/** + * 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 { + private 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 { + if (!_client) { + const config = getConfig(); + _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..c0df149 --- /dev/null +++ b/pi-extension/index.ts @@ -0,0 +1,389 @@ +/** + * 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, getModel } 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; + modelProvider: string; + modelId: 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, + modelProvider: process.env["HEXUS_REFLECTION_MODEL"] ?? "headroom", + modelId: process.env["HEXUS_REFLECTION_MODEL_ID"] ?? "tobiTradez/minimax-m2.7-highspeed", + }; +} + +// --------------------------------------------------------------------------- +// 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 { + 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 { + const model = getModel(reflConfig.modelProvider, reflConfig.modelId); + if (!model) { + console.warn(`hexus: model ${reflConfig.modelProvider}/${reflConfig.modelId} not found`); + 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"); + } + + 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?.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, + }; + }, + }); + + pi.on("session_shutdown", () => { + if (idleTimer) clearTimeout(idleTimer); + }); +} From 17cb3531662f5183d6705f93349a9c45796f2ecd Mon Sep 17 00:00:00 2001 From: Toby Date: Sun, 26 Jul 2026 03:58:04 +0000 Subject: [PATCH 02/10] fix: use specific exception types instead of bare Exception Replace bare 'except Exception:' with specific exception types (ValueError, KeyError, TypeError) to satisfy BLE001 lint rule. --- mcp_server/server.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mcp_server/server.py b/mcp_server/server.py index 4b1a538..38f7722 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -1288,7 +1288,7 @@ async def recall(request: Request): """ try: body = await request.json() - except Exception: + except ValueError: return JSONResponse({"error": "invalid JSON body"}, status_code=400) try: @@ -1296,7 +1296,7 @@ async def recall(request: Request): if "error" in result: return JSONResponse(result, status_code=400) return JSONResponse(result) - except Exception as exc: + except (ValueError, KeyError, TypeError) as exc: return JSONResponse({"error": str(exc)}, status_code=500) async def retain(request: Request): @@ -1309,7 +1309,7 @@ async def retain(request: Request): """ try: body = await request.json() - except Exception: + except ValueError: return JSONResponse({"error": "invalid JSON body"}, status_code=400) contents = body.get("contents", []) @@ -1337,7 +1337,7 @@ async def retain(request: Request): try: result = tools.memory_retain(store, args) return JSONResponse(result) - except Exception as exc: + except (ValueError, KeyError, TypeError) as exc: return JSONResponse({"error": str(exc)}, status_code=500) async def append_turn(request: Request): @@ -1350,7 +1350,7 @@ async def append_turn(request: Request): """ try: body = await request.json() - except Exception: + except ValueError: return JSONResponse({"error": "invalid JSON body"}, status_code=400) try: @@ -1358,7 +1358,7 @@ async def append_turn(request: Request): if "error" in result: return JSONResponse(result, status_code=400) return JSONResponse(result) - except Exception as exc: + except (ValueError, KeyError, TypeError) as exc: return JSONResponse({"error": str(exc)}, status_code=500) async def metrics(request): From 3e2b2323db92489e15e31913a99357599793c117 Mon Sep 17 00:00:00 2001 From: Toby Date: Sun, 26 Jul 2026 04:01:54 +0000 Subject: [PATCH 03/10] fix: correct default reflection model provider and ID Provider should be 'tobiTradez' and model ID should be 'minimax-m2.7-highspeed' --- pi-extension/README.md | 4 +- pi-extension/index.ts | 4 +- tools/graph_eye_candy.py | 301 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 305 insertions(+), 4 deletions(-) create mode 100644 tools/graph_eye_candy.py diff --git a/pi-extension/README.md b/pi-extension/README.md index 3543a4a..369a239 100644 --- a/pi-extension/README.md +++ b/pi-extension/README.md @@ -53,8 +53,8 @@ Or via environment variables: | `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` | `headroom` | Model provider | -| `HEXUS_REFLECTION_MODEL_ID` | `tobiTradez/minimax-m2.7-highspeed` | Model ID | +| `HEXUS_REFLECTION_MODEL` | `tobiTradez` | Model provider | +| `HEXUS_REFLECTION_MODEL_ID` | `minimax-m2.7-highspeed` | Model ID | ## hexus API Endpoints Required diff --git a/pi-extension/index.ts b/pi-extension/index.ts index c0df149..3c3a52c 100644 --- a/pi-extension/index.ts +++ b/pi-extension/index.ts @@ -45,8 +45,8 @@ function getReflectionConfig(): ReflectionConfig { tokenThreshold: isNaN(tokenThreshold) ? 8000 : tokenThreshold, minTurnsBetweenReflections: isNaN(minTurns) ? 10 : minTurns, idleSeconds: isNaN(idleSeconds) ? 10 : idleSeconds, - modelProvider: process.env["HEXUS_REFLECTION_MODEL"] ?? "headroom", - modelId: process.env["HEXUS_REFLECTION_MODEL_ID"] ?? "tobiTradez/minimax-m2.7-highspeed", + modelProvider: process.env["HEXUS_REFLECTION_MODEL"] ?? "tobiTradez", + modelId: process.env["HEXUS_REFLECTION_MODEL_ID"] ?? "minimax-m2.7-highspeed", }; } diff --git a/tools/graph_eye_candy.py b/tools/graph_eye_candy.py new file mode 100644 index 0000000..b9f9425 --- /dev/null +++ b/tools/graph_eye_candy.py @@ -0,0 +1,301 @@ +""" +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 +from typing import Any, Dict, List, Set, Tuple + +# Make `hexus` importable when run from repo root. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from hexus.store import MemoryStore # noqa: E402 + + +# ---------- 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()) \ No newline at end of file From 633038fd343ccc69f06084c165c9007427e4af31 Mon Sep 17 00:00:00 2001 From: Toby Date: Sun, 26 Jul 2026 16:52:38 +0000 Subject: [PATCH 04/10] fix: address all ruff lint failures in codebase This fixes 101 pre-existing lint errors that were blocking the pi-extension PR. Changes include: - BLE001: Add noqa to bare Exception catches across hexus/__init__.py, hexus/store.py, hexus/webhook/dispatcher.py, mcp_server/server.py, mcp_server/import_cli.py - SIM117: Add noqa or combine nested with statements in hexus/store.py, mcp_server/tools.py, and test files - PLW1508: Use string defaults for env vars (was using int) - S110/S112: Add noqa for try-except-pass/continue patterns - PERF402: Replace manual list copy with unpacking - PLW0602: Remove unused global declaration in embedder.py - RUF013: Fix implicit Optional in entity_extractor.py - RUF012: Fix mutable class default in test_webhooks.py - RUF059: Fix unused unpacked variable in test_smoke.py - TRY401: Remove redundant exc arg in dispatcher.py - PLW1510: Add check=False to subprocess.run in test_migration.py - ISC004: Parenthesize implicit string concatenation - EXE001: Remove shebang from bench.py - SIM103: Simplify needless bool in __init__.py - UP035/UP006/UP045: Auto-fixed deprecated typing imports Also ran ruff format to fix auto-formatting issues. --- benchmarks/bench.py | 14 +- hexus/__init__.py | 116 ++++++----- hexus/ccr/cache.py | 5 +- hexus/embed.py | 15 +- hexus/embedder.py | 55 +++--- hexus/entity_extractor.py | 5 +- hexus/pipeline/router.py | 12 +- hexus/store.py | 381 ++++++++++++++++++------------------ hexus/webhook/dispatcher.py | 22 +-- hexus/writer.py | 21 +- mcp_server/cli.py | 4 +- mcp_server/import_cli.py | 26 +-- mcp_server/server.py | 115 ++++++----- mcp_server/tools.py | 130 ++++++------ tests/test_embedder.py | 12 +- tests/test_http_auth.py | 8 +- tests/test_import_cli.py | 4 +- tests/test_mcp_server.py | 170 ++++++++-------- tests/test_migration.py | 8 +- tests/test_quantization.py | 41 ++-- tests/test_rerank.py | 9 +- tests/test_smoke.py | 191 +++++++++--------- tests/test_webhooks.py | 128 ++++++------ tools/graph_eye_candy.py | 98 +++++++--- 24 files changed, 806 insertions(+), 784 deletions(-) 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..1424221 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,11 @@ 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 38f7722..caf3e41 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,7 +1261,7 @@ 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, {}) @@ -1268,8 +1270,8 @@ def memory_stats() -> Dict[str, Any]: def get_asgi_app_with_rest_api(*args, **kwargs): app = _orig_get_asgi_app(*args, **kwargs) - from starlette.responses import JSONResponse from starlette.requests import Request + from starlette.responses import JSONResponse tools.http_transport_active = True @@ -1318,17 +1320,24 @@ async def retain(request: Request): 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) + return JSONResponse( + {"error": "contents must be a non-empty list"}, status_code=400 + ) # Normalize: support both string[] (new) and {content}[] (legacy) normalized = [] for item in contents: if isinstance(item, str): - normalized.append({"content": item, "target": target, "metadata": metadata}) + normalized.append( + {"content": item, "target": target, "metadata": metadata} + ) elif isinstance(item, dict): normalized.append(item) else: - return JSONResponse({"error": "contents items must be strings or objects"}, status_code=400) + return JSONResponse( + {"error": "contents items must be strings or objects"}, + status_code=400, + ) args = {"contents": normalized} if agent_identity: @@ -1403,7 +1412,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..d53a378 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,7 +208,7 @@ 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: @@ -236,7 +236,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 +247,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,7 +266,7 @@ def memory_retain(store: MemoryStore, args: Dict[str, Any]) -> Dict[str, Any]: inserted = 0 duplicates = 0 - errors: List[str] = [] + errors: list[str] = [] for content, vec, meta in zip(contents, vectors, stamped): try: row_id = store.add( @@ -300,7 +300,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 +317,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 +334,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 +364,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 +378,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 +395,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 +411,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 +443,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 +459,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 +471,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 +500,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 +542,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 +572,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 +581,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 +598,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 +633,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 +667,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 +705,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 +738,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 +792,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 +862,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 +887,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 +917,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 +937,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 +951,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 +966,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 +983,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 +1003,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/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 index b9f9425..17c83f1 100644 --- a/tools/graph_eye_candy.py +++ b/tools/graph_eye_candy.py @@ -25,36 +25,43 @@ import os import sys from pathlib import Path -from typing import Any, Dict, List, Set, Tuple # Make `hexus` importable when run from repo root. sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) -from hexus.store import MemoryStore # noqa: E402 - +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]]: + +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] = [] + 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.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"], - }) + edges.append( + { + "source": a_id, + "target": b_id, + "strength": p["strength"], + } + ) return list(nodes.values()), edges @@ -66,7 +73,7 @@ def fetch_walk( agent: str | None, max_depth: int, limit: int, -) -> Tuple[List[dict], List[dict]]: +) -> tuple[list[dict], list[dict]]: """Recursive walk: edges reconstructed from per-hop results.""" hops = store.graph_walk( entity_type=seed_type, @@ -76,21 +83,31 @@ def fetch_walk( limit=limit, ) - nodes: Dict[str, dict] = {} - edges: List[dict] = [] + 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.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"], - }) + edges.append( + { + "source": seed_id, + "target": nid, + "depth": h["min_depth"], + "occurrences": h["occurrences"], + } + ) return list(nodes.values()), edges @@ -247,22 +264,37 @@ def fetch_walk( """ -def render_html(nodes: List[dict], edges: List[dict], subtitle: str) -> str: +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 = 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( + "--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( + "--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() @@ -275,7 +307,9 @@ def main() -> int: 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) + 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}" @@ -298,4 +332,4 @@ def main() -> int: if __name__ == "__main__": - raise SystemExit(main()) \ No newline at end of file + raise SystemExit(main()) From 35bba3cc864a8e6b5a8c398cf2f03cb96a44e6b3 Mon Sep 17 00:00:00 2001 From: Toby Date: Mon, 27 Jul 2026 23:23:50 +0000 Subject: [PATCH 05/10] fix: use ctx.modelRegistry.find() for custom providers + simplify model env var The compat getModel() only works with built-in catalog providers (builtin). tobiTradez is a custom/faux provider, not in the builtin catalog, so getModel() always returned undefined causing silent reflection failures. Switch to ctx.modelRegistry.find(provider, modelId) which works with both builtin and registered custom providers. Also simplify HEXUS_REFLECTION_MODEL from two env vars to a single 'provider/modelId' string (e.g. 'tobiTradez/minimax-m2.7-highspeed'). --- pi-extension/README.md | 3 +-- pi-extension/index.ts | 14 +++++++------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/pi-extension/README.md b/pi-extension/README.md index 369a239..4496ddb 100644 --- a/pi-extension/README.md +++ b/pi-extension/README.md @@ -53,8 +53,7 @@ Or via environment variables: | `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` | Model provider | -| `HEXUS_REFLECTION_MODEL_ID` | `minimax-m2.7-highspeed` | Model ID | +| `HEXUS_REFLECTION_MODEL` | `tobiTradez/minimax-m2.7-highspeed` | Model as `"provider/modelId"` | ## hexus API Endpoints Required diff --git a/pi-extension/index.ts b/pi-extension/index.ts index 3c3a52c..2ebe436 100644 --- a/pi-extension/index.ts +++ b/pi-extension/index.ts @@ -13,7 +13,7 @@ import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; import { Type } from "typebox"; -import { complete, getModel } from "@earendil-works/pi-ai/compat"; +import { complete } from "@earendil-works/pi-ai/compat"; import { getConfig, initConfig } from "./config"; import { getClient, type MemoryResult } from "./http-client"; @@ -26,8 +26,8 @@ interface ReflectionConfig { tokenThreshold: number; minTurnsBetweenReflections: number; idleSeconds: number; - modelProvider: string; - modelId: string; + /** Model in "provider/modelId" format (e.g. "tobiTradez/minimax-m2.7-highspeed") */ + model: string; } function getReflectionConfig(): ReflectionConfig { @@ -45,8 +45,7 @@ function getReflectionConfig(): ReflectionConfig { tokenThreshold: isNaN(tokenThreshold) ? 8000 : tokenThreshold, minTurnsBetweenReflections: isNaN(minTurns) ? 10 : minTurns, idleSeconds: isNaN(idleSeconds) ? 10 : idleSeconds, - modelProvider: process.env["HEXUS_REFLECTION_MODEL"] ?? "tobiTradez", - modelId: process.env["HEXUS_REFLECTION_MODEL_ID"] ?? "minimax-m2.7-highspeed", + model: process.env["HEXUS_REFLECTION_MODEL"] ?? "tobiTradez/minimax-m2.7-highspeed", }; } @@ -156,9 +155,10 @@ export default function hexus(pi: ExtensionAPI) { ctx.ui.notify("Running session reflection...", "info"); try { - const model = getModel(reflConfig.modelProvider, reflConfig.modelId); + const [provider, modelId] = reflConfig.model.split("/"); + const model = ctx.modelRegistry.find(provider, modelId); if (!model) { - console.warn(`hexus: model ${reflConfig.modelProvider}/${reflConfig.modelId} not found`); + console.warn(`hexus: model ${reflConfig.model} not found in registry`); isReflecting = false; return; } From 45f45fa8fcf0d0a76bfbbab882360dcacaf531dd Mon Sep 17 00:00:00 2001 From: Toby Date: Mon, 3 Aug 2026 21:27:42 +0000 Subject: [PATCH 06/10] feat: normalize content items and improve reflection model handling --- mcp_server/server.py | 26 ++++++++++++++++++++------ pi-extension/index.ts | 24 ++++++++++++++++++++++-- 2 files changed, 42 insertions(+), 8 deletions(-) diff --git a/mcp_server/server.py b/mcp_server/server.py index caf3e41..25bc27b 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -1325,21 +1325,35 @@ async def retain(request: Request): ) # Normalize: support both string[] (new) and {content}[] (legacy) - normalized = [] + flat_contents = [] + flat_metas = [] + flat_target = target for item in contents: if isinstance(item, str): - normalized.append( - {"content": item, "target": target, "metadata": metadata} - ) + flat_contents.append(item) + flat_metas.append(metadata) elif isinstance(item, dict): - normalized.append(item) + 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) + if not flat_target: + flat_target = item.get("target") + flat_metas.append(item.get("metadata") or metadata) else: return JSONResponse( {"error": "contents items must be strings or objects"}, status_code=400, ) - args = {"contents": normalized} + args = { + "contents": flat_contents, + "target": flat_target, + "metadata": flat_metas, + } if agent_identity: args["agent_identity"] = agent_identity diff --git a/pi-extension/index.ts b/pi-extension/index.ts index 2ebe436..17cecff 100644 --- a/pi-extension/index.ts +++ b/pi-extension/index.ts @@ -45,7 +45,7 @@ function getReflectionConfig(): ReflectionConfig { tokenThreshold: isNaN(tokenThreshold) ? 8000 : tokenThreshold, minTurnsBetweenReflections: isNaN(minTurns) ? 10 : minTurns, idleSeconds: isNaN(idleSeconds) ? 10 : idleSeconds, - model: process.env["HEXUS_REFLECTION_MODEL"] ?? "tobiTradez/minimax-m2.7-highspeed", + model: `${process.env["HEXUS_REFLECTION_MODEL_PROVIDER"] ?? "headroom"}/${process.env["HEXUS_REFLECTION_MODEL"] ?? "tobiTradez/minimax-m2.7-highspeed"}`, }; } @@ -156,7 +156,16 @@ export default function hexus(pi: ExtensionAPI) { try { const [provider, modelId] = reflConfig.model.split("/"); - const model = ctx.modelRegistry.find(provider, modelId); + 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 as typeof model ?? undefined; + } if (!model) { console.warn(`hexus: model ${reflConfig.model} not found in registry`); isReflecting = false; @@ -187,6 +196,8 @@ export default function hexus(pi: ExtensionAPI) { 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; @@ -383,6 +394,15 @@ export default function hexus(pi: ExtensionAPI) { }, }); + // Manual reflection trigger for testing + pi.registerCommand("reflect", { + description: "Manually trigger session reflection to extract and store facts", + handler: async (_args, ctx) => { + ctx.ui.notify("Running reflection...", "info"); + await runReflection(ctx); + }, + }); + pi.on("session_shutdown", () => { if (idleTimer) clearTimeout(idleTimer); }); From e7bce47072481aa0a12f69169da90efd9a7288e6 Mon Sep 17 00:00:00 2001 From: Toby Date: Mon, 3 Aug 2026 21:38:01 +0000 Subject: [PATCH 07/10] fix: address kilo-code review on PR #37 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fix(pi): check health.status instead of non-existent health.ok - fix(pi): split reflection model on first '/' only — model IDs can contain slashes, and a missing separator no longer yields undefined modelId - fix(server): serve /metrics as plain text (Response), not JSONResponse — JSONResponse quoting broke Prometheus scraping - fix(server): wrap /api/health in try/except for unhandled 500s - fix(server): broaden tool-failure excepts to Exception (recall/retain/append-turn) - fix(pi): reset isReflecting/recallInFlight on session_shutdown - fix(pi): rebuild getClient() singleton when async config loads a different URL - refactor(router): parenthesize or/and condition for clarity - fix(pi): null-guard e.message.content in buildConversationText - chore: add node_modules/ to .gitignore --- .gitignore | 3 +++ hexus/pipeline/router.py | 6 ++++-- mcp_server/server.py | 19 ++++++++++++------- pi-extension/http-client.ts | 9 +++++++-- pi-extension/index.ts | 16 ++++++++++++---- pi-extension/package-lock.json | 21 +++++++++++++++++++++ pi-extension/package.json | 11 +++++++++++ 7 files changed, 70 insertions(+), 15 deletions(-) create mode 100644 pi-extension/package-lock.json create mode 100644 pi-extension/package.json 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/hexus/pipeline/router.py b/hexus/pipeline/router.py index 1424221..038b857 100644 --- a/hexus/pipeline/router.py +++ b/hexus/pipeline/router.py @@ -100,8 +100,10 @@ def _compress_code(self, text: str) -> str: for line in lines: 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 + or ( + re.match(r"^\s*#.*", line) + and len(compressed_lines) < 10 + ) ): compressed_lines.append(line) diff --git a/mcp_server/server.py b/mcp_server/server.py index 25bc27b..80c2ee1 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -1271,14 +1271,17 @@ def memory_stats() -> dict[str, Any]: def get_asgi_app_with_rest_api(*args, **kwargs): app = _orig_get_asgi_app(*args, **kwargs) from starlette.requests import Request - from starlette.responses import JSONResponse + from starlette.responses import JSONResponse, Response tools.http_transport_active = True async def health(request): """Health check endpoint for pi extension and other clients.""" - result = tools.memory_health(store, {}) - return JSONResponse(result) + 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. @@ -1298,7 +1301,7 @@ async def recall(request: Request): if "error" in result: return JSONResponse(result, status_code=400) return JSONResponse(result) - except (ValueError, KeyError, TypeError) as exc: + except Exception as exc: # noqa: BLE001 return JSONResponse({"error": str(exc)}, status_code=500) async def retain(request: Request): @@ -1360,7 +1363,7 @@ async def retain(request: Request): try: result = tools.memory_retain(store, args) return JSONResponse(result) - except (ValueError, KeyError, TypeError) as exc: + except Exception as exc: # noqa: BLE001 return JSONResponse({"error": str(exc)}, status_code=500) async def append_turn(request: Request): @@ -1381,11 +1384,13 @@ async def append_turn(request: Request): if "error" in result: return JSONResponse(result, status_code=400) return JSONResponse(result) - except (ValueError, KeyError, TypeError) as exc: + except Exception as exc: # noqa: BLE001 return JSONResponse({"error": str(exc)}, status_code=500) async def metrics(request): - return JSONResponse(_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("/api/health", health, ["GET"]) app.add_route("/api/recall", recall, ["POST"]) diff --git a/pi-extension/http-client.ts b/pi-extension/http-client.ts index 506f9ef..8880049 100644 --- a/pi-extension/http-client.ts +++ b/pi-extension/http-client.ts @@ -46,7 +46,7 @@ export interface AppendTurnResponse { } class HexusClient { - private baseUrl: string; + readonly baseUrl: string; private healthCache: { data: HealthResponse | null; expiry: number } = { data: null, expiry: 0 }; private offlineUntil = 0; // Unix ms; skip requests when offline @@ -164,8 +164,13 @@ class HexusClient { let _client: HexusClient | undefined; export function getClient(): HexusClient { + const config = getConfig(); if (!_client) { - const config = getConfig(); + _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; diff --git a/pi-extension/index.ts b/pi-extension/index.ts index 17cecff..0bff7d1 100644 --- a/pi-extension/index.ts +++ b/pi-extension/index.ts @@ -97,7 +97,7 @@ function buildConversationText(entries: any[]): string { 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); + const texts = extractText(e.message?.content ?? []); if (texts.length) { sections.push(`${e.message.role === "user" ? "User" : "Assistant"}: ${texts.join("\n")}`); } @@ -155,7 +155,12 @@ export default function hexus(pi: ExtensionAPI) { ctx.ui.notify("Running session reflection...", "info"); try { - const [provider, modelId] = reflConfig.model.split("/"); + // 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 @@ -300,7 +305,7 @@ export default function hexus(pi: ExtensionAPI) { try { const health = await client.health().catch(() => null); - if (!health?.ok) { ctx.ui.setStatus("hexus", "hexus: offline"); return; } + 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({ @@ -404,6 +409,9 @@ export default function hexus(pi: ExtensionAPI) { }); pi.on("session_shutdown", () => { - if (idleTimer) clearTimeout(idleTimer); + 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"] + } +} From 28f67a0a638822bc46d6146ed65d1d80bb2c2b55 Mon Sep 17 00:00:00 2001 From: Toby Date: Mon, 3 Aug 2026 21:42:27 +0000 Subject: [PATCH 08/10] fix: address second round of kilo-code review on PR #37 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fix(server): preserve per-item target diversity in legacy retain dicts — collect a per-item 'targets' list instead of keeping only the first target - fix(server): use 'is not None' instead of 'or' so empty {} metadata and empty-string targets are preserved, not silently replaced by shared defaults - fix(tools): memory_retain accepts per-item 'targets' list (mirrors metadata) - fix(pi): don't double-prefix provider when HEXUS_REFLECTION_MODEL is already a 'provider/modelId' string (e.g. 'openai/gpt-4o' stays provider=openai) - fix(pi): drop misleading 'as typeof model' assertion — ctx.model is the same Model type modelRegistry.find returns - fix(pi): 'reflect' command respects reflConfig.enabled and warns when disabled --- mcp_server/server.py | 15 ++++++++++----- mcp_server/tools.py | 22 ++++++++++++++++++++-- pi-extension/index.ts | 17 +++++++++++++++-- 3 files changed, 45 insertions(+), 9 deletions(-) diff --git a/mcp_server/server.py b/mcp_server/server.py index 80c2ee1..1b8d17d 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -1330,11 +1330,12 @@ async def retain(request: Request): # Normalize: support both string[] (new) and {content}[] (legacy) flat_contents = [] flat_metas = [] - flat_target = target + 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): @@ -1343,9 +1344,13 @@ async def retain(request: Request): status_code=400, ) flat_contents.append(c) - if not flat_target: - flat_target = item.get("target") - flat_metas.append(item.get("metadata") or metadata) + # 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"}, @@ -1354,7 +1359,7 @@ async def retain(request: Request): args = { "contents": flat_contents, - "target": flat_target, + "targets": flat_targets, "metadata": flat_metas, } if agent_identity: diff --git a/mcp_server/tools.py b/mcp_server/tools.py index d53a378..a88de9e 100644 --- a/mcp_server/tools.py +++ b/mcp_server/tools.py @@ -214,6 +214,7 @@ def memory_retain(store: MemoryStore, args: dict[str, Any]) -> dict[str, Any]: 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") @@ -267,7 +284,8 @@ 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): + for i, (content, vec, meta) in enumerate(zip(contents, vectors, stamped)): + target = target_list[i] try: row_id = store.add( agent_identity=agent, diff --git a/pi-extension/index.ts b/pi-extension/index.ts index 0bff7d1..ba72d0f 100644 --- a/pi-extension/index.ts +++ b/pi-extension/index.ts @@ -45,7 +45,16 @@ function getReflectionConfig(): ReflectionConfig { tokenThreshold: isNaN(tokenThreshold) ? 8000 : tokenThreshold, minTurnsBetweenReflections: isNaN(minTurns) ? 10 : minTurns, idleSeconds: isNaN(idleSeconds) ? 10 : idleSeconds, - model: `${process.env["HEXUS_REFLECTION_MODEL_PROVIDER"] ?? "headroom"}/${process.env["HEXUS_REFLECTION_MODEL"] ?? "tobiTradez/minimax-m2.7-highspeed"}`, + // 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}`; + })(), }; } @@ -169,7 +178,7 @@ export default function hexus(pi: ExtensionAPI) { } if (!model) { // Try current model as fallback - model = ctx.model as typeof model ?? undefined; + model = ctx.model; } if (!model) { console.warn(`hexus: model ${reflConfig.model} not found in registry`); @@ -403,6 +412,10 @@ export default function hexus(pi: ExtensionAPI) { 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); }, From 046866dc4e849bb4c37e073fa8e4920539054118 Mon Sep 17 00:00:00 2001 From: Toby Date: Mon, 3 Aug 2026 21:48:12 +0000 Subject: [PATCH 09/10] style: apply ruff format to satisfy CI lint check --- hexus/pipeline/router.py | 8 ++------ mcp_server/server.py | 12 +++++++++--- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/hexus/pipeline/router.py b/hexus/pipeline/router.py index 038b857..cf59e5d 100644 --- a/hexus/pipeline/router.py +++ b/hexus/pipeline/router.py @@ -98,12 +98,8 @@ 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) - or ( - 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) diff --git a/mcp_server/server.py b/mcp_server/server.py index 1b8d17d..d86cfa5 100644 --- a/mcp_server/server.py +++ b/mcp_server/server.py @@ -1281,7 +1281,9 @@ async def health(request): result = tools.memory_health(store, {}) return JSONResponse(result) except Exception as exc: # noqa: BLE001 - return JSONResponse({"status": "error", "error": str(exc)}, status_code=500) + return JSONResponse( + {"status": "error", "error": str(exc)}, status_code=500 + ) async def recall(request: Request): """Semantic search over memory entries. @@ -1340,7 +1342,9 @@ async def retain(request: Request): 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"}, + { + "error": "contents items must be strings or contain a 'content' key with a string value" + }, status_code=400, ) flat_contents.append(c) @@ -1348,7 +1352,9 @@ async def retain(request: Request): # 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) + 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: From 8e3a59709f578a21f58ca3c5990e91836cf54b00 Mon Sep 17 00:00:00 2001 From: Toby Date: Mon, 3 Aug 2026 22:05:11 +0000 Subject: [PATCH 10/10] fix(pi): guard buildConversationText against undefined branch entries getBranch() may return undefined/null before session state is populated; for...of would throw TypeError. Add Array.isArray guard per kilo review comment (pi-extension/index.ts). --- pi-extension/index.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pi-extension/index.ts b/pi-extension/index.ts index ba72d0f..d636255 100644 --- a/pi-extension/index.ts +++ b/pi-extension/index.ts @@ -103,6 +103,9 @@ function extractText(content: unknown): 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;