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:
Igor Lins e Silva
2026-04-07 17:19:53 -03:00
parent 71736a3f4f
commit a67b00d7c7
2 changed files with 25 additions and 4 deletions
+12 -4
View File
@@ -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
+13
View File
@@ -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.