refactor(mcp): replace eager _kg with lazy per-path cache (#1136)
Swap the module-level KnowledgeGraph singleton for a lazy, per-path cache keyed by the resolved sqlite path. Import no longer creates a sqlite file as a side effect, and MCP servers started with --palace now route KG calls to the correct tenant when MEMPALACE_PALACE_PATH changes between calls, matching the per-call behavior of _get_client() on the ChromaDB side. Default-path behavior is preserved: without --palace at startup, KG stays on DEFAULT_KG_PATH regardless of env var. The "no --palace but env var set" case is #540's scope and is not changed here.
This commit is contained in:
+40
-12
@@ -46,6 +46,7 @@ import argparse # noqa: E402 (deferred until after stdio protection above)
|
|||||||
import json # noqa: E402
|
import json # noqa: E402
|
||||||
import logging # noqa: E402
|
import logging # noqa: E402
|
||||||
import hashlib # noqa: E402
|
import hashlib # noqa: E402
|
||||||
|
import threading # noqa: E402
|
||||||
import time # noqa: E402
|
import time # noqa: E402
|
||||||
from datetime import date, datetime # noqa: E402
|
from datetime import date, datetime # noqa: E402
|
||||||
from pathlib import Path # noqa: E402
|
from pathlib import Path # noqa: E402
|
||||||
@@ -78,7 +79,7 @@ from .palace_graph import ( # noqa: E402
|
|||||||
follow_tunnels,
|
follow_tunnels,
|
||||||
)
|
)
|
||||||
|
|
||||||
from .knowledge_graph import KnowledgeGraph # noqa: E402
|
from .knowledge_graph import KnowledgeGraph, DEFAULT_KG_PATH # noqa: E402
|
||||||
|
|
||||||
logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stderr)
|
logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stderr)
|
||||||
logger = logging.getLogger("mempalace_mcp")
|
logger = logging.getLogger("mempalace_mcp")
|
||||||
@@ -103,12 +104,39 @@ if _args.palace:
|
|||||||
os.environ["MEMPALACE_PALACE_PATH"] = os.path.abspath(_args.palace)
|
os.environ["MEMPALACE_PALACE_PATH"] = os.path.abspath(_args.palace)
|
||||||
|
|
||||||
_config = MempalaceConfig()
|
_config = MempalaceConfig()
|
||||||
# Only override KG path when --palace is explicitly provided; otherwise use
|
|
||||||
# KnowledgeGraph's default (~/.mempalace/knowledge_graph.sqlite3).
|
# Lazy per-path KG cache. Import no longer creates the sqlite file as a side
|
||||||
if _args.palace:
|
# effect (see issue #1136). The path is resolved on each tool call so that a
|
||||||
_kg = KnowledgeGraph(db_path=os.path.join(_config.palace_path, "knowledge_graph.sqlite3"))
|
# multi-tenant host rotating MEMPALACE_PALACE_PATH between calls routes each
|
||||||
else:
|
# call to the correct KG file, matching the per-call behavior of _get_client()
|
||||||
_kg = KnowledgeGraph()
|
# on the ChromaDB side.
|
||||||
|
_kg_by_path: dict[str, KnowledgeGraph] = {}
|
||||||
|
_kg_cache_lock = threading.Lock()
|
||||||
|
|
||||||
|
# Whether --palace was given at startup. Controls default-path resolution:
|
||||||
|
# with the flag, KG follows _config.palace_path per call; without it, KG stays
|
||||||
|
# on DEFAULT_KG_PATH regardless of env var (issue #540's territory, out of
|
||||||
|
# scope here).
|
||||||
|
_palace_flag_given: bool = bool(_args.palace)
|
||||||
|
|
||||||
|
|
||||||
|
def _resolve_kg_path() -> str:
|
||||||
|
if _palace_flag_given:
|
||||||
|
return os.path.join(_config.palace_path, "knowledge_graph.sqlite3")
|
||||||
|
return DEFAULT_KG_PATH
|
||||||
|
|
||||||
|
|
||||||
|
def _get_kg() -> KnowledgeGraph:
|
||||||
|
path = os.path.abspath(_resolve_kg_path())
|
||||||
|
kg = _kg_by_path.get(path)
|
||||||
|
if kg is not None:
|
||||||
|
return kg
|
||||||
|
with _kg_cache_lock:
|
||||||
|
kg = _kg_by_path.get(path)
|
||||||
|
if kg is None:
|
||||||
|
kg = KnowledgeGraph(db_path=path)
|
||||||
|
_kg_by_path[path] = kg
|
||||||
|
return kg
|
||||||
|
|
||||||
|
|
||||||
_client_cache = None
|
_client_cache = None
|
||||||
@@ -1063,7 +1091,7 @@ def tool_kg_query(entity: str, as_of: str = None, direction: str = "both"):
|
|||||||
return {"error": str(e)}
|
return {"error": str(e)}
|
||||||
if direction not in ("outgoing", "incoming", "both"):
|
if direction not in ("outgoing", "incoming", "both"):
|
||||||
return {"error": "direction must be 'outgoing', 'incoming', or 'both'"}
|
return {"error": "direction must be 'outgoing', 'incoming', or 'both'"}
|
||||||
results = _kg.query_entity(entity, as_of=as_of, direction=direction)
|
results = _get_kg().query_entity(entity, as_of=as_of, direction=direction)
|
||||||
return {"entity": entity, "as_of": as_of, "facts": results, "count": len(results)}
|
return {"entity": entity, "as_of": as_of, "facts": results, "count": len(results)}
|
||||||
|
|
||||||
|
|
||||||
@@ -1108,7 +1136,7 @@ def tool_kg_add(
|
|||||||
"source_drawer_id": source_drawer_id,
|
"source_drawer_id": source_drawer_id,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
triple_id = _kg.add_triple(
|
triple_id = _get_kg().add_triple(
|
||||||
subject,
|
subject,
|
||||||
predicate,
|
predicate,
|
||||||
object,
|
object,
|
||||||
@@ -1147,7 +1175,7 @@ def tool_kg_invalidate(subject: str, predicate: str, object: str, ended: str = N
|
|||||||
"ended": resolved_ended,
|
"ended": resolved_ended,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
_kg.invalidate(subject, predicate, object, ended=resolved_ended)
|
_get_kg().invalidate(subject, predicate, object, ended=resolved_ended)
|
||||||
return {
|
return {
|
||||||
"success": True,
|
"success": True,
|
||||||
"fact": f"{subject} → {predicate} → {object}",
|
"fact": f"{subject} → {predicate} → {object}",
|
||||||
@@ -1162,13 +1190,13 @@ def tool_kg_timeline(entity: str = None):
|
|||||||
entity = sanitize_kg_value(entity, "entity")
|
entity = sanitize_kg_value(entity, "entity")
|
||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
return {"error": str(e)}
|
return {"error": str(e)}
|
||||||
results = _kg.timeline(entity)
|
results = _get_kg().timeline(entity)
|
||||||
return {"entity": entity or "all", "timeline": results, "count": len(results)}
|
return {"entity": entity or "all", "timeline": results, "count": len(results)}
|
||||||
|
|
||||||
|
|
||||||
def tool_kg_stats():
|
def tool_kg_stats():
|
||||||
"""Knowledge graph overview: entities, triples, relationship types."""
|
"""Knowledge graph overview: entities, triples, relationship types."""
|
||||||
return _kg.stats()
|
return _get_kg().stats()
|
||||||
|
|
||||||
|
|
||||||
# ==================== AGENT DIARY ====================
|
# ==================== AGENT DIARY ====================
|
||||||
|
|||||||
Reference in New Issue
Block a user