chore: add method get_session_facts for backends to improve performance

Author: scream4ik

memstate/backends/base.py

@@ -45,6 +45,11 @@ def remove_last_tx(self, count: int) -> None:
         """Removes the last N entries from the transaction log (LIFO)."""
         pass
 
+    @abstractmethod
+    def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
+        """Retrieve all facts belonging to a specific session."""
+        pass
+
     def close(self) -> None:
         """Cleanup resources (optional)."""
         pass
@@ -95,6 +100,11 @@ async def remove_last_tx(self, count: int) -> None:
         """Removes the last N entries from the transaction log (LIFO)."""
         pass
 
+    @abstractmethod
+    async def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
+        """Retrieve all facts belonging to a specific session."""
+        pass
+
     async def close(self) -> None:
         """Cleanup resources asynchronously."""
         pass

memstate/backends/inmemory.py

@@ -79,6 +79,9 @@ def remove_last_tx(self, count: int) -> None:
             else:
                 self._tx_log = self._tx_log[:-count]
 
+    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 close(self) -> None:
         pass
 
@@ -161,5 +164,8 @@ async def remove_last_tx(self, count: int) -> None:
             else:
                 self._tx_log = self._tx_log[:-count]
 
+    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 close(self) -> None:
         pass

memstate/backends/postgres.py

@@ -140,6 +140,12 @@ def remove_last_tx(self, count: int) -> None:
         with self._engine.begin() as conn:
             conn.execute(stmt)
 
+    def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
+        stmt = select(self._facts_table.c.doc).where(self._facts_table.c.doc["session_id"].astext == session_id)
+        with self._engine.connect() as conn:
+            rows = conn.execute(stmt).all()
+            return [r[0] for r in rows]
+
     def close(self) -> None:
         self._engine.dispose()
 
@@ -248,5 +254,11 @@ async def remove_last_tx(self, count: int) -> None:
         async with self._engine.begin() as conn:
             await conn.execute(stmt)
 
+    async def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
+        stmt = select(self._facts_table.c.doc).where(self._facts_table.c.doc["session_id"].astext == session_id)
+        async with self._engine.connect() as conn:
+            result = await conn.execute(stmt)
+            return [r[0] for r in result.all()]
+
     async def close(self) -> None:
         await self._engine.dispose()

memstate/backends/redis.py

@@ -149,6 +149,25 @@ def remove_last_tx(self, count: int) -> None:
             return
         self.r.ltrim(f"{self.prefix}tx_log", count, -1)
 
+    def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
+        key = f"{self.prefix}session:{session_id}"
+        ids = self.r.smembers(key)
+
+        if not ids:
+            return []
+
+        pipe = self.r.pipeline()
+        for i in ids:
+            pipe.get(self._key(i))
+        raw_docs = pipe.execute()
+
+        results = []
+        for raw_doc in raw_docs:
+            doc_str = self._to_str(raw_doc)
+            if doc_str:
+                results.append(json.loads(doc_str))
+        return results
+
     def close(self) -> None:
         if self._owns_client:
             self.r.close()
@@ -270,6 +289,24 @@ async def remove_last_tx(self, count: int) -> None:
             return
         await self.r.ltrim(f"{self.prefix}tx_log", count, -1)
 
+    async def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
+        key = f"{self.prefix}session:{session_id}"
+        ids = await self.r.smembers(key)
+
+        if not ids:
+            return []
+
+        async with self.r.pipeline() as pipe:
+            for i in ids:
+                pipe.get(self._key(i))
+            raw_docs = await pipe.execute()
+
+        results = []
+        for raw_doc in raw_docs:
+            if raw_doc:
+                results.append(json.loads(raw_doc))
+        return results
+
     async def close(self):
         if self._owns_client:
             await self.r.aclose()

memstate/backends/sqlite.py

@@ -155,6 +155,12 @@ def remove_last_tx(self, count: int) -> None:
             )
             self._conn.commit()
 
+    def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
+        with self._lock:
+            c = self._conn.cursor()
+            c.execute("SELECT data FROM facts WHERE json_extract(data, '$.session_id') = ?", (session_id,))
+            return [json.loads(row["data"]) for row in c.fetchall()]
+
     def close(self):
         if self._owns_connection:
             self._conn.close()
@@ -315,6 +321,14 @@ async def remove_last_tx(self, count: int) -> None:
             )
             await self._db.commit()
 
+    async def get_session_facts(self, session_id: str) -> list[dict[str, Any]]:
+        async with self._lock:
+            async with self._db.execute(
+                "SELECT data FROM facts WHERE json_extract(data, '$.session_id') = ?", (session_id,)
+            ) as cursor:
+                rows = await cursor.fetchall()
+                return [json.loads(row["data"]) for row in rows]
+
     async def close(self) -> None:
         if self._db:
             await self._db.close()

memstate/storage.py

@@ -242,13 +242,7 @@ def promote_session(
         reason: str | None = None,
     ) -> list[str]:
         with self._lock:
-            # Find all session facts (via query or backend index)
-            # RedisStorage and InMemory can search by session_id if we add it to the query.
-            # For MVP, we use a query with a filter (slow on large volumes, fast with indexes)
-
-            # In RedisStorage, it is better to create a separate get_session_facts method, but we use query
-            # We assume that storage stores session_id in json data
-            candidates = self.storage.query(json_filters={"session_id": session_id})
+            candidates = self.storage.get_session_facts(session_id)
 
             promoted = []
             for fact_dict in candidates:
@@ -492,7 +486,7 @@ async def promote_session(
         reason: str | None = None,
     ) -> list[str]:
         async with self._lock:
-            candidates = await self.storage.query(json_filters={"session_id": session_id})
+            candidates = await self.storage.get_session_facts(session_id)
 
             promoted = []
             for fact_dict in candidates:

tests/async/test_backends_async.py

@@ -165,3 +165,26 @@ async def test_remove_last_tx(storage):
 
     logs = await storage.get_tx_log(limit=10)
     assert len(logs) == 0
+
+
+async def test_get_session_facts(storage):
+    ts = datetime.now(timezone.utc).isoformat()
+
+    await storage.save({"id": "a1", "type": "msg", "session_id": "session_A", "payload": {"val": 1}, "ts": ts})
+    await storage.save({"id": "a2", "type": "msg", "session_id": "session_A", "payload": {"val": 2}, "ts": ts})
+
+    await storage.save({"id": "b1", "type": "msg", "session_id": "session_B", "payload": {"val": 3}, "ts": ts})
+
+    await storage.save({"id": "g1", "type": "config", "payload": {"val": 0}, "ts": ts})
+
+    facts_a = await storage.get_session_facts("session_A")
+    assert len(facts_a) == 2
+    ids_a = sorted([f["id"] for f in facts_a])
+    assert ids_a == ["a1", "a2"]
+
+    facts_b = await storage.get_session_facts("session_B")
+    assert len(facts_b) == 1
+    assert facts_b[0]["id"] == "b1"
+
+    facts_empty = await storage.get_session_facts("ghost_session")
+    assert facts_empty == []

tests/sync/test_backends.py

@@ -157,3 +157,26 @@ def test_remove_last_tx(storage):
 
     logs = storage.get_tx_log(limit=10)
     assert len(logs) == 0
+
+
+def test_get_session_facts(storage):
+    ts = datetime.now(timezone.utc).isoformat()
+
+    storage.save({"id": "a1", "type": "msg", "session_id": "session_A", "payload": {"val": 1}, "ts": ts})
+    storage.save({"id": "a2", "type": "msg", "session_id": "session_A", "payload": {"val": 2}, "ts": ts})
+
+    storage.save({"id": "b1", "type": "msg", "session_id": "session_B", "payload": {"val": 3}, "ts": ts})
+
+    storage.save({"id": "g1", "type": "config", "payload": {"val": 0}, "ts": ts})
+
+    facts_a = storage.get_session_facts("session_A")
+    assert len(facts_a) == 2
+    ids_a = sorted([f["id"] for f in facts_a])
+    assert ids_a == ["a1", "a2"]
+
+    facts_b = storage.get_session_facts("session_B")
+    assert len(facts_b) == 1
+    assert facts_b[0]["id"] == "b1"
+
+    facts_empty = storage.get_session_facts("ghost_session")
+    assert facts_empty == []