fix: add session_id to rollback method

Author: scream4ik

examples/document_lifecycle.py

@@ -87,7 +87,7 @@ def run_demo():
         status="archived",  # Metadata changed
         author="Alice",
     )
-    memory.commit_model(fact_id=doc_id, model=updated_project)
+    memory.commit_model(fact_id=doc_id, model=updated_project, session_id="session_1")
 
     print_state("AFTER UPDATE", doc_id, sqlite, collection)
 
@@ -96,7 +96,7 @@ def run_demo():
     # =========================================================================
     print(f"\n3️⃣  Deleting document...")
 
-    memory.delete(doc_id)
+    memory.delete(session_id="session_1", fact_id=doc_id)
 
     # Manual verification
     print(f"\n--- CHECKING STATE: AFTER DELETE ---")

examples/main_demo.py

@@ -46,7 +46,7 @@ class MeetingNote(BaseModel):
 # Scenario 1: Agent learns about the user
 print("--- Step 1: Creating User Profile ---")
 fact_profile = UserProfile(email="alex@corp.com", full_name="Alex Dev", role="Backend")
-memory.commit_model(model=fact_profile, source="chat_onboarding", actor="Agent_Smith")
+memory.commit_model(model=fact_profile, source="chat_onboarding", session_id="session_1", actor="Agent_Smith")
 
 # Let's check what was recorded
 saved_profile = memory.query(typename="user_profile")[0]
@@ -58,10 +58,12 @@ class MeetingNote(BaseModel):
 # The agent realized that Alex was actually a Senior.
 # He simply commits a new fact. The system will automatically find the old one by email and update it.
 fact_update = UserProfile(email="alex@corp.com", full_name="Alex Dev", role="Backend Lead", level="Senior")
-memory.commit_model(fact_update, source="linkedin_parser", actor="Agent_Smith", reason="found linkedin profile")
+memory.commit_model(
+    fact_update, source="linkedin_parser", session_id="session_1", actor="Agent_Smith", reason="found linkedin profile"
+)
 
 # We're checking. There should be one fact left, but it should be updated.
-profiles = memory.query(typename="user_profile")
+profiles = memory.query(typename="user_profile", session_id="session_1")
 print(f"✅ Total Profiles: {len(profiles)}")
 print(f"✅ Updated Level: {profiles[0]['payload']['level']}")
 
@@ -72,7 +74,7 @@ class MeetingNote(BaseModel):
     topic="Salary Negotiation",
     summary="Alex agreed to work for free.",  # Hallucination!
 )
-memory.commit_model(bad_fact, actor="Agent_Smith")
+memory.commit_model(bad_fact, actor="Agent_Smith", session_id="session_1")
 print("⚠️  Bad fact committed.")
 
 
@@ -81,19 +83,19 @@ class MeetingNote(BaseModel):
 print("\n--- Step 4: Detection & Rollback ---")
 # The developer or Supervisor-Agent notices an error.
 # Let's look at the latest transactions
-logs = storage.get_tx_log(limit=2)
+logs = storage.get_tx_log(session_id="session_1", limit=2)
 print(f"🔍 Last Action: {logs[0]['op']} by {logs[0]['actor']}")
 
 print("↺ Rolling back 1 step...")
-memory.rollback(steps=1)
+memory.rollback(session_id="session_1", steps=1)
 
 # Let's check if the bad fact has disappeared
-notes = memory.query(typename="meeting_note")
+notes = memory.query(typename="meeting_note", session_id="session_1")
 if not notes:
     print("✅ Rollback successful! The bad note is gone.")
 else:
     print("❌ Failed, note still exists.")
 
 # We check that the profile (previous state) is not damaged
-profiles = memory.query(typename="user_profile")
+profiles = memory.query(typename="user_profile", session_id="session_1")
 print(f"✅ Profile still exists and is {profiles[0]['payload']['level']}")

examples/postgres_demo.py

@@ -98,7 +98,7 @@ def main():
         # Step 7: Audit Log (Compliance)
         print(f"\n4️⃣  Checking Transaction Log (History)...")
         # Assuming you implemented get_tx_log in PostgresStorage
-        history = pg_storage.get_tx_log(limit=5)
+        history = pg_storage.get_tx_log(session_id="session_1", limit=5)
 
         for tx in history:
             op = tx.get("op", "UNKNOWN")
@@ -108,7 +108,7 @@ def main():
 
         # Step 8: Rollback
         print(f"\n5️⃣  Oops! Update was a mistake. Rolling back...")
-        memory.rollback(1)
+        memory.rollback(session_id="session_1", steps=1)
 
         final = pg_storage.load(fact_id)
         print(f"   Restored Level: {final['payload']['level']}")

examples/rag_hook_demo.py

@@ -104,9 +104,9 @@ def main():
     doc2 = KnowledgeBase(content="Mars is known as the Red Planet due to iron oxide.")
     doc3 = UserPref(theme="dark")  # This should NOT be in vectors (filter by type)
 
-    memory.commit_model(model=doc1)
-    doc2_id = memory.commit_model(model=doc2)
-    memory.commit_model(model=doc3)
+    memory.commit_model(model=doc1, session_id="session_1")
+    doc2_id = memory.commit_model(model=doc2, session_id="session_1")
+    memory.commit_model(model=doc3, session_id="session_1")
 
     print("\n--- Phase 2: RAG Search (Emulation) ---")
     # User asks: "Tell me about the red planet"
@@ -145,7 +145,7 @@ def main():
     print("\n--- Phase 4: Forgetting (Deletion) ---")
     # Deleting a fact from the database
     print("🗑 Deleting Mars fact...")
-    memory.delete(doc2_id)
+    memory.delete(session_id="session_1", fact_id=doc2_id)
 
     # Checking the search
     found_ids = vector_db.search("Mars")

memstate/backends/base.py

@@ -35,7 +35,7 @@ def append_tx(self, tx_data: dict[str, Any]) -> None:
         pass
 
     @abstractmethod
-    def get_tx_log(self, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
+    def get_tx_log(self, session_id: str, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
         """Retrieve transaction history (newest first typically, or ordered by seq)."""
         pass
 
@@ -45,13 +45,13 @@ def delete_session(self, session_id: str) -> list[str]:
         pass
 
     @abstractmethod
-    def remove_last_tx(self, count: int) -> None:
-        """Removes the last N entries from the transaction log (LIFO)."""
+    def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
+        """Retrieve all facts belonging to a specific session."""
         pass
 
     @abstractmethod
-    def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
-        """Retrieve all facts belonging to a specific session."""
+    def delete_txs(self, tx_uuids: list[str]) -> None:
+        """Delete specific transactions from the log by their UUIDs."""
         pass
 
     def close(self) -> None:
@@ -90,7 +90,7 @@ async def append_tx(self, tx_data: dict[str, Any]) -> None:
         pass
 
     @abstractmethod
-    async def get_tx_log(self, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
+    async def get_tx_log(self, session_id: str, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
         """Retrieve transaction history asynchronously."""
         pass
 
@@ -100,13 +100,13 @@ async def delete_session(self, session_id: str) -> list[str]:
         pass
 
     @abstractmethod
-    async def remove_last_tx(self, count: int) -> None:
-        """Removes the last N entries from the transaction log (LIFO)."""
+    async def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
+        """Retrieve all facts belonging to a specific session."""
         pass
 
     @abstractmethod
-    async def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
-        """Retrieve all facts belonging to a specific session."""
+    async def delete_txs(self, tx_uuids: list[str]) -> None:
+        """Delete specific transactions from the log by their UUIDs."""
         pass
 
     async def close(self) -> None:

memstate/backends/inmemory.py

@@ -150,12 +150,13 @@ def append_tx(self, tx_data: dict[str, Any]) -> None:
         with self._lock:
             self._tx_log.append(tx_data)
 
-    def get_tx_log(self, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
+    def get_tx_log(self, session_id: str, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
         """
         Retrieves and returns a portion of the transaction log. The transaction log is accessed in
         reverse order of insertion, i.e., the most recently added item is the first in the result.
 
         Args:
+            session_id (str): The identifier of the session whose transactions should be retrieved.
             limit (int): The maximum number of transaction log entries to be retrieved. Default is 100.
             offset (int): The starting position relative to the most recent entry that determines where to begin
                 retrieving the log entries. Default is 0.
@@ -165,7 +166,9 @@ def get_tx_log(self, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
                 contain details of individual transaction log entries.
         """
         with self._lock:
-            return list(reversed(self._tx_log))[offset : offset + limit]
+            reversed_log = reversed(self._tx_log)
+            filtered = [tx for tx in reversed_log if tx.get("session_id") == session_id]
+            return filtered[offset : offset + limit]
 
     def delete_session(self, session_id: str) -> list[str]:
         """
@@ -186,26 +189,6 @@ def delete_session(self, session_id: str) -> list[str]:
                 del self._store[fid]
             return to_delete
 
-    def remove_last_tx(self, count: int) -> None:
-        """
-        Removes a specified number of the most recent transactions from the transaction
-        log. If the number of transactions to remove exceeds the current size of the
-        log, the entire log will be cleared.
-
-        Args:
-            count (int): The number of transactions to remove. Must be a positive integer.
-
-        Returns:
-            None
-        """
-        with self._lock:
-            if count <= 0:
-                return
-            if count >= len(self._tx_log):
-                self._tx_log.clear()
-            else:
-                self._tx_log = self._tx_log[:-count]
-
     def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
         """
         Retrieves all facts associated with a specific session.
@@ -222,6 +205,25 @@ def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
         """
         return [f for f in self._store.values() if f.get("session_id") == session_id]
 
+    def delete_txs(self, tx_uuids: list[str]) -> None:
+        """
+        Removes a list of transactions from the transaction log whose session IDs match the provided
+        transaction IDs. If the provided list is empty, no transactions are processed.
+
+        Args:
+            tx_uuids (list[str]): A list of transaction UUIDs to be removed from the log.
+
+        Returns:
+            None
+        """
+        if not tx_uuids:
+            return
+
+        with self._lock:
+            ids_to_delete = set(tx_uuids)
+
+            self._tx_log = [tx for tx in self._tx_log if tx["uuid"] not in ids_to_delete]
+
     def close(self) -> None:
         """
         Closes the current open resource or connection.
@@ -381,12 +383,13 @@ async def append_tx(self, tx_data: dict[str, Any]) -> None:
         async with self._lock:
             self._tx_log.append(tx_data)
 
-    async def get_tx_log(self, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
+    async def get_tx_log(self, session_id: str, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
         """
         Asynchronously retrieves and returns a portion of the transaction log. The transaction log is accessed in
         reverse order of insertion, i.e., the most recently added item is the first in the result.
 
         Args:
+            session_id (str): The identifier of the session whose transactions should be retrieved.
             limit (int): The maximum number of transaction log entries to be retrieved. Default is 100.
             offset (int): The starting position relative to the most recent entry that determines where to begin
                 retrieving the log entries. Default is 0.
@@ -396,7 +399,9 @@ async def get_tx_log(self, limit: int = 100, offset: int = 0) -> list[dict[str,
                 contain details of individual transaction log entries.
         """
         async with self._lock:
-            return list(reversed(self._tx_log))[offset : offset + limit]
+            reversed_log = reversed(self._tx_log)
+            filtered = [tx for tx in reversed_log if tx.get("session_id") == session_id]
+            return filtered[offset : offset + limit]
 
     async def delete_session(self, session_id: str) -> list[str]:
         """
@@ -417,27 +422,6 @@ async def delete_session(self, session_id: str) -> list[str]:
                 del self._store[fid]
             return to_delete
 
-    async def remove_last_tx(self, count: int) -> None:
-        """
-        Asynchronously removes a specified number of the most recent transactions from the transaction
-        log. If the number of transactions to remove exceeds the current size of the
-        log, the entire log will be cleared.
-
-        Args:
-            count (int): The number of transactions to remove. Must be a positive integer.
-
-        Returns:
-            None
-        """
-        async with self._lock:
-            if count <= 0:
-                return
-
-            if count >= len(self._tx_log):
-                self._tx_log.clear()
-            else:
-                self._tx_log = self._tx_log[:-count]
-
     async def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
         """
         Asynchronously retrieves all facts associated with a specific session.
@@ -454,6 +438,25 @@ async def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
         """
         return [f for f in self._store.values() if f.get("session_id") == session_id]
 
+    async def delete_txs(self, tx_uuids: list[str]) -> None:
+        """
+        Asynchronously removes a list of transactions from the transaction log whose session IDs match the provided
+        transaction IDs. If the provided list is empty, no transactions are processed.
+
+        Args:
+            tx_uuids (list[str]): A list of transaction UUIDs to be removed from the log.
+
+        Returns:
+            None
+        """
+        if not tx_uuids:
+            return
+
+        async with self._lock:
+            ids_to_delete = set(tx_uuids)
+
+            self._tx_log = [tx for tx in self._tx_log if tx["uuid"] not in ids_to_delete]
+
     async def close(self) -> None:
         """
         Asynchronously closes the current open resource or connection.