perf: cache ChromaDB PersistentClient instead of re-instantiating per call
The MCP server previously created a new PersistentClient on every tool call via _get_collection(). This incurs HNSW index loading overhead on each request. Cache the client and collection at module level. The cache resets naturally on process restart (MCP runs as a subprocess). Also adds a _reset_mcp_cache fixture to conftest.py for test isolation. Includes test infrastructure from PR #131. 92 tests pass.
This commit is contained in:
+12
-4
@@ -39,13 +39,21 @@ logger = logging.getLogger("mempalace_mcp")
|
|||||||
_config = MempalaceConfig()
|
_config = MempalaceConfig()
|
||||||
|
|
||||||
|
|
||||||
|
_client_cache = None
|
||||||
|
_collection_cache = None
|
||||||
|
|
||||||
|
|
||||||
def _get_collection(create=False):
|
def _get_collection(create=False):
|
||||||
"""Return the ChromaDB collection, or None on failure."""
|
"""Return the ChromaDB collection, caching the client between calls."""
|
||||||
|
global _client_cache, _collection_cache
|
||||||
try:
|
try:
|
||||||
client = chromadb.PersistentClient(path=_config.palace_path)
|
if _client_cache is None:
|
||||||
|
_client_cache = chromadb.PersistentClient(path=_config.palace_path)
|
||||||
if create:
|
if create:
|
||||||
return client.get_or_create_collection(_config.collection_name)
|
_collection_cache = _client_cache.get_or_create_collection(_config.collection_name)
|
||||||
return client.get_collection(_config.collection_name)
|
elif _collection_cache is None:
|
||||||
|
_collection_cache = _client_cache.get_collection(_config.collection_name)
|
||||||
|
return _collection_cache
|
||||||
except Exception:
|
except Exception:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|||||||
@@ -34,6 +34,19 @@ from mempalace.config import MempalaceConfig # noqa: E402
|
|||||||
from mempalace.knowledge_graph import KnowledgeGraph # noqa: E402
|
from mempalace.knowledge_graph import KnowledgeGraph # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(autouse=True)
|
||||||
|
def _reset_mcp_cache():
|
||||||
|
"""Reset the MCP server's cached ChromaDB client/collection between tests."""
|
||||||
|
yield
|
||||||
|
try:
|
||||||
|
from mempalace import mcp_server
|
||||||
|
|
||||||
|
mcp_server._client_cache = None
|
||||||
|
mcp_server._collection_cache = None
|
||||||
|
except (ImportError, AttributeError):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(scope="session", autouse=True)
|
@pytest.fixture(scope="session", autouse=True)
|
||||||
def _isolate_home():
|
def _isolate_home():
|
||||||
"""Ensure HOME points to a temp dir for the entire test session.
|
"""Ensure HOME points to a temp dir for the entire test session.
|
||||||
|
|||||||
Reference in New Issue
Block a user