core[minor]: Adds an in-memory implementation of RecordManager

libs/community/langchain_community/indexes/memory_recordmanager.py

@@ -0,0 +1,118 @@
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Sequence
+
+from langchain_community.indexes.base import RecordManager
+
+
+class MemoryRecordManager(RecordManager):
+    data: List[Dict[str, Any]] = []
+
+    def __init__(self, namespace: str = ""):
+        super().__init__(namespace=namespace)
+
+    def create_schema(self) -> None:
+        pass
+
+    async def acreate_schema(self) -> None:
+        pass
+
+    def get_time(self) -> float:
+        return datetime.now().timestamp()
+
+    def update(
+        self,
+        keys: Sequence[str],
+        *,
+        group_ids: Optional[Sequence[Optional[str]]] = None,
+        time_at_least: Optional[float] = None,
+    ) -> None:
+        if group_ids is None:
+            group_ids = [None] * len(keys)
+        if len(keys) != len(group_ids):
+            raise ValueError(
+                f"Number of keys ({len(keys)}) does not match number of "
+                f"group_ids ({len(group_ids)})"
+            )
+
+        update_time = self.get_time()
+        if time_at_least and update_time < time_at_least:
+            # Safeguard against time sync issues
+            raise AssertionError(f"Time sync issue: {update_time} < {time_at_least}")
+
+        records_to_upsert = [
+            {
+                "key": key,
+                "namespace": self.namespace,
+                "updated_at": update_time,
+                "group_id": group_id,
+            }
+            for key, group_id in zip(keys, group_ids)
+        ]
+        self.delete_keys(keys)
+        self.data.extend(records_to_upsert)
+
+    def exists(self, keys: Sequence[str]) -> List[bool]:
+        return [
+            len(list(filter(lambda record: record["key"] == key, self.data))) == 1
+            for key in keys
+        ]
+
+    def list_keys(
+        self,
+        *,
+        before: Optional[float] = None,
+        after: Optional[float] = None,
+        group_ids: Optional[Sequence[str]] = None,
+        limit: Optional[int] = None,
+    ) -> List[str]:
+        keys = [
+            record["key"]
+            for record in filter(
+                lambda record: record["namespace"] == self.namespace
+                and (not group_ids or record["group_id"] in group_ids)
+                and (not before or record["updated_at"] < before)
+                and (not after or record["updated_at"] > after),
+                self.data,
+            )
+        ]
+        return keys[:limit]
+
+    def delete_keys(self, keys: Sequence[str]) -> None:
+        self.data = list(
+            filter(
+                lambda record: record["namespace"] != self.namespace
+                or record["key"] not in keys,
+                self.data,
+            )
+        )
+
+    # %% Async versions
+    async def aget_time(self) -> float:
+        return datetime.now().timestamp()
+
+    async def aupdate(
+        self,
+        keys: Sequence[str],
+        *,
+        group_ids: Optional[Sequence[Optional[str]]] = None,
+        time_at_least: Optional[float] = None,
+    ) -> None:
+        return self.update(keys=keys, group_ids=group_ids, time_at_least=time_at_least)
+
+    async def aexists(self, keys: Sequence[str]) -> List[bool]:
+        return self.exists(keys=keys)
+
+    async def alist_keys(
+        self,
+        *,
+        before: Optional[float] = None,
+        after: Optional[float] = None,
+        group_ids: Optional[Sequence[str]] = None,
+        limit: Optional[int] = None,
+    ) -> List[str]:
+        return self.list_keys(
+            before=before, after=after, group_ids=group_ids, limit=limit
+        )
+
+    async def adelete_keys(self, keys: Sequence[str]) -> None:
+        return self.delete_keys(keys=keys)

libs/community/tests/unit_tests/indexes/test_memory_record_manager.py

@@ -0,0 +1,132 @@
+import pytest
+import pytest_asyncio
+
+from langchain_community.indexes.memory_recordmanager import MemoryRecordManager
+
+
+@pytest.fixture()
+def manager() -> MemoryRecordManager:
+    """Initialize the test database and yield the TimestampedSet instance."""
+    # Initialize and yield the TimestampedSet instance
+    record_manager = MemoryRecordManager(namespace="kittens")
+    return record_manager
+
+
+@pytest_asyncio.fixture  # type: ignore
+async def amanager() -> MemoryRecordManager:
+    """Initialize the test database and yield the TimestampedSet instance."""
+    # Initialize and yield the TimestampedSet instance
+    record_manager = MemoryRecordManager(namespace="kittens")
+    return record_manager
+
+
+def test_update(manager: MemoryRecordManager) -> None:
+    """Test updating records in the database."""
+    # no keys should be present in the set
+    read_keys = manager.list_keys()
+    assert read_keys == []
+    # Insert records
+    keys = ["key1", "key2", "key3"]
+    manager.update(keys)
+    # Retrieve the records
+    read_keys = manager.list_keys()
+    assert read_keys == ["key1", "key2", "key3"]
+
+
+async def test_aupdate(amanager: MemoryRecordManager) -> None:
+    """Test updating records in the database."""
+    # no keys should be present in the set
+    read_keys = await amanager.alist_keys()
+    assert read_keys == []
+    # Insert records
+    keys = ["key1", "key2", "key3"]
+    await amanager.aupdate(keys)
+    # Retrieve the records
+    read_keys = await amanager.alist_keys()
+    assert read_keys == ["key1", "key2", "key3"]
+
+
+def test_update_with_group_ids(manager: MemoryRecordManager) -> None:
+    """Test updating records in the database."""
+    # no keys should be present in the set
+    read_keys = manager.list_keys()
+    assert read_keys == []
+    # Insert records
+    keys = ["key1", "key2", "key3"]
+    manager.update(keys)
+    # Retrieve the records
+    read_keys = manager.list_keys()
+    assert read_keys == ["key1", "key2", "key3"]
+
+
+async def test_aupdate_with_group_ids(amanager: MemoryRecordManager) -> None:
+    """Test updating records in the database."""
+    # no keys should be present in the set
+    read_keys = await amanager.alist_keys()
+    assert read_keys == []
+    # Insert records
+    keys = ["key1", "key2", "key3"]
+    await amanager.aupdate(keys)
+    # Retrieve the records
+    read_keys = await amanager.alist_keys()
+    assert read_keys == ["key1", "key2", "key3"]
+
+
+def test_exists(manager: MemoryRecordManager) -> None:
+    """Test checking if keys exist in the database."""
+    # Insert records
+    keys = ["key1", "key2", "key3"]
+    manager.update(keys)
+    # Check if the keys exist in the database
+    exists = manager.exists(keys)
+    assert len(exists) == len(keys)
+    assert exists == [True, True, True]
+
+    exists = manager.exists(["key1", "key4"])
+    assert len(exists) == 2
+    assert exists == [True, False]
+
+
+async def test_aexists(amanager: MemoryRecordManager) -> None:
+    """Test checking if keys exist in the database."""
+    # Insert records
+    keys = ["key1", "key2", "key3"]
+    await amanager.aupdate(keys)
+    # Check if the keys exist in the database
+    exists = await amanager.aexists(keys)
+    assert len(exists) == len(keys)
+    assert exists == [True, True, True]
+
+    exists = await amanager.aexists(["key1", "key4"])
+    assert len(exists) == 2
+    assert exists == [True, False]
+
+
+def test_delete_keys(manager: MemoryRecordManager) -> None:
+    """Test deleting keys from the database."""
+    # Insert records
+    keys = ["key1", "key2", "key3"]
+    manager.update(keys)
+
+    # Delete some keys
+    keys_to_delete = ["key1", "key2"]
+    manager.delete_keys(keys_to_delete)
+
+    # Check if the deleted keys are no longer in the database
+    remaining_keys = manager.list_keys()
+    assert remaining_keys == ["key3"]
+
+
+async def test_adelete_keys(amanager: MemoryRecordManager) -> None:
+    """Test deleting keys from the database."""
+    # Insert records
+    keys = ["key1", "key2", "key3"]
+    await amanager.aupdate(keys)
+
+    # Delete some keys
+    keys_to_delete = ["key1", "key2"]
+    await amanager.adelete_keys(keys_to_delete)
+
+    # Check if the deleted keys are no longer in the database
+    remaining_keys = await amanager.alist_keys()
+    assert remaining_keys == ["key3"]