Fix episodic configuration is not set correctly (MemMachine#791)

Author: malatewang

src/memmachine/common/configuration/episodic_config.py

@@ -232,8 +232,16 @@ def merge(self, other: Self) -> EpisodicMemoryConf:
         # ---- Step 4: update nested configuration in the base result ----
         return EpisodicMemoryConf(
             session_key=merged.session_key,
-            metrics_factory_id=merged.metrics_factory_id,
+            metrics_factory_id=merged.metrics_factory_id
+            if merged.metrics_factory_id is not None
+            else "prometheus",
             short_term_memory=stm_merged,
             long_term_memory=ltm_merged,
-            enabled=merged.enabled,
+            long_term_memory_enabled=True
+            if merged.long_term_memory_enabled is None and ltm_merged is not None
+            else merged.long_term_memory_enabled,
+            short_term_memory_enabled=True
+            if merged.short_term_memory_enabled is None and stm_merged is not None
+            else merged.short_term_memory_enabled,
+            enabled=True if merged.enabled is None else merged.enabled,
         )

src/memmachine/main/memmachine.py

@@ -7,15 +7,17 @@
 from enum import Enum
 from typing import Any, Final, Protocol, cast
 
-from pydantic import BaseModel, InstanceOf
+from pydantic import BaseModel, InstanceOf, ValidationError
 
 from memmachine.common.configuration import Configuration
 from memmachine.common.configuration.episodic_config import (
     EpisodicMemoryConf,
-    LongTermMemoryConf,
-    ShortTermMemoryConf,
+    EpisodicMemoryConfPartial,
+    LongTermMemoryConfPartial,
+    ShortTermMemoryConfPartial,
 )
 from memmachine.common.episode_store import Episode, EpisodeEntry, EpisodeIdT
+from memmachine.common.errors import ConfigurationError
 from memmachine.common.filter.filter_parser import (
     And as FilterAnd,
 )
@@ -74,17 +76,62 @@ def __init__(
     ) -> None:
         """Create a MemMachine using the provided configuration."""
         self._conf = conf
-
         if resources is not None:
             self._resources = resources
         else:
             self._resources = ResourceManagerImpl(conf)
+        self._initialize_default_episodic_configuration()
+        self._started = False
+
+    def _initialize_default_episodic_configuration(self) -> None:
+        # initialize the default value for episodic memory configuration
+        # Can not put the logic into the data type
+        default_prompt = "Based on the following episodes: {episodes}, and the previous summary: {summary}, please update the summary. Keep it under {max_length} characters."
+        if self._conf.episodic_memory is None:
+            self._conf.episodic_memory = EpisodicMemoryConfPartial()
+            self._conf.episodic_memory.enabled = False
+        if self._conf.episodic_memory.long_term_memory is None:
+            self._conf.episodic_memory.long_term_memory = LongTermMemoryConfPartial()
+            self._conf.episodic_memory.long_term_memory_enabled = False
+        if self._conf.episodic_memory.short_term_memory is None:
+            self._conf.episodic_memory.short_term_memory = ShortTermMemoryConfPartial()
+            self._conf.episodic_memory.short_term_memory_enabled = False
+        if self._conf.episodic_memory.long_term_memory.embedder is None:
+            self._conf.episodic_memory.long_term_memory.embedder = (
+                self._conf.default_long_term_memory_embedder
+            )
+        if self._conf.episodic_memory.long_term_memory.reranker is None:
+            self._conf.episodic_memory.long_term_memory.reranker = (
+                self._conf.default_long_term_memory_reranker
+            )
+        if self._conf.episodic_memory.short_term_memory.llm_model is None:
+            self._conf.episodic_memory.short_term_memory.llm_model = "gpt-4.1"
+        if self._conf.episodic_memory.short_term_memory.summary_prompt_system is None:
+            self._conf.episodic_memory.short_term_memory.summary_prompt_system = (
+                "You are a helpful assistant."
+            )
+        if self._conf.episodic_memory.short_term_memory.summary_prompt_user is None:
+            self._conf.episodic_memory.short_term_memory.summary_prompt_user = (
+                default_prompt
+            )
+        if self._conf.episodic_memory.long_term_memory.vector_graph_store is None:
+            self._conf.episodic_memory.long_term_memory.vector_graph_store = (
+                "default_store"
+            )
 
     async def start(self) -> None:
+        if self._started:
+            return
+        self._started = True
+
         semantic_service = await self._resources.get_semantic_service()
         await semantic_service.start()
 
     async def stop(self) -> None:
+        if not self._started:
+            return
+        self._started = False
+
         semantic_service = await self._resources.get_semantic_service()
         await semantic_service.stop()
 
@@ -93,75 +140,39 @@ async def stop(self) -> None:
     def _with_default_episodic_memory_conf(
         self,
         *,
-        embedder_name: str | None = None,
-        reranker_name: str | None = None,
+        user_conf: EpisodicMemoryConfPartial | None = None,
         session_key: str,
     ) -> EpisodicMemoryConf:
         # Get default prompts from config, with fallbacks
-        short_term = self._conf.episodic_memory.short_term_memory
-        summary_prompt_system = (
-            short_term.summary_prompt_system
-            if short_term and short_term.summary_prompt_system
-            else "You are a helpful assistant."
-        )
-        summary_prompt_user = (
-            short_term.summary_prompt_user
-            if short_term and short_term.summary_prompt_user
-            else "Based on the following episodes: {episodes}, and the previous summary: {summary}, please update the summary. Keep it under {max_length} characters."
-        )
-
-        # Get default embedder and reranker from config
-        long_term = self._conf.episodic_memory.long_term_memory
-
-        if not embedder_name:
-            embedder_name = self._conf.default_long_term_memory_embedder
-        if not reranker_name:
-            reranker_name = self._conf.default_long_term_memory_reranker
-
-        self._conf.check_reranker(reranker_name)
-        self._conf.check_embedder(embedder_name)
-
-        target_vector_store = (
-            long_term.vector_graph_store
-            if long_term and long_term.vector_graph_store
-            else "default_store"
-        )
-
-        target_short_llm_model = (
-            short_term.llm_model if short_term and short_term.llm_model else "gpt-4.1"
-        )
-
-        return EpisodicMemoryConf(
-            session_key=session_key,
-            long_term_memory=LongTermMemoryConf(
-                session_id=session_key,
-                vector_graph_store=target_vector_store,
-                embedder=embedder_name,
-                reranker=reranker_name,
-            ),
-            short_term_memory=ShortTermMemoryConf(
-                session_key=session_key,
-                llm_model=target_short_llm_model,
-                summary_prompt_system=summary_prompt_system,
-                summary_prompt_user=summary_prompt_user,
-            ),
-            long_term_memory_enabled=True,
-            short_term_memory_enabled=True,
-            enabled=True,
-        )
+        try:
+            if user_conf is None:
+                user_conf = EpisodicMemoryConfPartial()
+            user_conf.session_key = session_key
+            episodic_conf = user_conf.merge(self._conf.episodic_memory)
+            if episodic_conf.long_term_memory is not None:
+                if episodic_conf.long_term_memory.embedder is not None:
+                    self._conf.check_embedder(episodic_conf.long_term_memory.embedder)
+                if episodic_conf.long_term_memory.reranker is not None:
+                    self._conf.check_reranker(episodic_conf.long_term_memory.reranker)
+        except ValidationError as e:
+            logger.exception(
+                "Faield to merge configuration: %s, %s",
+                str(user_conf),
+                str(self._conf.episodic_memory),
+            )
+            raise ConfigurationError("Failed to merge configuration") from e
+        return episodic_conf
 
     async def create_session(
         self,
         session_key: str,
         *,
         description: str = "",
-        embedder_name: str | None = None,
-        reranker_name: str | None = None,
+        user_conf: EpisodicMemoryConfPartial | None = None,
     ) -> SessionDataManager.SessionInfo:
         """Create a new session."""
         episodic_memory_conf = self._with_default_episodic_memory_conf(
-            embedder_name=embedder_name,
-            reranker_name=reranker_name,
+            user_conf=user_conf,
             session_key=session_key,
         )
 

src/memmachine/server/api_v2/router.py

@@ -23,6 +23,10 @@
     SearchMemoriesSpec,
     SearchResult,
 )
+from memmachine.common.configuration.episodic_config import (
+    EpisodicMemoryConfPartial,
+    LongTermMemoryConfPartial,
+)
 from memmachine.common.errors import (
     ConfigurationError,
     InvalidArgumentError,
@@ -50,11 +54,16 @@ async def create_project(
         project_id=spec.project_id,
     )
     try:
+        user_conf = EpisodicMemoryConfPartial(
+            long_term_memory=LongTermMemoryConfPartial(
+                embedder=spec.config.embedder if spec.config.embedder else None,
+                reranker=spec.config.reranker if spec.config.reranker else None,
+            )
+        )
         session = await memmachine.create_session(
             session_key=session_data.session_key,
             description=spec.description,
-            embedder_name=spec.config.embedder,
-            reranker_name=spec.config.reranker,
+            user_conf=user_conf,
         )
     except InvalidArgumentError as e:
         raise RestError(code=422, message="invalid argument: " + str(e)) from e

tests/memmachine/main/test_memmachine_mock.py

@@ -10,9 +10,9 @@
 
 from memmachine.common.configuration import (
     Configuration,
-    EpisodicMemoryConfPartial,
 )
 from memmachine.common.configuration.episodic_config import (
+    EpisodicMemoryConfPartial,
     LongTermMemoryConfPartial,
     ShortTermMemoryConfPartial,
 )
@@ -39,8 +39,9 @@ def session_key(self) -> str:  # pragma: no cover - trivial accessor
         return self._session_key
 
 
-@pytest.fixture
-def minimal_conf() -> Configuration:
+def _minimal_conf(
+    short_memory_enabled: bool = True, long_term_memory_enabled: bool = True
+) -> Configuration:
     """Provide the minimal subset of configuration accessed in tests."""
     mock_rerankers = MagicMock()
     mock_rerankers.contains_reranker.return_value = True
@@ -65,12 +66,24 @@ def minimal_conf() -> Configuration:
             embedder="default-embedder",
             reranker="default-reranker",
         ),
+        short_term_memory_enabled=short_memory_enabled,
+        long_term_memory_enabled=long_term_memory_enabled,
     )
     ret.default_long_term_memory_embedder = "default-embedder"
     ret.default_long_term_memory_reranker = "default-reranker"
     return ret
 
 
+@pytest.fixture
+def minimal_conf() -> Configuration:
+    return _minimal_conf()
+
+
+@pytest.fixture
+def minimal_conf_factory():
+    return _minimal_conf
+
+
 @pytest.fixture
 def patched_resource_manager(monkeypatch):
     """Replace :class:`ResourceManagerImpl` with a controllable double."""
@@ -122,6 +135,52 @@ def test_with_default_episodic_memory_conf_uses_fallbacks(
     )
 
 
+def test_with_default_short_conf_enable_status(
+    minimal_conf_factory, patched_resource_manager
+):
+    min_conf = minimal_conf_factory(
+        short_memory_enabled=False, long_term_memory_enabled=True
+    )
+    memmachine = MemMachine(min_conf, patched_resource_manager)
+    conf = memmachine._with_default_episodic_memory_conf(session_key="session-2")
+    assert min_conf.episodic_memory.short_term_memory_enabled is False
+    assert min_conf.episodic_memory.long_term_memory_enabled is True
+    assert conf.short_term_memory_enabled is False
+    assert conf.long_term_memory_enabled is True
+    user_conf = EpisodicMemoryConfPartial(
+        short_term_memory_enabled=True,
+        long_term_memory_enabled=False,
+    )
+    conf = memmachine._with_default_episodic_memory_conf(
+        session_key="session-2", user_conf=user_conf
+    )
+    assert conf.short_term_memory_enabled is True
+    assert conf.long_term_memory_enabled is False
+
+
+def test_with_default_long_conf_enable_status(
+    minimal_conf_factory, patched_resource_manager
+):
+    min_conf = minimal_conf_factory(
+        short_memory_enabled=True, long_term_memory_enabled=False
+    )
+    memmachine = MemMachine(min_conf, patched_resource_manager)
+    conf = memmachine._with_default_episodic_memory_conf(session_key="session-2")
+    assert min_conf.episodic_memory.short_term_memory_enabled is True
+    assert min_conf.episodic_memory.long_term_memory_enabled is False
+    assert conf.short_term_memory_enabled is True
+    assert conf.long_term_memory_enabled is False
+    user_conf = EpisodicMemoryConfPartial(
+        short_term_memory_enabled=False,
+        long_term_memory_enabled=True,
+    )
+    conf = memmachine._with_default_episodic_memory_conf(
+        session_key="session-2", user_conf=user_conf
+    )
+    assert conf.short_term_memory_enabled is False
+    assert conf.long_term_memory_enabled is True
+
+
 @pytest.mark.asyncio
 async def test_create_session_passes_generated_config(
     minimal_conf, patched_resource_manager
@@ -133,11 +192,16 @@ async def test_create_session_passes_generated_config(
 
     memmachine = MemMachine(minimal_conf, patched_resource_manager)
 
+    user_conf = EpisodicMemoryConfPartial(
+        long_term_memory=LongTermMemoryConfPartial(
+            embedder="custom-embed",
+            reranker="custom-reranker",
+        )
+    )
     await memmachine.create_session(
         "alpha",
         description="demo",
-        embedder_name="custom-embed",
-        reranker_name="custom-reranker",
+        user_conf=user_conf,
     )
 
     session_manager.create_new_session.assert_awaited_once()

tests/memmachine/server/api_v2/test_router.py

@@ -60,8 +60,8 @@ def test_create_project(client, mock_memmachine):
     call_args = mock_memmachine.create_session.call_args[1]
     assert call_args["session_key"] == "test_org/test_proj"
     assert call_args["description"] == "A test project"
-    assert call_args["embedder_name"] == "openai"
-    assert call_args["reranker_name"] == "cohere"
+    assert call_args["user_conf"].long_term_memory.embedder == "openai"
+    assert call_args["user_conf"].long_term_memory.reranker == "cohere"
 
     mock_memmachine.create_session.reset_mock()
     mock_memmachine.create_session.side_effect = InvalidArgumentError(