edsl 0.1.37.dev6__py3-none-any.whl → 0.1.38__py3-none-any.whl

edsl/data/CacheHandler.py

@@ -1,149 +1,149 @@
-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
+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()

edsl/data/RemoteCacheSync.py

@@ -1,97 +1,78 @@
-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
+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."
+        )