Teknium
603 lines
29 KiB
Python
603 lines
29 KiB
Python
"""RetainDB memory plugin — MemoryProvider interface.
|
|
|
|
Cross-session memory via the RetainDB cloud API: durable SQLite write-behind queue, semantic
|
|
search + profile, context overlay, dialectic/agent self-model prefetch, shared file store tools.
|
|
|
|
Config (env vars, or config.yaml ``memory.retaindb`` for the non-secret ones): RETAINDB_API_KEY (required),
|
|
RETAINDB_BASE_URL (default https://api.retaindb.com), RETAINDB_PROJECT (optional; defaults to "default").
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import queue
|
|
import re
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
from contextlib import suppress
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Dict, List
|
|
from urllib.parse import quote
|
|
|
|
from agent.memory_provider import MemoryProvider
|
|
from agent.secret_scope import get_secret
|
|
from agent.file_safety import raise_if_read_blocked
|
|
from tools.registry import tool_error
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_DEFAULT_BASE_URL = "https://api.retaindb.com"
|
|
_ASYNC_SHUTDOWN = object()
|
|
_TEXT_EXTS = (".txt", ".md", ".json", ".csv", ".yaml", ".yml", ".xml", ".html")
|
|
|
|
|
|
def _load_retaindb_config() -> Dict[str, Any]:
|
|
"""``memory.retaindb`` block from config.yaml (empty on error): Dashboard-persisted base_url/project; api_key stays in scoped secrets."""
|
|
try:
|
|
from hermes_cli.config import load_config_readonly
|
|
|
|
provider_config = load_config_readonly().get("memory", {}).get("retaindb", {})
|
|
return dict(provider_config) if isinstance(provider_config, dict) else {}
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def _config_str(value: Any) -> str:
|
|
"""Stripped string for a config value, else ``""``."""
|
|
return value.strip() if isinstance(value, str) else ""
|
|
|
|
|
|
def _q(s: str) -> str:
|
|
return quote(s, safe="")
|
|
|
|
|
|
# ── Tool schemas ─────────────────────────────────────────────────────────────
|
|
|
|
def _schema(name: str, description: str, properties: dict | None = None, required: tuple = ()) -> dict:
|
|
return {
|
|
"name": name,
|
|
"description": description,
|
|
"parameters": {"type": "object", "properties": properties or {}, "required": list(required)},
|
|
}
|
|
|
|
|
|
def _prop(type_: str, description: str, **extra) -> dict:
|
|
return {"type": type_, **extra, "description": description}
|
|
|
|
|
|
def _s(description: str, **extra) -> dict:
|
|
return _prop("string", description, **extra)
|
|
|
|
|
|
PROFILE_SCHEMA = _schema(
|
|
"retaindb_profile", "Get the user's stable profile — preferences, facts, and patterns recalled from long-term memory.")
|
|
SEARCH_SCHEMA = _schema(
|
|
"retaindb_search", "Semantic search across stored memories. Returns ranked results with relevance scores.",
|
|
{"query": _s("What to search for."), "top_k": _prop("integer", "Max results (default: 8, max: 20).")}, ("query",))
|
|
CONTEXT_SCHEMA = _schema(
|
|
"retaindb_context", "Synthesized context block — what matters most for the current task, pulled from long-term memory.",
|
|
{"query": _s("Current task or question.")}, ("query",))
|
|
REMEMBER_SCHEMA = _schema(
|
|
"retaindb_remember", "Persist an explicit fact, preference, or decision to long-term memory.",
|
|
{"content": _s("The fact to remember."),
|
|
"memory_type": _s("Category (default: factual).", enum=["factual", "preference", "goal", "instruction", "event", "opinion"]),
|
|
"importance": _prop("number", "Importance 0-1 (default: 0.7).")}, ("content",))
|
|
FORGET_SCHEMA = _schema("retaindb_forget", "Delete a specific memory by ID.", {"memory_id": _s("Memory ID to delete.")}, ("memory_id",))
|
|
FILE_UPLOAD_SCHEMA = _schema(
|
|
"retaindb_upload_file", "Upload a file to the shared RetainDB file store. Returns an rdb:// URI any agent can reference.",
|
|
{"local_path": _s("Local file path to upload."), "remote_path": _s("Destination path, e.g. /reports/q1.pdf"),
|
|
"scope": _s("Access scope (default: PROJECT).", enum=["USER", "PROJECT", "ORG"]),
|
|
"ingest": _prop("boolean", "Also extract memories from file after upload (default: false).")}, ("local_path",))
|
|
FILE_LIST_SCHEMA = _schema(
|
|
"retaindb_list_files", "List files in the shared file store.",
|
|
{"prefix": _s("Path prefix to filter by, e.g. /reports/"), "limit": _prop("integer", "Max results (default: 50).")})
|
|
FILE_READ_SCHEMA = _schema(
|
|
"retaindb_read_file", "Read the text content of a stored file by its file ID.",
|
|
{"file_id": _s("File ID returned from upload or list.")}, ("file_id",))
|
|
FILE_INGEST_SCHEMA = _schema(
|
|
"retaindb_ingest_file", "Chunk, embed, and extract memories from a stored file. Makes its contents searchable.",
|
|
{"file_id": _s("File ID to ingest.")}, ("file_id",))
|
|
FILE_DELETE_SCHEMA = _schema("retaindb_delete_file", "Delete a stored file.", {"file_id": _s("File ID to delete.")}, ("file_id",))
|
|
_SCHEMAS = (
|
|
PROFILE_SCHEMA, SEARCH_SCHEMA, CONTEXT_SCHEMA, REMEMBER_SCHEMA, FORGET_SCHEMA,
|
|
FILE_UPLOAD_SCHEMA, FILE_LIST_SCHEMA, FILE_READ_SCHEMA, FILE_INGEST_SCHEMA, FILE_DELETE_SCHEMA,
|
|
)
|
|
|
|
|
|
# ── HTTP client ──────────────────────────────────────────────────────────────
|
|
|
|
class _Client:
|
|
def __init__(self, api_key: str, base_url: str, project: str):
|
|
self.api_key = api_key
|
|
self.base_url = re.sub(r"/+$", "", base_url)
|
|
self.project = project
|
|
|
|
def _headers(self, path: str, json_body: bool = True) -> dict:
|
|
token = self.api_key.replace("Bearer ", "").strip()
|
|
return {
|
|
"Authorization": f"Bearer {token}", "x-sdk-runtime": "hermes-plugin",
|
|
**({"Content-Type": "application/json"} if json_body else {}),
|
|
# memory/context routes also accept the key as X-API-Key
|
|
**({"X-API-Key": token} if path.startswith(("/v1/memory", "/v1/context")) else {}),
|
|
}
|
|
|
|
def request(self, method: str, path: str, *, params=None, json_body=None, timeout: float = 8.0) -> Any:
|
|
import requests
|
|
method = method.upper()
|
|
resp = requests.request(
|
|
method, f"{self.base_url}{path}", params=params, json=json_body if method not in {"GET", "DELETE"} else None,
|
|
headers=self._headers(path), timeout=timeout,
|
|
)
|
|
try:
|
|
payload = resp.json()
|
|
except Exception:
|
|
payload = resp.text
|
|
if not resp.ok:
|
|
msg = str(payload.get("message") or payload.get("error") or "") if isinstance(payload, dict) else ""
|
|
raise RuntimeError(f"RetainDB {method} {path} failed ({resp.status_code}): {msg or payload}")
|
|
return payload
|
|
|
|
@staticmethod
|
|
def _with_fallback(primary: Callable[[], dict], fallback: Callable[[], dict]) -> dict:
|
|
"""Try the current API route; on any error retry via the legacy route."""
|
|
try:
|
|
return primary()
|
|
except Exception:
|
|
return fallback()
|
|
|
|
def _scoped(self, user_id: str, session_id: str, **extra) -> dict:
|
|
return {"project": self.project, "user_id": user_id, "session_id": session_id, **extra}
|
|
|
|
# Memory
|
|
|
|
def query_context(self, user_id: str, session_id: str, query: str, max_tokens: int = 1200) -> dict:
|
|
body = self._scoped(user_id, session_id, query=query, include_memories=True, max_tokens=max_tokens)
|
|
return self.request("POST", "/v1/context/query", json_body=body)
|
|
|
|
def search(self, user_id: str, session_id: str, query: str, top_k: int = 8) -> dict:
|
|
body = self._scoped(user_id, session_id, query=query, top_k=top_k, include_pending=True)
|
|
return self.request("POST", "/v1/memory/search", json_body=body)
|
|
|
|
def get_profile(self, user_id: str) -> dict:
|
|
return self._with_fallback(
|
|
lambda: self.request("GET", f"/v1/memory/profile/{_q(user_id)}", params={"project": self.project, "include_pending": "true"}),
|
|
lambda: self.request("GET", "/v1/memories", params={"project": self.project, "user_id": user_id, "limit": "200"}),
|
|
)
|
|
|
|
def add_memory(self, user_id: str, session_id: str, content: str, memory_type: str = "factual", importance: float = 0.7) -> dict:
|
|
body = self._scoped(user_id, session_id, content=content, memory_type=memory_type, importance=importance)
|
|
return self._with_fallback(
|
|
lambda: self.request("POST", "/v1/memory", json_body={**body, "write_mode": "sync"}, timeout=5.0),
|
|
lambda: self.request("POST", "/v1/memories", json_body=body, timeout=5.0),
|
|
)
|
|
|
|
def delete_memory(self, memory_id: str) -> dict:
|
|
return self._with_fallback(
|
|
lambda: self.request("DELETE", f"/v1/memory/{_q(memory_id)}", timeout=5.0),
|
|
lambda: self.request("DELETE", f"/v1/memories/{_q(memory_id)}", timeout=5.0),
|
|
)
|
|
|
|
def ingest_session(self, user_id: str, session_id: str, messages: list, timeout: float = 15.0) -> dict:
|
|
body = self._scoped(user_id, session_id, messages=messages, write_mode="sync")
|
|
return self.request("POST", "/v1/memory/ingest/session", json_body=body, timeout=timeout)
|
|
|
|
def ask_user(self, user_id: str, query: str, reasoning_level: str = "low") -> dict:
|
|
body = {"project": self.project, "query": query, "reasoning_level": reasoning_level}
|
|
return self.request("POST", f"/v1/memory/profile/{_q(user_id)}/ask", json_body=body, timeout=8.0)
|
|
|
|
def get_agent_model(self, agent_id: str) -> dict:
|
|
return self.request("GET", f"/v1/memory/agent/{_q(agent_id)}/model", params={"project": self.project}, timeout=4.0)
|
|
|
|
def seed_agent_identity(self, agent_id: str, content: str, source: str = "soul_md") -> dict:
|
|
body = {"project": self.project, "content": content, "source": source}
|
|
return self.request("POST", f"/v1/memory/agent/{_q(agent_id)}/seed", json_body=body, timeout=20.0)
|
|
|
|
# Files
|
|
|
|
def _raw(self, method: str, path: str, **kwargs) -> Any:
|
|
"""Non-JSON request (multipart upload / binary download); raises on HTTP error."""
|
|
import requests
|
|
resp = requests.request(method, f"{self.base_url}{path}", headers=self._headers(path, json_body=False), timeout=30, **kwargs)
|
|
resp.raise_for_status()
|
|
return resp
|
|
|
|
def upload_file(self, data: bytes, filename: str, remote_path: str, mime_type: str, scope: str, project_id: str | None) -> dict:
|
|
import io
|
|
fields = {"path": remote_path, "scope": scope.upper(), **({"project_id": project_id} if project_id else {})}
|
|
return self._raw("POST", "/v1/files", files={"file": (filename, io.BytesIO(data), mime_type)}, data=fields).json()
|
|
|
|
def list_files(self, prefix: str | None = None, limit: int = 50) -> dict:
|
|
return self.request("GET", "/v1/files", params={"limit": limit, **({"prefix": prefix} if prefix else {})})
|
|
|
|
def get_file(self, file_id: str) -> dict:
|
|
return self.request("GET", f"/v1/files/{_q(file_id)}")
|
|
|
|
def read_file_content(self, file_id: str) -> bytes:
|
|
return self._raw("GET", f"/v1/files/{_q(file_id)}/content", allow_redirects=True).content
|
|
|
|
def ingest_file(self, file_id: str, user_id: str | None = None, agent_id: str | None = None) -> dict:
|
|
body = {k: v for k, v in (("user_id", user_id), ("agent_id", agent_id)) if v}
|
|
return self.request("POST", f"/v1/files/{_q(file_id)}/ingest", json_body=body, timeout=60.0)
|
|
|
|
def delete_file(self, file_id: str) -> dict:
|
|
return self.request("DELETE", f"/v1/files/{_q(file_id)}", timeout=5.0)
|
|
|
|
|
|
# ── Durable write-behind queue ───────────────────────────────────────────────
|
|
|
|
class _WriteQueue:
|
|
"""SQLite-backed async write queue. Survives crashes — pending rows replay on startup."""
|
|
|
|
def __init__(self, client: _Client, db_path: Path):
|
|
self._client = client
|
|
self._db_path = db_path
|
|
self._q: queue.Queue = queue.Queue()
|
|
self._thread = threading.Thread(target=self._loop, name="retaindb-writer", daemon=True)
|
|
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._local = threading.local() # one cached connection per thread
|
|
self._connections: set[sqlite3.Connection] = set()
|
|
self._connections_lock = threading.Lock()
|
|
self._shutdown_lock = threading.Lock()
|
|
self._shutdown = False
|
|
conn = self._execute(
|
|
"CREATE TABLE IF NOT EXISTS pending (id INTEGER PRIMARY KEY AUTOINCREMENT, user_id TEXT, "
|
|
"session_id TEXT, messages_json TEXT, created_at TEXT, last_error TEXT)"
|
|
).connection
|
|
self._thread.start()
|
|
# Replay any rows left from a previous crash
|
|
for row_id, user_id, session_id, msgs_json in conn.execute(
|
|
"SELECT id, user_id, session_id, messages_json FROM pending ORDER BY id ASC LIMIT 200"
|
|
).fetchall():
|
|
self._q.put((row_id, user_id, session_id, json.loads(msgs_json)))
|
|
|
|
def _get_conn(self) -> sqlite3.Connection:
|
|
"""Return a cached connection for the current thread."""
|
|
conn = getattr(self._local, "conn", None)
|
|
if conn is None:
|
|
conn = sqlite3.connect(str(self._db_path), timeout=30, check_same_thread=False)
|
|
conn.row_factory = sqlite3.Row
|
|
self._local.conn = conn
|
|
with self._connections_lock:
|
|
self._connections.add(conn)
|
|
return conn
|
|
|
|
def _execute(self, sql: str, params: tuple = ()) -> sqlite3.Cursor:
|
|
"""Execute + commit on this thread's connection."""
|
|
cur = self._get_conn().execute(sql, params)
|
|
cur.connection.commit()
|
|
return cur
|
|
|
|
def _close_thread_conn(self) -> None:
|
|
conn = getattr(self._local, "conn", None)
|
|
if conn is None:
|
|
return
|
|
self._local.conn = None
|
|
with self._connections_lock:
|
|
self._connections.discard(conn)
|
|
with suppress(Exception):
|
|
conn.close()
|
|
|
|
def enqueue(self, user_id: str, session_id: str, messages: list) -> None:
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
with self._shutdown_lock:
|
|
if self._shutdown:
|
|
return
|
|
cur = self._execute(
|
|
"INSERT INTO pending (user_id, session_id, messages_json, created_at) VALUES (?,?,?,?)",
|
|
(user_id, session_id, json.dumps(messages, ensure_ascii=False), now),
|
|
)
|
|
self._q.put((cur.lastrowid, user_id, session_id, messages))
|
|
|
|
def _flush_row(self, row_id: int, user_id: str, session_id: str, messages: list) -> None:
|
|
try:
|
|
self._client.ingest_session(user_id, session_id, messages)
|
|
self._execute("DELETE FROM pending WHERE id = ?", (row_id,))
|
|
except Exception as exc:
|
|
logger.warning("RetainDB ingest failed (will retry): %s", exc)
|
|
self._execute("UPDATE pending SET last_error = ? WHERE id = ?", (str(exc), row_id))
|
|
time.sleep(2)
|
|
|
|
def _loop(self) -> None:
|
|
try:
|
|
while (item := self._q.get()) is not _ASYNC_SHUTDOWN:
|
|
try:
|
|
self._flush_row(*item)
|
|
except Exception as exc:
|
|
logger.error("RetainDB writer error: %s", exc)
|
|
finally:
|
|
self._close_thread_conn() # sqlite3 connections must close on their owning thread
|
|
|
|
def shutdown(self) -> None:
|
|
with self._shutdown_lock:
|
|
if self._shutdown:
|
|
return
|
|
self._shutdown = True
|
|
self._q.put(_ASYNC_SHUTDOWN)
|
|
self._close_thread_conn() # caller thread owns the connection opened in __init__
|
|
self._thread.join(timeout=10)
|
|
if not self._thread.is_alive():
|
|
# Executor workers that already exited may have left tracked handles;
|
|
# check_same_thread=False lets shutdown close them deterministically.
|
|
with self._connections_lock:
|
|
connections, self._connections = list(self._connections), set()
|
|
for conn in connections:
|
|
with suppress(Exception):
|
|
conn.close()
|
|
|
|
|
|
# ── Overlay formatter ────────────────────────────────────────────────────────
|
|
|
|
def _compact(s: str) -> str:
|
|
return re.sub(r"\s+", " ", str(s or "")).strip()[:320]
|
|
|
|
|
|
def _norm(s: str) -> str:
|
|
return re.sub(r"[^a-z0-9 ]", "", _compact(s).lower())
|
|
|
|
|
|
def _build_overlay(profile: dict, query_result: dict, local_entries: list[str] | None = None) -> str:
|
|
"""Profile + query memories (5 each, compacted, deduped against each other and *local_entries*)."""
|
|
seen = {_norm(e) for e in (local_entries or []) if _norm(e)}
|
|
|
|
def _dedupe(items) -> list[str]:
|
|
out: list[str] = []
|
|
for m in list(items or [])[:5]:
|
|
c = _compact((m or {}).get("content") or "")
|
|
if c and _norm(c) not in seen:
|
|
seen.add(_norm(c))
|
|
out.append(c)
|
|
return out
|
|
|
|
profile_items = _dedupe((profile or {}).get("memories"))
|
|
query_items = _dedupe((query_result or {}).get("results"))
|
|
if not profile_items and not query_items:
|
|
return ""
|
|
return "\n".join(
|
|
["[RetainDB Context]", "Profile:"] + ([f"- {i}" for i in profile_items] or ["- None"])
|
|
+ ["Relevant memories:"] + ([f"- {i}" for i in query_items] or ["- None"])
|
|
)
|
|
|
|
|
|
# ── Provider ─────────────────────────────────────────────────────────────────
|
|
|
|
# Agent self-model keys -> prefetch line formatter, in display order.
|
|
_AGENT_MODEL_FIELDS = (
|
|
("persona", lambda v: f"Persona: {v}"),
|
|
("persistent_instructions", lambda v: "Instructions:\n" + "\n".join(f"- {i}" for i in v)),
|
|
("working_style", lambda v: f"Working style: {v}"),
|
|
)
|
|
|
|
|
|
class RetainDBMemoryProvider(MemoryProvider):
|
|
"""RetainDB cloud memory — durable queue, semantic search, dialectic synthesis, shared files."""
|
|
|
|
def __init__(self):
|
|
self._client: _Client | None = None
|
|
self._queue: _WriteQueue | None = None
|
|
self._user_id, self._session_id, self._agent_id = "default", "", "hermes"
|
|
self._lock = threading.Lock()
|
|
# Prefetch caches + thread tracking (prevents accumulation on rapid calls)
|
|
self._context_result = self._dialectic_result = ""
|
|
self._agent_model: dict = {}
|
|
self._prefetch_threads: list[threading.Thread] = []
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "retaindb"
|
|
|
|
def is_available(self) -> bool:
|
|
return bool(get_secret("RETAINDB_API_KEY"))
|
|
|
|
def get_config_schema(self) -> List[Dict[str, Any]]:
|
|
return [
|
|
{"key": "api_key", "description": "RetainDB API key", "secret": True, "required": True, "env_var": "RETAINDB_API_KEY", "url": "https://retaindb.com"},
|
|
{"key": "base_url", "description": "API endpoint", "default": _DEFAULT_BASE_URL},
|
|
{"key": "project", "description": "Project identifier (optional — uses 'default' project if not set)", "default": ""},
|
|
]
|
|
|
|
def initialize(self, session_id: str, **kwargs) -> None:
|
|
# Non-secret fields resolve env -> config.yaml (written by the Dashboard) -> default.
|
|
provider_config = _load_retaindb_config()
|
|
base_url = re.sub(r"/+$", "", os.environ.get("RETAINDB_BASE_URL") or _config_str(provider_config.get("base_url")) or _DEFAULT_BASE_URL)
|
|
# Project: RETAINDB_PROJECT > config.yaml > hermes-<profile> > "default" (API auto-creates "default").
|
|
project = os.environ.get("RETAINDB_PROJECT") or _config_str(provider_config.get("project"))
|
|
if not project:
|
|
profile_name = os.path.basename(str(kwargs.get("hermes_home", "")))
|
|
project = f"hermes-{profile_name}" if profile_name not in {"", ".hermes"} else "default"
|
|
|
|
self._client = _Client(get_secret("RETAINDB_API_KEY", "") or "", base_url, project)
|
|
self._session_id = session_id
|
|
self._user_id = kwargs.get("user_id", "default") or "default"
|
|
self._agent_id = kwargs.get("agent_id", "hermes") or "hermes"
|
|
|
|
from hermes_constants import get_hermes_home
|
|
hermes_home_path = get_hermes_home()
|
|
self._queue = _WriteQueue(self._client, hermes_home_path / "retaindb_queue.db")
|
|
# Seed agent identity from SOUL.md in background
|
|
soul_path = hermes_home_path / "SOUL.md"
|
|
soul_content = soul_path.read_text(encoding="utf-8", errors="replace").strip() if soul_path.exists() else ""
|
|
if soul_content:
|
|
threading.Thread(target=self._seed_soul, args=(soul_content,), name="retaindb-soul-seed", daemon=True).start()
|
|
|
|
def _seed_soul(self, content: str) -> None:
|
|
try:
|
|
self._client.seed_agent_identity(self._agent_id, content, source="soul_md")
|
|
except Exception as exc:
|
|
logger.debug("RetainDB soul seed failed: %s", exc)
|
|
|
|
def system_prompt_block(self) -> str:
|
|
project = self._client.project if self._client else "retaindb"
|
|
return (
|
|
f"# RetainDB Memory\nActive. Project: {project}.\n"
|
|
"Use retaindb_search to find memories, retaindb_remember to store facts, "
|
|
"retaindb_profile for a user overview, retaindb_context for current-task context."
|
|
)
|
|
|
|
# Background prefetch (fires at turn-end, consumed next turn-start)
|
|
|
|
def queue_prefetch(self, query: str, *, session_id: str = "") -> None:
|
|
"""Fire context + dialectic + agent model prefetches in background."""
|
|
if not self._client:
|
|
return
|
|
# Wait for the previous batch so threads don't accumulate on rapid turns.
|
|
for t in self._prefetch_threads:
|
|
t.join(timeout=2.0)
|
|
if any(t.is_alive() for t in self._prefetch_threads):
|
|
logger.debug("RetainDB prefetch still running; skipping new batch")
|
|
return
|
|
jobs = (
|
|
("retaindb-ctx", "context", lambda: ("_context_result", self._context_overlay(query)["context"])),
|
|
("retaindb-dialectic", "dialectic", lambda: self._fetch_dialectic(query)),
|
|
("retaindb-agent-model", "agent model", self._fetch_agent_model),
|
|
)
|
|
threads = [threading.Thread(target=self._store, args=(label, fetch), name=name, daemon=True) for name, label, fetch in jobs]
|
|
self._prefetch_threads = threads
|
|
for t in threads:
|
|
t.start()
|
|
|
|
def _context_overlay(self, query: str) -> dict:
|
|
query_result = self._client.query_context(self._user_id, self._session_id, query)
|
|
profile = self._client.get_profile(self._user_id)
|
|
return {"context": _build_overlay(profile, query_result), "raw": query_result}
|
|
|
|
def _fetch_dialectic(self, query: str) -> tuple[str, str | None]:
|
|
result = self._client.ask_user(self._user_id, query, reasoning_level=self._reasoning_level(query))
|
|
return "_dialectic_result", str(result.get("answer") or "") or None
|
|
|
|
def _fetch_agent_model(self) -> tuple[str, dict | None]:
|
|
model = self._client.get_agent_model(self._agent_id)
|
|
return "_agent_model", model if model.get("memory_count", 0) > 0 else None
|
|
|
|
def _store(self, label: str, fetch: Callable[[], tuple[str, Any]]) -> None:
|
|
"""Run one prefetch job; store (attr, value) under the lock unless value is None; log failures at debug."""
|
|
try:
|
|
attr, value = fetch()
|
|
if value is not None:
|
|
with self._lock:
|
|
setattr(self, attr, value)
|
|
except Exception as exc:
|
|
logger.debug("RetainDB %s prefetch failed: %s", label, exc)
|
|
|
|
@staticmethod
|
|
def _reasoning_level(query: str) -> str:
|
|
n = len(query)
|
|
return "low" if n < 120 else "medium" if n < 400 else "high"
|
|
|
|
def prefetch(self, query: str, *, session_id: str = "") -> str:
|
|
"""Consume prefetched results and return them as a context block."""
|
|
with self._lock:
|
|
context, dialectic, agent_model = self._context_result, self._dialectic_result, self._agent_model
|
|
self._context_result = self._dialectic_result = ""
|
|
self._agent_model = {}
|
|
parts = [context] if context else []
|
|
if dialectic:
|
|
parts.append(f"[RetainDB User Synthesis]\n{dialectic}")
|
|
if agent_model.get("memory_count", 0) > 0:
|
|
model_lines = [fmt(agent_model[k]) for k, fmt in _AGENT_MODEL_FIELDS if agent_model.get(k)]
|
|
if model_lines:
|
|
parts.append("[RetainDB Agent Self-Model]\n" + "\n".join(model_lines))
|
|
return "\n\n".join(parts)
|
|
|
|
def sync_turn(self, user_content: str, assistant_content: str, *, session_id: str = "") -> None:
|
|
"""Queue turn for async ingest. Returns immediately."""
|
|
if not self._queue or not user_content:
|
|
return
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
self._queue.enqueue(self._user_id, session_id or self._session_id, [
|
|
{"role": "user", "content": user_content, "timestamp": now},
|
|
{"role": "assistant", "content": assistant_content, "timestamp": now},
|
|
])
|
|
|
|
def get_tool_schemas(self) -> List[Dict[str, Any]]:
|
|
return list(_SCHEMAS)
|
|
|
|
def handle_tool_call(self, tool_name: str, args: dict, **kwargs) -> str:
|
|
if not self._client:
|
|
return tool_error("RetainDB not initialized")
|
|
try:
|
|
return json.dumps(self._dispatch(tool_name, args))
|
|
except Exception as exc:
|
|
return tool_error(str(exc))
|
|
|
|
def _dispatch(self, tool_name: str, args: dict) -> Any:
|
|
entry = _TOOLS.get(tool_name)
|
|
if entry is None:
|
|
return {"error": f"Unknown tool: {tool_name}"}
|
|
required, handler = entry
|
|
value = args.get(required, "") if required else None
|
|
if required and not value:
|
|
return {"error": f"{required} is required"}
|
|
return handler(self, args, value)
|
|
|
|
def _tool_upload_file(self, args: dict, local_path: str) -> Any:
|
|
path_obj = Path(local_path)
|
|
if not path_obj.exists():
|
|
return {"error": f"File not found: {local_path}"}
|
|
try:
|
|
raise_if_read_blocked(str(path_obj))
|
|
except ValueError as exc:
|
|
return {"error": str(exc)}
|
|
import mimetypes
|
|
mime = mimetypes.guess_type(path_obj.name)[0] or "application/octet-stream"
|
|
result = self._client.upload_file(path_obj.read_bytes(), path_obj.name, args.get("remote_path") or f"/{path_obj.name}",
|
|
mime, args.get("scope", "PROJECT"), None)
|
|
if args.get("ingest") and result.get("file", {}).get("id"):
|
|
result["ingest"] = self._ingest(result["file"]["id"])
|
|
return result
|
|
|
|
def _tool_read_file(self, args: dict, file_id: str) -> Any:
|
|
file_info = self._client.get_file(file_id).get("file") or {}
|
|
mime = (file_info.get("mime_type") or "").lower()
|
|
raw = self._client.read_file_content(file_id)
|
|
out = {"file_id": file_id, "rdb_uri": file_info.get("rdb_uri"), "name": file_info.get("name")}
|
|
if not (mime.startswith("text/") or file_info.get("name", "").endswith(_TEXT_EXTS)):
|
|
return {**out, "content": None, "note": "Binary file — use retaindb_ingest_file to extract text into memory."}
|
|
text = raw.decode("utf-8", errors="replace")
|
|
return {**out, "content": text[:32000], "truncated": len(text) > 32000}
|
|
|
|
def _ingest(self, file_id: str) -> Any:
|
|
return self._client.ingest_file(file_id, user_id=self._user_id, agent_id=self._agent_id)
|
|
|
|
def on_memory_write(self, action: str, target: str, content: str) -> None:
|
|
"""Mirror built-in memory writes to RetainDB."""
|
|
if action != "add" or not content or not self._client:
|
|
return
|
|
try:
|
|
memory_type = "preference" if target == "user" else "factual"
|
|
self._client.add_memory(self._user_id, self._session_id, content, memory_type=memory_type)
|
|
except Exception as exc:
|
|
logger.debug("RetainDB memory mirror failed: %s", exc)
|
|
|
|
def shutdown(self) -> None:
|
|
for t in self._prefetch_threads:
|
|
t.join(timeout=3.0)
|
|
self._prefetch_threads = []
|
|
queue_obj, self._queue, self._client = self._queue, None, None
|
|
if queue_obj:
|
|
queue_obj.shutdown()
|
|
|
|
|
|
# tool name -> (required arg or None, handler(provider, args, required_value)); missing arg -> "<arg> is required"
|
|
_TOOLS: Dict[str, tuple[str | None, Callable[..., Any]]] = {
|
|
"retaindb_profile": (None, lambda p, a, _: p._client.get_profile(p._user_id)),
|
|
"retaindb_search": ("query", lambda p, a, q: p._client.search(p._user_id, p._session_id, q, top_k=min(int(a.get("top_k", 8)), 20))),
|
|
"retaindb_context": ("query", lambda p, a, q: p._context_overlay(q)),
|
|
"retaindb_remember": ("content", lambda p, a, c: p._client.add_memory(
|
|
p._user_id, p._session_id, c, memory_type=a.get("memory_type", "factual"), importance=float(a.get("importance", 0.7)))),
|
|
"retaindb_forget": ("memory_id", lambda p, a, m: p._client.delete_memory(m)),
|
|
"retaindb_upload_file": ("local_path", RetainDBMemoryProvider._tool_upload_file),
|
|
"retaindb_list_files": (None, lambda p, a, _: p._client.list_files(prefix=a.get("prefix"), limit=int(a.get("limit", 50)))),
|
|
"retaindb_read_file": ("file_id", RetainDBMemoryProvider._tool_read_file),
|
|
"retaindb_ingest_file": ("file_id", lambda p, a, f: p._ingest(f)),
|
|
"retaindb_delete_file": ("file_id", lambda p, a, f: p._client.delete_file(f)),
|
|
}
|
|
|
|
|
|
def register(ctx) -> None:
|
|
"""Register RetainDB as a memory provider plugin."""
|
|
ctx.register_memory_provider(RetainDBMemoryProvider())
|