edsl 0.1.38.dev3__py3-none-any.whl → 0.1.39__py3-none-any.whl

edsl/data/CacheHandler.py

@@ -1,149 +1,168 @@
-from __future__ import annotations
-import ast
-import json
-import os
-import shutil
-import sqlite3
-from edsl.config import CONFIG
-from edsl.data.Cache import Cache
-from edsl.data.CacheEntry import CacheEntry
-from edsl.data.SQLiteDict import SQLiteDict
-
-from edsl.config import CONFIG
-
-
-def set_session_cache(cache: Cache) -> None:
-    """
-    Set the session cache.
-    """
-    CONFIG.EDSL_SESSION_CACHE = cache
-
-
-def unset_session_cache() -> None:
-    """
-    Unset the session cache.
-    """
-    if hasattr(CONFIG, "EDSL_SESSION_CACHE"):
-        del CONFIG.EDSL_SESSION_CACHE
-
-
-class CacheHandler:
-    """
-    This CacheHandler figures out what caches are available and does migrations, as needed.
-    """
-
-    CACHE_PATH = CONFIG.get("EDSL_DATABASE_PATH")
-
-    def __init__(self, test: bool = False):
-        self.test = test
-        self.create_cache_directory()
-        self.cache = self.gen_cache()
-        old_data = self.from_old_sqlite_cache()
-        self.cache.add_from_dict(old_data)
-
-    def create_cache_directory(self, notify=False) -> None:
-        """
-        Create the cache directory if one is required and it does not exist.
-        """
-        path = self.CACHE_PATH.replace("sqlite:///", "")
-        dir_path = os.path.dirname(path)
-        if dir_path and not os.path.exists(dir_path):
-            os.makedirs(dir_path)
-            if notify:
-                print(f"Created cache directory: {dir_path}")
-
-    def gen_cache(self) -> Cache:
-        """
-        Generate a Cache object.
-        """
-        if self.test:
-            return Cache(data={})
-
-        if hasattr(CONFIG, "EDSL_SESSION_CACHE"):
-            return CONFIG.EDSL_SESSION_CACHE
-
-        cache = Cache(data=SQLiteDict(self.CACHE_PATH))
-        return cache
-
-    def from_old_sqlite_cache(
-        self, path: str = "edsl_cache.db"
-    ) -> dict[str, CacheEntry]:
-        """
-        Convert an old-style cache to the new format.
-        - NB: Not worth converting to sqlalchemy - this is a one-time operation.
-        """
-        old_data = {}
-        if not os.path.exists(os.path.join(os.getcwd(), path)):
-            return old_data
-        try:
-            conn = sqlite3.connect(path)
-            with conn:
-                cur = conn.cursor()
-                table_name = "responses"
-                cur.execute(f"PRAGMA table_info({table_name})")
-                columns = cur.fetchall()
-                schema = {column[1]: column[2] for column in columns}
-                data = cur.execute(f"SELECT * FROM {table_name}").fetchall()
-                for row in data:
-                    entry = self._parse_old_cache_entry(row, schema)
-                    old_data[entry.key] = entry
-            print(
-                f"Found old cache at {path} with {len(old_data)} entries.\n"
-                f"We will convert this to the new cache format.\n"
-                f"The old cache is backed up to {path}.bak"
-            )
-            shutil.copy(path, f"{path}.bak")
-            os.remove(path)
-        except sqlite3.OperationalError:
-            print("Found an old Cache but could not convert it to new format.")
-
-        return old_data
-
-    def _parse_old_cache_entry(self, row: tuple, schema) -> CacheEntry:
-        """
-        Parse an old cache entry.
-        """
-        entry_dict = {k: row[i] for i, k in enumerate(schema.keys())}
-        _ = entry_dict.pop("id")
-        entry_dict["user_prompt"] = entry_dict.pop("prompt")
-        parameters = entry_dict["parameters"]
-        entry_dict["parameters"] = ast.literal_eval(parameters)
-        entry = CacheEntry(**entry_dict)
-        return entry
-
-    def get_cache(self) -> Cache:
-        return self.cache
-
-    ###############
-    # NOT IN USE
-    ###############
-    def from_sqlite(uri="new_edsl_cache.db") -> dict[str, CacheEntry]:
-        """
-        Read in a new-style sqlite cache and return a dictionary of dictionaries.
-        """
-        conn = sqlite3.connect(uri)
-        with conn:
-            cur = conn.cursor()
-            data = cur.execute("SELECT key, value FROM data").fetchall()
-            newdata = {}
-            for _, value in data:
-                entry = CacheEntry.from_dict(json.loads(value))
-                newdata[entry.key] = entry
-        return newdata
-
-    def from_jsonl(filename="edsl_cache.jsonl") -> dict[str, CacheEntry]:
-        """Read in a jsonl file and return a dictionary of CacheEntry objects."""
-        with open(filename, "a+") as f:
-            f.seek(0)
-            lines = f.readlines()
-        newdata = {}
-        for line in lines:
-            d = json.loads(line)
-            key = list(d.keys())[0]
-            value = list(d.values())[0]
-            newdata[key] = CacheEntry.from_dict(value)
-        return newdata
-
-
-if __name__ == "__main__":
-    ch = CacheHandler()
+from __future__ import annotations
+import ast
+import json
+import os
+import shutil
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.data.Cache import Cache
+    from edsl.data.CacheEntry import CacheEntry
+
+
+def set_session_cache(cache: "Cache") -> None:
+    """
+    Set the session cache.
+    """
+    from edsl.config import CONFIG
+
+    CONFIG.EDSL_SESSION_CACHE = cache
+
+
+def unset_session_cache() -> None:
+    """
+    Unset the session cache.
+    """
+    from edsl.config import CONFIG
+
+    if hasattr(CONFIG, "EDSL_SESSION_CACHE"):
+        del CONFIG.EDSL_SESSION_CACHE
+
+
+class CacheHandler:
+    """
+    This CacheHandler figures out what caches are available and does migrations, as needed.
+    """
+
+    @property
+    def CACHE_PATH(self):
+        from edsl.config import CONFIG
+
+        return CONFIG.get("EDSL_DATABASE_PATH")
+
+    def __init__(self, test: bool = False):
+        self.test = test
+        self.create_cache_directory()
+        self.cache = self.gen_cache()
+        old_data = self.from_old_sqlite_cache()
+        self.cache.add_from_dict(old_data)
+
+    def create_cache_directory(self, notify=False) -> None:
+        """
+        Create the cache directory if one is required and it does not exist.
+        """
+        path = self.CACHE_PATH.replace("sqlite:///", "")
+        dir_path = os.path.dirname(path)
+        if dir_path and not os.path.exists(dir_path):
+            os.makedirs(dir_path)
+            if notify:
+                print(f"Created cache directory: {dir_path}")
+
+    def gen_cache(self) -> "Cache":
+        """
+        Generate a Cache object.
+        """
+        from edsl.data.Cache import Cache
+
+        if self.test:
+            return Cache(data={})
+
+        from edsl.config import CONFIG
+
+        if hasattr(CONFIG, "EDSL_SESSION_CACHE"):
+            return CONFIG.EDSL_SESSION_CACHE
+
+        from edsl.data.SQLiteDict import SQLiteDict
+
+        cache = Cache(data=SQLiteDict(self.CACHE_PATH))
+        return cache
+
+    def from_old_sqlite_cache(
+        self, path: str = "edsl_cache.db"
+    ) -> dict[str, CacheEntry]:
+        """
+        Convert an old-style cache to the new format.
+        - NB: Not worth converting to sqlalchemy - this is a one-time operation.
+        """
+        old_data = {}
+        if not os.path.exists(os.path.join(os.getcwd(), path)):
+            return old_data
+        try:
+            import sqlite3
+
+            conn = sqlite3.connect(path)
+            with conn:
+                cur = conn.cursor()
+                table_name = "responses"
+                cur.execute(f"PRAGMA table_info({table_name})")
+                columns = cur.fetchall()
+                schema = {column[1]: column[2] for column in columns}
+                data = cur.execute(f"SELECT * FROM {table_name}").fetchall()
+                for row in data:
+                    entry = self._parse_old_cache_entry(row, schema)
+                    old_data[entry.key] = entry
+            print(
+                f"Found old cache at {path} with {len(old_data)} entries.\n"
+                f"We will convert this to the new cache format.\n"
+                f"The old cache is backed up to {path}.bak"
+            )
+            shutil.copy(path, f"{path}.bak")
+            os.remove(path)
+        except sqlite3.OperationalError:
+            print("Found an old Cache but could not convert it to new format.")
+
+        return old_data
+
+    def _parse_old_cache_entry(self, row: tuple, schema) -> CacheEntry:
+        """
+        Parse an old cache entry.
+        """
+        entry_dict = {k: row[i] for i, k in enumerate(schema.keys())}
+        _ = entry_dict.pop("id")
+        entry_dict["user_prompt"] = entry_dict.pop("prompt")
+        parameters = entry_dict["parameters"]
+        entry_dict["parameters"] = ast.literal_eval(parameters)
+        from edsl.data.CacheEntry import CacheEntry
+
+        entry = CacheEntry(**entry_dict)
+        return entry
+
+    def get_cache(self) -> Cache:
+        return self.cache
+
+    ###############
+    # NOT IN USE
+    ###############
+    def from_sqlite(uri="new_edsl_cache.db") -> dict[str, "CacheEntry"]:
+        """
+        Read in a new-style sqlite cache and return a dictionary of dictionaries.
+        """
+        conn = sqlite3.connect(uri)
+        with conn:
+            cur = conn.cursor()
+            data = cur.execute("SELECT key, value FROM data").fetchall()
+            newdata = {}
+            for _, value in data:
+                entry = CacheEntry.from_dict(json.loads(value))
+                newdata[entry.key] = entry
+        return newdata
+
+    def from_jsonl(filename="edsl_cache.jsonl") -> dict[str, "CacheEntry"]:
+        """Read in a jsonl file and return a dictionary of CacheEntry objects."""
+        with open(filename, "a+") as f:
+            f.seek(0)
+            lines = f.readlines()
+        newdata = {}
+        for line in lines:
+            d = json.loads(line)
+            key = list(d.keys())[0]
+            value = list(d.values())[0]
+            newdata[key] = CacheEntry.from_dict(value)
+        return newdata
+
+
+if __name__ == "__main__":
+    # ch = CacheHandler()
+    import doctest
+
+    doctest.testmod()

edsl/data/RemoteCacheSync.py

@@ -1,97 +1,186 @@
-class RemoteCacheSync:
-    def __init__(
-        self, coop, cache, output_func, remote_cache=True, remote_cache_description=""
-    ):
-        self.coop = coop
-        self.cache = cache
-        self._output = output_func
-        self.remote_cache = remote_cache
-        self.old_entry_keys = []
-        self.new_cache_entries = []
-        self.remote_cache_description = remote_cache_description
-
-    def __enter__(self):
-        if self.remote_cache:
-            self._sync_from_remote()
-            self.old_entry_keys = list(self.cache.keys())
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        if self.remote_cache:
-            self._sync_to_remote()
-        return False  # Propagate exceptions
-
-    def _sync_from_remote(self):
-        cache_difference = self.coop.remote_cache_get_diff(self.cache.keys())
-        client_missing_cacheentries = cache_difference.get(
-            "client_missing_cacheentries", []
-        )
-        missing_entry_count = len(client_missing_cacheentries)
-
-        if missing_entry_count > 0:
-            self._output(
-                f"Updating local cache with {missing_entry_count:,} new "
-                f"{'entry' if missing_entry_count == 1 else 'entries'} from remote..."
-            )
-            self.cache.add_from_dict(
-                {entry.key: entry for entry in client_missing_cacheentries}
-            )
-            self._output("Local cache updated!")
-        else:
-            self._output("No new entries to add to local cache.")
-
-    def _sync_to_remote(self):
-        cache_difference = self.coop.remote_cache_get_diff(self.cache.keys())
-        server_missing_cacheentry_keys = cache_difference.get(
-            "server_missing_cacheentry_keys", []
-        )
-        server_missing_cacheentries = [
-            entry
-            for key in server_missing_cacheentry_keys
-            if (entry := self.cache.data.get(key)) is not None
-        ]
-
-        new_cache_entries = [
-            entry
-            for entry in self.cache.values()
-            if entry.key not in self.old_entry_keys
-        ]
-        server_missing_cacheentries.extend(new_cache_entries)
-        new_entry_count = len(server_missing_cacheentries)
-
-        if new_entry_count > 0:
-            self._output(
-                f"Updating remote cache with {new_entry_count:,} new "
-                f"{'entry' if new_entry_count == 1 else 'entries'}..."
-            )
-            self.coop.remote_cache_create_many(
-                server_missing_cacheentries,
-                visibility="private",
-                description=self.remote_cache_description,
-            )
-            self._output("Remote cache updated!")
-        else:
-            self._output("No new entries to add to remote cache.")
-
-        self._output(
-            f"There are {len(self.cache.keys()):,} entries in the local cache."
-        )
-
-
-# # Usage example
-# def run_job(self, n, progress_bar, cache, stop_on_exception, sidecar_model, print_exceptions, raise_validation_errors, use_remote_cache=True):
-#     with RemoteCacheSync(self.coop, cache, self._output, remote_cache=use_remote_cache):
-#         self._output("Running job...")
-#         results = self._run_local(
-#             n=n,
-#             progress_bar=progress_bar,
-#             cache=cache,
-#             stop_on_exception=stop_on_exception,
-#             sidecar_model=sidecar_model,
-#             print_exceptions=print_exceptions,
-#             raise_validation_errors=raise_validation_errors,
-#         )
-#         self._output("Job completed!")
-
-#     results.cache = cache.new_entries_cache()
-#     return results
+from typing import List, Dict, Any, Optional, TYPE_CHECKING, Callable
+from dataclasses import dataclass
+from contextlib import AbstractContextManager
+from collections import UserList
+
+if TYPE_CHECKING:
+    from .Cache import Cache
+    from edsl.coop.coop import Coop
+    from .CacheEntry import CacheEntry
+
+from logging import Logger
+
+
+class CacheKeyList(UserList):
+    def __init__(self, data: List[str]):
+        super().__init__(data)
+        self.data = data
+
+    def __repr__(self):
+        import reprlib
+
+        keys_repr = reprlib.repr(self.data)
+        return f"CacheKeyList({keys_repr})"
+
+
+class CacheEntriesList(UserList):
+    def __init__(self, data: List["CacheEntry"]):
+        super().__init__(data)
+        self.data = data
+
+    def __repr__(self):
+        import reprlib
+
+        entries_repr = reprlib.repr(self.data)
+        return f"CacheEntries({entries_repr})"
+
+    def to_cache(self) -> "Cache":
+        from edsl.data.Cache import Cache
+
+        return Cache({entry.key: entry for entry in self.data})
+
+
+@dataclass
+class CacheDifference:
+    client_missing_entries: CacheEntriesList
+    server_missing_keys: List[str]
+
+    def __repr__(self):
+        """Returns a string representation of the CacheDifference object."""
+        import reprlib
+
+        missing_entries_repr = reprlib.repr(self.client_missing_entries)
+        missing_keys_repr = reprlib.repr(self.server_missing_keys)
+        return f"CacheDifference(client_missing_entries={missing_entries_repr}, server_missing_keys={missing_keys_repr})"
+
+
+class RemoteCacheSync(AbstractContextManager):
+    """Synchronizes a local cache with a remote cache.
+
+    Handles bidirectional synchronization:
+    - Downloads missing entries from remote to local cache
+    - Uploads new local entries to remote cache
+    """
+
+    def __init__(
+        self,
+        coop: "Coop",
+        cache: "Cache",
+        output_func: Callable,
+        remote_cache: bool = True,
+        remote_cache_description: str = "",
+    ):
+        """
+        Initializes a RemoteCacheSync object.
+
+        :param coop: Coop object for interacting with the remote cache
+        :param cache: Cache object for local cache
+        :param output_func: Function for outputting messages
+        :param remote_cache: Whether to enable remote cache synchronization
+        :param remote_cache_description: Description for remote cache entries
+
+        """
+        self.coop = coop
+        self.cache = cache
+        self._output = output_func
+        self.remote_cache_enabled = remote_cache
+        self.remote_cache_description = remote_cache_description
+        self.initial_cache_keys = []
+
+    def __enter__(self) -> "RemoteCacheSync":
+        if self.remote_cache_enabled:
+            self._sync_from_remote()
+            self.initial_cache_keys = list(self.cache.keys())
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.remote_cache_enabled:
+            self._sync_to_remote()
+        return False  # Propagate exceptions
+
+    def _get_cache_difference(self) -> CacheDifference:
+        """Retrieves differences between local and remote caches."""
+        diff = self.coop.remote_cache_get_diff(self.cache.keys())
+        return CacheDifference(
+            client_missing_entries=diff.get("client_missing_cacheentries", []),
+            server_missing_keys=diff.get("server_missing_cacheentry_keys", []),
+        )
+
+    def _sync_from_remote(self) -> None:
+        """Downloads missing entries from remote cache to local cache."""
+        diff: CacheDifference = self._get_cache_difference()
+        missing_count = len(diff.client_missing_entries)
+
+        if missing_count == 0:
+            self._output("No new entries to add to local cache.")
+            return
+
+        self._output(
+            f"Updating local cache with {missing_count:,} new "
+            f"{'entry' if missing_count == 1 else 'entries'} from remote..."
+        )
+
+        self.cache.add_from_dict(
+            {entry.key: entry for entry in diff.client_missing_entries}
+        )
+        self._output("Local cache updated!")
+
+    def _get_entries_to_upload(self, diff: CacheDifference) -> CacheEntriesList:
+        """Determines which entries need to be uploaded to remote cache."""
+        # Get entries for keys missing from server
+        server_missing_entries = CacheEntriesList(
+            [
+                entry
+                for key in diff.server_missing_keys
+                if (entry := self.cache.data.get(key)) is not None
+            ]
+        )
+
+        # Get newly added entries since sync started
+        new_entries = CacheEntriesList(
+            [
+                entry
+                for entry in self.cache.values()
+                if entry.key not in self.initial_cache_keys
+            ]
+        )
+
+        return server_missing_entries + new_entries
+
+    def _sync_to_remote(self) -> None:
+        """Uploads new local entries to remote cache."""
+        diff: CacheDifference = self._get_cache_difference()
+        entries_to_upload: CacheEntriesList = self._get_entries_to_upload(diff)
+        upload_count = len(entries_to_upload)
+
+        if upload_count > 0:
+            self._output(
+                f"Updating remote cache with {upload_count:,} new "
+                f"{'entry' if upload_count == 1 else 'entries'}..."
+            )
+
+            self.coop.remote_cache_create_many(
+                entries_to_upload,
+                visibility="private",
+                description=self.remote_cache_description,
+            )
+            self._output("Remote cache updated!")
+        else:
+            self._output("No new entries to add to remote cache.")
+
+        self._output(
+            f"There are {len(self.cache.keys()):,} entries in the local cache."
+        )
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
+
+    from edsl.coop.coop import Coop
+    from edsl.data.Cache import Cache
+    from edsl.data.CacheEntry import CacheEntry
+
+    r = RemoteCacheSync(Coop(), Cache(), print)
+    diff = r._get_cache_difference()