feat(integrations): add Hermes provider with shared memory utilities
This commit is contained in:
kartik-mem0
@@ -11,6 +11,25 @@ You can run Mem0 in three ways:
|
||||
- **Self-hosted server mode**: point the plugin at a Mem0 server you run yourself (the Docker-shipped server). The plugin only talks HTTP to your server.
|
||||
- **OSS mode**: run Mem0 in-process with your own LLM, embedder, and vector store. No Mem0 server required.
|
||||
|
||||
## Requirements and compatibility
|
||||
|
||||
Use Python 3.11+ and a recent Hermes release. The host contract was checked against
|
||||
Hermes `v0.21.3` (`v2026.9.14`) and main at
|
||||
`c62bd9f2078a946108f1c9d9b24bf118963277ef`. Earlier host versions have not been validated.
|
||||
|
||||
<Warning>
|
||||
Both checked Hermes revisions still bundle Mem0. Hermes gives bundled memory providers
|
||||
precedence over a same-named user plugin. Installing or enabling this plugin does not
|
||||
replace that bundled copy. Use the isolated preview instructions in the
|
||||
[plugin README](https://github.com/mem0ai/mem0/tree/main/integrations/hermes-plugin#local-worktree-preview)
|
||||
to try this version in a separate Hermes checkout. Once Hermes removes its bundled Mem0
|
||||
provider, the user-installed copy loads directly.
|
||||
</Warning>
|
||||
|
||||
Keep `memory.provider: mem0`, `$HERMES_HOME/mem0.json`, and your existing `MEM0_*` variables.
|
||||
There is no memory migration. An explicitly configured user ID continues to share memories
|
||||
across hosts; otherwise the provider uses the gateway user ID, then `hermes-user`.
|
||||
|
||||
## How It Works
|
||||
|
||||
Hermes runs a built-in memory system (file-based `MEMORY.md` and `USER.md`) alongside one external provider. When Mem0 is active, it works additively with the built-in system at two points in every conversation turn.
|
||||
@@ -21,7 +40,9 @@ When you send a message, Hermes searches your stored memories for the current qu
|
||||
|
||||
### 2. Background fact extraction (sync)
|
||||
|
||||
Once the model finishes, Hermes sends the `(user message, assistant response)` pair to Mem0 in a background thread. Mem0 extracts facts automatically (for example, "user prefers Python" or "user works at Acme Corp"), so you never have to tell it what to remember. Each write is tagged with the gateway channel it came from.
|
||||
Once the model finishes, Hermes sends the `(user message, assistant response)` pair to Mem0 in a background thread. Mem0 extracts facts automatically (for example, "user prefers Python" or "user works at Acme Corp"), so you never have to tell it what to remember. Each write is tagged with the gateway channel it came from and the Hermes session as top-level `run_id`. Recall remains scoped to the user across sessions.
|
||||
|
||||
The plugin uses the shared `agent-plugin-core` redaction and token-aware batching. Full non-empty user and assistant text is preserved after known-secret redaction; oversized messages are split across requests instead of truncated. Raw tool results and the full historical transcript are not captured. Pattern-based redaction cannot recognize every possible secret.
|
||||
|
||||
## Agent Tools
|
||||
|
||||
@@ -30,7 +51,7 @@ When Mem0 is active, the model gets four tools it can call during a conversation
|
||||
| Tool | Description | Parameters |
|
||||
|------|-------------|------------|
|
||||
| `mem0_search` | Semantic search by meaning, ranked by relevance | `query` (required), `top_k` (default 10, max 50), `rerank` (default `false`, Platform mode only) |
|
||||
| `mem0_add` | Store a fact verbatim, with no LLM extraction | `content` (required) |
|
||||
| `mem0_add` | Store a fact after known-secret redaction, with no LLM extraction | `content` (required) |
|
||||
| `mem0_update` | Update a memory's text by ID | `memory_id`, `text` (both required) |
|
||||
| `mem0_delete` | Delete a memory by ID | `memory_id` (required) |
|
||||
|
||||
@@ -43,7 +64,27 @@ curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scri
|
||||
source ~/.bashrc
|
||||
```
|
||||
|
||||
The `mem0ai` package is installed automatically when you enable the Mem0 provider, so there is no manual pip step. OSS providers may need extra packages (for example `qdrant-client`, `psycopg2-binary`, or `ollama`), which the setup flow installs for you when you pick them.
|
||||
After this integration is published to the Mem0 repository:
|
||||
|
||||
```bash
|
||||
hermes plugins install mem0ai/mem0/integrations/hermes-plugin
|
||||
hermes plugins enable mem0
|
||||
hermes memory setup
|
||||
hermes memory status
|
||||
```
|
||||
|
||||
For unpublished worktree changes, use the README's local installation instructions instead.
|
||||
The bundled-provider precedence described above applies to both installation methods.
|
||||
|
||||
Recent Hermes development versions install the plugin's declared dependencies automatically.
|
||||
On Hermes `v0.21.3`, install them into the environment that runs Hermes:
|
||||
|
||||
```bash
|
||||
HERMES_REPO=/path/to/hermes-agent
|
||||
uv pip install --python "$HERMES_REPO/.venv/bin/python" 'mem0ai>=2.0.10,<3' 'httpx>=0.27,<1'
|
||||
```
|
||||
|
||||
OSS providers may need extra packages such as `qdrant-client`, `psycopg2-binary`, or `ollama`, which the setup flow installs when you select them.
|
||||
|
||||
## Platform Setup
|
||||
|
||||
@@ -196,9 +237,10 @@ Behavioral settings live in `~/.hermes/mem0.json` and are written for you by `he
|
||||
| `mode` | `platform` | `platform` (Mem0 Cloud) or `oss` (self-managed, in-process). Self-hosted server routing is set via `host` |
|
||||
| `host` | none | Self-hosted Mem0 server URL. When set, the plugin talks HTTP to your server instead of the cloud |
|
||||
| `api_key` | none | Mem0 Platform API key, or the admin key of a self-hosted server. Stored in `.env` as `MEM0_API_KEY` |
|
||||
| `user_id` | `hermes-user` | Identifier that scopes memories. See cross-channel behavior below |
|
||||
| `user_id` | Gateway user ID, then `hermes-user` | Identifier that scopes memories. See cross-channel behavior below |
|
||||
| `agent_id` | `hermes` | Agent identifier attached to writes |
|
||||
| `rerank` | `false` | Rerank search results for relevance (Platform mode only) |
|
||||
| `sync_max_chars` | Uncapped for platform/HTTP; `450` for OSS | Positive maximum per chunk; longer text is split without dropping its tail. `0` disables the character cap |
|
||||
|
||||
### Cross-channel memories
|
||||
|
||||
@@ -213,7 +255,8 @@ Either way, every write is tagged with `metadata.channel` (for example `telegram
|
||||
|
||||
- **Circuit breaker**: if Mem0 fails five times in a row, Hermes pauses calls for two minutes, then retries. The agent keeps working without memory during that window. Expected client errors, like a 404 on a missing memory id, do not count toward tripping the breaker.
|
||||
- **Non-blocking**: fact extraction runs in a background daemon thread, and current-turn recall waits at most 3 seconds, so a slow or failed call never blocks your conversation.
|
||||
- **Thread-safe**: the client uses lazy initialization with locking, and the background sync and recall threads are guarded so concurrent gateway messages cannot produce duplicate memories.
|
||||
- **Capture queue**: overlapping turns are queued in memory instead of skipped while an earlier turn is syncing. Failed requests are logged without durable retries. Shutdown waits at most five seconds; a process exit can lose pending turns.
|
||||
- **Existing collections**: an OSS embedding-dimension mismatch fails initialization without deleting vectors. Restore the original embedding configuration or select a new collection name.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
@@ -246,7 +289,7 @@ curl http://localhost:11434/api/tags
|
||||
|
||||
### Memories not appearing
|
||||
|
||||
- `mem0_add` stores text verbatim with no extraction. Ordinary conversation turns are extracted automatically by the background sync.
|
||||
- `mem0_add` stores text after known-secret redaction with no extraction. Ordinary conversation turns are extracted automatically by the background sync.
|
||||
- Search is semantic, so try a broader query.
|
||||
- Confirm `user_id` is the same across sessions (check `~/.hermes/mem0.json`).
|
||||
|
||||
|
||||
+1
-1
@@ -256,7 +256,7 @@ If the user is on a pre-current major (Python < 2, TS < 3, or a Platform call st
|
||||
- [Agno](https://docs.mem0.ai/integrations/agno) [Platform]: Use when the user is on Agno.
|
||||
- [Camel AI](https://docs.mem0.ai/integrations/camel-ai) [Both]: Use when the user is on Camel AI.
|
||||
- [ChatDev](https://docs.mem0.ai/integrations/chatdev) [Platform]: Use when the user is on ChatDev.
|
||||
- [Hermes](https://docs.mem0.ai/integrations/hermes) [Both]: Use when the user is on Hermes.
|
||||
- [Hermes Agent](https://docs.mem0.ai/integrations/hermes) [Both]: Use when adding the native Mem0 memory provider to Hermes Agent, preserving existing configuration, or checking bundled-provider precedence.
|
||||
- [Pi Agent](https://docs.mem0.ai/integrations/pi-agent) [Platform]: Use when adding automatic capture, prompt recall, scoped memory, and six memory commands to Pi Agent.
|
||||
- [DeepSeek Harness](https://docs.mem0.ai/integrations/deepseek-plugin) [Platform]: Use when adding automatic recall, completed-turn capture, and native search/add tools to DeepSeek Harness.
|
||||
- [OpenAI Agents SDK](https://docs.mem0.ai/integrations/openai-agents-sdk) [Platform]: Use when the user is on the OpenAI Agents SDK.
|
||||
|
||||
@@ -7,6 +7,7 @@ Agent and editor integrations. Most packages are self-contained; coding-agent pl
|
||||
| `vercel-ai-sdk/` | `@mem0/vercel-ai-provider` | tsup (CJS+ESM) | ESLint + Prettier | jest + vitest (edge/node) |
|
||||
| `openclaw/` | `@mem0/openclaw-mem0` | tsup (ESM) | none | vitest |
|
||||
| `agent-plugin-core/` | Shared Python/TypeScript behavior, skill templates, builds, and conformance | Python build script | ruff + tsc | pytest + node:test |
|
||||
| `hermes-plugin/` | Native Hermes memory provider; generated shared redaction and batching | Python build script | ruff | pytest with `--confcutdir=integrations/hermes-plugin/tests` |
|
||||
| `mem0-agent-plugin/` | One portable Agent Plugins v1 package | Python | ruff | shared conformance |
|
||||
| `claude-code-plugin/`, `cursor-plugin/`, `codex-plugin/`, `kimi-plugin/`, `antigravity-plugin/` | Self-contained native plugins generated from the shared Python core | Python | ruff | pytest |
|
||||
| `opencode-plugin/` | `@mem0/opencode-plugin` (Bun/TypeScript) | tsup (via Bun) | tsc | bun test |
|
||||
@@ -43,6 +44,7 @@ Run the type check after every TypeScript change: `pnpm run typecheck` or `tsc -
|
||||
|
||||
- **`vercel-ai-sdk/`** wraps the Vercel AI SDK through a `createMem0` provider. Integrations for AI-SDK repos go through this wrapper, not raw `MemoryClient`.
|
||||
- **`agent-plugin-core/`** owns the shared Python memory runtime, TypeScript lifecycle utilities, skill templates, builds, and conformance runner. Claude Code is the behavioral source of truth. Native manifests and adapters live in sibling plugin directories; do not hand-edit their generated `core/` or `skills/` trees. Build and validation details are in [`agent-plugin-core/README.md`](agent-plugin-core/README.md).
|
||||
- **`hermes-plugin/`** preserves the upstream Hermes memory-provider API, setup, and cloud/HTTP/OSS backends. Only `core/message_utils.py` is generated; no MCP server or generic skills are installed. Its offline tests use a separate pytest invocation with `--confcutdir=integrations/hermes-plugin/tests` to avoid importing the host entry point during collection.
|
||||
- **`opencode-plugin/`** is a Bun/TypeScript plugin for OpenCode (`@mem0/opencode-plugin` on npm). It registers Mem0 memory tools as an OpenCode plugin with its own skills and telemetry.
|
||||
- **`openclaw/`**, **`pi-agent-plugin/`**, **`deepseek-plugin/`** are editor and agent plugins with the same shape. `deepseek-plugin/` registers Mem0 search/add tools as a native DeepSeek Harness (Cordis) plugin.
|
||||
- **`n8n-nodes-mem0/`** is an n8n community node: add, search, get, update, delete.
|
||||
|
||||
@@ -18,10 +18,11 @@ integrations/
|
||||
├── cursor-plugin/ # Native Cursor package and adapter
|
||||
├── codex-plugin/ # Native Codex package and adapter
|
||||
├── kimi-plugin/ # Native Kimi package and adapter
|
||||
└── antigravity-plugin/ # Native Antigravity package and adapter
|
||||
├── antigravity-plugin/ # Native Antigravity package and adapter
|
||||
└── hermes-plugin/ # Native Hermes provider; shared message helpers only
|
||||
```
|
||||
|
||||
Each native directory owns only its manifest, native hooks or adapter, tests, and `plugin-build.json`. Its `core/` and `skills/` directories are generated from this module. They are committed because clients install a self-contained plugin directory and the Agent Plugins specification forbids package files from resolving outside the plugin root.
|
||||
Each native directory owns its manifest, native hooks or adapter, tests, and `plugin-build.json`. Hermes also retains its upstream setup wizard and backend adapters. Its `core/` and `skills/` directories are generated from this module. They are committed because clients install a self-contained plugin directory and the Agent Plugins specification forbids package files from resolving outside the plugin root.
|
||||
|
||||
Sidekick belongs only to Claude Code. Its agent definition is in `claude-code-plugin/agents/sidekick.md`; its hooks are in `claude-code-plugin/adapters/claude/hook.py`. Other plugins must not register Sidekick. The shared core handles memory and native subagent tracking for Claude Code and Codex.
|
||||
|
||||
@@ -29,7 +30,7 @@ TypeScript integrations (`openclaw`, `opencode-plugin`, `pi-agent-plugin`, and `
|
||||
|
||||
## Shared memory behavior
|
||||
|
||||
The six Python packages use the same `search_memories` MCP tool and six skill templates. Native hooks collect conversations and flush them to Mem0 in the background. The portable package uses the Agent Plugins v1 layout so compatible hosts can load its MCP server and skills. It has no lifecycle hooks or flush worker; its bundled `remember` skill assumes automatic capture and cannot save a memory on its own.
|
||||
The six MCP-based Python packages use the same `search_memories` MCP tool and six skill templates. Native hooks collect conversations and flush them to Mem0 in the background. The portable package uses the Agent Plugins v1 layout so compatible hosts can load its MCP server and skills. It has no lifecycle hooks or flush worker; its bundled `remember` skill assumes automatic capture and cannot save a memory on its own.
|
||||
|
||||
Python search accepts `query`, `top_k`, `category`, `scope`, and optional `run_id`:
|
||||
|
||||
@@ -49,6 +50,8 @@ TypeScript hosts reuse redaction and lifecycle utilities but retain their own to
|
||||
|
||||
For installation, follow the host guides: [Claude Code](../../docs/integrations/claude-code.mdx), [Cursor](../../docs/integrations/cursor.mdx), [Codex](../../docs/integrations/codex.mdx), [Kimi](../../docs/integrations/kimi.mdx), and [Antigravity](../../docs/integrations/antigravity.mdx).
|
||||
|
||||
Hermes bundles only `python/message_utils.py`, sharing the same secret redaction and lossless token batching. It retains Hermes user-wide recall, four native tools, three backend modes, and legacy configuration. Completed turns are queued in memory, with session `run_id` on writes; network failures are logged and are not durably retried. Explicit `sync_max_chars` values split text instead of truncating it. It does not install the generic MCP skills, repository scopes, or Sidekick.
|
||||
|
||||
## Build and verify
|
||||
|
||||
From the repository root:
|
||||
@@ -58,7 +61,7 @@ python3.11 -m venv /tmp/mem0-agent-plugins
|
||||
/tmp/mem0-agent-plugins/bin/pip install \
|
||||
-r integrations/agent-plugin-core/requirements-dev.txt
|
||||
|
||||
for host in claude-code cursor codex kimi antigravity; do
|
||||
for host in claude-code cursor codex kimi antigravity hermes; do
|
||||
/tmp/mem0-agent-plugins/bin/python \
|
||||
integrations/agent-plugin-core/build/build.py "$host" \
|
||||
--kind native --check
|
||||
@@ -99,7 +102,7 @@ Do not put a real key in source files, command history shared with others, or pu
|
||||
For another native Python host:
|
||||
|
||||
1. Add `integrations/<host>-plugin/` with its native manifest and the smallest adapter that translates host events.
|
||||
2. Add `plugin-build.json` declaring the plugin-root variable and runtime files.
|
||||
2. Add `plugin-build.json` declaring the plugin-root variable and runtime files. Native providers can select `native.pythonFiles` and set `native.skills: false` when they expose host-native tools instead of MCP.
|
||||
3. Add one adapter contract test.
|
||||
4. Register the host in `build/build.py` and `conformance/run.py`.
|
||||
5. Run `--sync`, `--check`, and the conformance command above.
|
||||
@@ -117,6 +120,7 @@ For a TypeScript host, import the shared lifecycle modules directly and keep onl
|
||||
| Codex | Native prompt and final-response fields | Structured failure indicators when present; otherwise unknown | Parent context; native agent ID |
|
||||
| Kimi | Prompt hooks and completed v2 wire output | Native success/failure hooks | No plugin subagent hooks or agent declaration |
|
||||
| Antigravity | Incremental completed transcript messages, including later prompts | Native tool errors | No plugin subagent hooks or agent declaration |
|
||||
| Hermes | Completed turns via `sync_turn`, split without truncation | No separate tool-evidence capture | No plugin subagent hooks or agent declaration |
|
||||
| Portable v1 | Explicit memory skills | No native lifecycle hooks | No native subagent declaration |
|
||||
|
||||
Python status uses `subagent_runs` and `last_subagent`. Legacy SQLite names and event handling keep existing records and running workers compatible.
|
||||
|
||||
@@ -25,6 +25,7 @@ NATIVE_PLUGINS = {
|
||||
"codex": INTEGRATIONS_ROOT / "codex-plugin",
|
||||
"kimi": INTEGRATIONS_ROOT / "kimi-plugin",
|
||||
"antigravity": INTEGRATIONS_ROOT / "antigravity-plugin",
|
||||
"hermes": INTEGRATIONS_ROOT / "hermes-plugin",
|
||||
}
|
||||
PROTECTED_OUTPUTS = {
|
||||
REPOSITORY_ROOT,
|
||||
@@ -88,13 +89,26 @@ def _bundle_python(
|
||||
*,
|
||||
plugin_data: str = "",
|
||||
portable: bool = False,
|
||||
python_files: list[str] | None = None,
|
||||
skills: bool = True,
|
||||
) -> None:
|
||||
core = staged / "core"
|
||||
core.mkdir()
|
||||
for source in sorted((CORE_ROOT / "python").glob("*.py")):
|
||||
if portable and source.name in {"flush_worker.py", "hook_runner.py"}:
|
||||
continue
|
||||
shutil.copy2(source, core / source.name)
|
||||
if python_files is None:
|
||||
python_files = sorted(source.name for source in (CORE_ROOT / "python").glob("*.py"))
|
||||
if not isinstance(python_files, list) or any(
|
||||
not isinstance(name, str) or Path(name).name != name or Path(name).suffix != ".py" for name in python_files
|
||||
):
|
||||
raise ValueError("pythonFiles must be a list of Python source filenames")
|
||||
_copy_declared_files(
|
||||
core,
|
||||
CORE_ROOT / "python",
|
||||
{name: name for name in python_files if not portable or name not in {"flush_worker.py", "hook_runner.py"}},
|
||||
)
|
||||
if not isinstance(skills, bool):
|
||||
raise ValueError("native skills must be a boolean")
|
||||
if not skills:
|
||||
return
|
||||
|
||||
values = {
|
||||
"PLUGIN_ROOT": plugin_root,
|
||||
@@ -150,6 +164,8 @@ def _build_native(host: str, source_root: Path, staged: Path, descriptor: dict)
|
||||
host,
|
||||
native["pluginRoot"],
|
||||
plugin_data=str(native.get("pluginData") or ""),
|
||||
python_files=native.get("pythonFiles"),
|
||||
skills=native.get("skills", True),
|
||||
)
|
||||
_copy_declared_files(staged, source_root, native.get("files", {}))
|
||||
|
||||
@@ -221,7 +237,10 @@ def sync_generated(host: str, kind: str) -> Path:
|
||||
with tempfile.TemporaryDirectory(prefix=f"mem0-sync-{host}-") as temporary:
|
||||
generated = build(host, kind, Path(temporary) / "bundle")
|
||||
for directory in ("core", "skills"):
|
||||
replace_output(generated / directory, target / directory)
|
||||
if (generated / directory).is_dir():
|
||||
replace_output(generated / directory, target / directory)
|
||||
elif (target / directory).exists():
|
||||
shutil.rmtree(target / directory)
|
||||
return target
|
||||
|
||||
|
||||
|
||||
@@ -13,10 +13,9 @@ import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
CORE_ROOT = Path(__file__).resolve().parents[1]
|
||||
REPOSITORY_ROOT = CORE_ROOT.parents[1]
|
||||
PYTHON_HOSTS = ("claude-code", "cursor", "codex", "kimi", "antigravity")
|
||||
PYTHON_HOSTS = ("claude-code", "cursor", "codex", "kimi", "antigravity", "hermes")
|
||||
GROUPS = (
|
||||
"python-bundles",
|
||||
"python-tests",
|
||||
@@ -30,9 +29,14 @@ LIVE_GROUP = "live-platform"
|
||||
|
||||
sys.path.insert(0, str(CORE_ROOT))
|
||||
sys.path.insert(0, str(CORE_ROOT / "python"))
|
||||
from memory_core import redact # noqa: E402
|
||||
from build.build import build # noqa: E402
|
||||
from conformance.artifacts import TYPESCRIPT_ARTIFACTS, verify_artifact as _typescript_artifact_check # noqa: E402
|
||||
from conformance.artifacts import ( # noqa: E402
|
||||
TYPESCRIPT_ARTIFACTS,
|
||||
)
|
||||
from conformance.artifacts import ( # noqa: E402
|
||||
verify_artifact as _typescript_artifact_check,
|
||||
)
|
||||
from message_utils import redact # noqa: E402
|
||||
|
||||
|
||||
def _package_directories() -> dict[str, Path]:
|
||||
@@ -61,6 +65,14 @@ def _runtime_commands() -> dict[str, list[list[str]]]:
|
||||
"-q",
|
||||
"--ignore=integrations/agent-plugin-core/tests/test_conformance.py",
|
||||
"--ignore=integrations/claude-code-plugin/tests/integration",
|
||||
],
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pytest",
|
||||
"integrations/hermes-plugin/tests",
|
||||
"--confcutdir=integrations/hermes-plugin/tests",
|
||||
"-q",
|
||||
]
|
||||
],
|
||||
"typescript-core": [["pnpm", "test"], ["pnpm", "typecheck"]],
|
||||
|
||||
@@ -11,7 +11,6 @@ from __future__ import annotations
|
||||
import functools
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
@@ -27,6 +26,14 @@ from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
import telemetry
|
||||
from message_utils import MAX_EXTRACTION_INPUT_TOKENS as MAX_EXTRACTION_INPUT_TOKENS
|
||||
from message_utils import SECRET_PATTERNS as SECRET_PATTERNS
|
||||
from message_utils import _estimated_tokens as _estimated_tokens
|
||||
from message_utils import _is_agent_assignment as _is_agent_assignment
|
||||
from message_utils import _is_agent_response as _is_agent_response
|
||||
from message_utils import _message_tokens as _message_tokens
|
||||
from message_utils import extraction_message_batches as extraction_message_batches
|
||||
from message_utils import redact as redact
|
||||
|
||||
DEFAULT_API_URL = "https://api.mem0.ai"
|
||||
PLUGIN_VERSION = "0.3.1"
|
||||
@@ -66,7 +73,6 @@ CHECKPOINT_EXCHANGES = 5
|
||||
CHECKPOINT_MESSAGES = 10
|
||||
CHECKPOINT_SOURCE_CHARS = 40000
|
||||
DEFAULT_MAX_CONTEXT_CHARS = 4000
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
MAX_FLUSH_ATTEMPTS = 5
|
||||
FORGET_PAGE_SIZE = 100
|
||||
FORGET_MAX_PAGES = 50
|
||||
@@ -139,48 +145,11 @@ BUILD_COMMAND_RE = re.compile(
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"
|
||||
),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r'|(?:access|refresh|session)[_-]?token|token|authorization|credential'
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = (
|
||||
value
|
||||
if isinstance(value, str)
|
||||
else json.dumps(value, ensure_ascii=False, default=str)
|
||||
)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def bounded(value: Any, limit: int) -> str:
|
||||
text = redact(value).strip()
|
||||
if len(text) <= limit:
|
||||
@@ -1706,100 +1675,6 @@ def build_extraction_messages(structured: dict[str, Any]) -> list[dict[str, str]
|
||||
return messages
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if (
|
||||
_is_agent_assignment(message)
|
||||
and index + 1 < len(exchange)
|
||||
and _is_agent_response(exchange[index + 1])
|
||||
):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
|
||||
|
||||
def _request_json(
|
||||
url: str, key: str, payload: dict[str, Any], timeout: float
|
||||
) -> tuple[dict[str, Any] | list[Any], int, int]:
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Shared, host-independent redaction and lossless extraction batching."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r"|(?:access|refresh|session)[_-]?token|token|authorization|credential"
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, default=str)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if _is_agent_assignment(message) and index + 1 < len(exchange) and _is_agent_response(exchange[index + 1]):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
@@ -1,3 +1,4 @@
|
||||
jsonschema>=4.23,<5
|
||||
pytest>=8,<10
|
||||
skills-ref==0.1.1
|
||||
httpx>=0.27,<1
|
||||
|
||||
@@ -15,6 +15,7 @@ from build.build import ( # noqa: E402
|
||||
bundle_drift,
|
||||
render_template,
|
||||
replace_output,
|
||||
sync_generated,
|
||||
)
|
||||
from build.validate import validate_bundle # noqa: E402
|
||||
|
||||
@@ -113,6 +114,7 @@ def test_native_control_skills_select_the_host_store(host: str, tmp_path: Path)
|
||||
("codex", "native"),
|
||||
("kimi", "native"),
|
||||
("antigravity", "native"),
|
||||
("hermes", "native"),
|
||||
],
|
||||
)
|
||||
def test_installable_plugin_directories_are_current(host: str, kind: str) -> None:
|
||||
@@ -133,3 +135,91 @@ def test_marketplaces_keep_public_names_and_reference_real_plugins() -> None:
|
||||
assert [plugin["name"] for plugin in codex_marketplace["plugins"]] == ["mem0"]
|
||||
codex = codex_marketplace["plugins"][0]
|
||||
assert codex["source"]["path"] == "./integrations/codex-plugin"
|
||||
|
||||
|
||||
def test_native_bundle_can_select_runtime_without_skills(tmp_path: Path, monkeypatch) -> None:
|
||||
from build import build as builder
|
||||
|
||||
source = tmp_path / "plugin"
|
||||
source.mkdir()
|
||||
(source / "__init__.py").write_text("# Native plugin adapter\n", encoding="utf-8")
|
||||
(source / "plugin-build.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"native": {
|
||||
"pluginRoot": "${PLUGIN_ROOT}",
|
||||
"pythonFiles": ["message_utils.py"],
|
||||
"skills": False,
|
||||
"files": {"__init__.py": "__init__.py"},
|
||||
}
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.setitem(builder.NATIVE_PLUGINS, "test-host", source)
|
||||
stale_skill = source / "skills" / "remember" / "SKILL.md"
|
||||
stale_skill.parent.mkdir(parents=True)
|
||||
stale_skill.write_text("stale generated skill", encoding="utf-8")
|
||||
|
||||
root = build("test-host", "native", tmp_path / "output")
|
||||
|
||||
assert {path.name for path in (root / "core").iterdir()} == {"message_utils.py"}
|
||||
assert not (root / "skills").exists()
|
||||
assert (root / "__init__.py").read_text(encoding="utf-8") == "# Native plugin adapter\n"
|
||||
sync_generated("test-host", "native")
|
||||
assert bundle_drift("test-host", "native") == []
|
||||
assert not (source / "skills").exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("files", ["message_utils.py", ["../README.md"], ["/tmp/source.py"], [7], ["missing.py"]])
|
||||
def test_native_runtime_selection_rejects_invalid_sources(files, tmp_path: Path) -> None:
|
||||
from build.build import _build_native
|
||||
|
||||
with pytest.raises(ValueError, match="pythonFiles|native source file"):
|
||||
_build_native(
|
||||
"test",
|
||||
tmp_path,
|
||||
tmp_path,
|
||||
{
|
||||
"native": {
|
||||
"pluginRoot": "${PLUGIN_ROOT}",
|
||||
"pythonFiles": files,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_native_runtime_selection_rejects_escaping_symlink(tmp_path: Path, monkeypatch) -> None:
|
||||
from build import build as builder
|
||||
|
||||
source = tmp_path / "shared" / "python"
|
||||
source.mkdir(parents=True)
|
||||
secret = tmp_path / "outside.py"
|
||||
secret.write_text("secret", encoding="utf-8")
|
||||
(source / "message_utils.py").symlink_to(secret)
|
||||
monkeypatch.setattr(builder, "CORE_ROOT", source.parent)
|
||||
staged = tmp_path / "bundle"
|
||||
staged.mkdir()
|
||||
|
||||
with pytest.raises(ValueError, match="inside their roots"):
|
||||
builder._build_native(
|
||||
"test",
|
||||
tmp_path,
|
||||
staged,
|
||||
{
|
||||
"native": {
|
||||
"pluginRoot": "${PLUGIN_ROOT}",
|
||||
"pythonFiles": ["message_utils.py"],
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_hermes_bundle_uses_only_host_independent_runtime(tmp_path: Path) -> None:
|
||||
root = build("hermes", "native", tmp_path / "hermes")
|
||||
|
||||
assert {path.name for path in (root / "core").iterdir()} == {"message_utils.py"}
|
||||
assert (root / "__init__.py").is_file()
|
||||
assert (root / "plugin.yaml").is_file()
|
||||
assert not (root / "skills").exists()
|
||||
assert not (root / "agents").exists()
|
||||
|
||||
@@ -10,10 +10,9 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from conformance import run as conformance_run # noqa: E402
|
||||
from conformance.run import _command_check # noqa: E402
|
||||
|
||||
|
||||
PLUGIN_ROOT = Path(__file__).resolve().parents[1]
|
||||
RUNNER = PLUGIN_ROOT / "conformance" / "run.py"
|
||||
PYTHON_HOSTS = {"claude-code", "cursor", "codex", "kimi", "antigravity"}
|
||||
PYTHON_HOSTS = {"claude-code", "cursor", "codex", "kimi", "antigravity", "hermes"}
|
||||
|
||||
|
||||
def test_python_bundle_conformance_builds_every_host(tmp_path: Path) -> None:
|
||||
@@ -78,6 +77,12 @@ def test_conformance_plan_covers_every_runtime(tmp_path: Path) -> None:
|
||||
"deepseek",
|
||||
}
|
||||
assert all(entry["status"] == "planned" for entry in payload["checks"])
|
||||
python_commands = [entry["command"] for entry in payload["checks"] if entry["group"] == "python-tests"]
|
||||
assert any(
|
||||
"integrations/hermes-plugin/tests" in command
|
||||
and "--confcutdir=integrations/hermes-plugin/tests" in command
|
||||
for command in python_commands
|
||||
)
|
||||
assert {
|
||||
entry["group"]
|
||||
for entry in payload["checks"]
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
CORE = Path(__file__).resolve().parents[1] / "python"
|
||||
|
||||
|
||||
def test_message_helpers_are_standalone_and_keep_legacy_exports() -> None:
|
||||
result = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
import message_utils
|
||||
assert "memory_core" not in sys.modules
|
||||
assert "telemetry" not in sys.modules
|
||||
text = "prefix " + "日本語🙂" * 250 + ' {"password": "secret value"}'
|
||||
redacted = message_utils.redact(text)
|
||||
assert "secret value" not in redacted
|
||||
messages = [{"role": "user", "content": redacted}]
|
||||
batches = message_utils.extraction_message_batches(messages, max_tokens=100)
|
||||
assert len(batches) > 1
|
||||
assert all(message_utils._message_tokens(batch) <= 100 for batch in batches)
|
||||
assert "".join(m["content"] for batch in batches for m in batch) == redacted
|
||||
import memory_core
|
||||
for name in ("redact", "SECRET_PATTERNS", "extraction_message_batches", "_message_tokens"):
|
||||
assert getattr(memory_core, name) is getattr(message_utils, name)
|
||||
""",
|
||||
],
|
||||
cwd=CORE,
|
||||
text=True,
|
||||
capture_output=True,
|
||||
check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
@@ -11,7 +11,6 @@ from __future__ import annotations
|
||||
import functools
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
@@ -27,6 +26,14 @@ from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
import telemetry
|
||||
from message_utils import MAX_EXTRACTION_INPUT_TOKENS as MAX_EXTRACTION_INPUT_TOKENS
|
||||
from message_utils import SECRET_PATTERNS as SECRET_PATTERNS
|
||||
from message_utils import _estimated_tokens as _estimated_tokens
|
||||
from message_utils import _is_agent_assignment as _is_agent_assignment
|
||||
from message_utils import _is_agent_response as _is_agent_response
|
||||
from message_utils import _message_tokens as _message_tokens
|
||||
from message_utils import extraction_message_batches as extraction_message_batches
|
||||
from message_utils import redact as redact
|
||||
|
||||
DEFAULT_API_URL = "https://api.mem0.ai"
|
||||
PLUGIN_VERSION = "0.3.1"
|
||||
@@ -66,7 +73,6 @@ CHECKPOINT_EXCHANGES = 5
|
||||
CHECKPOINT_MESSAGES = 10
|
||||
CHECKPOINT_SOURCE_CHARS = 40000
|
||||
DEFAULT_MAX_CONTEXT_CHARS = 4000
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
MAX_FLUSH_ATTEMPTS = 5
|
||||
FORGET_PAGE_SIZE = 100
|
||||
FORGET_MAX_PAGES = 50
|
||||
@@ -139,48 +145,11 @@ BUILD_COMMAND_RE = re.compile(
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"
|
||||
),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r'|(?:access|refresh|session)[_-]?token|token|authorization|credential'
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = (
|
||||
value
|
||||
if isinstance(value, str)
|
||||
else json.dumps(value, ensure_ascii=False, default=str)
|
||||
)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def bounded(value: Any, limit: int) -> str:
|
||||
text = redact(value).strip()
|
||||
if len(text) <= limit:
|
||||
@@ -1706,100 +1675,6 @@ def build_extraction_messages(structured: dict[str, Any]) -> list[dict[str, str]
|
||||
return messages
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if (
|
||||
_is_agent_assignment(message)
|
||||
and index + 1 < len(exchange)
|
||||
and _is_agent_response(exchange[index + 1])
|
||||
):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
|
||||
|
||||
def _request_json(
|
||||
url: str, key: str, payload: dict[str, Any], timeout: float
|
||||
) -> tuple[dict[str, Any] | list[Any], int, int]:
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Shared, host-independent redaction and lossless extraction batching."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r"|(?:access|refresh|session)[_-]?token|token|authorization|credential"
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, default=str)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if _is_agent_assignment(message) and index + 1 < len(exchange) and _is_agent_response(exchange[index + 1]):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
@@ -11,7 +11,6 @@ from __future__ import annotations
|
||||
import functools
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
@@ -27,6 +26,14 @@ from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
import telemetry
|
||||
from message_utils import MAX_EXTRACTION_INPUT_TOKENS as MAX_EXTRACTION_INPUT_TOKENS
|
||||
from message_utils import SECRET_PATTERNS as SECRET_PATTERNS
|
||||
from message_utils import _estimated_tokens as _estimated_tokens
|
||||
from message_utils import _is_agent_assignment as _is_agent_assignment
|
||||
from message_utils import _is_agent_response as _is_agent_response
|
||||
from message_utils import _message_tokens as _message_tokens
|
||||
from message_utils import extraction_message_batches as extraction_message_batches
|
||||
from message_utils import redact as redact
|
||||
|
||||
DEFAULT_API_URL = "https://api.mem0.ai"
|
||||
PLUGIN_VERSION = "0.3.1"
|
||||
@@ -66,7 +73,6 @@ CHECKPOINT_EXCHANGES = 5
|
||||
CHECKPOINT_MESSAGES = 10
|
||||
CHECKPOINT_SOURCE_CHARS = 40000
|
||||
DEFAULT_MAX_CONTEXT_CHARS = 4000
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
MAX_FLUSH_ATTEMPTS = 5
|
||||
FORGET_PAGE_SIZE = 100
|
||||
FORGET_MAX_PAGES = 50
|
||||
@@ -139,48 +145,11 @@ BUILD_COMMAND_RE = re.compile(
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"
|
||||
),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r'|(?:access|refresh|session)[_-]?token|token|authorization|credential'
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = (
|
||||
value
|
||||
if isinstance(value, str)
|
||||
else json.dumps(value, ensure_ascii=False, default=str)
|
||||
)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def bounded(value: Any, limit: int) -> str:
|
||||
text = redact(value).strip()
|
||||
if len(text) <= limit:
|
||||
@@ -1706,100 +1675,6 @@ def build_extraction_messages(structured: dict[str, Any]) -> list[dict[str, str]
|
||||
return messages
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if (
|
||||
_is_agent_assignment(message)
|
||||
and index + 1 < len(exchange)
|
||||
and _is_agent_response(exchange[index + 1])
|
||||
):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
|
||||
|
||||
def _request_json(
|
||||
url: str, key: str, payload: dict[str, Any], timeout: float
|
||||
) -> tuple[dict[str, Any] | list[Any], int, int]:
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Shared, host-independent redaction and lossless extraction batching."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r"|(?:access|refresh|session)[_-]?token|token|authorization|credential"
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, default=str)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if _is_agent_assignment(message) and index + 1 < len(exchange) and _is_agent_response(exchange[index + 1]):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
@@ -11,7 +11,6 @@ from __future__ import annotations
|
||||
import functools
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
@@ -27,6 +26,14 @@ from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
import telemetry
|
||||
from message_utils import MAX_EXTRACTION_INPUT_TOKENS as MAX_EXTRACTION_INPUT_TOKENS
|
||||
from message_utils import SECRET_PATTERNS as SECRET_PATTERNS
|
||||
from message_utils import _estimated_tokens as _estimated_tokens
|
||||
from message_utils import _is_agent_assignment as _is_agent_assignment
|
||||
from message_utils import _is_agent_response as _is_agent_response
|
||||
from message_utils import _message_tokens as _message_tokens
|
||||
from message_utils import extraction_message_batches as extraction_message_batches
|
||||
from message_utils import redact as redact
|
||||
|
||||
DEFAULT_API_URL = "https://api.mem0.ai"
|
||||
PLUGIN_VERSION = "0.3.1"
|
||||
@@ -66,7 +73,6 @@ CHECKPOINT_EXCHANGES = 5
|
||||
CHECKPOINT_MESSAGES = 10
|
||||
CHECKPOINT_SOURCE_CHARS = 40000
|
||||
DEFAULT_MAX_CONTEXT_CHARS = 4000
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
MAX_FLUSH_ATTEMPTS = 5
|
||||
FORGET_PAGE_SIZE = 100
|
||||
FORGET_MAX_PAGES = 50
|
||||
@@ -139,48 +145,11 @@ BUILD_COMMAND_RE = re.compile(
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"
|
||||
),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r'|(?:access|refresh|session)[_-]?token|token|authorization|credential'
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = (
|
||||
value
|
||||
if isinstance(value, str)
|
||||
else json.dumps(value, ensure_ascii=False, default=str)
|
||||
)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def bounded(value: Any, limit: int) -> str:
|
||||
text = redact(value).strip()
|
||||
if len(text) <= limit:
|
||||
@@ -1706,100 +1675,6 @@ def build_extraction_messages(structured: dict[str, Any]) -> list[dict[str, str]
|
||||
return messages
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if (
|
||||
_is_agent_assignment(message)
|
||||
and index + 1 < len(exchange)
|
||||
and _is_agent_response(exchange[index + 1])
|
||||
):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
|
||||
|
||||
def _request_json(
|
||||
url: str, key: str, payload: dict[str, Any], timeout: float
|
||||
) -> tuple[dict[str, Any] | list[Any], int, int]:
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Shared, host-independent redaction and lossless extraction batching."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r"|(?:access|refresh|session)[_-]?token|token|authorization|credential"
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, default=str)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if _is_agent_assignment(message) and index + 1 < len(exchange) and _is_agent_response(exchange[index + 1]):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
@@ -11,7 +11,6 @@ from __future__ import annotations
|
||||
import functools
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
@@ -27,6 +26,14 @@ from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
import telemetry
|
||||
from message_utils import MAX_EXTRACTION_INPUT_TOKENS as MAX_EXTRACTION_INPUT_TOKENS
|
||||
from message_utils import SECRET_PATTERNS as SECRET_PATTERNS
|
||||
from message_utils import _estimated_tokens as _estimated_tokens
|
||||
from message_utils import _is_agent_assignment as _is_agent_assignment
|
||||
from message_utils import _is_agent_response as _is_agent_response
|
||||
from message_utils import _message_tokens as _message_tokens
|
||||
from message_utils import extraction_message_batches as extraction_message_batches
|
||||
from message_utils import redact as redact
|
||||
|
||||
DEFAULT_API_URL = "https://api.mem0.ai"
|
||||
PLUGIN_VERSION = "0.3.1"
|
||||
@@ -66,7 +73,6 @@ CHECKPOINT_EXCHANGES = 5
|
||||
CHECKPOINT_MESSAGES = 10
|
||||
CHECKPOINT_SOURCE_CHARS = 40000
|
||||
DEFAULT_MAX_CONTEXT_CHARS = 4000
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
MAX_FLUSH_ATTEMPTS = 5
|
||||
FORGET_PAGE_SIZE = 100
|
||||
FORGET_MAX_PAGES = 50
|
||||
@@ -139,48 +145,11 @@ BUILD_COMMAND_RE = re.compile(
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"
|
||||
),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r'|(?:access|refresh|session)[_-]?token|token|authorization|credential'
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = (
|
||||
value
|
||||
if isinstance(value, str)
|
||||
else json.dumps(value, ensure_ascii=False, default=str)
|
||||
)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def bounded(value: Any, limit: int) -> str:
|
||||
text = redact(value).strip()
|
||||
if len(text) <= limit:
|
||||
@@ -1706,100 +1675,6 @@ def build_extraction_messages(structured: dict[str, Any]) -> list[dict[str, str]
|
||||
return messages
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if (
|
||||
_is_agent_assignment(message)
|
||||
and index + 1 < len(exchange)
|
||||
and _is_agent_response(exchange[index + 1])
|
||||
):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
|
||||
|
||||
def _request_json(
|
||||
url: str, key: str, payload: dict[str, Any], timeout: float
|
||||
) -> tuple[dict[str, Any] | list[Any], int, int]:
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Shared, host-independent redaction and lossless extraction batching."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r"|(?:access|refresh|session)[_-]?token|token|authorization|credential"
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, default=str)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if _is_agent_assignment(message) and index + 1 < len(exchange) and _is_agent_response(exchange[index + 1]):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
@@ -0,0 +1,64 @@
|
||||
# Source and compatibility notes
|
||||
|
||||
Imported from [NousResearch/hermes-plugin-mem0](https://github.com/NousResearch/hermes-plugin-mem0)
|
||||
version 1.3.0 at commit `3fc36950b2b7c19cdd81c6de99f10d2cbed850af`.
|
||||
The original native provider MIT license is retained in `LICENSE`. Generated shared core
|
||||
is Apache-2.0, with its license in `LICENSE-APACHE-2.0`; package metadata records both.
|
||||
This Mem0-owned port is version 1.4.0.
|
||||
The source was a handoff of Hermes' bundled provider, with package-relative imports and
|
||||
standalone dependency declarations; its historical authors remain in the upstream repositories.
|
||||
|
||||
## Changes in this port
|
||||
|
||||
- Generates only shared `core/message_utils.py` through `plugin-build.json`.
|
||||
- Uses shared secret redaction and token-aware batching for non-empty completed-turn text.
|
||||
- Preserves full message text after redaction; explicit `sync_max_chars` and the OSS default of 450
|
||||
limit chunk size instead of discarding the tail. Platform/HTTP has no default character cap.
|
||||
- Attaches Hermes session IDs as top-level Mem0 `run_id` on writes; existing user recall
|
||||
remains unfiltered by session. The queue is in-memory, with bounded shutdown, logged failures,
|
||||
and no durable retry.
|
||||
- Keeps `mem0`, `memory.provider`, `mem0.json`, environment fallbacks, all four tool names,
|
||||
identity precedence, SDK backends, and the three setup modes.
|
||||
|
||||
- Refuses OSS collection dimension mismatches without deleting existing vectors.
|
||||
|
||||
## Hermes contract review
|
||||
|
||||
Reviewed the [official memory-provider documentation](https://hermes-agent.nousresearch.com/docs/developer-guide/memory-provider-plugin),
|
||||
release [v0.21.3 / v2026.9.14](https://github.com/NousResearch/hermes-agent/releases/tag/v2026.9.14)
|
||||
(commit `345cd2b057a452236de401d3534b8502a7465e8d`), and main at
|
||||
`c62bd9f2078a946108f1c9d9b24bf118963277ef` on September 18, 2026.
|
||||
|
||||
- These Hermes revisions still bundle Mem0. Bundled providers take precedence over user
|
||||
directories. See README for an isolated preview; installation alone does not replace core.
|
||||
- The installer supports `owner/repo/subdirectory` and copies only that subtree. All runtime
|
||||
imports must stay within this directory or use Hermes/declared Python dependencies.
|
||||
- Main installs dependencies from `[project].dependencies` and reapplies them after updates;
|
||||
v0.21.3 does not. Install dependencies explicitly in the Hermes environment on that release.
|
||||
- `register(ctx)` registers a `MemoryProvider` instance. Hermes dispatches tool schemas and
|
||||
lifecycle hooks; no general-plugin hook registration is required. Background work retains
|
||||
context variables through `spawn_context_thread`, including profile and secret scope.
|
||||
- Hermes supports the existing `sync_turn(user, assistant, *, session_id="")` signature.
|
||||
New optional `messages`/`turn_author` arguments are passed only to providers accepting them.
|
||||
- `post_setup(hermes_home, config)` takes over setup and activation. The copied wizard already
|
||||
owns its prompt helpers; it still uses Hermes config, curses and credential utilities.
|
||||
- The handoff's original notes mentioned `config_schema.py`, but neither its actual directory
|
||||
nor the checked Hermes Mem0 directory contains it. This port preserves CLI setup and does
|
||||
not claim a provider-specific dashboard configuration panel.
|
||||
|
||||
Older Hermes releases and live cloud/OSS services require additional validation; preserving
|
||||
configuration and tool names does not imply compatibility with every historical host API.
|
||||
|
||||
Actual external-loader and `MemoryManager` smoke checks passed on both pinned host
|
||||
revisions: initialization, tool dispatch, optional sync arguments, full-turn tail capture,
|
||||
session ID forwarding and shutdown. SDK backends were mocked; no live service was contacted.
|
||||
|
||||
Reproduce the host smoke from the repository root (use the Hermes environment or a test
|
||||
venv with its imported dependencies installed):
|
||||
|
||||
```bash
|
||||
HERMES_SOURCE=/path/to/hermes-agent python integrations/hermes-plugin/tests/smoke_hermes.py
|
||||
```
|
||||
|
||||
The check confirms bundled precedence, then isolates external discovery by replacing the
|
||||
bundled search root with an empty temporary directory in the test process only.
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2025 Nous Research
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [2023] [Taranjeet Singh]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,273 @@
|
||||
# Mem0 for Hermes Agent
|
||||
|
||||
Native Hermes memory provider, version **1.4.0**, based on the Nous Research
|
||||
[1.3.0 handoff](https://github.com/NousResearch/hermes-plugin-mem0/tree/3fc36950b2b7c19cdd81c6de99f10d2cbed850af).
|
||||
Keeps platform, self-hosted HTTP, and in-process OSS modes and all four existing tools.
|
||||
See [HANDOFF.md](HANDOFF.md) for provenance and compatibility details.
|
||||
|
||||
## Install and activate
|
||||
|
||||
Requires Python 3.11+ and the Hermes memory-provider API present in **v0.21.3
|
||||
(v2026.9.14)**. Prefer the latest Hermes release. Earlier Hermes versions have not
|
||||
been validated. Install dependencies into the Python environment that runs Hermes:
|
||||
|
||||
```bash
|
||||
# Replace this with your Hermes checkout path.
|
||||
HERMES_REPO=/path/to/hermes-agent
|
||||
uv pip install --python "$HERMES_REPO/.venv/bin/python" 'mem0ai>=2.0.10,<3' 'httpx>=0.27,<1'
|
||||
```
|
||||
|
||||
Recent Hermes development versions install `pyproject.toml` dependencies automatically;
|
||||
v0.21.3 needs the explicit command above. Optional OSS providers can require additional
|
||||
packages installed by the setup wizard.
|
||||
|
||||
### Local worktree preview
|
||||
|
||||
From the Mem0 worktree root, copy the complete directory to the active Hermes profile.
|
||||
This command refuses to overwrite an existing user plugin:
|
||||
|
||||
```bash
|
||||
python3 - <<'PYINSTALL'
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
home = Path(os.environ.get("HERMES_HOME", "~/.hermes")).expanduser()
|
||||
shutil.copytree("integrations/hermes-plugin", home / "plugins" / "mem0")
|
||||
PYINSTALL
|
||||
```
|
||||
|
||||
**Installing is not enough on Hermes versions that bundle Mem0.** Both v0.21.3 and
|
||||
Hermes main checked on September 18, 2026 still contain `plugins/memory/mem0`.
|
||||
Hermes always selects that bundled provider ahead of a same-named user plugin;
|
||||
`hermes plugins enable mem0` does not change this precedence.
|
||||
|
||||
To preview this version without changing your normal Hermes checkout, make a separate
|
||||
Hermes worktree and replace its bundled copy. From the Mem0 worktree root:
|
||||
|
||||
```bash
|
||||
MEM0_WORKTREE="$PWD"
|
||||
HERMES_PREVIEW="${TMPDIR:-/tmp}/hermes-mem0-preview"
|
||||
git -C "$HERMES_REPO" worktree add --detach "$HERMES_PREVIEW" HEAD
|
||||
mv "$HERMES_PREVIEW/plugins/memory/mem0" "$HERMES_PREVIEW/.mem0-bundled-backup"
|
||||
cp -R "$MEM0_WORKTREE/integrations/hermes-plugin" "$HERMES_PREVIEW/plugins/memory/mem0"
|
||||
cd "$HERMES_PREVIEW"
|
||||
"$HERMES_REPO/.venv/bin/python" -m hermes_cli.main memory setup
|
||||
"$HERMES_REPO/.venv/bin/python" -m hermes_cli.main chat
|
||||
```
|
||||
|
||||
Use a fresh `HERMES_HOME` for a separate test profile, or your existing profile to retain
|
||||
its configuration and memory identity. The preview uses your existing Hermes environment;
|
||||
the original provider remains in the normal checkout and in the preview backup.
|
||||
Updating Hermes can restore its bundled copy, so recheck which provider is selected.
|
||||
Once Hermes no longer bundles Mem0, the user plugin copy loads directly.
|
||||
|
||||
### After this integration is published
|
||||
|
||||
Only after `integrations/hermes-plugin` is available on the remote branch:
|
||||
|
||||
```bash
|
||||
hermes plugins install mem0ai/mem0/integrations/hermes-plugin
|
||||
hermes plugins enable mem0
|
||||
hermes memory setup
|
||||
hermes memory status
|
||||
```
|
||||
|
||||
The bundled-provider precedence above still applies. This remote command does not install
|
||||
unpublished worktree changes. Keep the name `mem0`, existing `mem0.json`, `MEM0_*` variables,
|
||||
and `memory.provider: mem0`; no memory migration or new user ID is required.
|
||||
|
||||
## Capture and recall
|
||||
|
||||
Automatic capture processes the completed user/assistant turn with the shared Mem0
|
||||
message preparation: known-secret redaction and token-aware batching. Empty text is skipped.
|
||||
It preserves full completed-turn text after redaction instead of dropping everything after 450 characters.
|
||||
Raw tool results and the full historical transcript are not captured. Platform mode sends
|
||||
prepared turn text to Mem0 Cloud for extraction; HTTP and OSS modes use their configured
|
||||
server or model providers. Pattern-based redaction cannot recognize every possible secret.
|
||||
|
||||
An explicit positive `sync_max_chars` still limits each message chunk; long messages are
|
||||
split across chunks without losing their remaining text. The default is no character cap
|
||||
for platform/HTTP and 450 characters per chunk for OSS. Batches also obey the shared token
|
||||
budget. Writes attach the Hermes session as Mem0’s top-level `run_id`. Recall does not filter by
|
||||
`run_id`, so the existing user identity continues to recall across sessions. Automatic
|
||||
capture uses an in-memory queue: process crashes can lose queued turns, and failed backend
|
||||
batches are logged without durable retries. Shutdown drains the queue for a bounded time;
|
||||
a stuck backend can leave queued turns unfinished.
|
||||
|
||||
Only `core/message_utils.py` is generated from `agent-plugin-core`; Hermes owns provider
|
||||
lifecycle, setup, tools, and SDK backends. This package does not add MCP or coding-agent skills.
|
||||
|
||||
## Config
|
||||
|
||||
Behavioral settings live in `$HERMES_HOME/mem0.json` (set them via `hermes memory setup`). Store credentials in the active profile’s `$HERMES_HOME/.env`.
|
||||
|
||||
| Key | Default | Description |
|
||||
|-----|---------|-------------|
|
||||
| `mode` | `platform` | `platform` (Mem0 Cloud) or `oss` (self-managed, in-process) |
|
||||
| `host` | — | Self-hosted Mem0 server URL (the Docker dashboard). When set, connects over HTTP with `X-API-Key`. Don't combine with `mode: oss` |
|
||||
| `user_id` | Gateway user ID, else `hermes-user` | Explicit config or `MEM0_USER_ID` preserves the same identity across hosts. The legacy `hermes-user` placeholder permits gateway fallback. |
|
||||
| `agent_id` | `hermes` | Agent identifier |
|
||||
| `rerank` | `false` | Rerank search results for relevance (platform mode only) |
|
||||
| `sync_max_chars` | Uncapped (platform/HTTP), `450` (OSS) | Positive character limit per chunk; longer text is split, not discarded. Increase for larger OSS embedding windows. |
|
||||
|
||||
The plugin has three connection modes:
|
||||
|
||||
- **Platform** — Mem0's hosted cloud (`api.mem0.ai`). Set `MEM0_API_KEY`. (default)
|
||||
- **Self-hosted dashboard** — a Mem0 server you run yourself via Docker. Set `host`. See below.
|
||||
- **OSS** — run Mem0 in-process with your own LLM + vector store. Set `mode: oss`. See below.
|
||||
|
||||
## Self-Hosted Dashboard (Server) Mode
|
||||
|
||||
Connect the plugin to a standalone Mem0 server you run yourself — the Docker-shipped Mem0 dashboard/server with its own REST API. Unlike OSS mode (which runs `mem0ai` in-process with your own vector store), here the plugin just talks HTTP to your server.
|
||||
|
||||
1. Run the Mem0 server (FastAPI + pgvector) from its Docker image and note its URL and `ADMIN_API_KEY`.
|
||||
2. Point the plugin at it — via the setup wizard:
|
||||
```bash
|
||||
hermes memory setup # select "mem0" → "Self-hosted server"
|
||||
# Or non-interactive:
|
||||
hermes memory setup mem0 --mode selfhosted --host http://localhost:8888 --api-key your-admin-api-key
|
||||
```
|
||||
or via env vars:
|
||||
```bash
|
||||
echo "MEM0_HOST=http://localhost:8888" >> ~/.hermes/.env
|
||||
echo "MEM0_API_KEY=your-admin-api-key" >> ~/.hermes/.env
|
||||
```
|
||||
or in `$HERMES_HOME/mem0.json`:
|
||||
```json
|
||||
{
|
||||
"host": "http://localhost:8888",
|
||||
"api_key": "your-admin-api-key"
|
||||
}
|
||||
```
|
||||
3. Start a fresh Hermes session and call `mem0_search` — it connects to your server.
|
||||
|
||||
The plugin authenticates with `X-API-Key` and uses the server's `/search` and `/memories` routes. `api_key` is optional — omit it only for servers running with `AUTH_DISABLED`.
|
||||
|
||||
> Setting `host` routes to the self-hosted server automatically. Don't set `mode: oss` — OSS takes precedence and ignores `host`.
|
||||
|
||||
## OSS (Self-Hosted) Mode
|
||||
|
||||
Run Mem0 locally with your own LLM, embedder, and vector store. This is the in-process SDK mode. To instead connect to a Mem0 server you run via Docker, see [Self-Hosted Dashboard (Server) Mode](#self-hosted-dashboard-server-mode) above.
|
||||
|
||||
### Interactive Setup
|
||||
|
||||
```bash
|
||||
hermes memory setup
|
||||
# Select "mem0" → "Open Source (self-hosted)"
|
||||
# Follow prompts for LLM, embedder, and vector store
|
||||
```
|
||||
|
||||
### Agent-Driven Setup (Flags)
|
||||
|
||||
```bash
|
||||
hermes memory setup mem0 --mode oss \
|
||||
--oss-llm openai --oss-llm-key sk-... \
|
||||
--oss-vector qdrant
|
||||
```
|
||||
|
||||
### Supported Providers
|
||||
|
||||
| Component | Providers |
|
||||
|-----------|-----------|
|
||||
| LLM | openai, ollama |
|
||||
| Embedder | openai, ollama |
|
||||
| Vector Store | qdrant (local/server), pgvector |
|
||||
|
||||
### Flags Reference
|
||||
|
||||
| Flag | Description |
|
||||
|------|-------------|
|
||||
| `--mode` | `platform`, `selfhosted`, or `oss` |
|
||||
| `--oss-llm` | LLM provider (default: openai) |
|
||||
| `--oss-llm-key` | LLM API key |
|
||||
| `--oss-embedder` | Embedder provider (default: openai) |
|
||||
| `--oss-vector` | Vector store (default: qdrant) |
|
||||
| `--oss-vector-path` | Qdrant local path |
|
||||
| `--user-id` | User identifier |
|
||||
|
||||
## Switching Modes
|
||||
|
||||
### Platform to OSS
|
||||
|
||||
```bash
|
||||
hermes memory setup mem0 --mode oss --oss-llm-key sk-...
|
||||
```
|
||||
|
||||
Or edit `$HERMES_HOME/mem0.json` directly:
|
||||
```json
|
||||
{
|
||||
"mode": "oss",
|
||||
"oss": {
|
||||
"llm": {"provider": "openai", "config": {"model": "gpt-5-mini", "is_reasoning_model": true}},
|
||||
"embedder": {"provider": "openai", "config": {"model": "text-embedding-3-small"}},
|
||||
"vector_store": {"provider": "qdrant", "config": {"path": "~/.hermes/mem0_qdrant"}}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### OSS to Platform
|
||||
|
||||
```bash
|
||||
hermes memory setup mem0 --mode platform --api-key sk-...
|
||||
```
|
||||
|
||||
### Dry Run (preview without writing)
|
||||
|
||||
```bash
|
||||
hermes memory setup mem0 --mode oss --oss-llm-key sk-... --dry-run
|
||||
```
|
||||
|
||||
## Tools
|
||||
|
||||
| Tool | Description |
|
||||
|------|-------------|
|
||||
| `mem0_search` | Semantic search by meaning |
|
||||
| `mem0_add` | Store an explicit fact after known-secret redaction (no LLM extraction) |
|
||||
| `mem0_update` | Update a memory's text by ID |
|
||||
| `mem0_delete` | Delete a memory by ID |
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### "Mem0 temporarily unavailable"
|
||||
|
||||
Circuit breaker tripped after 5 consecutive failures. Resets after 2 minutes.
|
||||
|
||||
- **Platform mode**: Check API key and internet connectivity.
|
||||
- **OSS mode**: Check that your vector store (qdrant/pgvector) is running.
|
||||
|
||||
### OSS: Qdrant connection refused
|
||||
|
||||
```bash
|
||||
# If using local Qdrant, check the storage path is writable:
|
||||
ls -la ~/.hermes/mem0_qdrant
|
||||
|
||||
# If using Qdrant server, check it's reachable:
|
||||
curl http://localhost:6333/healthz
|
||||
```
|
||||
|
||||
### OSS: PGVector connection refused
|
||||
|
||||
```bash
|
||||
# Verify PostgreSQL is running and accepting connections:
|
||||
pg_isready -h localhost -p 5432
|
||||
```
|
||||
|
||||
### OSS: Ollama not reachable
|
||||
|
||||
```bash
|
||||
# Check Ollama is running:
|
||||
curl http://localhost:11434/api/tags
|
||||
```
|
||||
|
||||
### Memories not appearing
|
||||
|
||||
- `mem0_add` stores a redacted fact without extraction. Completed turns are extracted automatically.
|
||||
- Search uses semantic matching — try broader queries.
|
||||
- Check `user_id` matches between sessions (`$HERMES_HOME/mem0.json`).
|
||||
|
||||
### Existing OSS collection has different embedding dimensions
|
||||
|
||||
Initialization fails with a configuration error instead of deleting or recreating the
|
||||
collection. Restore the original embedding model/dimensions, or explicitly configure a new
|
||||
collection name. Existing vectors remain untouched.
|
||||
@@ -0,0 +1,663 @@
|
||||
"""Mem0 memory plugin — MemoryProvider interface.
|
||||
|
||||
Server-side fact extraction and semantic search via the Mem0 Platform API (cloud), a
|
||||
self-hosted Mem0 server (MEM0_HOST, HTTP), or OSS Memory. Secrets live in $HERMES_HOME/.env
|
||||
(MEM0_API_KEY, MEM0_HOST); settings in $HERMES_HOME/mem0.json via `hermes memory setup`:
|
||||
mode ("platform"|"oss"), host, user_id (canonical id across gateways; unset → gateway-native
|
||||
id), agent_id. MEM0_* env vars remain a fallback.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
from contextlib import suppress
|
||||
from contextvars import copy_context
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from agent.memory_provider import MemoryProvider, spawn_context_thread
|
||||
from agent.secret_scope import get_secret
|
||||
from tools.registry import tool_error
|
||||
from utils import atomic_json_write, read_json_or_empty
|
||||
|
||||
from .core.message_utils import extraction_message_batches, redact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Circuit breaker: after _BREAKER_THRESHOLD consecutive failures, pause API
|
||||
# calls for _BREAKER_COOLDOWN_SECS to avoid hammering a down server.
|
||||
_BREAKER_THRESHOLD, _BREAKER_COOLDOWN_SECS, _PREFETCH_WAIT_SECS = 5, 120, 3
|
||||
_SHUTDOWN_WAIT_SECS = 5.0
|
||||
_CLIENT_ERROR_TYPES = ("MemoryNotFoundError", "ValidationError")
|
||||
# Placeholder user_id. initialize() treats it as "no operator-configured user_id"
|
||||
# so legacy mem0.json files written by the wizard don't override gateway-native ids.
|
||||
_DEFAULT_USER_ID = "hermes-user"
|
||||
|
||||
# Legacy OSS embedding windows remain configurable; split instead of discarding the tail.
|
||||
_SYNC_MSG_MAX_CHARS = 450
|
||||
|
||||
|
||||
# Sentence ends recognized when trimming a synced message. Deliberately unordered:
|
||||
# the LAST boundary of ANY kind wins, so one CJK stop early in a mixed-script turn
|
||||
# cannot outrank a Latin stop near the end of the window. ``".\n"`` is not listed —
|
||||
# its index can never exceed the bare ``"."`` it starts with.
|
||||
_SYNC_SENTENCE_ENDS = ("。", "！", "？", ".", "!", "?")
|
||||
|
||||
|
||||
def _truncate_for_sync(text: str, max_len: int = _SYNC_MSG_MAX_CHARS) -> str:
|
||||
"""Cap a synced message at its last sentence boundary within ``max_len``.
|
||||
|
||||
Short messages pass through unchanged; long ones keep the last complete
|
||||
sentence inside the window so fact extraction still sees coherent statements,
|
||||
with a hard cut as fallback when no boundary exists (or one only appears in
|
||||
the first third of the window, which usually means unsegmented input).
|
||||
"""
|
||||
if len(text) <= max_len:
|
||||
return text
|
||||
window = text[:max_len]
|
||||
cut = max(window.rfind(sep) for sep in _SYNC_SENTENCE_ENDS)
|
||||
if cut > max_len // 3:
|
||||
return text[: cut + 1]
|
||||
return text[:max_len]
|
||||
|
||||
|
||||
def _is_client_error(exc: Exception) -> bool:
|
||||
"""True for user-caused errors (bad ID, not found) that should NOT trip circuit breaker."""
|
||||
err_str = str(exc).lower()
|
||||
return type(exc).__name__ in _CLIENT_ERROR_TYPES or any(s in err_str for s in ("404", "not found", "valid uuid"))
|
||||
|
||||
|
||||
def _load_config() -> dict:
|
||||
"""Env vars provide defaults; $HERMES_HOME/mem0.json overrides individual keys.
|
||||
Layering avoids a silent failure when the JSON file exists but lacks fields
|
||||
like ``api_key`` that the user set in ``.env``."""
|
||||
from hermes_constants import get_hermes_home
|
||||
|
||||
# Identity (user/agent id), host and mode are .env values like the key: read them through the
|
||||
# profile scope too, or a secondary profile's memories land in the default profile's account.
|
||||
# A scope-less multiplex caller raises here on purpose — that is a spawn-site bug, and
|
||||
# swallowing it would silently route the turn's memories to the default profile.
|
||||
config = {
|
||||
"mode": get_secret("MEM0_MODE", "") or "platform",
|
||||
"host": get_secret("MEM0_HOST", "") or "",
|
||||
"agent_id": get_secret("MEM0_AGENT_ID", "") or "hermes",
|
||||
"oss": {},
|
||||
}
|
||||
if user_id := get_secret(
|
||||
"MEM0_USER_ID", ""
|
||||
): # only when explicitly configured, so initialize() can fall back to the gateway-native id
|
||||
config["user_id"] = user_id
|
||||
file_cfg = read_json_or_empty(get_hermes_home() / "mem0.json")
|
||||
config.update({k: v for k, v in file_cfg.items() if v is not None and v != ""})
|
||||
# MEM0_API_KEY authenticates the Platform and self-hosted HTTP backends; pure OSS mode builds its
|
||||
# backend from the local ``oss`` config and has no platform credential to resolve, so a profile
|
||||
# scope WITHOUT the key must still load an OSS config (#99121 as it stands today: the caller is
|
||||
# scoped, the scope is just empty). Decided after mem0.json overrode the env defaults because
|
||||
# the file may be what selects ``oss``. Scope-less callers already raised above.
|
||||
if config.get("mode", "platform") == "oss":
|
||||
config.setdefault("api_key", "")
|
||||
elif not config.get("api_key"):
|
||||
config["api_key"] = get_secret("MEM0_API_KEY", "")
|
||||
return config
|
||||
|
||||
|
||||
def _schema(name: str, description: str, properties: dict[str, tuple[str, str]], required: list[str]) -> dict:
|
||||
props = {k: {"type": t, "description": d} for k, (t, d) in properties.items()}
|
||||
return {
|
||||
"name": name,
|
||||
"description": description,
|
||||
"parameters": {"type": "object", "properties": props, "required": required},
|
||||
}
|
||||
|
||||
|
||||
TOOL_SCHEMAS = [
|
||||
_schema(
|
||||
"mem0_search",
|
||||
"Search the user's memories by meaning; returns facts ranked by relevance. Use this before answering any question that may depend on what you know about the user (preferences, facts, history, people, projects, past decisions). For multi-part or multi-hop questions, call it several times — vary the wording and run follow-up searches on what earlier results reveal; one search is rarely enough.",
|
||||
{
|
||||
"query": ("string", "What to search for."),
|
||||
"top_k": ("integer", "Max results (default: 10, max: 50)."),
|
||||
"rerank": ("boolean", "Rerank results for relevance (default: false, platform mode only)."),
|
||||
},
|
||||
["query"],
|
||||
),
|
||||
_schema(
|
||||
"mem0_add",
|
||||
"Store a durable fact about the user, verbatim (no LLM extraction). Call this the moment the user states a lasting preference, correction, decision, or personal detail worth recalling on future turns — don't wait to be asked to remember. Skip transient chit-chat and facts you've already stored.",
|
||||
{"content": ("string", "The fact to store.")},
|
||||
["content"],
|
||||
),
|
||||
_schema(
|
||||
"mem0_update",
|
||||
"Replace the text of an existing memory by its ID (take the ID from a mem0_search result). Use when a stored fact has changed or was wrong — correct it in place instead of adding a duplicate.",
|
||||
{"memory_id": ("string", "Memory UUID to update."), "text": ("string", "New text content.")},
|
||||
["memory_id", "text"],
|
||||
),
|
||||
_schema(
|
||||
"mem0_delete",
|
||||
"Delete a memory by its ID (take the ID from a mem0_search result). Use when a stored fact is obsolete or the user asks you to forget it; prefer mem0_update if the fact merely changed.",
|
||||
{"memory_id": ("string", "Memory UUID to delete.")},
|
||||
["memory_id"],
|
||||
),
|
||||
]
|
||||
|
||||
_PROMPT_BODY = (
|
||||
"You have persistent memory of this user from past conversations. You should call mem0_search before answering anything that could depend on prior context (the user's preferences, facts, history, people, projects, or earlier decisions) — do not rely on the chat window alone, and do not assume you have no memory.\n"
|
||||
"For multi-part or multi-hop questions, run several searches with different wording/angles and follow-up searches on what the first results surface; one search is rarely enough. Keep searching until you have every fact the question needs before you answer.\n"
|
||||
"Tools: mem0_search to find memories, mem0_add to store facts, mem0_update and mem0_delete to manage by ID."
|
||||
)
|
||||
|
||||
|
||||
class Mem0MemoryProvider(MemoryProvider):
|
||||
"""Mem0 memory with server-side extraction and semantic search (platform, self-hosted or OSS)."""
|
||||
|
||||
def __init__(self):
|
||||
self._config = self._backend = self._sync_thread = self._prefetch_thread = None
|
||||
self._mode, self._api_key, self._host, self._user_id, self._agent_id = (
|
||||
"platform",
|
||||
"",
|
||||
"",
|
||||
_DEFAULT_USER_ID,
|
||||
"hermes",
|
||||
)
|
||||
self._rerank_default, self._channel = False, "cli" # channel = gateway name (cli/telegram/discord/...)
|
||||
self._sync_max_chars = 0
|
||||
self._session_id = ""
|
||||
self._sync_queue = deque()
|
||||
self._closed = False
|
||||
self._prefetch_query = self._prefetch_result = ""
|
||||
self._prefetch_done = self._atexit_registered = False
|
||||
self._prefetch_threads = set()
|
||||
self._consecutive_failures, self._breaker_open_until = 0, 0.0 # circuit breaker state
|
||||
self._breaker_lock, self._sync_lock, self._prefetch_lock = threading.Lock(), threading.Lock(), threading.Lock()
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "mem0"
|
||||
|
||||
def is_available(self) -> bool:
|
||||
cfg = _load_config()
|
||||
if cfg.get("mode", "platform") == "oss":
|
||||
return bool(cfg.get("oss", {}).get("vector_store"))
|
||||
return bool(
|
||||
cfg.get("api_key") or cfg.get("host")
|
||||
) # platform needs a key; self-hosted a host (key optional with AUTH_DISABLED)
|
||||
|
||||
def save_config(self, values, hermes_home):
|
||||
"""Merge-write config to $HERMES_HOME/mem0.json."""
|
||||
config_path = Path(hermes_home) / "mem0.json"
|
||||
atomic_json_write(config_path, {**read_json_or_empty(config_path), **values}, mode=0o600)
|
||||
|
||||
def get_config_schema(self):
|
||||
cfg = _load_config()
|
||||
api_key_required = cfg.get("mode", "platform") != "oss" and not cfg.get("host")
|
||||
return [
|
||||
{
|
||||
"key": "api_key",
|
||||
"description": "Mem0 Platform API key",
|
||||
"secret": True,
|
||||
"required": api_key_required,
|
||||
"env_var": "MEM0_API_KEY",
|
||||
"url": "https://app.mem0.ai",
|
||||
},
|
||||
{
|
||||
"key": "host",
|
||||
"description": "Self-hosted Mem0 server URL (leave blank for cloud)",
|
||||
"required": False,
|
||||
"env_var": "MEM0_HOST",
|
||||
},
|
||||
{"key": "user_id", "description": "User identifier", "default": "hermes-user"},
|
||||
{"key": "agent_id", "description": "Agent identifier", "default": "hermes"},
|
||||
{
|
||||
"key": "rerank",
|
||||
"description": "Enable reranking for recall",
|
||||
"default": "false",
|
||||
"choices": ["true", "false"],
|
||||
},
|
||||
]
|
||||
|
||||
def post_setup(self, hermes_home: str, config: dict) -> None:
|
||||
from ._setup import post_setup
|
||||
|
||||
post_setup(hermes_home, config)
|
||||
|
||||
def _oss_hint(self, template: str, default: str = "vector store") -> str:
|
||||
"""OSS-only hint; ``{vs}`` is the configured vector-store provider. "" in other modes."""
|
||||
return (
|
||||
template.format(vs=self._config.get("oss", {}).get("vector_store", {}).get("provider", default))
|
||||
if self._mode == "oss"
|
||||
else ""
|
||||
)
|
||||
|
||||
def _create_backend(self):
|
||||
try:
|
||||
from . import _backend
|
||||
|
||||
if self._mode == "oss":
|
||||
return _backend.OSSBackend(self._config.get("oss", {}))
|
||||
return (
|
||||
_backend.SelfHostedBackend(self._api_key, self._host)
|
||||
if self._host
|
||||
else _backend.PlatformBackend(self._api_key)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Mem0 backend failed to initialize (%s mode): %s", self._mode, redact(str(e)))
|
||||
self._init_error = redact(str(e))
|
||||
return None
|
||||
|
||||
def _is_breaker_open(self) -> bool:
|
||||
"""True while the breaker is tripped; an expired cooldown resets the failure count."""
|
||||
with self._breaker_lock:
|
||||
if self._consecutive_failures >= _BREAKER_THRESHOLD and time.monotonic() < self._breaker_open_until:
|
||||
return True
|
||||
if self._consecutive_failures >= _BREAKER_THRESHOLD:
|
||||
self._consecutive_failures = 0
|
||||
return False
|
||||
|
||||
def _format_error(self, prefix: str, exc: Exception) -> str:
|
||||
msg = f"{prefix}: {redact(str(exc))}"
|
||||
if any(s in str(exc).lower() for s in ("connection", "refused", "timeout")):
|
||||
msg += self._oss_hint(" (check that {vs} is running)")
|
||||
return msg
|
||||
|
||||
def _record_success(self):
|
||||
with self._breaker_lock:
|
||||
self._consecutive_failures = 0
|
||||
|
||||
def _record_failure(self):
|
||||
with self._breaker_lock:
|
||||
self._consecutive_failures = count = self._consecutive_failures + 1
|
||||
if count >= _BREAKER_THRESHOLD:
|
||||
self._breaker_open_until = time.monotonic() + _BREAKER_COOLDOWN_SECS
|
||||
if count >= _BREAKER_THRESHOLD:
|
||||
hint = self._oss_hint(" Check that your {vs} vector store is running and reachable.", "unknown")
|
||||
logger.warning(
|
||||
"Mem0 circuit breaker tripped after %d consecutive failures. Pausing API calls for %ds.%s",
|
||||
count,
|
||||
_BREAKER_COOLDOWN_SECS,
|
||||
hint,
|
||||
)
|
||||
|
||||
def _try(self, call, log, msg: str):
|
||||
"""Background-path wrapper: run ``call`` under the breaker; on error log ``msg`` and return None."""
|
||||
try:
|
||||
result = call()
|
||||
except Exception as e:
|
||||
self._record_failure()
|
||||
log(msg, redact(str(e)))
|
||||
return None
|
||||
self._record_success()
|
||||
return result
|
||||
|
||||
def initialize(self, session_id: str, **kwargs) -> None:
|
||||
self._session_id = session_id
|
||||
self._config = cfg = _load_config()
|
||||
self._mode, self._api_key, self._host, self._agent_id = (
|
||||
cfg.get("mode", "platform"),
|
||||
cfg.get("api_key", ""),
|
||||
cfg.get("host", ""),
|
||||
cfg.get("agent_id", "hermes"),
|
||||
)
|
||||
# user_id precedence: operator-configured (env/mem0.json) > gateway-native id (kwargs) > _DEFAULT_USER_ID.
|
||||
# The literal placeholder counts as unset so wizard users still get gateway-native ids.
|
||||
configured = cfg.get("user_id")
|
||||
self._user_id = (
|
||||
(None if configured == _DEFAULT_USER_ID else configured) or kwargs.get("user_id") or _DEFAULT_USER_ID
|
||||
)
|
||||
# Persisted rerank preference: default for mem0_search when the model omits ``rerank``. Platform-only.
|
||||
_rr = cfg.get("rerank", False)
|
||||
self._rerank_default = _rr.lower() in ("true", "1", "yes") if isinstance(_rr, str) else bool(_rr)
|
||||
self._channel = kwargs.get("platform") or "cli"
|
||||
default_cap = _SYNC_MSG_MAX_CHARS if self._mode == "oss" else 0
|
||||
self._sync_max_chars = int(cfg.get("sync_max_chars", default_cap))
|
||||
if self._sync_max_chars < 0:
|
||||
raise ValueError("sync_max_chars must be zero (unlimited) or positive")
|
||||
self._backend = self._create_backend()
|
||||
if self._backend and not self._atexit_registered:
|
||||
atexit.register(self.shutdown)
|
||||
self._atexit_registered = True
|
||||
|
||||
def _search(self, query: str, top_k: int = 10, rerank: bool = False, backend=None) -> list:
|
||||
# Scoped to user_id only — by design — so recall surfaces memories from any gateway/agent under this
|
||||
# principal; writes attach agent_id and metadata.channel so narrower views remain possible at query time.
|
||||
return (backend or self._backend).search(
|
||||
redact(query), filters={"user_id": self._user_id}, top_k=top_k, rerank=rerank
|
||||
)
|
||||
|
||||
def _add(self, messages: list, infer: bool, *, session_id: str = ""):
|
||||
"""Send messages already redacted at the capture or explicit-tool boundary."""
|
||||
metadata = {"channel": self._channel} if self._channel else {}
|
||||
return self._backend.add(
|
||||
messages,
|
||||
user_id=self._user_id,
|
||||
agent_id=self._agent_id,
|
||||
infer=infer,
|
||||
metadata=metadata,
|
||||
run_id=session_id or self._session_id,
|
||||
)
|
||||
|
||||
def system_prompt_block(self) -> str:
|
||||
# Mirror _create_backend precedence (oss > host > platform). Rerank is a Mem0 Platform feature only.
|
||||
mode_label = (
|
||||
"OSS (self-hosted)"
|
||||
if self._mode == "oss"
|
||||
else "self-hosted (HTTP API)"
|
||||
if self._host
|
||||
else "platform (cloud API)"
|
||||
)
|
||||
rerank_note = " Rerank is available on search." if (self._mode == "platform" and not self._host) else ""
|
||||
return f"# Mem0 Memory\nActive. Mode: {mode_label}. User: {self._user_id}.\n{_PROMPT_BODY}{rerank_note}"
|
||||
|
||||
def on_session_switch(self, new_session_id: str, **kwargs) -> None:
|
||||
self._session_id = new_session_id
|
||||
with self._prefetch_lock:
|
||||
self._prefetch_query = self._prefetch_result = ""
|
||||
self._prefetch_done = False
|
||||
|
||||
def on_turn_start(self, turn_number: int, message: str, **kwargs) -> None:
|
||||
self._start_prefetch(message)
|
||||
|
||||
def _consume_prefetch_result(self, query: str) -> str | None:
|
||||
"""Pop the finished prefetch body for ``query`` (None if absent or still running)."""
|
||||
with self._prefetch_lock:
|
||||
if self._prefetch_query != query or not self._prefetch_done:
|
||||
return None
|
||||
result, self._prefetch_result, self._prefetch_done = self._prefetch_result, "", False
|
||||
return result
|
||||
|
||||
def _start_prefetch(self, query: str) -> None:
|
||||
backend = self._backend
|
||||
if not query or backend is None or self._is_breaker_open():
|
||||
return
|
||||
|
||||
def _run():
|
||||
try:
|
||||
results = self._try(
|
||||
lambda: self._search(query, rerank=self._rerank_default, backend=backend),
|
||||
logger.debug,
|
||||
"Mem0 prefetch failed: %s",
|
||||
)
|
||||
lines = [redact(r.get("memory", "")) for r in (results or []) if r.get("memory")]
|
||||
body = "## Mem0 Memory\n" + "\n".join(f"- {line}" for line in lines) if lines else ""
|
||||
with self._prefetch_lock:
|
||||
if self._prefetch_query == query:
|
||||
self._prefetch_result, self._prefetch_done = body, True
|
||||
finally:
|
||||
with self._prefetch_lock:
|
||||
self._prefetch_threads.discard(threading.current_thread())
|
||||
self._close_if_idle()
|
||||
|
||||
with self._prefetch_lock:
|
||||
if self._closed:
|
||||
return
|
||||
# Same query already answered or still in flight: don't restart it.
|
||||
if self._prefetch_query == query and (
|
||||
self._prefetch_done or (self._prefetch_thread and self._prefetch_thread.is_alive())
|
||||
):
|
||||
return
|
||||
self._prefetch_query, self._prefetch_result, self._prefetch_done = query, "", False
|
||||
self._prefetch_thread = t = spawn_context_thread(_run, name="mem0-prefetch")
|
||||
self._prefetch_threads.add(t)
|
||||
t.start()
|
||||
|
||||
def prefetch(self, query: str, *, session_id: str = "") -> str:
|
||||
"""Recall memories for the CURRENT question with a short hot-path wait."""
|
||||
if (cached := self._consume_prefetch_result(query)) is not None:
|
||||
return cached
|
||||
self._start_prefetch(query)
|
||||
with self._prefetch_lock:
|
||||
thread = self._prefetch_thread if self._prefetch_query == query else None
|
||||
if thread:
|
||||
thread.join(timeout=_PREFETCH_WAIT_SECS)
|
||||
return (
|
||||
self._consume_prefetch_result(query) or ""
|
||||
) # slow backend: skip injection; mem0_search remains the backstop
|
||||
|
||||
def _sync_batches(self, messages: list) -> list:
|
||||
if self._sync_max_chars and any(len(m["content"]) > self._sync_max_chars for m in messages):
|
||||
chunks = []
|
||||
for message in messages:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
part = _truncate_for_sync(remaining, self._sync_max_chars)
|
||||
chunks.extend(extraction_message_batches([{**message, "content": part}]))
|
||||
remaining = remaining[len(part) :]
|
||||
return chunks
|
||||
return extraction_message_batches(messages)
|
||||
|
||||
def _drain_sync(self) -> None:
|
||||
while True:
|
||||
with self._sync_lock:
|
||||
if not self._sync_queue:
|
||||
self._sync_thread = None
|
||||
break
|
||||
context, messages, session_id = self._sync_queue.popleft()
|
||||
context.run(self._sync_messages, messages, session_id)
|
||||
self._close_if_idle()
|
||||
|
||||
def _sync_messages(self, messages: list, session_id: str) -> None:
|
||||
for batch in self._sync_batches(messages):
|
||||
if self._is_breaker_open():
|
||||
logger.warning(
|
||||
"Mem0 turn not synced for session %s: circuit breaker open; capture is best effort", session_id
|
||||
)
|
||||
return
|
||||
self._try(
|
||||
lambda: self._add(batch, infer=True, session_id=session_id),
|
||||
logger.warning,
|
||||
"Mem0 sync failed: %s",
|
||||
)
|
||||
|
||||
def sync_turn(self, user_content: str, assistant_content: str, *, session_id: str = "") -> None:
|
||||
"""Queue full, redacted turns without blocking the agent or dropping busy turns."""
|
||||
if self._backend is None or self._is_breaker_open():
|
||||
return
|
||||
messages = [
|
||||
{"role": role, "content": redact(content)}
|
||||
for role, content in (("user", user_content), ("assistant", assistant_content))
|
||||
if content
|
||||
]
|
||||
if not messages:
|
||||
return
|
||||
with self._sync_lock:
|
||||
if self._closed:
|
||||
return
|
||||
# ponytail: in-memory queue; use a durable spool if crash recovery is required.
|
||||
self._sync_queue.append((copy_context(), messages, session_id or self._session_id))
|
||||
if self._sync_thread is None:
|
||||
self._sync_thread = spawn_context_thread(self._drain_sync, name="mem0-sync")
|
||||
self._sync_thread.start()
|
||||
|
||||
def get_tool_schemas(self) -> List[Dict[str, Any]]:
|
||||
return list(TOOL_SCHEMAS)
|
||||
|
||||
# -- tool handlers: (required params, error label, body, client-error policy) ---
|
||||
# Client errors (bad ID / not found) never trip the breaker, except for mem0_add
|
||||
# where they count as failures; update/delete answer them with "Memory not found".
|
||||
|
||||
def _tool_search(self, args: dict) -> str:
|
||||
top_k = max(1, min(int(args.get("top_k", 10)), 50))
|
||||
rerank_raw = args.get("rerank", self._rerank_default)
|
||||
rerank = rerank_raw.lower() not in ("false", "0", "no") if isinstance(rerank_raw, str) else bool(rerank_raw)
|
||||
results = self._search(args["query"], top_k, rerank)
|
||||
if not results:
|
||||
return json.dumps({"result": "No relevant memories found."})
|
||||
items = [
|
||||
{"id": r.get("id"), "memory": redact(r.get("memory", "")), "score": r.get("score", 0)} for r in results
|
||||
]
|
||||
return json.dumps({"results": items, "count": len(items)})
|
||||
|
||||
def _tool_add(self, args: dict) -> str:
|
||||
result = self._add([{"role": "user", "content": redact(args["content"])}], infer=False)
|
||||
event_id = result.get("event_id") if isinstance(result, dict) else None
|
||||
# Cloud add is async (server-side extraction); OSS and self-hosted store synchronously.
|
||||
msg = "Fact stored." if (self._mode == "oss" or self._host) else "Fact queued for storage."
|
||||
return json.dumps({"result": msg, "event_id": event_id})
|
||||
|
||||
_TOOL_HANDLERS = {
|
||||
"mem0_search": (("query",), "Search failed", _tool_search, "skip"),
|
||||
"mem0_add": (("content",), "Failed to store", _tool_add, "count"),
|
||||
"mem0_update": (
|
||||
("memory_id", "text"),
|
||||
"Update failed",
|
||||
lambda self, a: json.dumps(self._backend.update(a["memory_id"], redact(a["text"]))),
|
||||
"not_found",
|
||||
),
|
||||
"mem0_delete": (
|
||||
("memory_id",),
|
||||
"Delete failed",
|
||||
lambda self, a: json.dumps(self._backend.delete(a["memory_id"])),
|
||||
"not_found",
|
||||
),
|
||||
}
|
||||
|
||||
def handle_tool_call(self, tool_name: str, args: dict, **kwargs) -> str:
|
||||
if self._backend is None:
|
||||
err = getattr(self, "_init_error", "unknown error")
|
||||
return json.dumps(
|
||||
{
|
||||
"error": f"Mem0 backend not initialized: {err}.{self._oss_hint(' Check that {vs} is running and reachable.')}"
|
||||
}
|
||||
)
|
||||
if self._is_breaker_open():
|
||||
return json.dumps(
|
||||
{
|
||||
"error": f"Mem0 temporarily unavailable (multiple consecutive failures). Will retry automatically.{self._oss_hint(' Check that your {vs} is running.')}"
|
||||
}
|
||||
)
|
||||
if tool_name not in self._TOOL_HANDLERS:
|
||||
return tool_error(f"Unknown tool: {tool_name}")
|
||||
required, label, body, on_client_error = self._TOOL_HANDLERS[tool_name]
|
||||
if not isinstance(args, dict):
|
||||
return tool_error("Tool arguments must be an object")
|
||||
if missing := next((k for k in required if not isinstance(args.get(k), str) or not args[k].strip()), None):
|
||||
return tool_error(f"Missing or invalid required parameter: {missing}")
|
||||
if tool_name == "mem0_search":
|
||||
try:
|
||||
int(args.get("top_k", 10))
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return tool_error("top_k must be an integer")
|
||||
try:
|
||||
result = body(self, args)
|
||||
except Exception as e:
|
||||
client = _is_client_error(e)
|
||||
if client and on_client_error == "not_found":
|
||||
return tool_error(f"Memory not found: {args['memory_id']}")
|
||||
if not client or on_client_error == "count":
|
||||
self._record_failure()
|
||||
return tool_error(self._format_error(label, e))
|
||||
self._record_success()
|
||||
return result
|
||||
|
||||
def _shutdown_backend(self):
|
||||
with suppress(Exception):
|
||||
if self._backend:
|
||||
self._backend.close()
|
||||
self._backend = None
|
||||
|
||||
def _close_if_idle(self) -> None:
|
||||
with self._sync_lock, self._prefetch_lock:
|
||||
if self._closed and self._sync_thread is None and not self._prefetch_threads:
|
||||
self._shutdown_backend()
|
||||
|
||||
def shutdown(self) -> None:
|
||||
with self._sync_lock, self._prefetch_lock:
|
||||
self._closed = True
|
||||
threads = [*self._prefetch_threads, self._sync_thread]
|
||||
deadline = time.monotonic() + _SHUTDOWN_WAIT_SECS
|
||||
for thread in threads:
|
||||
if thread and thread.is_alive():
|
||||
thread.join(timeout=max(0, deadline - time.monotonic()))
|
||||
if any(thread and thread.is_alive() for thread in threads):
|
||||
logger.warning("Mem0 shutdown timed out; pending capture may be lost if the process exits")
|
||||
# Active workers close their backend when finished, never underneath a request.
|
||||
self._close_if_idle()
|
||||
|
||||
|
||||
def register(ctx) -> None:
|
||||
"""Register Mem0 as a memory provider plugin."""
|
||||
ctx.register_memory_provider(Mem0MemoryProvider())
|
||||
|
||||
|
||||
# Compatibility names retained for callers of the original bundled provider.
|
||||
|
||||
ADD_SCHEMA = {
|
||||
"name": "mem0_add",
|
||||
"description": (
|
||||
"Store a durable fact about the user, verbatim (no LLM extraction). "
|
||||
"Call this the moment the user states a lasting preference, correction, "
|
||||
"decision, or personal detail worth recalling on future turns — don't "
|
||||
"wait to be asked to remember. Skip transient chit-chat and facts you've "
|
||||
"already stored."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"content": {"type": "string", "description": "The fact to store."},
|
||||
},
|
||||
"required": ["content"],
|
||||
},
|
||||
}
|
||||
|
||||
DELETE_SCHEMA = {
|
||||
"name": "mem0_delete",
|
||||
"description": (
|
||||
"Delete a memory by its ID (take the ID from a mem0_search "
|
||||
"result). Use when a stored fact is obsolete or the user asks you to "
|
||||
"forget it; prefer mem0_update if the fact merely changed."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"memory_id": {"type": "string", "description": "Memory UUID to delete."},
|
||||
},
|
||||
"required": ["memory_id"],
|
||||
},
|
||||
}
|
||||
|
||||
SEARCH_SCHEMA = {
|
||||
"name": "mem0_search",
|
||||
"description": (
|
||||
"Search the user's memories by meaning; returns facts ranked by "
|
||||
"relevance. Use this before answering any question that may depend on "
|
||||
"what you know about the user (preferences, facts, history, people, "
|
||||
"projects, past decisions). For multi-part or multi-hop questions, "
|
||||
"call it several times — vary the wording and run follow-up searches "
|
||||
"on what earlier results reveal; one search is rarely enough."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "What to search for."},
|
||||
"top_k": {"type": "integer", "description": "Max results (default: 10, max: 50)."},
|
||||
"rerank": {
|
||||
"type": "boolean",
|
||||
"description": "Rerank results for relevance (default: false, platform mode only).",
|
||||
},
|
||||
},
|
||||
"required": ["query"],
|
||||
},
|
||||
}
|
||||
|
||||
UPDATE_SCHEMA = {
|
||||
"name": "mem0_update",
|
||||
"description": (
|
||||
"Replace the text of an existing memory by its ID (take the ID from a "
|
||||
"mem0_search result). Use when a stored fact has changed "
|
||||
"or was wrong — correct it in place instead of adding a duplicate."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"memory_id": {"type": "string", "description": "Memory UUID to update."},
|
||||
"text": {"type": "string", "description": "New text content."},
|
||||
},
|
||||
"required": ["memory_id", "text"],
|
||||
},
|
||||
}
|
||||
# ---- END PLUGIN-COMPAT ----
|
||||
@@ -0,0 +1,298 @@
|
||||
"""Backend abstraction for Mem0 Platform and OSS modes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from contextlib import closing, suppress
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _add_kwargs(user_id: str, agent_id: str, infer: bool, metadata: dict | None, run_id: str = "") -> dict[str, Any]:
|
||||
return {
|
||||
"user_id": user_id,
|
||||
"agent_id": agent_id,
|
||||
"infer": infer,
|
||||
**({"metadata": metadata} if metadata else {}),
|
||||
**({"run_id": run_id} if run_id else {}),
|
||||
}
|
||||
|
||||
|
||||
def _unwrap_results(response: Any) -> list:
|
||||
"""Normalize API response — extract results list from dict or pass through."""
|
||||
return response.get("results", []) if isinstance(response, dict) else response if isinstance(response, list) else []
|
||||
|
||||
|
||||
class Mem0Backend(ABC):
|
||||
"""Unified interface over Platform (MemoryClient), self-hosted (HTTP) and OSS (Memory) backends.
|
||||
update()/delete() are template methods: subclasses implement raw ``_update``/``_delete``."""
|
||||
|
||||
@abstractmethod
|
||||
def search(self, query: str, *, filters: dict, top_k: int = 10, rerank: bool = False) -> list[dict]: ...
|
||||
@abstractmethod
|
||||
def add(
|
||||
self,
|
||||
messages: list,
|
||||
*,
|
||||
user_id: str,
|
||||
agent_id: str,
|
||||
infer: bool = False,
|
||||
metadata: dict | None = None,
|
||||
run_id: str = "",
|
||||
) -> dict: ...
|
||||
@abstractmethod
|
||||
def _update(self, memory_id: str, text: str) -> None: ...
|
||||
@abstractmethod
|
||||
def _delete(self, memory_id: str) -> None: ...
|
||||
|
||||
def update(self, memory_id: str, text: str) -> dict:
|
||||
self._update(memory_id, text)
|
||||
return {"result": "Memory updated.", "memory_id": memory_id}
|
||||
|
||||
def delete(self, memory_id: str) -> dict:
|
||||
self._delete(memory_id)
|
||||
return {"result": "Memory deleted.", "memory_id": memory_id}
|
||||
|
||||
def close(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class PlatformBackend(Mem0Backend):
|
||||
"""Wraps mem0.MemoryClient for Mem0 Platform (cloud API)."""
|
||||
|
||||
def __init__(self, api_key: str):
|
||||
from mem0 import MemoryClient
|
||||
|
||||
self._client = MemoryClient(api_key=api_key)
|
||||
|
||||
def search(self, query: str, *, filters: dict, top_k: int = 10, rerank: bool = False) -> list[dict]:
|
||||
return _unwrap_results(self._client.search(query, filters=filters, top_k=top_k, rerank=rerank))
|
||||
|
||||
def add(
|
||||
self,
|
||||
messages: list,
|
||||
*,
|
||||
user_id: str,
|
||||
agent_id: str,
|
||||
infer: bool = False,
|
||||
metadata: dict | None = None,
|
||||
run_id: str = "",
|
||||
) -> dict:
|
||||
return self._client.add(messages, **_add_kwargs(user_id, agent_id, infer, metadata, run_id))
|
||||
|
||||
def _update(self, memory_id: str, text: str) -> None:
|
||||
self._client.update(memory_id=memory_id, text=text)
|
||||
|
||||
def _delete(self, memory_id: str) -> None:
|
||||
self._client.delete(memory_id=memory_id)
|
||||
|
||||
|
||||
class SelfHostedBackend(Mem0Backend):
|
||||
"""Direct HTTP backend for a self-hosted Mem0 server (the FastAPI ``server/``).
|
||||
mem0.MemoryClient is hardwired to the cloud API (``Authorization: Token``, ``GET /v1/ping/`` in ``__init__``),
|
||||
so this speaks the server's real contract: ``X-API-Key`` auth and the ``/memories`` / ``/search`` routes."""
|
||||
|
||||
def __init__(self, api_key: str, host: str, transport=None):
|
||||
import httpx
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
**({"X-API-Key": api_key} if api_key else {}),
|
||||
} # key omitted only for AUTH_DISABLED servers
|
||||
# Connect-level retries keep one dropped SYN from counting toward the breaker. ``transport`` is injectable for tests.
|
||||
self._client = httpx.Client(
|
||||
base_url=host.rstrip("/"),
|
||||
headers=headers,
|
||||
timeout=30.0,
|
||||
transport=transport or httpx.HTTPTransport(retries=2),
|
||||
)
|
||||
|
||||
def _json(self, method: str, path: str, **kwargs) -> Any:
|
||||
resp = self._client.request(method, path, **kwargs)
|
||||
resp.raise_for_status()
|
||||
return resp.json() if resp.content else {}
|
||||
|
||||
def search(self, query: str, *, filters: dict, top_k: int = 10, rerank: bool = False) -> list[dict]:
|
||||
# rerank is platform-only; the self-hosted /search ignores it. user_id belongs in filters (top-level is deprecated).
|
||||
return _unwrap_results(
|
||||
self._json(
|
||||
"POST", "/search", json={"query": query, "top_k": top_k, **({"filters": filters} if filters else {})}
|
||||
)
|
||||
)
|
||||
|
||||
def add(
|
||||
self,
|
||||
messages: list,
|
||||
*,
|
||||
user_id: str,
|
||||
agent_id: str,
|
||||
infer: bool = False,
|
||||
metadata: dict | None = None,
|
||||
run_id: str = "",
|
||||
) -> dict:
|
||||
return self._json(
|
||||
"POST", "/memories", json={"messages": messages, **_add_kwargs(user_id, agent_id, infer, metadata, run_id)}
|
||||
)
|
||||
|
||||
def _update(self, memory_id: str, text: str) -> None:
|
||||
self._json("PUT", f"/memories/{memory_id}", json={"text": text})
|
||||
|
||||
def _delete(self, memory_id: str) -> None:
|
||||
self._json("DELETE", f"/memories/{memory_id}")
|
||||
|
||||
def close(self) -> None:
|
||||
with suppress(Exception):
|
||||
self._client.close()
|
||||
|
||||
|
||||
_DIRECT_OPENAI_PROVIDER = "hermes_openai"
|
||||
_DIRECT_OPENAI_CLASS_PATH = f"{__package__}._openai_llm.DirectOpenAILLM"
|
||||
|
||||
|
||||
def _register_direct_openai_provider() -> None:
|
||||
"""Register Hermes' OpenAI-only Mem0 LLM provider once per factory."""
|
||||
from mem0.configs.llms.openai import OpenAIConfig
|
||||
from mem0.utils.factory import LlmFactory
|
||||
|
||||
provider_map = getattr(LlmFactory, "provider_to_class", None)
|
||||
register_provider = getattr(LlmFactory, "register_provider", None)
|
||||
if not isinstance(provider_map, dict) or not callable(register_provider):
|
||||
raise RuntimeError(
|
||||
"mem0 LlmFactory does not support the provider registration required for the Hermes OpenAI OSS backend"
|
||||
)
|
||||
if provider_map.get(_DIRECT_OPENAI_PROVIDER) != (_DIRECT_OPENAI_CLASS_PATH, OpenAIConfig):
|
||||
register_provider(_DIRECT_OPENAI_PROVIDER, _DIRECT_OPENAI_CLASS_PATH, OpenAIConfig)
|
||||
|
||||
|
||||
class OSSBackend(Mem0Backend):
|
||||
"""Wraps mem0.Memory for self-hosted (OSS) mode."""
|
||||
|
||||
def __init__(self, oss_config: dict):
|
||||
import os
|
||||
|
||||
from mem0 import Memory
|
||||
|
||||
from ._oss_providers import EMBEDDER_PROVIDERS, KNOWN_DIMS, LLM_PROVIDERS
|
||||
|
||||
def _provider_block(name: str, registry: dict) -> dict:
|
||||
"""Copy of oss_config[name] with the legacy ``api_base`` key mapped to the provider's canonical base-URL key."""
|
||||
block = dict(oss_config[name])
|
||||
provider_config = dict(block.get("config", {}))
|
||||
legacy_base = provider_config.pop("api_base", None)
|
||||
canonical_key = registry.get(str(block.get("provider") or "").strip().lower(), {}).get("base_url_key")
|
||||
if legacy_base and canonical_key:
|
||||
provider_config.setdefault(canonical_key, legacy_base)
|
||||
block["config"] = provider_config
|
||||
return block
|
||||
|
||||
vector_store = dict(oss_config["vector_store"])
|
||||
vs_config = dict(vector_store.get("config", {}))
|
||||
if "path" in vs_config:
|
||||
vs_config["path"] = os.path.expanduser(vs_config["path"])
|
||||
embedder_config = oss_config.get("embedder", {}).get("config", {})
|
||||
dims = embedder_config.get("embedding_dims") or KNOWN_DIMS.get(embedder_config.get("model", ""))
|
||||
if dims:
|
||||
vs_config["embedding_model_dims"] = dims
|
||||
self._recreate_collection_if_dims_changed(vector_store.get("provider", "qdrant"), vs_config, dims)
|
||||
vector_store["config"] = vs_config
|
||||
config = {
|
||||
"vector_store": vector_store,
|
||||
"llm": _provider_block("llm", LLM_PROVIDERS),
|
||||
"embedder": _provider_block("embedder", EMBEDDER_PROVIDERS),
|
||||
"version": "v1.1",
|
||||
}
|
||||
if str(config["llm"].get("provider") or "").strip().lower() == "openai":
|
||||
# mem0 validates LlmConfig.provider before its factory lookup: build the supported OpenAI config, then swap the provider.
|
||||
_register_direct_openai_provider()
|
||||
from mem0.configs.base import MemoryConfig
|
||||
|
||||
memory_config = MemoryConfig(**config)
|
||||
try:
|
||||
memory_config.llm.provider = _DIRECT_OPENAI_PROVIDER
|
||||
except (AttributeError, TypeError) as exc:
|
||||
raise RuntimeError(
|
||||
"mem0 MemoryConfig does not expose a mutable llm.provider for the Hermes OpenAI OSS backend"
|
||||
) from exc
|
||||
self._memory = Memory(memory_config)
|
||||
else:
|
||||
self._memory = Memory.from_config(config)
|
||||
|
||||
@staticmethod
|
||||
def _recreate_collection_if_dims_changed(provider: str, vs_config: dict, expected_dims: int) -> None:
|
||||
"""Reject dimension changes without deleting existing memories (legacy method name)."""
|
||||
collection_name = vs_config.get("collection_name", "mem0")
|
||||
current_dims = None
|
||||
with suppress(Exception):
|
||||
if provider == "qdrant":
|
||||
from qdrant_client import QdrantClient
|
||||
|
||||
path, url = vs_config.get("path"), vs_config.get("url")
|
||||
if path:
|
||||
client = QdrantClient(path=path)
|
||||
elif url:
|
||||
client = QdrantClient(url=url, api_key=vs_config.get("api_key"))
|
||||
else:
|
||||
return
|
||||
with closing(client):
|
||||
if not client.collection_exists(collection_name):
|
||||
return
|
||||
vectors = client.get_collection(collection_name).config.params.vectors
|
||||
# Named-vector collections expose a dict; unnamed expose an object with .size.
|
||||
if isinstance(vectors, dict):
|
||||
vectors = next(iter(vectors.values()), None)
|
||||
current_dims = getattr(vectors, "size", None)
|
||||
elif provider == "pgvector":
|
||||
import psycopg2
|
||||
|
||||
conn_params = {
|
||||
k: vs_config[k]
|
||||
for k in ("host", "port", "user", "password", "dbname", "sslmode")
|
||||
if vs_config.get(k)
|
||||
}
|
||||
with closing(psycopg2.connect(**conn_params)) as conn:
|
||||
conn.autocommit = True
|
||||
with closing(conn.cursor()) as cur:
|
||||
cur.execute(
|
||||
"SELECT atttypmod FROM pg_attribute WHERE attrelid = %s::regclass AND attname = 'vector'",
|
||||
(collection_name,),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
current_dims = row[0] if row and row[0] > 0 else None
|
||||
if current_dims is not None and current_dims != expected_dims:
|
||||
raise ValueError(
|
||||
f"Collection {collection_name!r} has {current_dims} embedding dimensions, but {expected_dims} are configured. "
|
||||
"Existing memories were preserved. Restore the previous embedder or choose a new collection_name "
|
||||
"and migrate your memories explicitly."
|
||||
)
|
||||
|
||||
def search(self, query: str, *, filters: dict, top_k: int = 10, rerank: bool = False) -> list[dict]:
|
||||
return _unwrap_results(self._memory.search(query, filters=filters, top_k=top_k))
|
||||
|
||||
def add(
|
||||
self,
|
||||
messages: list,
|
||||
*,
|
||||
user_id: str,
|
||||
agent_id: str,
|
||||
infer: bool = False,
|
||||
metadata: dict | None = None,
|
||||
run_id: str = "",
|
||||
) -> dict:
|
||||
return self._memory.add(messages, **_add_kwargs(user_id, agent_id, infer, metadata, run_id))
|
||||
|
||||
def _update(self, memory_id: str, text: str) -> None:
|
||||
self._memory.update(memory_id, data=text)
|
||||
|
||||
def _delete(self, memory_id: str) -> None:
|
||||
self._memory.delete(memory_id)
|
||||
|
||||
def close(self):
|
||||
with suppress(Exception):
|
||||
telemetry = getattr(self._memory, "telemetry", None)
|
||||
if telemetry and hasattr(telemetry, "posthog"):
|
||||
with suppress(Exception):
|
||||
telemetry.posthog.shutdown()
|
||||
vs = getattr(self._memory, "vector_store", None)
|
||||
# Memory, then its vector store, then the store's raw client; the first failure aborts the chain.
|
||||
for obj in filter(None, (self._memory, vs, getattr(vs, "client", None))):
|
||||
if hasattr(obj, "close"):
|
||||
obj.close()
|
||||
@@ -0,0 +1,88 @@
|
||||
"""OpenAI-only LLM adapter for Mem0 OSS mode."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
from mem0.configs.llms.base import BaseLlmConfig
|
||||
from mem0.configs.llms.openai import OpenAIConfig
|
||||
from mem0.llms.base import LLMBase
|
||||
from mem0.llms.openai import OpenAILLM
|
||||
|
||||
# BaseLlmConfig fields copied into OpenAIConfig; the last two may be absent on older mem0.
|
||||
_COPIED_FIELDS = (
|
||||
"model",
|
||||
"temperature",
|
||||
"api_key",
|
||||
"max_tokens",
|
||||
"top_p",
|
||||
"top_k",
|
||||
"enable_vision",
|
||||
"vision_details",
|
||||
"http_client_proxies",
|
||||
)
|
||||
_OPTIONAL_FIELDS = ("reasoning_effort", "is_reasoning_model")
|
||||
|
||||
|
||||
class DirectOpenAILLM(OpenAILLM):
|
||||
"""Use OpenAI credentials and requests regardless of router environment."""
|
||||
|
||||
def __init__(self, config: Optional[Union[BaseLlmConfig, OpenAIConfig, Dict]] = None):
|
||||
if config is None:
|
||||
config = OpenAIConfig()
|
||||
elif isinstance(config, dict):
|
||||
config = OpenAIConfig(**config)
|
||||
elif isinstance(config, BaseLlmConfig) and not isinstance(config, OpenAIConfig):
|
||||
fields = {k: getattr(config, k) for k in _COPIED_FIELDS}
|
||||
fields.update({k: getattr(config, k, None) for k in _OPTIONAL_FIELDS})
|
||||
config = OpenAIConfig(**fields)
|
||||
if not config.model:
|
||||
config.model = "gpt-5-mini"
|
||||
# Configs predating the setup marker: keep the default model reasoning-safe
|
||||
# without overriding an explicit user choice.
|
||||
if config.model == "gpt-5-mini" and config.is_reasoning_model is None:
|
||||
config.is_reasoning_model = True
|
||||
# Bypass OpenAILLM.__init__ (it picks OpenRouter when OPENROUTER_API_KEY is
|
||||
# set); LLMBase still owns validation and supported-parameter filtering.
|
||||
LLMBase.__init__(self, config)
|
||||
# OPENAI_API_KEY / OPENAI_BASE_URL are profile credentials: read them through the secret
|
||||
# scope, never raw os.environ, or a multiplexed secondary's memory extraction runs on the
|
||||
# default profile's OpenAI account (and its proxy).
|
||||
from agent.secret_scope import get_secret
|
||||
|
||||
api_key = self.config.api_key or get_secret("OPENAI_API_KEY", "")
|
||||
if not api_key:
|
||||
raise ValueError("OpenAI API key is required for the Hermes Mem0 OSS provider")
|
||||
from openai import OpenAI
|
||||
|
||||
self.client = OpenAI(
|
||||
api_key=api_key,
|
||||
base_url=self.config.openai_base_url or get_secret("OPENAI_BASE_URL", "") or "https://api.openai.com/v1",
|
||||
)
|
||||
|
||||
def generate_response(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
response_format=None,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
tool_choice: str = "auto",
|
||||
**kwargs,
|
||||
):
|
||||
params = self._get_supported_params(messages=messages, **kwargs)
|
||||
params.update({"model": self.config.model, "messages": messages})
|
||||
# No OpenRouter-only fields; ``store`` is opt-in so OpenAI-compatible endpoints never receive unknown fields.
|
||||
if self.config.store is not None:
|
||||
params["store"] = self.config.store
|
||||
if response_format:
|
||||
params["response_format"] = response_format
|
||||
if tools:
|
||||
params["tools"], params["tool_choice"] = tools, tool_choice
|
||||
response = self.client.chat.completions.create(**params)
|
||||
parsed_response = self._parse_response(response, tools)
|
||||
if self.config.response_callback:
|
||||
try:
|
||||
self.config.response_callback(self, response, params)
|
||||
except Exception:
|
||||
logging.error("Error running Mem0 OpenAI response callback")
|
||||
return parsed_response
|
||||
@@ -0,0 +1,98 @@
|
||||
"""OSS provider definitions for LLM, embedder, and vector store."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from hermes_constants import get_hermes_home
|
||||
|
||||
LLM_PROVIDERS: dict[str, dict[str, Any]] = {
|
||||
"openai": {
|
||||
"label": "OpenAI",
|
||||
"needs_key": True,
|
||||
"env_var": "OPENAI_API_KEY",
|
||||
"default_model": "gpt-5-mini",
|
||||
"base_url_key": "openai_base_url",
|
||||
},
|
||||
"ollama": {
|
||||
"label": "Ollama (local)",
|
||||
"needs_key": False,
|
||||
"default_model": "llama3.1:8b",
|
||||
"default_url": "http://localhost:11434",
|
||||
"base_url_key": "ollama_base_url",
|
||||
"pip_dep": "ollama",
|
||||
},
|
||||
}
|
||||
|
||||
EMBEDDER_PROVIDERS: dict[str, dict[str, Any]] = {
|
||||
"openai": {
|
||||
"label": "OpenAI",
|
||||
"needs_key": True,
|
||||
"env_var": "OPENAI_API_KEY",
|
||||
"default_model": "text-embedding-3-small",
|
||||
"base_url_key": "openai_base_url",
|
||||
"dims": 1536,
|
||||
},
|
||||
"ollama": {
|
||||
"label": "Ollama (local)",
|
||||
"needs_key": False,
|
||||
"default_model": "nomic-embed-text",
|
||||
"default_url": "http://localhost:11434",
|
||||
"base_url_key": "ollama_base_url",
|
||||
"dims": 768,
|
||||
"pip_dep": "ollama",
|
||||
},
|
||||
}
|
||||
|
||||
VECTOR_PROVIDERS: dict[str, dict[str, Any]] = {
|
||||
# Resolved lazily (see ``vector_default_config``): the profile home is a ContextVar at call time,
|
||||
# not an import-time constant, and ``~/.hermes`` is wrong on Windows and under profiles.
|
||||
"qdrant": {
|
||||
"label": "Qdrant",
|
||||
"default_config": {"path": lambda: str(get_hermes_home() / "mem0_qdrant")},
|
||||
"pip_dep": "qdrant-client",
|
||||
},
|
||||
"pgvector": {
|
||||
"label": "PGVector",
|
||||
"default_config": {
|
||||
"host": "localhost",
|
||||
"port": 5432,
|
||||
"user": os.getenv("USER", "postgres"),
|
||||
"dbname": "postgres",
|
||||
},
|
||||
"pip_dep": "psycopg2-binary",
|
||||
},
|
||||
}
|
||||
|
||||
KNOWN_DIMS: dict[str, int] = {
|
||||
"text-embedding-3-small": 1536,
|
||||
"text-embedding-3-large": 3072,
|
||||
"text-embedding-ada-002": 1536,
|
||||
"nomic-embed-text": 768,
|
||||
}
|
||||
|
||||
|
||||
def vector_default_config(provider_id: str) -> dict[str, Any]:
|
||||
"""A vector store's ``default_config`` with callable defaults resolved for the active profile."""
|
||||
return {k: (v() if callable(v) else v) for k, v in VECTOR_PROVIDERS[provider_id]["default_config"].items()}
|
||||
|
||||
|
||||
SECTION_REGISTRIES = (("llm", LLM_PROVIDERS), ("embedder", EMBEDDER_PROVIDERS), ("vector_store", VECTOR_PROVIDERS))
|
||||
|
||||
|
||||
def validate_oss_config(oss_config: dict) -> list[str]:
|
||||
"""Validate an OSS config dict. Returns list of error strings (empty = valid)."""
|
||||
errors: list[str] = []
|
||||
for section, registry in SECTION_REGISTRIES:
|
||||
block = oss_config.get(section)
|
||||
if not block or not isinstance(block, dict):
|
||||
errors.append(f"Missing required section: {section}")
|
||||
elif block.get("provider", "") not in registry:
|
||||
errors.append(
|
||||
f"Unknown {section} provider '{block.get('provider', '')}'. Valid: {', '.join(registry.keys())}"
|
||||
)
|
||||
vs = oss_config.get("vector_store", {})
|
||||
if vs.get("provider") == "pgvector" and not vs.get("config", {}).get("user"):
|
||||
errors.append("PGVector requires 'user' in vector_store.config")
|
||||
return errors
|
||||
@@ -0,0 +1,741 @@
|
||||
"""Setup wizard for Mem0 plugin — interactive and flag-based modes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import getpass
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from hermes_constants import get_hermes_home # noqa: F401 — patched by tests
|
||||
|
||||
from ._oss_providers import (
|
||||
EMBEDDER_PROVIDERS,
|
||||
KNOWN_DIMS,
|
||||
LLM_PROVIDERS,
|
||||
SECTION_REGISTRIES,
|
||||
VECTOR_PROVIDERS,
|
||||
validate_oss_config,
|
||||
vector_default_config,
|
||||
)
|
||||
|
||||
_OLLAMA_URL = "http://localhost:11434"
|
||||
_PGVECTOR_CONTAINER, _PGVECTOR_IMAGE, _PGVECTOR_PASSWORD = "hermes-pgvector", "pgvector/pgvector:pg17", "hermes"
|
||||
|
||||
|
||||
def _curses_select(title: str, items: list[tuple[str, str]], default: int = 0) -> int:
|
||||
from hermes_cli.curses_ui import curses_radiolist
|
||||
|
||||
return curses_radiolist(
|
||||
title,
|
||||
[f"{label} {desc}" if desc else label for label, desc in items],
|
||||
selected=default,
|
||||
cancel_returns=default,
|
||||
)
|
||||
|
||||
|
||||
def _prompt(label: str, default: str | None = None, secret: bool = False) -> str:
|
||||
"""Prompt for a value with optional default and secret masking."""
|
||||
sys.stdout.write(f" {label}{f' [{default}]' if default else ''}: ")
|
||||
sys.stdout.flush()
|
||||
val = getpass.getpass(prompt="") if secret and sys.stdin.isatty() else sys.stdin.readline().strip()
|
||||
return val or (default or "")
|
||||
|
||||
|
||||
def _input(label: str, default: str) -> str:
|
||||
return input(f" {label} [{default}]: ").strip() or default
|
||||
|
||||
|
||||
def _masked(secret: str) -> str:
|
||||
return f"...{secret[-4:]}" if len(secret) > 4 else "set"
|
||||
|
||||
|
||||
def _http_get(url: str, path: str, timeout: int):
|
||||
return urllib.request.urlopen(urllib.request.Request(f"{url.rstrip('/')}{path}", method="GET"), timeout=timeout)
|
||||
|
||||
|
||||
def _prompt_api_key(label: str, env_var: str, hermes_home: str) -> str:
|
||||
"""Prompt for API key, showing masked existing value if found."""
|
||||
existing = os.environ.get(env_var, "")
|
||||
if not existing:
|
||||
from agent.secret_scope import load_env_file
|
||||
|
||||
existing = load_env_file(Path(hermes_home) / ".env").get(env_var, "")
|
||||
hint = f" (current: {_masked(existing)}, blank to keep)" if existing else ""
|
||||
return getpass.getpass(f" {label} API key{hint}: ").strip()
|
||||
|
||||
|
||||
def _api_key_writes(
|
||||
flags: dict, label: str, *, url: str | None = None, fresh_label: str | None = None
|
||||
) -> dict[str, str]:
|
||||
"""MEM0_API_KEY for .env: from --api-key, else prompt (masking any key already in the environment)."""
|
||||
if flags.get("api_key"):
|
||||
return {"MEM0_API_KEY": flags["api_key"]}
|
||||
existing = os.environ.get("MEM0_API_KEY", "")
|
||||
if url and not existing:
|
||||
print(f" Get yours at {url}")
|
||||
val = _prompt(
|
||||
f"{label} (current: {_masked(existing)}, blank to keep)" if existing else fresh_label or label, secret=True
|
||||
)
|
||||
return {"MEM0_API_KEY": val} if val else {}
|
||||
|
||||
|
||||
def _print_dry_run(summary: str, env_writes: dict, check=None) -> None:
|
||||
print(f"\n [dry-run] Would save config: {summary}")
|
||||
if env_writes:
|
||||
print(" [dry-run] Would write API key to .env")
|
||||
if check:
|
||||
check()
|
||||
print(" [dry-run] No files written.\n")
|
||||
|
||||
|
||||
# --oss-vector-<key> flags accepted per vector store (also the pgvector key order).
|
||||
_VECTOR_FLAG_KEYS = {"qdrant": ("path", "url"), "pgvector": ("host", "port", "user", "password", "dbname")}
|
||||
_FLAG_KEYS = (
|
||||
"mode",
|
||||
"api_key",
|
||||
"host",
|
||||
*(f"oss_{s}{k}" for s in ("llm", "embedder") for k in ("", "_key", "_model", "_url")),
|
||||
"oss_vector",
|
||||
*(f"oss_vector_{k}" for ks in _VECTOR_FLAG_KEYS.values() for k in ks),
|
||||
"user_id",
|
||||
)
|
||||
_FLAG_DEFAULTS = {"oss_llm": "openai", "oss_embedder": "openai", "oss_vector": "qdrant"}
|
||||
|
||||
|
||||
def parse_flags(argv: list[str] | None = None) -> dict[str, str]:
|
||||
args = argv if argv is not None else sys.argv[1:]
|
||||
flags: dict[str, Any] = {**{k: _FLAG_DEFAULTS.get(k, "") for k in _FLAG_KEYS}, "dry_run": False}
|
||||
flag_map = {"--" + k.replace("_", "-"): k for k in _FLAG_KEYS}
|
||||
i = 0
|
||||
while i < len(args):
|
||||
if args[i] == "--dry-run":
|
||||
flags["dry_run"] = True
|
||||
elif args[i] in flag_map and i + 1 < len(args):
|
||||
flags[flag_map[args[i]]] = args[i + 1]
|
||||
i += 1
|
||||
i += 1
|
||||
return flags
|
||||
|
||||
|
||||
def _model_block(flags: dict, registry: dict, prefix: str) -> tuple[str, dict, dict[str, Any]]:
|
||||
"""Resolve (provider_id, provider_def, config) for an LLM/embedder section from flags."""
|
||||
pid = flags.get(prefix, "openai")
|
||||
pdef = registry[pid]
|
||||
cfg: dict[str, Any] = {"model": flags.get(f"{prefix}_model") or pdef["default_model"]}
|
||||
url = flags.get(f"{prefix}_url") or pdef.get("default_url")
|
||||
if url and pdef.get("base_url_key"):
|
||||
cfg[pdef["base_url_key"]] = url
|
||||
return pid, pdef, cfg
|
||||
|
||||
|
||||
def build_oss_config(flags: dict[str, str]) -> tuple[dict, dict[str, str]]:
|
||||
"""Build (oss_config for mem0.json, env_writes of secrets for .env) from parsed flags."""
|
||||
llm_id, llm_def, llm_config = _model_block(flags, LLM_PROVIDERS, "oss_llm")
|
||||
if llm_id == "openai" and llm_config["model"] == "gpt-5-mini":
|
||||
llm_config["is_reasoning_model"] = True
|
||||
embedder_id, embedder_def, embedder_config = _model_block(flags, EMBEDDER_PROVIDERS, "oss_embedder")
|
||||
dims = KNOWN_DIMS.get(embedder_config["model"])
|
||||
if dims:
|
||||
embedder_config["embedding_dims"] = dims
|
||||
vector_id = flags.get("oss_vector", "qdrant")
|
||||
vector_config = vector_default_config(vector_id)
|
||||
for key in _VECTOR_FLAG_KEYS.get(vector_id, ()):
|
||||
if val := flags.get(f"oss_vector_{key}"):
|
||||
vector_config[key] = int(val) if key == "port" else val
|
||||
if "url" in vector_config:
|
||||
vector_config.pop("path", None) # a remote Qdrant URL replaces local storage
|
||||
oss_config = {
|
||||
"llm": {"provider": llm_id, "config": llm_config},
|
||||
"embedder": {"provider": embedder_id, "config": embedder_config},
|
||||
"vector_store": {"provider": vector_id, "config": vector_config},
|
||||
}
|
||||
# An embedder sharing the LLM's provider reuses the LLM key when no embedder key was given.
|
||||
llm_key = flags.get("oss_llm_key") if llm_def.get("needs_key") else ""
|
||||
emb_key = (
|
||||
(flags.get("oss_embedder_key") or (flags.get("oss_llm_key") if embedder_id == llm_id else ""))
|
||||
if embedder_def.get("needs_key")
|
||||
else ""
|
||||
)
|
||||
env_writes = {d["env_var"]: k for d, k in ((llm_def, llm_key), (embedder_def, emb_key)) if k}
|
||||
return oss_config, env_writes
|
||||
|
||||
|
||||
def _write_env(env_path: Path, env_writes: dict[str, str]) -> None:
|
||||
env_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
# utf-8-sig like the canonical .env readers: a BOM'd first line would miss the key match and get duplicated.
|
||||
existing_lines = env_path.read_text(encoding="utf-8-sig").splitlines() if env_path.exists() else []
|
||||
keys = [
|
||||
line.split("=", 1)[0].strip() if "=" in line and not line.startswith("#") else None for line in existing_lines
|
||||
]
|
||||
new_lines = [f"{k}={env_writes[k]}" if k in env_writes else line for k, line in zip(keys, existing_lines)]
|
||||
new_lines += [f"{k}={v}" for k, v in env_writes.items() if k not in keys]
|
||||
fd, temporary = tempfile.mkstemp(prefix=".mem0-env-", dir=env_path.parent)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
||||
stream.write("\n".join(new_lines) + "\n")
|
||||
os.replace(temporary, env_path)
|
||||
finally:
|
||||
Path(temporary).unlink(missing_ok=True)
|
||||
|
||||
|
||||
def _activate_provider(config: dict) -> None:
|
||||
"""Point config.yaml's memory.provider at mem0."""
|
||||
from hermes_cli.config import save_config
|
||||
|
||||
config["memory"]["provider"] = "mem0"
|
||||
save_config(config)
|
||||
|
||||
|
||||
def _persist_provider_config(
|
||||
hermes_home: str,
|
||||
config: dict,
|
||||
provider_config: dict,
|
||||
env_writes: dict[str, str],
|
||||
label: str,
|
||||
key_line: str,
|
||||
server: str | None = None,
|
||||
) -> None:
|
||||
"""Shared platform/self-hosted tail: activate, write mem0.json (0600), then .env, then a saved summary."""
|
||||
_activate_provider(config)
|
||||
from . import Mem0MemoryProvider
|
||||
|
||||
Mem0MemoryProvider().save_config(provider_config, hermes_home)
|
||||
if env_writes:
|
||||
_write_env(Path(hermes_home) / ".env", env_writes)
|
||||
if server:
|
||||
_check_selfhosted_server(server)
|
||||
print(
|
||||
"\n".join(
|
||||
[
|
||||
"",
|
||||
f" Memory provider: {label}",
|
||||
*([f" Server: {server}"] if server else []),
|
||||
" Activation saved to config.yaml",
|
||||
" Provider config saved",
|
||||
*([f" {key_line}"] if env_writes else []),
|
||||
"",
|
||||
" Start a new session to activate.",
|
||||
"",
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _setup_platform(hermes_home: str, config: dict, flags: dict[str, str]) -> None:
|
||||
"""Platform mode setup — prompts for API key (secret -> .env), user/agent ids and rerank (-> mem0.json)."""
|
||||
from utils import read_json_or_empty
|
||||
|
||||
provider_config = read_json_or_empty(Path(hermes_home) / "mem0.json")
|
||||
print("\n Configuring mem0:\n")
|
||||
env_writes = _api_key_writes(flags, "Mem0 Platform API key", url="https://app.mem0.ai")
|
||||
for key, desc, default in (
|
||||
("user_id", "User identifier", "hermes-user"),
|
||||
("agent_id", "Agent identifier", "hermes"),
|
||||
):
|
||||
if val := _prompt(desc, default=str(provider_config.get(key) or default)):
|
||||
provider_config[key] = val
|
||||
choices = ["true", "false"]
|
||||
current = str(provider_config.get("rerank", "false") or "").lower()
|
||||
provider_config["rerank"] = choices[
|
||||
_curses_select(
|
||||
" Enable reranking for recall",
|
||||
[(c, "") for c in choices],
|
||||
default=choices.index(current) if current in choices else 0,
|
||||
)
|
||||
]
|
||||
if flags.get("dry_run"):
|
||||
_print_dry_run(str(provider_config), env_writes)
|
||||
return
|
||||
# Routing checks ``host`` before platform, so clear a stale self-hosted host. "" rather than
|
||||
# pop(): save_config merges into the existing mem0.json, so a popped key would survive.
|
||||
provider_config.update(mode="platform", host="")
|
||||
# _load_config() also seeds ``host`` from MEM0_HOST (.env); the file clear can't help there, so warn.
|
||||
if os.environ.get("MEM0_HOST", "").strip():
|
||||
print(
|
||||
f"\n ⚠ MEM0_HOST is set in your environment ({os.environ['MEM0_HOST']}). It overrides platform mode — remove it from ~/.hermes/.env (or unset it) or Hermes will keep routing to the self-hosted server."
|
||||
)
|
||||
_persist_provider_config(hermes_home, config, provider_config, env_writes, "mem0", "API keys saved to .env")
|
||||
|
||||
|
||||
def _check_selfhosted_server(host: str) -> None:
|
||||
"""Best-effort reachability check for a self-hosted Mem0 server (non-fatal)."""
|
||||
try:
|
||||
_http_get(host, "/docs", 5)
|
||||
print(f" ✓ Mem0 server reachable at {host}")
|
||||
except urllib.error.HTTPError:
|
||||
# Any HTTP response (401/403/404) still means something is listening.
|
||||
print(f" ✓ Mem0 server responding at {host}")
|
||||
except Exception:
|
||||
print(f" ⚠ Could not reach {host} — check the URL and that the server is running.")
|
||||
|
||||
|
||||
def _setup_selfhosted(hermes_home: str, config: dict, flags: dict[str, str]) -> None:
|
||||
"""Self-hosted mode — point at an existing Mem0 server: URL -> mem0.json, key -> .env (MEM0_API_KEY)."""
|
||||
from utils import read_json_or_empty
|
||||
|
||||
provider_config = read_json_or_empty(Path(hermes_home) / "mem0.json")
|
||||
print("\n Configuring mem0 (self-hosted server):\n")
|
||||
host = flags.get("host") or _prompt(
|
||||
"Mem0 server URL (e.g. http://localhost:8888)", default=provider_config.get("host") or None
|
||||
)
|
||||
if not host:
|
||||
print(" Error: a server URL is required for self-hosted mode.", file=sys.stderr)
|
||||
return
|
||||
host = host.rstrip("/")
|
||||
env_writes = _api_key_writes(flags, "Server API key", fresh_label="Server API key (blank if AUTH_DISABLED)")
|
||||
user_id = flags.get("user_id") or _prompt(
|
||||
"User identifier", default=provider_config.get("user_id") or "hermes-user"
|
||||
)
|
||||
agent_id = _prompt("Agent identifier", default=provider_config.get("agent_id") or "hermes")
|
||||
if flags.get("dry_run"):
|
||||
_print_dry_run(
|
||||
f"host={host}, user_id={user_id}, agent_id={agent_id}", env_writes, lambda: _check_selfhosted_server(host)
|
||||
)
|
||||
return
|
||||
provider_config.update(
|
||||
mode="platform", host=host, user_id=user_id, agent_id=agent_id
|
||||
) # routing: oss > host > platform
|
||||
_persist_provider_config(
|
||||
hermes_home, config, provider_config, env_writes, "mem0 (self-hosted)", "API key saved to .env", server=host
|
||||
)
|
||||
|
||||
|
||||
def _print_oss_summary(oss_config: dict, env_writes: dict, dry_run: bool = False) -> None:
|
||||
llm, emb = oss_config["llm"], oss_config["embedder"]
|
||||
w = 0 if dry_run else 9 # final summary column-aligns the labels
|
||||
lines = [
|
||||
"",
|
||||
" [dry-run] OSS config would be:" if dry_run else " ✓ Mem0 configured (OSS mode)",
|
||||
f" {'LLM:':<{w}} {llm['provider']} ({llm['config'].get('model', '')})",
|
||||
f" {'Embedder:':<{w}} {emb['provider']} ({emb['config'].get('model', '')})",
|
||||
f" {'Vector:':<{w}} {oss_config['vector_store']['provider']}",
|
||||
]
|
||||
if dry_run:
|
||||
lines += [f" Env vars: {', '.join(env_writes.keys())}"] if env_writes else []
|
||||
else:
|
||||
lines += [
|
||||
*([" API keys saved to .env"] if env_writes else []),
|
||||
" Config saved to mem0.json",
|
||||
" Provider set in config.yaml",
|
||||
"",
|
||||
" Start a new session to activate.",
|
||||
"",
|
||||
]
|
||||
print("\n".join(lines))
|
||||
|
||||
|
||||
def _finish_oss(
|
||||
hermes_home: str,
|
||||
config: dict,
|
||||
oss_config: dict,
|
||||
env_writes: dict[str, str],
|
||||
user_id: str,
|
||||
agent_id: str,
|
||||
pgvector_config: dict | None = None,
|
||||
) -> None:
|
||||
"""Shared OSS tail: write secrets + mem0.json, install deps, activate, check, summarize."""
|
||||
from . import Mem0MemoryProvider
|
||||
|
||||
if env_writes:
|
||||
_write_env(Path(hermes_home) / ".env", env_writes)
|
||||
Mem0MemoryProvider().save_config(
|
||||
{"mode": "oss", "user_id": user_id, "agent_id": agent_id, "oss": oss_config}, hermes_home
|
||||
)
|
||||
_install_provider_deps(
|
||||
oss_config["llm"]["provider"], oss_config["embedder"]["provider"], oss_config["vector_store"]["provider"]
|
||||
)
|
||||
if pgvector_config:
|
||||
_ensure_pgvector_extension(pgvector_config)
|
||||
_activate_provider(config)
|
||||
_run_connectivity_checks(oss_config)
|
||||
_print_oss_summary(oss_config, env_writes)
|
||||
|
||||
|
||||
def _setup_oss(hermes_home: str, config: dict, flags: dict[str, str]) -> None:
|
||||
"""OSS mode — non-interactive when --mode was given, otherwise curses pickers."""
|
||||
if not flags.get("_mode_from_flag"):
|
||||
_setup_oss_interactive(hermes_home, config)
|
||||
return
|
||||
oss_config, env_writes = build_oss_config(flags)
|
||||
if errors := validate_oss_config(oss_config):
|
||||
print("".join(f" Error: {e}\n" for e in errors), end="", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
if flags.get("dry_run"):
|
||||
_print_oss_summary(oss_config, env_writes, dry_run=True)
|
||||
_run_connectivity_checks(oss_config)
|
||||
print(" [dry-run] No files written.\n")
|
||||
return
|
||||
_finish_oss(
|
||||
hermes_home, config, oss_config, env_writes, flags.get("user_id") or os.getenv("USER", "hermes-user"), "hermes"
|
||||
)
|
||||
|
||||
|
||||
def _docker(*args: str, timeout: int, **kwargs) -> subprocess.CompletedProcess:
|
||||
return subprocess.run(["docker", *args], capture_output=True, timeout=timeout, stdin=subprocess.DEVNULL, **kwargs)
|
||||
|
||||
|
||||
def _pg_ready(host: str, port: int, wait: int) -> bool:
|
||||
"""Wait up to ``wait`` seconds for the port, then report whether PostgreSQL answers."""
|
||||
_wait_for_port(host, port, timeout=wait)
|
||||
return _check_pgvector(host, port)[0]
|
||||
|
||||
|
||||
def _ensure_pgvector(host: str = "localhost", port: int = 5432) -> dict | None:
|
||||
"""Ensure pgvector is reachable, offering Docker if not; returns the started container's vector_config, else None."""
|
||||
if _check_pgvector(host, port)[0]:
|
||||
print(f" ✓ PostgreSQL reachable at {host}:{port}")
|
||||
return None
|
||||
print(f" PostgreSQL not reachable at {host}:{port}")
|
||||
if not shutil.which("docker"):
|
||||
print(" Docker not found. Install Docker to auto-start pgvector,\n or run PostgreSQL with pgvector manually.")
|
||||
return None
|
||||
with suppress(Exception): # restart our own container if it exists but is stopped
|
||||
result = _docker(
|
||||
"inspect",
|
||||
_PGVECTOR_CONTAINER,
|
||||
"--format",
|
||||
"{{.State.Status}}",
|
||||
timeout=10,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
)
|
||||
if result.returncode == 0 and "exited" in result.stdout:
|
||||
print(f" Found stopped container '{_PGVECTOR_CONTAINER}', restarting...")
|
||||
_docker("start", _PGVECTOR_CONTAINER, timeout=15)
|
||||
if _pg_ready(host, port, 15):
|
||||
print(" ✓ PostgreSQL container restarted")
|
||||
return None
|
||||
if input(" Start pgvector via Docker? [Y/n]: ").strip().lower() not in ("", "y", "yes"):
|
||||
print(" Skipping Docker setup. Make sure PostgreSQL with pgvector is running.")
|
||||
return None
|
||||
try:
|
||||
print(f" Pulling {_PGVECTOR_IMAGE}...")
|
||||
_docker("pull", _PGVECTOR_IMAGE, timeout=120)
|
||||
_docker("rm", "-f", _PGVECTOR_CONTAINER, timeout=10) # remove existing container if present
|
||||
print(f" Starting container '{_PGVECTOR_CONTAINER}' on port {port}...")
|
||||
_docker(
|
||||
"run",
|
||||
"-d",
|
||||
"--name",
|
||||
_PGVECTOR_CONTAINER,
|
||||
"-e",
|
||||
f"POSTGRES_PASSWORD={_PGVECTOR_PASSWORD}",
|
||||
"-p",
|
||||
f"{port}:5432",
|
||||
_PGVECTOR_IMAGE,
|
||||
timeout=30,
|
||||
check=True,
|
||||
)
|
||||
if _pg_ready(host, port, 20):
|
||||
print(f" ✓ pgvector running on {host}:{port}")
|
||||
else:
|
||||
print(
|
||||
" Warning: Container started but PostgreSQL not yet accepting connections.\n It may need a few more seconds. Config will be saved; retry later."
|
||||
)
|
||||
return {"host": host, "port": port, "user": "postgres", "password": _PGVECTOR_PASSWORD, "dbname": "postgres"}
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f" Failed to start Docker container: {e}")
|
||||
except Exception as e:
|
||||
print(f" Docker error: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _ensure_ollama(models: list[str]) -> bool:
|
||||
"""Ensure Ollama is running and ``models`` are pulled; False when the user must handle it manually."""
|
||||
ollama_bin = shutil.which("ollama")
|
||||
if not (ok := _check_ollama(_OLLAMA_URL)[0]):
|
||||
if not ollama_bin:
|
||||
print(
|
||||
" Ollama not found. Install it:\n curl -fsSL https://ollama.com/install.sh | sh\n Or on macOS: brew install ollama"
|
||||
)
|
||||
return False
|
||||
print(" Ollama installed but not running. Starting...")
|
||||
try:
|
||||
subprocess.Popen(
|
||||
[ollama_bin, "serve"], stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
|
||||
)
|
||||
_wait_for_port("localhost", 11434, timeout=10)
|
||||
if ok := _check_ollama(_OLLAMA_URL)[0]:
|
||||
print(" ✓ Ollama started")
|
||||
except Exception as e:
|
||||
print(f" Could not start Ollama: {e}")
|
||||
if not ok:
|
||||
print(" Warning: Ollama not reachable. Models cannot be pulled.")
|
||||
return False
|
||||
for model in models:
|
||||
try:
|
||||
names = [
|
||||
m.get("name", "") for m in json.loads(_http_get(_OLLAMA_URL, "/api/tags", 5).read()).get("models", [])
|
||||
]
|
||||
except Exception:
|
||||
names = []
|
||||
if any(model in n or model.split(":")[0] in n for n in names):
|
||||
print(f" ✓ Model '{model}' available")
|
||||
continue
|
||||
print(f" Pulling '{model}'... (this may take a few minutes)")
|
||||
try:
|
||||
subprocess.run([ollama_bin or "ollama", "pull", model], timeout=600, stdin=subprocess.DEVNULL)
|
||||
print(f" ✓ Model '{model}' pulled")
|
||||
except Exception as e:
|
||||
print(f" Warning: Could not pull '{model}': {e}\n Run manually: ollama pull {model}")
|
||||
return True
|
||||
|
||||
|
||||
def _ensure_pgvector_extension(pg_config: dict) -> None:
|
||||
try:
|
||||
import psycopg2
|
||||
except ImportError:
|
||||
return
|
||||
defaults = {"host": "localhost", "port": 5432, "user": "postgres", "dbname": "postgres"}
|
||||
try:
|
||||
conn = psycopg2.connect(
|
||||
**(defaults | {k: v for k, v in pg_config.items() if k in defaults or (k == "password" and v)})
|
||||
)
|
||||
conn.autocommit = True
|
||||
conn.cursor().execute("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
conn.close()
|
||||
print(" ✓ pgvector extension enabled")
|
||||
except Exception as e:
|
||||
print(f" Warning: Could not enable pgvector extension: {e}")
|
||||
|
||||
|
||||
def _wait_for_port(host: str, port: int, timeout: int = 15) -> None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
socket.create_connection((host, port), timeout=1).close()
|
||||
return
|
||||
except OSError:
|
||||
time.sleep(0.5)
|
||||
|
||||
|
||||
# Picker descriptions: LLM/embedder show model (+ URL); vector stores by provider id (default: the id itself).
|
||||
_VECTOR_DESCRIPTIONS = {
|
||||
"qdrant": lambda cfg: cfg.get("path", "local storage"),
|
||||
"pgvector": lambda cfg: f"{cfg.get('host', 'localhost')}:{cfg.get('port', 5432)}",
|
||||
}
|
||||
|
||||
|
||||
def _configure_model_provider(
|
||||
kind: str, registry: dict, hermes_home: str, env_writes: dict[str, str], llm: tuple[str, dict] | None = None
|
||||
) -> tuple[str, dict, str, str | None]:
|
||||
"""Pick an LLM/embedder provider, collect its key, and (for Ollama) model + URL -> (id, definition, model, url).
|
||||
For the embedder (``llm`` given), a provider shared with the LLM reuses the LLM key instead of prompting again."""
|
||||
items = [
|
||||
(
|
||||
v["label"],
|
||||
f"{v.get('default_model', '')} ({v['default_url']})"
|
||||
if v.get("default_url")
|
||||
else v.get("default_model", ""),
|
||||
)
|
||||
for v in registry.values()
|
||||
]
|
||||
pid = list(registry)[_curses_select(f"{kind} Provider", items, 0)]
|
||||
pdef = registry[pid]
|
||||
model, url = pdef["default_model"], pdef.get("default_url")
|
||||
if pdef["needs_key"]:
|
||||
if llm is None or pid != llm[0]:
|
||||
if key := _prompt_api_key(
|
||||
pdef["label"] if llm is None else f"{pdef['label']} embedder", pdef["env_var"], hermes_home
|
||||
):
|
||||
env_writes[pdef["env_var"]] = key
|
||||
elif llm[1].get("env_var") in env_writes:
|
||||
env_writes[pdef["env_var"]] = env_writes[llm[1]["env_var"]]
|
||||
if pid == "ollama":
|
||||
model = _input(f"{kind} model", pdef["default_model"])
|
||||
url = _input("Ollama URL", pdef["default_url"])
|
||||
return pid, pdef, model, url
|
||||
|
||||
|
||||
def _setup_oss_interactive(hermes_home: str, config: dict) -> None:
|
||||
env_writes: dict[str, str] = {}
|
||||
llm_id, llm_def, llm_model, llm_url = _configure_model_provider("LLM", LLM_PROVIDERS, hermes_home, env_writes)
|
||||
embedder_id, _, embedder_model, embedder_url = _configure_model_provider(
|
||||
"Embedder", EMBEDDER_PROVIDERS, hermes_home, env_writes, llm=(llm_id, llm_def)
|
||||
)
|
||||
vector_items = [
|
||||
(v["label"], _VECTOR_DESCRIPTIONS.get(pid, lambda cfg: pid)(vector_default_config(pid)))
|
||||
for pid, v in VECTOR_PROVIDERS.items()
|
||||
]
|
||||
vector_id = list(VECTOR_PROVIDERS)[_curses_select("Vector Store", vector_items, 0)]
|
||||
# Auto-setup: ensure Ollama is running and models are pulled; ensure pgvector is reachable (offer Docker if not).
|
||||
ollama_models = [m for pid, m in ((llm_id, llm_model), (embedder_id, embedder_model)) if pid == "ollama"]
|
||||
if ollama_models:
|
||||
_ensure_ollama(ollama_models)
|
||||
pgvector_config = _ensure_pgvector() if vector_id == "pgvector" else None
|
||||
if (
|
||||
vector_id == "pgvector" and not pgvector_config
|
||||
): # native PostgreSQL: prompt for connection details (user first, historical order)
|
||||
pg = {
|
||||
k: _input(f"PostgreSQL {label}", d)
|
||||
for k, label, d in (
|
||||
("user", "user", os.getenv("USER", "postgres")),
|
||||
("host", "host", "localhost"),
|
||||
("port", "port", "5432"),
|
||||
("dbname", "database", "postgres"),
|
||||
)
|
||||
}
|
||||
pg_password = getpass.getpass(" PostgreSQL password (blank if none): ").strip()
|
||||
pgvector_config = {**pg, "port": int(pg["port"]), **({"password": pg_password} if pg_password else {})}
|
||||
user_id = _input("User ID", os.getenv("USER", "hermes-user"))
|
||||
agent_id = _input("Agent ID", "hermes")
|
||||
flags = {
|
||||
"oss_llm": llm_id,
|
||||
"oss_llm_model": llm_model,
|
||||
"oss_llm_url": llm_url or "",
|
||||
"oss_llm_key": env_writes.get(llm_def["env_var"], "") if llm_def.get("env_var") else "",
|
||||
"oss_embedder": embedder_id,
|
||||
"oss_embedder_model": embedder_model,
|
||||
"oss_embedder_url": embedder_url or "",
|
||||
"oss_vector": vector_id,
|
||||
"user_id": user_id,
|
||||
}
|
||||
flags.update({f"oss_vector_{key}": str(val) for key, val in (pgvector_config or {}).items() if val})
|
||||
oss_config, _ = build_oss_config(flags)
|
||||
_finish_oss(hermes_home, config, oss_config, env_writes, user_id, agent_id, pgvector_config)
|
||||
|
||||
|
||||
def _install_provider_deps(llm_id: str, embedder_id: str, vector_id: str) -> None:
|
||||
deps = {
|
||||
registry[pid]["pip_dep"]
|
||||
for (_, registry), pid in zip(SECTION_REGISTRIES, (llm_id, embedder_id, vector_id))
|
||||
if registry.get(pid, {}).get("pip_dep")
|
||||
}
|
||||
for dep in sorted(deps):
|
||||
print(f" Installing {dep}...")
|
||||
try:
|
||||
# Environment-aware install: sealed hosted venvs redirect to the durable data-volume target instead of /opt/hermes.
|
||||
from tools.lazy_deps import install_specs
|
||||
|
||||
outcome = install_specs([dep], timeout=60)
|
||||
except Exception:
|
||||
outcome = None
|
||||
print(
|
||||
f" ✓ Installed {dep}"
|
||||
if outcome is not None and outcome.ok
|
||||
else f" Warning: cannot install {dep}: {outcome.reason}"
|
||||
if outcome is not None and outcome.blocked
|
||||
else f" Warning: Could not install {dep}. Install manually: uv pip install {dep}"
|
||||
)
|
||||
if deps:
|
||||
import importlib
|
||||
|
||||
importlib.invalidate_caches()
|
||||
|
||||
|
||||
def _probe(fn, ok: str, fail: str, exc=Exception) -> tuple[bool, str]:
|
||||
"""Run ``fn``; (True, ok) on success, (False, "fail: <error>") on ``exc``."""
|
||||
try:
|
||||
fn()
|
||||
return True, ok
|
||||
except exc as e:
|
||||
return False, f"{fail}: {e}"
|
||||
|
||||
|
||||
def _check_qdrant_path(path: str) -> tuple[bool, str]:
|
||||
"""Check that qdrant local storage parent dir is writable."""
|
||||
parent = Path(path).expanduser().parent
|
||||
return _probe(
|
||||
lambda: parent.mkdir(parents=True, exist_ok=True),
|
||||
f"Directory writable: {parent}",
|
||||
f"Cannot write to {parent}",
|
||||
OSError,
|
||||
)
|
||||
|
||||
|
||||
def _check_ollama(url: str) -> tuple[bool, str]:
|
||||
return _probe(lambda: _http_get(url, "/api/tags", 3), "Ollama reachable", f"Ollama not reachable at {url}")
|
||||
|
||||
|
||||
def _check_pgvector(host: str, port: int) -> tuple[bool, str]:
|
||||
return _probe(
|
||||
lambda: socket.create_connection((host, port), timeout=3).close(),
|
||||
f"PGVector reachable at {host}:{port}",
|
||||
f"PGVector not reachable at {host}:{port}",
|
||||
)
|
||||
|
||||
|
||||
def _warn_unless(check: tuple[bool, str]) -> None:
|
||||
ok, msg = check
|
||||
if not ok:
|
||||
print(f" Warning: {msg}")
|
||||
|
||||
|
||||
def _run_connectivity_checks(oss_config: dict) -> None:
|
||||
vs = oss_config.get("vector_store", {})
|
||||
cfg = vs.get("config", {})
|
||||
if vs.get("provider") == "qdrant":
|
||||
path, url = cfg.get("path"), cfg.get("url")
|
||||
if path:
|
||||
_warn_unless(_check_qdrant_path(path))
|
||||
elif url:
|
||||
_warn_unless(
|
||||
_probe(lambda: _http_get(url, "/healthz", 3), "Qdrant reachable", f"Qdrant not reachable at {url}")
|
||||
)
|
||||
elif vs.get("provider") == "pgvector":
|
||||
_warn_unless(_check_pgvector(cfg.get("host", "localhost"), cfg.get("port", 5432)))
|
||||
llm = oss_config.get("llm", {})
|
||||
if llm.get("provider") == "ollama":
|
||||
_warn_unless(_check_ollama(llm.get("config", {}).get("ollama_base_url", _OLLAMA_URL)))
|
||||
|
||||
|
||||
_MODE_HANDLERS = {
|
||||
"oss": _setup_oss,
|
||||
"selfhosted": _setup_selfhosted,
|
||||
"self-hosted": _setup_selfhosted,
|
||||
"platform": _setup_platform,
|
||||
}
|
||||
# Interactive picker order: Platform, Self-hosted server, Open Source.
|
||||
_MODE_ITEMS = [
|
||||
("Platform", "Mem0 Cloud API (lightweight, just needs an API key)"),
|
||||
("Self-hosted server", "Connect to an existing self-hosted Mem0 server (Docker/FastAPI)"),
|
||||
("Open Source", "Run Mem0 locally (self-hosted LLM + vector store)"),
|
||||
]
|
||||
_MODE_PICKER = (_setup_platform, _setup_selfhosted, _setup_oss)
|
||||
|
||||
|
||||
def post_setup(hermes_home: str, config: dict) -> None:
|
||||
"""Entry point for `hermes memory setup`: routes on --mode (platform / selfhosted / oss), else shows a picker.
|
||||
OSS is non-interactive only when the mode came from the flag."""
|
||||
with suppress(ImportError): # mem0ai must meet the minimum version from plugin.yaml
|
||||
import mem0
|
||||
|
||||
installed_ver = getattr(mem0, "__version__", None)
|
||||
if installed_ver and tuple(int(x) for x in installed_ver.split(".")[:3]) < (2, 0, 10):
|
||||
print(
|
||||
f"\n ⚠ mem0ai {installed_ver} installed but >=2.0.10 required.\n Run: uv pip install --python {sys.executable} 'mem0ai>=2.0.10'"
|
||||
)
|
||||
flags = parse_flags(sys.argv[1:])
|
||||
handler = _MODE_HANDLERS.get(flags["mode"])
|
||||
flags["_mode_from_flag"] = handler is not None
|
||||
if handler is None:
|
||||
handler = _MODE_PICKER[_curses_select(" Select mode", _MODE_ITEMS, 0)]
|
||||
handler(hermes_home, config, flags)
|
||||
|
||||
|
||||
# Compatibility name retained for callers of the original bundled provider.
|
||||
|
||||
|
||||
def has_oss_flags() -> bool:
|
||||
"""Check if OSS-related flags are present in sys.argv."""
|
||||
flags = parse_flags(sys.argv[1:])
|
||||
if flags["mode"] == "oss":
|
||||
return True
|
||||
if any(flags.get(k) for k in ("oss_llm_key", "oss_vector_path", "oss_vector_url")):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# ---- END PLUGIN-COMPAT ----
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Shared, host-independent redaction and lossless extraction batching."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r"|(?:access|refresh|session)[_-]?token|token|authorization|credential"
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, default=str)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if _is_agent_assignment(message) and index + 1 < len(exchange) and _is_agent_response(exchange[index + 1]):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
@@ -0,0 +1,25 @@
|
||||
{
|
||||
"id": "mem0",
|
||||
"version": "1.4.0",
|
||||
"homepage": "https://docs.mem0.ai/integrations/hermes",
|
||||
"native": {
|
||||
"pluginRoot": "",
|
||||
"pythonFiles": [
|
||||
"message_utils.py"
|
||||
],
|
||||
"skills": false,
|
||||
"files": {
|
||||
"plugin.yaml": "plugin.yaml",
|
||||
"pyproject.toml": "pyproject.toml",
|
||||
"__init__.py": "__init__.py",
|
||||
"_backend.py": "_backend.py",
|
||||
"_setup.py": "_setup.py",
|
||||
"_openai_llm.py": "_openai_llm.py",
|
||||
"_oss_providers.py": "_oss_providers.py",
|
||||
"README.md": "README.md",
|
||||
"HANDOFF.md": "HANDOFF.md",
|
||||
"LICENSE": "LICENSE",
|
||||
"LICENSE-APACHE-2.0": "LICENSE-APACHE-2.0"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
name: mem0
|
||||
version: 1.4.0
|
||||
description: "Mem0 — server-side LLM fact extraction with semantic search, automatic deduplication, and opt-in reranking (platform mode)."
|
||||
pip_dependencies:
|
||||
- mem0ai>=2.0.10,<3
|
||||
@@ -0,0 +1,17 @@
|
||||
[project]
|
||||
name = "hermes-plugin-mem0"
|
||||
version = "1.4.0"
|
||||
description = "Hermes Agent memory provider plugin: mem0"
|
||||
requires-python = ">=3.11"
|
||||
license = { text = "MIT AND Apache-2.0" }
|
||||
# Installed into the Hermes venv by `hermes plugins install` / `enable` and re-applied after
|
||||
# every `hermes update` (hermes-agent#113851). Keep upper bounds: Hermes pins its own deps exactly
|
||||
# and refuses a plugin whose requirements cannot resolve against them.
|
||||
dependencies = [
|
||||
"mem0ai>=2.0.10,<3",
|
||||
"httpx>=0.27,<1",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
postgres = ['psycopg2-binary>=2.9,<3']
|
||||
qdrant = ['qdrant-client>=1.9,<2']
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Offline contract smoke against a real Hermes checkout.
|
||||
|
||||
HERMES_SOURCE=/path/to/hermes-agent python tests/smoke_hermes.py
|
||||
Host modules are real; only the Mem0 backend is mocked. Uses a temporary profile.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
from unittest.mock import Mock
|
||||
|
||||
hermes_source = pathlib.Path(os.environ["HERMES_SOURCE"]).resolve()
|
||||
plugin_source = pathlib.Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(hermes_source))
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
os.environ["HERMES_HOME"] = d
|
||||
os.environ["MEM0_API_KEY"] = "test-key"
|
||||
import plugins.memory as pm
|
||||
from agent.memory_manager import MemoryManager
|
||||
|
||||
destination = pathlib.Path(d) / "plugins" / "mem0"
|
||||
shutil.copytree(plugin_source, destination)
|
||||
assert pm.find_provider_dir("mem0") == hermes_source / "plugins" / "memory" / "mem0"
|
||||
pm._MEMORY_PLUGINS_DIR = pathlib.Path(d) / "empty-bundled"
|
||||
pm._MEMORY_PLUGINS_DIR.mkdir()
|
||||
provider = pm.load_memory_provider("mem0", register_skills=False)
|
||||
assert provider is not None
|
||||
assert provider.__class__.__module__.startswith("_hermes_user_memory.")
|
||||
backend = Mock()
|
||||
backend.search.return_value = []
|
||||
backend.add.return_value = {}
|
||||
provider._create_backend = lambda: backend
|
||||
manager = MemoryManager()
|
||||
manager.add_provider(provider)
|
||||
manager.initialize_all("s1", platform="cli", user_id="alice")
|
||||
result = manager.handle_tool_call("mem0_search", {"query": "favorite programming language"})
|
||||
assert json.loads(result)["result"] == "No relevant memories found."
|
||||
manager.sync_all(
|
||||
"I prefer Python for tooling. " * 100 + "UNIQUE_END_MARKER",
|
||||
"I will remember that.",
|
||||
session_id="s2",
|
||||
messages=[{"role": "tool", "content": "NOT_FOR_CAPTURE"}],
|
||||
turn_author={"id": "alice"},
|
||||
)
|
||||
manager.shutdown_all()
|
||||
assert backend.add.call_count >= 1
|
||||
sent = str(backend.add.call_args_list)
|
||||
assert "UNIQUE_END_MARKER" in sent and "NOT_FOR_CAPTURE" not in sent
|
||||
assert backend.add.call_args.kwargs["run_id"] == "s2"
|
||||
print("PASS:", hermes_source, "real external loader + MemoryManager lifecycle/tools/full capture/session ID")
|
||||
@@ -0,0 +1,86 @@
|
||||
"""Offline backend contracts; no Hermes install, credentials, or database needed."""
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import httpx
|
||||
|
||||
spec = importlib.util.spec_from_file_location("hermes_backend_test._backend", Path(__file__).parents[1] / "_backend.py")
|
||||
backend = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(backend)
|
||||
|
||||
|
||||
class BackendTests(unittest.TestCase):
|
||||
def test_cloud_and_oss_preserve_scope_and_accept_optional_session(self):
|
||||
for backend_type, client_attr in ((backend.PlatformBackend, "_client"), (backend.OSSBackend, "_memory")):
|
||||
with self.subTest(backend=backend_type.__name__):
|
||||
instance = backend_type.__new__(backend_type)
|
||||
client = Mock()
|
||||
setattr(instance, client_attr, client)
|
||||
client.search.return_value = {"results": [{"id": "old", "memory": "Legacy memory"}]}
|
||||
instance.add([{"role": "user", "content": "fact"}], user_id="user", agent_id="hermes")
|
||||
self.assertNotIn("run_id", client.add.call_args.kwargs)
|
||||
instance.add([], user_id="user", agent_id="hermes", run_id="session", metadata={"channel": "cli"})
|
||||
self.assertEqual(client.add.call_args.kwargs["run_id"], "session")
|
||||
self.assertEqual(client.add.call_args.kwargs["metadata"], {"channel": "cli"})
|
||||
self.assertEqual(instance.search("fact", filters={"user_id": "user"}, top_k=3)[0]["id"], "old")
|
||||
self.assertEqual(client.search.call_args.kwargs["filters"], {"user_id": "user"})
|
||||
self.assertEqual(client.search.call_args.kwargs["top_k"], 3)
|
||||
|
||||
def test_selfhosted_http_contract(self):
|
||||
requests = []
|
||||
|
||||
def respond(request):
|
||||
requests.append(request)
|
||||
return httpx.Response(200, json={"results": [{"id": "legacy", "memory": "fact"}]})
|
||||
|
||||
instance = backend.SelfHostedBackend(
|
||||
"test-key", "http://localhost:8888/", transport=httpx.MockTransport(respond)
|
||||
)
|
||||
try:
|
||||
instance.add([], user_id="user", agent_id="hermes", run_id="session", infer=True)
|
||||
self.assertEqual(json.loads(requests[-1].content)["run_id"], "session")
|
||||
self.assertEqual(requests[-1].headers["X-API-Key"], "test-key")
|
||||
instance.search("fact", filters={"user_id": "user"}, top_k=7, rerank=True)
|
||||
self.assertEqual(str(requests[-1].url), "http://localhost:8888/search")
|
||||
self.assertEqual(
|
||||
json.loads(requests[-1].content), {"query": "fact", "filters": {"user_id": "user"}, "top_k": 7}
|
||||
)
|
||||
instance.update("legacy", "new fact")
|
||||
self.assertEqual((requests[-1].method, json.loads(requests[-1].content)), ("PUT", {"text": "new fact"}))
|
||||
instance.delete("legacy")
|
||||
self.assertEqual(requests[-1].method, "DELETE")
|
||||
finally:
|
||||
instance.close()
|
||||
|
||||
def test_qdrant_dimension_change_never_deletes_memories(self):
|
||||
client = Mock()
|
||||
client.collection_exists.return_value = True
|
||||
client.get_collection.return_value.config.params.vectors = SimpleNamespace(size=1536)
|
||||
with patch.dict(sys.modules, {"qdrant_client": SimpleNamespace(QdrantClient=Mock(return_value=client))}):
|
||||
with self.assertRaisesRegex(ValueError, "1536.*768"):
|
||||
backend.OSSBackend._recreate_collection_if_dims_changed("qdrant", {"path": "/unused"}, 768)
|
||||
client.delete_collection.assert_not_called()
|
||||
client.close.assert_called_once()
|
||||
|
||||
def test_pgvector_dimension_change_never_drops_table(self):
|
||||
cursor = Mock()
|
||||
cursor.fetchone.return_value = (1536,)
|
||||
connection = Mock()
|
||||
connection.cursor.return_value = cursor
|
||||
driver = SimpleNamespace(connect=Mock(return_value=connection), sql=Mock())
|
||||
with patch.dict(sys.modules, {"psycopg2": driver}):
|
||||
with self.assertRaisesRegex(ValueError, "1536.*768"):
|
||||
backend.OSSBackend._recreate_collection_if_dims_changed("pgvector", {"user": "test"}, 768)
|
||||
self.assertEqual(cursor.execute.call_count, 1)
|
||||
self.assertTrue(cursor.execute.call_args.args[0].startswith("SELECT"))
|
||||
connection.close.assert_called_once()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,225 @@
|
||||
"""Offline contracts for the native Hermes provider and legacy configuration."""
|
||||
|
||||
import contextvars
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def plugin(monkeypatch, tmp_path):
|
||||
def spawn(target, *, name):
|
||||
context = contextvars.copy_context()
|
||||
return threading.Thread(target=context.run, args=(target,), name=name)
|
||||
|
||||
modules = {
|
||||
"agent": {},
|
||||
"agent.memory_provider": {"MemoryProvider": object, "spawn_context_thread": spawn},
|
||||
"agent.secret_scope": {"get_secret": lambda key, default="": default},
|
||||
"tools": {},
|
||||
"tools.registry": {"tool_error": lambda text: json.dumps({"error": text})},
|
||||
"utils": {
|
||||
"read_json_or_empty": lambda p: json.loads(p.read_text()) if p.exists() else {},
|
||||
"atomic_json_write": lambda p, value, **kw: p.write_text(json.dumps(value)),
|
||||
},
|
||||
"hermes_constants": {"get_hermes_home": lambda: tmp_path},
|
||||
}
|
||||
for name, values in modules.items():
|
||||
module = types.ModuleType(name)
|
||||
module.__dict__.update(values)
|
||||
monkeypatch.setitem(sys.modules, name, module)
|
||||
name = "hermes_test_mem0"
|
||||
spec = importlib.util.spec_from_file_location(name, ROOT / "__init__.py", submodule_search_locations=[str(ROOT)])
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, name, module)
|
||||
spec.loader.exec_module(module)
|
||||
monkeypatch.setattr(module.atexit, "register", lambda *args: None)
|
||||
return module
|
||||
|
||||
|
||||
def provider(plugin, monkeypatch, config=None, **identity):
|
||||
backend = Mock()
|
||||
backend.add.return_value = {"event_id": "event-1"}
|
||||
backend.search.return_value = []
|
||||
monkeypatch.setattr(plugin, "_load_config", lambda: config or {})
|
||||
monkeypatch.setattr(plugin.Mem0MemoryProvider, "_create_backend", lambda self: backend)
|
||||
instance = plugin.Mem0MemoryProvider()
|
||||
instance.initialize("session-1", **identity)
|
||||
return instance, backend
|
||||
|
||||
|
||||
def test_long_turns_preserve_tail_and_redact_before_chunking(plugin, monkeypatch):
|
||||
instance, backend = provider(plugin, monkeypatch, {"sync_max_chars": 450})
|
||||
text = "Opening. " + "x" * 1100 + " api_key=secret-value lasting preference at the end."
|
||||
instance.sync_turn(text, "Noted.", session_id="turn-session")
|
||||
instance.shutdown()
|
||||
messages = [m for call in backend.add.call_args_list for m in call.args[0]]
|
||||
assert "".join(m["content"] for m in messages if m["role"] == "user") == plugin.redact(text)
|
||||
assert all(len(m["content"]) <= 450 for m in messages)
|
||||
assert "secret-value" not in repr(backend.add.call_args_list)
|
||||
assert all(call.kwargs["run_id"] == "turn-session" for call in backend.add.call_args_list)
|
||||
|
||||
|
||||
def test_busy_capture_queues_each_turn_with_its_profile_context(plugin, monkeypatch):
|
||||
instance, backend = provider(plugin, monkeypatch)
|
||||
started, release = threading.Event(), threading.Event()
|
||||
profile = contextvars.ContextVar("profile", default="first")
|
||||
seen = []
|
||||
|
||||
def add(messages, **kwargs):
|
||||
if not seen:
|
||||
started.set()
|
||||
assert release.wait(3)
|
||||
seen.append((messages[0]["content"], profile.get()))
|
||||
return {}
|
||||
|
||||
backend.add.side_effect = add
|
||||
instance.sync_turn("first turn", "reply")
|
||||
assert started.wait(3)
|
||||
token = profile.set("second")
|
||||
try:
|
||||
instance.sync_turn("second turn", "reply")
|
||||
finally:
|
||||
profile.reset(token)
|
||||
release.set()
|
||||
instance.shutdown()
|
||||
assert seen == [("first turn", "first"), ("second turn", "second")]
|
||||
backend.close.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"configured,gateway,expected",
|
||||
[
|
||||
(None, "telegram-42", "telegram-42"),
|
||||
("hermes-user", "telegram-42", "telegram-42"),
|
||||
("existing-account", "telegram-42", "existing-account"),
|
||||
(None, None, "hermes-user"),
|
||||
],
|
||||
)
|
||||
def test_legacy_identity_and_tool_contract(plugin, monkeypatch, configured, gateway, expected):
|
||||
instance, backend = provider(plugin, monkeypatch, {"user_id": configured}, user_id=gateway)
|
||||
assert instance.name == "mem0"
|
||||
assert {s["name"] for s in instance.get_tool_schemas()} == {"mem0_search", "mem0_add", "mem0_update", "mem0_delete"}
|
||||
instance.handle_tool_call("mem0_search", {"query": "api_key=secret-value"})
|
||||
assert backend.search.call_args.kwargs["filters"] == {"user_id": expected}
|
||||
assert backend.search.call_args.args[0] == "api_key=[REDACTED]"
|
||||
instance.handle_tool_call("mem0_add", {"content": "password=secret-value"})
|
||||
assert backend.add.call_args.kwargs["infer"] is False
|
||||
assert backend.add.call_args.args[0][0]["content"] == "password=[REDACTED]"
|
||||
instance.handle_tool_call("mem0_update", {"memory_id": "old-id", "text": "password=secret-value"})
|
||||
backend.update.assert_called_once_with("old-id", "password=[REDACTED]")
|
||||
instance.shutdown()
|
||||
|
||||
|
||||
def test_config_keeps_legacy_file_over_env_precedence(plugin, monkeypatch, tmp_path):
|
||||
env = {"MEM0_API_KEY": "env-key", "MEM0_USER_ID": "env-user", "MEM0_HOST": "http://localhost:8888"}
|
||||
monkeypatch.setattr(plugin, "get_secret", lambda key, default="": env.get(key, default))
|
||||
(tmp_path / "mem0.json").write_text(json.dumps({"user_id": "existing-user", "rerank": True}))
|
||||
cfg = plugin._load_config()
|
||||
assert (cfg["api_key"], cfg["user_id"], cfg["host"], cfg["rerank"]) == (
|
||||
"env-key",
|
||||
"existing-user",
|
||||
"http://localhost:8888",
|
||||
True,
|
||||
)
|
||||
|
||||
|
||||
def test_prefetch_respects_rerank_and_redacts_recalled_context(plugin, monkeypatch):
|
||||
instance, backend = provider(plugin, monkeypatch, {"rerank": True})
|
||||
backend.search.return_value = [{"memory": "password=old-secret"}]
|
||||
assert "old-secret" not in instance.prefetch("preference")
|
||||
assert backend.search.call_args.kwargs["rerank"] is True
|
||||
instance.shutdown()
|
||||
|
||||
|
||||
def test_session_switch_updates_writes_without_narrowing_recall(plugin, monkeypatch):
|
||||
instance, backend = provider(plugin, monkeypatch)
|
||||
instance.on_session_switch("resumed-session")
|
||||
instance.handle_tool_call("mem0_add", {"content": "fact"})
|
||||
assert backend.add.call_args.kwargs["run_id"] == "resumed-session"
|
||||
instance.handle_tool_call("mem0_search", {"query": "fact"})
|
||||
assert backend.search.call_args.kwargs["filters"] == {"user_id": "hermes-user"}
|
||||
instance.shutdown()
|
||||
|
||||
|
||||
def test_invalid_tool_input_never_reaches_backend_or_trips_breaker(plugin, monkeypatch):
|
||||
instance, backend = provider(plugin, monkeypatch)
|
||||
for args in ({"query": []}, {"query": "fact", "top_k": "not-a-number"}, None):
|
||||
assert "error" in json.loads(instance.handle_tool_call("mem0_search", args))
|
||||
backend.search.assert_not_called()
|
||||
assert instance._consecutive_failures == 0
|
||||
instance.shutdown()
|
||||
|
||||
|
||||
def test_setup_writes_private_env_and_preserves_existing_values(plugin, tmp_path):
|
||||
import importlib
|
||||
|
||||
setup = importlib.import_module(f"{plugin.__name__}._setup")
|
||||
path = tmp_path / ".env"
|
||||
path.write_text("EXISTING=value\nMEM0_API_KEY=old\n")
|
||||
setup._write_env(path, {"MEM0_API_KEY": "new"})
|
||||
assert path.read_text() == "EXISTING=value\nMEM0_API_KEY=new\n"
|
||||
assert path.stat().st_mode & 0o777 == 0o600
|
||||
|
||||
|
||||
def test_redaction_marker_split_at_chunk_boundary_is_lossless(plugin, monkeypatch):
|
||||
instance, backend = provider(plugin, monkeypatch, {"sync_max_chars": 450})
|
||||
text = "x" * 438 + " password=secret-value tail"
|
||||
instance.sync_turn(text, "")
|
||||
instance.shutdown()
|
||||
stored = "".join(m["content"] for call in backend.add.call_args_list for m in call.args[0])
|
||||
assert stored == plugin.redact(text)
|
||||
|
||||
|
||||
def test_sync_queue_stops_network_calls_when_breaker_opens(plugin, monkeypatch, caplog):
|
||||
instance, backend = provider(plugin, monkeypatch)
|
||||
started, release = threading.Event(), threading.Event()
|
||||
|
||||
def unavailable(messages, **kwargs):
|
||||
started.set()
|
||||
assert release.wait(3)
|
||||
raise RuntimeError("server unavailable")
|
||||
|
||||
backend.add.side_effect = unavailable
|
||||
instance.sync_turn("first turn", "reply")
|
||||
assert started.wait(3)
|
||||
try:
|
||||
for number in range(9):
|
||||
instance.sync_turn(f"queued turn {number}", "reply")
|
||||
finally:
|
||||
release.set()
|
||||
instance.shutdown()
|
||||
assert backend.add.call_count == plugin._BREAKER_THRESHOLD
|
||||
assert any("not synced" in record.message.lower() for record in caplog.records)
|
||||
|
||||
|
||||
def test_shutdown_is_bounded_and_defers_close_until_sync_finishes(plugin, monkeypatch, caplog):
|
||||
instance, backend = provider(plugin, monkeypatch)
|
||||
started, release = threading.Event(), threading.Event()
|
||||
monkeypatch.setattr(plugin, "_SHUTDOWN_WAIT_SECS", 0.01)
|
||||
|
||||
def add(*args, **kwargs):
|
||||
started.set()
|
||||
assert release.wait(3)
|
||||
return {}
|
||||
|
||||
backend.add.side_effect = add
|
||||
instance.sync_turn("fact", "reply")
|
||||
assert started.wait(3)
|
||||
worker = instance._sync_thread
|
||||
try:
|
||||
instance.shutdown()
|
||||
backend.close.assert_not_called()
|
||||
assert "shutdown timed out" in caplog.text
|
||||
finally:
|
||||
release.set()
|
||||
worker.join(timeout=3)
|
||||
backend.close.assert_called_once()
|
||||
@@ -11,7 +11,6 @@ from __future__ import annotations
|
||||
import functools
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
@@ -27,6 +26,14 @@ from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
import telemetry
|
||||
from message_utils import MAX_EXTRACTION_INPUT_TOKENS as MAX_EXTRACTION_INPUT_TOKENS
|
||||
from message_utils import SECRET_PATTERNS as SECRET_PATTERNS
|
||||
from message_utils import _estimated_tokens as _estimated_tokens
|
||||
from message_utils import _is_agent_assignment as _is_agent_assignment
|
||||
from message_utils import _is_agent_response as _is_agent_response
|
||||
from message_utils import _message_tokens as _message_tokens
|
||||
from message_utils import extraction_message_batches as extraction_message_batches
|
||||
from message_utils import redact as redact
|
||||
|
||||
DEFAULT_API_URL = "https://api.mem0.ai"
|
||||
PLUGIN_VERSION = "0.3.1"
|
||||
@@ -66,7 +73,6 @@ CHECKPOINT_EXCHANGES = 5
|
||||
CHECKPOINT_MESSAGES = 10
|
||||
CHECKPOINT_SOURCE_CHARS = 40000
|
||||
DEFAULT_MAX_CONTEXT_CHARS = 4000
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
MAX_FLUSH_ATTEMPTS = 5
|
||||
FORGET_PAGE_SIZE = 100
|
||||
FORGET_MAX_PAGES = 50
|
||||
@@ -139,48 +145,11 @@ BUILD_COMMAND_RE = re.compile(
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"
|
||||
),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r'|(?:access|refresh|session)[_-]?token|token|authorization|credential'
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = (
|
||||
value
|
||||
if isinstance(value, str)
|
||||
else json.dumps(value, ensure_ascii=False, default=str)
|
||||
)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def bounded(value: Any, limit: int) -> str:
|
||||
text = redact(value).strip()
|
||||
if len(text) <= limit:
|
||||
@@ -1706,100 +1675,6 @@ def build_extraction_messages(structured: dict[str, Any]) -> list[dict[str, str]
|
||||
return messages
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if (
|
||||
_is_agent_assignment(message)
|
||||
and index + 1 < len(exchange)
|
||||
and _is_agent_response(exchange[index + 1])
|
||||
):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
|
||||
|
||||
def _request_json(
|
||||
url: str, key: str, payload: dict[str, Any], timeout: float
|
||||
) -> tuple[dict[str, Any] | list[Any], int, int]:
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Shared, host-independent redaction and lossless extraction batching."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r"|(?:access|refresh|session)[_-]?token|token|authorization|credential"
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, default=str)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if _is_agent_assignment(message) and index + 1 < len(exchange) and _is_agent_response(exchange[index + 1]):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
@@ -11,7 +11,6 @@ from __future__ import annotations
|
||||
import functools
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
@@ -27,6 +26,14 @@ from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
import telemetry
|
||||
from message_utils import MAX_EXTRACTION_INPUT_TOKENS as MAX_EXTRACTION_INPUT_TOKENS
|
||||
from message_utils import SECRET_PATTERNS as SECRET_PATTERNS
|
||||
from message_utils import _estimated_tokens as _estimated_tokens
|
||||
from message_utils import _is_agent_assignment as _is_agent_assignment
|
||||
from message_utils import _is_agent_response as _is_agent_response
|
||||
from message_utils import _message_tokens as _message_tokens
|
||||
from message_utils import extraction_message_batches as extraction_message_batches
|
||||
from message_utils import redact as redact
|
||||
|
||||
DEFAULT_API_URL = "https://api.mem0.ai"
|
||||
PLUGIN_VERSION = "0.3.1"
|
||||
@@ -66,7 +73,6 @@ CHECKPOINT_EXCHANGES = 5
|
||||
CHECKPOINT_MESSAGES = 10
|
||||
CHECKPOINT_SOURCE_CHARS = 40000
|
||||
DEFAULT_MAX_CONTEXT_CHARS = 4000
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
MAX_FLUSH_ATTEMPTS = 5
|
||||
FORGET_PAGE_SIZE = 100
|
||||
FORGET_MAX_PAGES = 50
|
||||
@@ -139,48 +145,11 @@ BUILD_COMMAND_RE = re.compile(
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"
|
||||
),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r'|(?:access|refresh|session)[_-]?token|token|authorization|credential'
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = (
|
||||
value
|
||||
if isinstance(value, str)
|
||||
else json.dumps(value, ensure_ascii=False, default=str)
|
||||
)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def bounded(value: Any, limit: int) -> str:
|
||||
text = redact(value).strip()
|
||||
if len(text) <= limit:
|
||||
@@ -1706,100 +1675,6 @@ def build_extraction_messages(structured: dict[str, Any]) -> list[dict[str, str]
|
||||
return messages
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get(
|
||||
"content", ""
|
||||
).startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if (
|
||||
_is_agent_assignment(message)
|
||||
and index + 1 < len(exchange)
|
||||
and _is_agent_response(exchange[index + 1])
|
||||
):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
|
||||
|
||||
def _request_json(
|
||||
url: str, key: str, payload: dict[str, Any], timeout: float
|
||||
) -> tuple[dict[str, Any] | list[Any], int, int]:
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Shared, host-independent redaction and lossless extraction batching."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
MAX_EXTRACTION_INPUT_TOKENS = 24000
|
||||
|
||||
SECRET_PATTERNS = [
|
||||
re.compile(r"(?i)(authorization\s*[:=]\s*(?:bearer|token)\s+)[^\s\"']+"),
|
||||
re.compile(r"(?i)((?:api[_-]?key|secret[_-]?access[_-]?key|session[_-]?token)\s*[:=]\s*)[^\s\"']+"),
|
||||
re.compile(
|
||||
r"(?i)((?:access[_-]?token|refresh[_-]?token|password|credential)"
|
||||
r"\s*[:=]\s*)[^\s&\"']+"
|
||||
),
|
||||
re.compile(r"\b(?:sk|m0|mem0_sk|psk)-[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(r"\b(?:ASIA|AKIA)[A-Z0-9]{12,}\b"),
|
||||
re.compile(r"\b(?:ghp_|github_pat_|xox[baprs]-)[A-Za-z0-9_\-]{12,}\b"),
|
||||
re.compile(
|
||||
r"-----BEGIN [^-]*PRIVATE KEY-----.*?-----END [^-]*PRIVATE KEY-----",
|
||||
re.DOTALL,
|
||||
),
|
||||
re.compile(
|
||||
r'(?i)("(?:api[_-]?key|password|secret(?:[_-]?access[_-]?key)?'
|
||||
r"|(?:access|refresh|session)[_-]?token|token|authorization|credential"
|
||||
r')"\s*:\s*")(?:\\.|[^"\\])*'
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def redact(value: Any) -> str:
|
||||
text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, default=str)
|
||||
for pattern in SECRET_PATTERNS:
|
||||
if pattern.groups:
|
||||
text = pattern.sub(r"\1[REDACTED]", text)
|
||||
else:
|
||||
text = pattern.sub("[REDACTED]", text)
|
||||
return text
|
||||
|
||||
|
||||
def _estimated_tokens(value: str) -> int:
|
||||
"""Conservatively estimate tokens without adding a tokenizer dependency."""
|
||||
ascii_chars = sum(ord(char) < 128 for char in value)
|
||||
return math.ceil((ascii_chars * 0.4) + (len(value) - ascii_chars))
|
||||
|
||||
|
||||
def _message_tokens(messages: list[dict[str, str]]) -> int:
|
||||
return _estimated_tokens(json.dumps(messages, ensure_ascii=False))
|
||||
|
||||
|
||||
def _is_agent_assignment(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent assignment (")
|
||||
|
||||
|
||||
def _is_agent_response(message: dict[str, str]) -> bool:
|
||||
return message.get("role") == "assistant" and message.get("content", "").startswith("Subagent response (")
|
||||
|
||||
|
||||
def extraction_message_batches(
|
||||
messages: list[dict[str, str]],
|
||||
*,
|
||||
max_tokens: int = MAX_EXTRACTION_INPUT_TOKENS,
|
||||
) -> list[list[dict[str, str]]]:
|
||||
"""Keep exchanges together when possible; split oversized messages to enforce the request budget."""
|
||||
if not messages or _message_tokens(messages) <= max_tokens:
|
||||
return [messages]
|
||||
|
||||
exchanges: list[list[dict[str, str]]] = []
|
||||
exchange: list[dict[str, str]] = []
|
||||
for message in messages:
|
||||
if message.get("role") == "user" and exchange:
|
||||
exchanges.append(exchange)
|
||||
exchange = []
|
||||
exchange.append(message)
|
||||
if exchange:
|
||||
exchanges.append(exchange)
|
||||
|
||||
units: list[list[dict[str, str]]] = []
|
||||
for exchange in exchanges:
|
||||
if _message_tokens(exchange) <= max_tokens:
|
||||
units.append(exchange)
|
||||
continue
|
||||
index = 0
|
||||
while index < len(exchange):
|
||||
message = exchange[index]
|
||||
if _is_agent_assignment(message) and index + 1 < len(exchange) and _is_agent_response(exchange[index + 1]):
|
||||
units.append(exchange[index : index + 2])
|
||||
index += 2
|
||||
else:
|
||||
units.append([message])
|
||||
index += 1
|
||||
|
||||
bounded_units: list[list[dict[str, str]]] = []
|
||||
for unit in units:
|
||||
if _message_tokens(unit) <= max_tokens:
|
||||
bounded_units.append(unit)
|
||||
continue
|
||||
for message in unit:
|
||||
remaining = message["content"]
|
||||
while remaining:
|
||||
low, high = 0, len(remaining)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if _message_tokens([{**message, "content": remaining[:middle]}]) <= max_tokens:
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if low == 0:
|
||||
raise ValueError("Extraction token budget cannot fit a message")
|
||||
bounded_units.append([{**message, "content": remaining[:low]}])
|
||||
remaining = remaining[low:]
|
||||
|
||||
batches: list[list[dict[str, str]]] = []
|
||||
batch: list[dict[str, str]] = []
|
||||
for unit in bounded_units:
|
||||
candidate = [*batch, *unit]
|
||||
if batch and _message_tokens(candidate) > max_tokens:
|
||||
batches.append(batch)
|
||||
batch = list(unit)
|
||||
else:
|
||||
batch = candidate
|
||||
if batch:
|
||||
batches.append(batch)
|
||||
return batches
|
||||
Reference in New Issue
Block a user