edsl 0.1.15__py3-none-any.whl → 0.1.40__py3-none-any.whl

edsl/data/CacheEntry.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+import json
+import datetime
+import hashlib
+from typing import Optional
+from uuid import uuid4
+
+from edsl.utilities.decorators import remove_edsl_version
+
+from edsl.Base import RepresentationMixin
+
+
+class CacheEntry(RepresentationMixin):
+    """
+    A Class to represent a cache entry.
+    """
+
+    key_fields = ["model", "parameters", "system_prompt", "user_prompt", "iteration"]
+    all_fields = key_fields + ["timestamp", "output"]
+
+    def __init__(
+        self,
+        *,
+        model: str,
+        parameters: dict,
+        system_prompt: str,
+        user_prompt: str,
+        iteration: Optional[int] = None,
+        output: str,
+        timestamp: Optional[int] = None,
+    ):
+        self.model = model
+        self.parameters = parameters
+        self.system_prompt = system_prompt
+        self.user_prompt = user_prompt
+        self.output = output
+        self.iteration = iteration or 0
+        self.timestamp = timestamp or int(
+            datetime.datetime.now(datetime.timezone.utc).timestamp()
+        )
+        self._check_types()
+
+    def _check_types(self):
+        """
+        Checks if the types of the fields are correct.
+        """
+        if not isinstance(self.model, str):
+            raise TypeError("`model` should be a string.")
+        if not isinstance(self.parameters, dict):
+            raise TypeError("`parameters` should be a dictionary.")
+        if not isinstance(self.system_prompt, str):
+            raise TypeError("`system_prompt` should be a string.")
+        if not isinstance(self.user_prompt, str):
+            raise TypeError("`user_prompt` should be a string")
+        if not isinstance(self.output, str):
+            raise TypeError("`output` should be a string")
+        if not isinstance(self.iteration, int):
+            raise TypeError("`iteration` should be an integer")
+        # TODO: should probably be float
+        if not isinstance(self.timestamp, int):
+            raise TypeError(f"`timestamp` should be an integer")
+
+    @classmethod
+    def gen_key(
+        self, *, model, parameters, system_prompt, user_prompt, iteration
+    ) -> str:
+        """
+        Generates a key for the cache entry.
+        - Treats single and double quotes as the same.
+
+        TODO: add more robustness.
+        """
+        long_key = f"{model}{json.dumps(parameters, sort_keys=True)}{system_prompt}{user_prompt}{iteration}"
+        return hashlib.md5(long_key.encode()).hexdigest()
+
+    @property
+    def key(self) -> str:
+        """
+        Returns the key for the cache entry.
+        - The key is a hash of the key fields.
+        """
+        d = {k: value for k, value in self.__dict__.items() if k in self.key_fields}
+        return self.gen_key(**d)
+
+    def to_dict(self, add_edsl_version=True) -> dict:
+        """
+        Returns a dictionary representation of a CacheEntry.
+        """
+        d = {
+            "model": self.model,
+            "parameters": self.parameters,
+            "system_prompt": self.system_prompt,
+            "user_prompt": self.user_prompt,
+            "output": self.output,
+            "iteration": self.iteration,
+            "timestamp": self.timestamp,
+        }
+        # if add_edsl_version:
+        #     from edsl import __version__
+
+        #     d["edsl_version"] = __version__
+        #     d["edsl_class_name"] = self.__class__.__name__
+        return d
+
+    def keys(self):
+        return list(self.to_dict().keys())
+
+    def values(self):
+        return list(self.to_dict().values())
+
+    def __getitem__(self, key):
+        """
+        Returns the value of a field.
+        """
+        return getattr(self, key)
+
+    @classmethod
+    def from_dict(cls, data: dict) -> CacheEntry:
+        """
+        Initializes a CacheEntry object from its dictionary representation.
+        """
+        return cls(**data)
+
+    def __eq__(self, other: CacheEntry) -> bool:
+        """
+        Checks if two CacheEntry objects are equal.
+        - Does not include timestamp in the comparison.
+        """
+        if not isinstance(other, CacheEntry):
+            return False
+        for field in self.all_fields:
+            if getattr(self, field) != getattr(other, field) and field != "timestamp":
+                return False
+        return True
+
+    def __repr__(self) -> str:
+        """
+        Returns a string representation of a CacheEntry.
+        """
+        return (
+            f"CacheEntry(model={repr(self.model)}, "
+            f"parameters={self.parameters}, "
+            f"system_prompt={repr(self.system_prompt)}, "
+            f"user_prompt={repr(self.user_prompt)}, "
+            f"output={repr(self.output)}, "
+            f"iteration={self.iteration}, "
+            f"timestamp={self.timestamp})"
+        )
+
+    @classmethod
+    def example(cls, randomize: bool = False) -> CacheEntry:
+        """
+        Returns an example CacheEntry instance.
+
+        :param randomize: If True, adds a random string to the system prompt.
+        """
+        # if random, create a uuid
+        addition = "" if not randomize else str(uuid4())
+        return CacheEntry(
+            model="gpt-3.5-turbo",
+            parameters={"temperature": 0.5},
+            system_prompt=f"The quick brown fox jumps over the lazy dog.{addition}",
+            user_prompt="What does the fox say?",
+            output="The fox says 'hello'",
+            iteration=1,
+            timestamp=int(datetime.datetime.now(datetime.timezone.utc).timestamp()),
+        )
+
+    @classmethod
+    def example_dict(cls) -> dict:
+        """
+        Returns an example dictionary with a single CacheEntry.
+        - This will be useful one level up, in the Cache class.
+        """
+        cache_entry = cls.example()
+        return {cache_entry.key: cache_entry}
+
+    @classmethod
+    def fetch_input_example(cls) -> dict:
+        """
+        Creates an example input for a 'fetch' operation.
+        - This will be useful one level up, in the Cache class.
+        """
+        input = cls.example().to_dict()
+        _ = input.pop("timestamp")
+        _ = input.pop("output")
+        return input
+
+    @classmethod
+    def store_input_example(cls) -> dict:
+        """
+        Creates an example input for a 'store' operation.
+        - This will be useful one level up, in the Cache class.
+        """
+        input = cls.example().to_dict()
+        _ = input.pop("timestamp")
+        input["response"] = input.pop("output")
+        return input
+
+
+def main():
+    from edsl.data.CacheEntry import CacheEntry
+
+    # an example of how a cache entry looks
+    cache_entry = CacheEntry.example()
+    cache_entry
+
+    # .key property returns the hash of the cache entry
+    cache_entry.key
+    # to dict / from dict
+    cache_entry.to_dict()
+    CacheEntry.from_dict(cache_entry.to_dict())
+    # TODO: this will be false because equality includes timestamp
+    CacheEntry.from_dict(cache_entry.to_dict()) == CacheEntry.example()
+    # equality by checking values
+    cache_entry == CacheEntry.example()
+    # equality by checking keys
+    cache_entry.key == CacheEntry.example().key
+    # evalable repr
+    eval(repr(cache_entry)) == cache_entry
+    # not sure what these are useful for yet
+    cache_entry.example_dict()
+    cache_entry.fetch_input_example()
+    cache_entry.store_input_example()
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

edsl/data/CacheHandler.py

@@ -0,0 +1,168 @@
+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

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