Create non-existent projects when adding semantic memories

Author: mwqgithub

src/memmachine/common/session_manager/session_data_manager_sql_impl.py

@@ -20,7 +20,9 @@
     update,
 )
 from sqlalchemy.dialects.postgresql import JSONB
-from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, async_sessionmaker
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker
+from sqlalchemy.engine import Connection
 from sqlalchemy.orm import (
     DeclarativeBase,
     Mapped,
@@ -94,7 +96,7 @@ def __init__(self, engine: AsyncEngine, schema: str | None = None) -> None:
     async def create_tables(self) -> None:
         """Create the necessary tables in the database."""
 
-        def _check_migration_needed(conn: AsyncConnection) -> bool:
+        def _check_migration_needed(conn: Connection) -> bool:
             inspector = inspect(conn)
             schema = self.SessionConfig.__table__.schema
             table_name = "sessions"
@@ -196,15 +198,6 @@ async def create_new_session(
             param_data = param.__dict__
 
         async with self._async_session() as dbsession:
-            # Query for an existing session with the same ID
-            sessions = await dbsession.execute(
-                select(self.SessionConfig).where(
-                    self.SessionConfig.session_key == session_key,
-                ),
-            )
-            session = sessions.first()
-            if session is not None:
-                raise SessionAlreadyExistsError(session_key)
             # create a new entry
             new_session = self.SessionConfig(
                 session_key=session_key,
@@ -215,7 +208,11 @@ async def create_new_session(
                 user_metadata=metadata,
             )
             dbsession.add(new_session)
-            await dbsession.commit()
+            try:
+                await dbsession.commit()
+            except IntegrityError as exc:
+                await dbsession.rollback()
+                raise SessionAlreadyExistsError(session_key) from exc
 
     async def delete_session(self, session_key: str) -> None:
         """Delete a session and its related data from the database."""

src/memmachine/main/memmachine.py

@@ -17,7 +17,11 @@
     ShortTermMemoryConfPartial,
 )
 from memmachine.common.episode_store import Episode, EpisodeEntry, EpisodeIdT
-from memmachine.common.errors import ConfigurationError, SessionNotFoundError
+from memmachine.common.errors import (
+    ConfigurationError,
+    SessionAlreadyExistsError,
+    SessionNotFoundError,
+)
 from memmachine.common.filter.filter_parser import (
     And as FilterAnd,
 )
@@ -162,6 +166,7 @@ async def create_session(
         *,
         description: str = "",
         user_conf: EpisodicMemoryConfPartial | None = None,
+        exist_ok: bool = False,
     ) -> SessionDataManager.SessionInfo:
         """Create a new session."""
         episodic_memory_conf = self._with_default_episodic_memory_conf(
@@ -170,13 +175,18 @@ async def create_session(
         )
 
         session_data_manager = await self._resources.get_session_data_manager()
-        await session_data_manager.create_new_session(
-            session_key=session_key,
-            configuration={},
-            param=episodic_memory_conf,
-            description=description,
-            metadata={},
-        )
+        try:
+            await session_data_manager.create_new_session(
+                session_key=session_key,
+                configuration={},
+                param=episodic_memory_conf,
+                description=description,
+                metadata={},
+            )
+        except SessionAlreadyExistsError:
+            if not exist_ok:
+                raise
+
         ret = await self.get_session(session_key=session_key)
         if ret is None:
             raise RuntimeError(f"Failed to create session {session_key}")
@@ -267,19 +277,17 @@ async def add_episodes(
         )
         episode_ids = [e.uid for e in episodes]
 
+        if await self.get_session(session_data.session_key) is None:
+            await self.create_session(session_data.session_key, exist_ok=True)
+
         tasks = []
 
         if MemoryType.Episodic in target_memories:
             episodic_memory_manager = (
                 await self._resources.get_episodic_memory_manager()
             )
-            async with episodic_memory_manager.open_or_create_episodic_memory(
+            async with episodic_memory_manager.open_episodic_memory(
                 session_key=session_data.session_key,
-                description="",
-                episodic_memory_config=self._with_default_episodic_memory_conf(
-                    session_key=session_data.session_key
-                ),
-                metadata={},
             ) as episodic_session:
                 tasks.append(episodic_session.add_memory_episodes(episodes))
 

tests/memmachine/main/test_memmachine_mock.py

@@ -392,6 +392,47 @@ async def test_add_episodes_dispatches_to_all_memories(
     )
 
 
+@pytest.mark.asyncio
+async def test_add_semantic_memory_will_create_nonexistent_session(
+    minimal_conf, patched_resource_manager
+):
+    memmachine = MemMachine(minimal_conf, patched_resource_manager)
+    session = DummySessionData("new-session")
+
+    entries = [
+        EpisodeEntry(content="hello", producer_id="user", producer_role="assistant"),
+    ]
+
+    mock_session_manager = AsyncMock()
+    patched_resource_manager.get_session_data_manager = AsyncMock(
+        return_value=mock_session_manager
+    )
+    mock_session_manager.get_session_info.return_value = None
+
+    semantic_manager = MagicMock()
+    semantic_manager.add_message = AsyncMock()
+    patched_resource_manager.get_semantic_session_manager = AsyncMock(
+        return_value=semantic_manager
+    )
+
+    async def _create_session_side_effect(*args, **kwargs):
+        mock_session_manager.get_session_info.return_value = session
+
+    mock_session_manager.create_new_session = AsyncMock(
+        side_effect=_create_session_side_effect
+    )
+
+    await memmachine.add_episodes(
+        session, entries, target_memories=[MemoryType.Semantic]
+    )
+
+    mock_session_manager.create_new_session.assert_awaited_once()
+    call_args = mock_session_manager.create_new_session.await_args[1]
+    assert call_args["session_key"] == session.session_key
+
+    semantic_manager.add_message.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_add_episodes_skips_memories_not_requested(
     minimal_conf, patched_resource_manager