Add Memory Gateway agent plugin
This commit is contained in:
187
plugins/memory-gateway-agent/scripts/cleanup_test_memories.py
Normal file
187
plugins/memory-gateway-agent/scripts/cleanup_test_memories.py
Normal file
@ -0,0 +1,187 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
PLUGIN_ROOT = Path(__file__).resolve().parents[1]
|
||||
if str(PLUGIN_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(PLUGIN_ROOT))
|
||||
|
||||
from memory_gateway_plugin.output import dumps_safe, short_id, summarize_data
|
||||
|
||||
|
||||
USER_ID = "test_user_memory_gateway_plugin"
|
||||
AGENT_ID = "test_hermes_memory_gateway_plugin"
|
||||
WORKSPACE_ID = "test_workspace_memory_gateway_plugin"
|
||||
SESSION_ID = "test_session_memory_gateway_plugin_001"
|
||||
|
||||
|
||||
def _assert_test_user(user_id: str) -> None:
|
||||
if not user_id.startswith("test_user_"):
|
||||
raise ValueError("cleanup_refuses_non_test_user")
|
||||
|
||||
|
||||
def _gateway_url() -> str:
|
||||
return os.environ.get("MEMORY_GATEWAY_URL", "http://127.0.0.1:1934").rstrip("/")
|
||||
|
||||
|
||||
def _api_key() -> str:
|
||||
return os.environ.get("MEMORY_GATEWAY_API_KEY", "")
|
||||
|
||||
|
||||
def _request(method: str, endpoint: str, payload: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if _api_key():
|
||||
headers["X-API-Key"] = _api_key()
|
||||
body = None if payload is None else json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||||
req = urllib.request.Request(_gateway_url() + endpoint, data=body, headers=headers, method=method)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=15) as response:
|
||||
raw = response.read().decode("utf-8")
|
||||
return {"ok": True, "status_code": getattr(response, "status", 200), "data": json.loads(raw) if raw else {}}
|
||||
except urllib.error.HTTPError as exc:
|
||||
return {"ok": False, "status_code": exc.code, "error": str(exc.reason)[:300]}
|
||||
except Exception as exc:
|
||||
return {"ok": False, "status_code": None, "error": str(exc)[:300]}
|
||||
|
||||
|
||||
def _search_candidates() -> dict[str, Any]:
|
||||
return _request(
|
||||
"POST",
|
||||
"/v1/memory/search",
|
||||
{
|
||||
"query": "integration_test plugin safe_to_delete Memory Gateway plugin",
|
||||
"user_id": USER_ID,
|
||||
"agent_id": AGENT_ID,
|
||||
"workspace_id": WORKSPACE_ID,
|
||||
"session_id": SESSION_ID,
|
||||
"tags": ["integration_test"],
|
||||
"limit": 100,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _audit_candidate_ids() -> list[str]:
|
||||
result = _request("GET", "/v1/audit?limit=1000")
|
||||
if not result.get("ok"):
|
||||
return []
|
||||
ids: list[str] = []
|
||||
for row in result.get("data") or []:
|
||||
if row.get("actor_user_id") != USER_ID:
|
||||
continue
|
||||
if row.get("actor_agent_id") not in {AGENT_ID, None, ""}:
|
||||
continue
|
||||
if row.get("target_type") == "memory" and row.get("action") in {"upsert_memory", "feedback:incorrect", "feedback:duplicate", "feedback:outdated"}:
|
||||
target_id = row.get("target_id")
|
||||
if target_id and target_id not in ids:
|
||||
ids.append(target_id)
|
||||
return ids
|
||||
|
||||
|
||||
def _memory_from_result(item: dict[str, Any]) -> dict[str, Any] | None:
|
||||
memory = item.get("memory")
|
||||
if isinstance(memory, dict) and memory.get("id"):
|
||||
return memory
|
||||
return None
|
||||
|
||||
|
||||
def _is_cleanup_candidate(memory: dict[str, Any]) -> bool:
|
||||
if memory.get("user_id") != USER_ID:
|
||||
return False
|
||||
tags = set(memory.get("tags") or [])
|
||||
return bool(tags.intersection({"integration_test", "safe_to_delete", "plugin"}))
|
||||
|
||||
|
||||
def _delete_memory(memory_id: str) -> dict[str, Any]:
|
||||
query = urllib.parse.urlencode({"user_id": USER_ID, "agent_id": AGENT_ID, "workspace_id": WORKSPACE_ID, "session_id": SESSION_ID})
|
||||
return _request("DELETE", f"/v1/memory/{urllib.parse.quote(memory_id)}?{query}")
|
||||
|
||||
|
||||
def _feedback_memory(memory_id: str) -> dict[str, Any]:
|
||||
return _request(
|
||||
"POST",
|
||||
f"/v1/memory/{urllib.parse.quote(memory_id)}/feedback",
|
||||
{
|
||||
"user_id": USER_ID,
|
||||
"agent_id": AGENT_ID,
|
||||
"workspace_id": WORKSPACE_ID,
|
||||
"session_id": SESSION_ID,
|
||||
"feedback": "incorrect",
|
||||
"comment": "cleanup marker for integration test memory",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def run(user_id: str = USER_ID) -> dict[str, Any]:
|
||||
_assert_test_user(user_id)
|
||||
if user_id != USER_ID:
|
||||
return {"ok": False, "error": "script_is_scoped_to_fixed_test_user", "user_id": user_id}
|
||||
|
||||
search = _search_candidates()
|
||||
if not search.get("ok"):
|
||||
return {"ok": False, "search": {"ok": False, "status_code": search.get("status_code"), "error": search.get("error")}, "deleted": 0, "feedback_marked": 0, "skipped": 0}
|
||||
|
||||
rows = (search.get("data") or {}).get("results") or []
|
||||
memory_ids = _audit_candidate_ids()
|
||||
deleted = 0
|
||||
feedback_marked = 0
|
||||
skipped = 0
|
||||
unable: list[dict[str, Any]] = []
|
||||
touched: list[str] = []
|
||||
|
||||
for item in rows:
|
||||
memory = _memory_from_result(item)
|
||||
if not memory or not _is_cleanup_candidate(memory):
|
||||
skipped += 1
|
||||
continue
|
||||
memory_id = memory["id"]
|
||||
if memory_id not in memory_ids:
|
||||
memory_ids.append(memory_id)
|
||||
|
||||
for memory_id in memory_ids:
|
||||
deletion = _delete_memory(memory_id)
|
||||
if deletion.get("ok"):
|
||||
deleted += 1
|
||||
touched.append(short_id(memory_id))
|
||||
continue
|
||||
if deletion.get("status_code") == 404:
|
||||
skipped += 1
|
||||
continue
|
||||
feedback = _feedback_memory(memory_id)
|
||||
if feedback.get("ok"):
|
||||
feedback_marked += 1
|
||||
touched.append(short_id(memory_id))
|
||||
else:
|
||||
unable.append({"memory_id": short_id(memory_id), "delete_status": deletion.get("status_code"), "feedback_status": feedback.get("status_code"), "reason": feedback.get("error") or deletion.get("error")})
|
||||
|
||||
return {
|
||||
"ok": not unable,
|
||||
"search": {"ok": True, "status_code": search.get("status_code"), "data": summarize_data(search.get("data"))},
|
||||
"deleted": deleted,
|
||||
"feedback_marked": feedback_marked,
|
||||
"skipped": skipped,
|
||||
"unable_count": len(unable),
|
||||
"unable": unable,
|
||||
"touched_memory_ids": touched,
|
||||
"limitation": "search API returns local MemoryRecord rows plus OpenViking context; cleanup only deletes local MemoryRecord rows for the fixed test user.",
|
||||
}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
try:
|
||||
result = run(os.environ.get("MEMORY_GATEWAY_CLEANUP_USER_ID", USER_ID))
|
||||
except ValueError as exc:
|
||||
result = {"ok": False, "error": str(exc), "deleted": 0, "feedback_marked": 0, "skipped": 0}
|
||||
print(dumps_safe(result))
|
||||
return 0 if result.get("ok") else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user