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

edsl/data/Cache.py

@@ -1,109 +1,580 @@
-from edsl import Config
-import sqlite3
-from collections import UserList
+"""
+The `Cache` class is used to store responses from a language model.
+"""
+
+from __future__ import annotations
+import json
+import os
+import warnings
+from typing import Optional, Union, TYPE_CHECKING
 from edsl.Base import Base
-from rich.table import Table
 
-config = Config()
+from edsl.utilities.remove_edsl_version import remove_edsl_version
+from edsl.exceptions.cache import CacheError
 
-path = config.get("EDSL_DATABASE_PATH")
 
+class Cache(Base):
+    """
+    A class that represents a cache of responses from a language model.
 
-class Cache(Base, UserList):
-    def __init__(self, data=None, schema=None):
-        self.db_path = config.get("EDSL_DATABASE_PATH")[len("sqlite:///") :]
-        self.table_name = "responses"
+    :param data: The data to initialize the cache with.
+    :param immediate_write: Whether to write to the cache immediately after storing a new entry.
 
-        if data is None or schema is None:
-            result = self.load_data()
-            self.data = result["data"]
-            self.schema = result["schema"]
-        else:
+    Deprecated:
+
+    :param method: The method of storage to use for the cache.
+    """
+
+    __documentation__ = "https://docs.expectedparrot.com/en/latest/data.html"
+
+    data = {}
+
+    def __init__(
+        self,
+        *,
+        filename: Optional[str] = None,
+        data: Optional[Union["SQLiteDict", dict]] = None,
+        immediate_write: bool = True,
+        method=None,
+        verbose=False,
+    ):
+        """
+        Create two dictionaries to store the cache data.
+
+        :param filename: The name of the file to read/write the cache from/to.
+        :param data: The data to initialize the cache with.
+        :param immediate_write: Whether to write to the cache immediately after storing a new entry.
+        :param method: The method of storage to use for the cache.
+
+        """
+
+        # self.data_at_init = data or {}
+        self.fetched_data = {}
+        self.immediate_write = immediate_write
+        self.method = method
+        self.new_entries = {}
+        self.new_entries_to_write_later = {}
+        self.coop = None
+        self.verbose = verbose
+
+        self.filename = filename
+        if filename and data:
+            raise CacheError("Cannot provide both filename and data")
+        if filename is None and data is None:
+            data = {}
+        if data is not None:
             self.data = data
-            self.schema = schema
-
-    def _connect(self, db_path=None):
-        return sqlite3.connect(db_path if db_path else self.db_path)
-
-    def load_data(self):
-        with self._connect() as conn:
-            cur = conn.cursor()
-            cur.execute(f"PRAGMA table_info({self.table_name})")
-            schema = cur.fetchall()
-            cur.execute(f"SELECT * FROM {self.table_name}")
-            return {
-                "data": [
-                    dict(zip([col[1] for col in schema], row)) for row in cur.fetchall()
-                ],
-                "schema": schema,
-            }
+        if filename is not None:
+            self.data = {}
+            if filename.endswith(".jsonl"):
+                if os.path.exists(filename):
+                    self.add_from_jsonl(filename)
+                else:
+                    print(
+                        f"File {filename} not found, but will write to this location."
+                    )
+            elif filename.endswith(".db"):
+                if os.path.exists(filename):
+                    self.add_from_sqlite(filename)
+            else:
+                raise CacheError("Invalid file extension. Must be .jsonl or .db")
+
+        self._perform_checks()
+
+    def code(sefl):
+        pass
+        # raise NotImplementedError("This method is not implemented yet.")
+
+    def keys(self):
+        """
+        >>> from edsl import Cache
+        >>> Cache.example().keys()
+        ['5513286eb6967abc0511211f0402587d']
+        """
+        return list(self.data.keys())
+
+    def values(self):
+        """
+        >>> from edsl import Cache
+        >>> Cache.example().values()
+        [CacheEntry(...)]
+        """
+        return list(self.data.values())
+
+    def items(self):
+        return zip(self.keys(), self.values())
+
+    def new_entries_cache(self) -> Cache:
+        """Return a new Cache object with the new entries."""
+        return Cache(data={**self.new_entries, **self.fetched_data})
+
+    def _perform_checks(self):
+        """Perform checks on the cache."""
+        from edsl.data.CacheEntry import CacheEntry
+
+        if any(not isinstance(value, CacheEntry) for value in self.data.values()):
+            raise CacheError("Not all values are CacheEntry instances")
+        if self.method is not None:
+            warnings.warn("Argument `method` is deprecated", DeprecationWarning)
+
+    ####################
+    # READ/WRITE
+    ####################
+    def fetch(
+        self,
+        *,
+        model: str,
+        parameters: dict,
+        system_prompt: str,
+        user_prompt: str,
+        iteration: int,
+    ) -> tuple(Union[None, str], str):
+        """
+        Fetch a value (LLM output) from the cache.
+
+        :param model: The name of the language model.
+        :param parameters: The model parameters.
+        :param system_prompt: The system prompt.
+        :param user_prompt: The user prompt.
+        :param iteration: The iteration number.
+
+        Return None if the response is not found.
+
+        >>> c = Cache()
+        >>> c.fetch(model="gpt-3", parameters="default", system_prompt="Hello", user_prompt="Hi", iteration=1)[0] is None
+        True
+
+
+        """
+        from edsl.data.CacheEntry import CacheEntry
 
-    def save_data_to_new_db(self, new_db_path):
-        with self._connect(new_db_path) as new_conn:
-            new_cur = new_conn.cursor()
-            # Create table in new database
-            columns = ", ".join([f"{col[1]} {col[2]}" for col in self.schema])
-            new_cur.execute(f"CREATE TABLE IF NOT EXISTS {self.table_name} ({columns})")
-
-            # Insert data into new table
-            placeholders = ", ".join(["?"] * len(self.schema))
-            new_cur.executemany(
-                f"INSERT INTO {self.table_name} ({', '.join([col[1] for col in self.schema])}) VALUES ({placeholders})",
-                [tuple(row[col[1]] for col in self.schema) for row in self.data],
-            )
-            new_conn.commit()
-
-    def _table_row(self, row):
-        table = Table()
-        table.add_column("Key")
-        table.add_column("Value")
-        for key, value in row.items():
-            table.add_row(str(key), str(value))
-        return table
-
-    def rich_print(self):
-        table = Table(title="Cache")
-        table.add_column("Entry", style="bold")
-        for index, row in enumerate(self.data):
-            table.add_row(self._table_row(row))
-        return table
+        key = CacheEntry.gen_key(
+            model=model,
+            parameters=parameters,
+            system_prompt=system_prompt,
+            user_prompt=user_prompt,
+            iteration=iteration,
+        )
+        entry = self.data.get(key, None)
+        if entry is not None:
+            if self.verbose:
+                print(f"Cache hit for key: {key}")
+            self.fetched_data[key] = entry
+        else:
+            if self.verbose:
+                print(f"Cache miss for key: {key}")
+        return None if entry is None else entry.output, key
+
+    def store(
+        self,
+        model: str,
+        parameters: str,
+        system_prompt: str,
+        user_prompt: str,
+        response: dict,
+        iteration: int,
+    ) -> str:
+        """
+        Add a new key-value pair to the cache.
+
+        * Key is a hash of the input parameters.
+        * Output is the response from the language model.
+
+        How it works:
+
+        * The key-value pair is added to `self.new_entries`
+        * If `immediate_write` is True , the key-value pair is added to `self.data`
+        * If `immediate_write` is False, the key-value pair is added to `self.new_entries_to_write_later`
+
+        >>> from edsl import Cache, Model, Question
+        >>> m = Model("test")
+        >>> c = Cache()
+        >>> len(c)
+        0
+        >>> results = Question.example("free_text").by(m).run(cache = c, disable_remote_cache = True, disable_remote_inference = True)
+        >>> len(c)
+        1
+        """
+        from edsl.data.CacheEntry import CacheEntry
+
+        entry = CacheEntry(
+            model=model,
+            parameters=parameters,
+            system_prompt=system_prompt,
+            user_prompt=user_prompt,
+            output=json.dumps(response),
+            iteration=iteration,
+        )
+        key = entry.key
+        self.new_entries[key] = entry
+        if self.immediate_write:
+            self.data[key] = entry
+        else:
+            self.new_entries_to_write_later[key] = entry
+        return key
+
+    def add_from_dict(
+        self, new_data: dict[str, "CacheEntry"], write_now: Optional[bool] = True
+    ) -> None:
+        """
+        Add entries to the cache from a dictionary.
+
+        :param write_now: Whether to write to the cache immediately (similar to `immediate_write`).
+        """
+        from edsl.data.CacheEntry import CacheEntry
+
+        for key, value in new_data.items():
+            if key in self.data:
+                if value != self.data[key]:
+                    raise CacheError("Mismatch in values")
+            if not isinstance(value, CacheEntry):
+                raise CacheError(f"Wrong type - the observed type is {type(value)}")
+
+        self.new_entries.update(new_data)
+        if write_now:
+            self.data.update(new_data)
+        else:
+            self.new_entries_to_write_later.update(new_data)
+
+    def add_from_jsonl(self, filename: str, write_now: Optional[bool] = True) -> None:
+        """
+        Add entries to the cache from a JSONL.
+
+        :param write_now: Whether to write to the cache immediately (similar to `immediate_write`).
+        """
+        from edsl.data.CacheEntry import CacheEntry
+
+        with open(filename, "a+") as f:
+            f.seek(0)
+            lines = f.readlines()
+        new_data = {}
+        for line in lines:
+            d = json.loads(line)
+            key = list(d.keys())[0]
+            value = list(d.values())[0]
+            new_data[key] = CacheEntry(**value)
+        self.add_from_dict(new_data=new_data, write_now=write_now)
+
+    def add_from_sqlite(self, db_path: str, write_now: Optional[bool] = True):
+        """
+        Add entries to the cache from an SQLite database.
+
+        :param write_now: Whether to write to the cache immediately (similar to `immediate_write`).
+        """
+        from edsl.data.SQLiteDict import SQLiteDict
+
+        db = SQLiteDict(db_path)
+        new_data = {}
+        for key, value in db.items():
+            new_data[key] = CacheEntry(**value)
+        self.add_from_dict(new_data=new_data, write_now=write_now)
 
     @classmethod
-    def example(cls):
-        data0 = """{'id': 1, 'model': 'gpt-3.5-turbo'}"""
-        schema = [
-            (0, "id", "INTEGER", 1, None, 1),
-            (1, "model", "VARCHAR(100)", 1, None, 0),
-            (2, "parameters", "TEXT", 1, None, 0),
-            (3, "system_prompt", "TEXT", 1, None, 0),
-            (4, "prompt", "TEXT", 1, None, 0),
-            (5, "output", "TEXT", 1, None, 0),
-        ]
-        return cls(data=[eval(data0)], schema=schema)
+    def from_sqlite_db(cls, db_path: str) -> Cache:
+        """
+        Construct a Cache from a SQLite database.
+        """
+        from edsl.data.SQLiteDict import SQLiteDict
 
-    def __repr__(self):
-        return f"Cache(data={self.data}, schema={self.schema})"
+        return cls(data=SQLiteDict(db_path))
+
+    @classmethod
+    def from_local_cache(cls) -> Cache:
+        """
+        Construct a Cache from a local cache file.
+        """
+        from edsl.config import CONFIG
+
+        CACHE_PATH = CONFIG.get("EDSL_DATABASE_PATH")
+        path = CACHE_PATH.replace("sqlite:///", "")
+        # db_path = os.path.join(os.path.dirname(path), "data.db")
+        return cls.from_sqlite_db(path)
+
+    @classmethod
+    def from_jsonl(cls, jsonlfile: str, db_path: Optional[str] = None) -> Cache:
+        """
+        Construct a Cache from a JSONL file.
+
+        :param jsonlfile: The path to the JSONL file of cache entries.
+        :param db_path: The path to the SQLite database used to store the cache.
 
-    def code():
+        * If `db_path` is None, the cache will be stored in memory, as a dictionary.
+        * If `db_path` is provided, the cache will be stored in an SQLite database.
+        """
+        # if a file doesn't exist at jsonfile, throw an error
+        from edsl.data.SQLiteDict import SQLiteDict
+
+        if not os.path.exists(jsonlfile):
+            raise FileNotFoundError(f"File {jsonlfile} not found")
+
+        if db_path is None:
+            data = {}
+        else:
+            data = SQLiteDict(db_path)
+
+        cache = Cache(data=data)
+        cache.add_from_jsonl(jsonlfile)
+        return cache
+
+    def write_sqlite_db(self, db_path: str) -> None:
+        """
+        Write the cache to an SQLite database.
+        """
+        ## TODO: Check to make sure not over-writing (?)
+        ## Should be added to SQLiteDict constructor (?)
+        from edsl.data.SQLiteDict import SQLiteDict
+
+        new_data = SQLiteDict(db_path)
+        for key, value in self.data.items():
+            new_data[key] = value
+
+    def write(self, filename: Optional[str] = None) -> None:
+        """
+        Write the cache to a file at the specified location.
+        """
+        if filename is None:
+            filename = self.filename
+        if filename.endswith(".jsonl"):
+            self.write_jsonl(filename)
+        elif filename.endswith(".db"):
+            self.write_sqlite_db(filename)
+        else:
+            raise CacheError("Invalid file extension. Must be .jsonl or .db")
+
+    def write_jsonl(self, filename: str) -> None:
+        """
+        Write the cache to a JSONL file.
+        """
+        path = os.path.join(os.getcwd(), filename)
+        with open(path, "w") as f:
+            for key, value in self.data.items():
+                f.write(json.dumps({key: value.to_dict()}) + "\n")
+
+    def to_scenario_list(self):
+        from edsl.scenarios.ScenarioList import ScenarioList
+        from edsl.scenarios.Scenario import Scenario
+
+        scenarios = []
+        for key, value in self.data.items():
+            new_d = value.to_dict()
+            new_d["cache_key"] = key
+            s = Scenario(new_d)
+            scenarios.append(s)
+        return ScenarioList(scenarios)
+
+    def __floordiv__(self, other: "Cache") -> "Cache":
+        """
+        Return a new Cache containing entries that are in self but not in other.
+        Uses // operator as alternative to subtraction.
+
+        :param other: Another Cache object to compare against
+        :return: A new Cache object containing unique entries
+
+        >>> from edsl.data.CacheEntry import CacheEntry
+        >>> ce1 = CacheEntry.example(randomize = True)
+        >>> ce2 = CacheEntry.example(randomize = True)
+        >>> ce2 = CacheEntry.example(randomize = True)
+        >>> c1 = Cache(data={ce1.key: ce1, ce2.key: ce2})
+        >>> c2 = Cache(data={ce1.key: ce1})
+        >>> c3 = c1 // c2
+        >>> len(c3)
+        1
+        >>> c3.data[ce2.key] == ce2
+        True
+        """
+        if not isinstance(other, Cache):
+            raise CacheError("Can only compare two caches")
+
+        diff_data = {k: v for k, v in self.data.items() if k not in other.data}
+        return Cache(data=diff_data, immediate_write=self.immediate_write)
+
+    @classmethod
+    def from_url(cls, db_path=None) -> Cache:
+        """
+        Construct a Cache object from a remote.
+        """
+        # ...do something here
+        # return Cache(data=db)
         pass
 
-    def add_row(self, row_data):
-        if set(row_data.keys()) != set([col[1] for col in self.schema]):
-            raise ValueError("Row keys do not match table schema")
-        self.data.append(row_data)
+    def __enter__(self):
+        """
+        Run when a context is entered.
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """
+        Run when a context is exited.
+        """
+        for key, entry in self.new_entries_to_write_later.items():
+            self.data[key] = entry
+
+        if self.filename:
+            self.write(self.filename)
+
+    def __hash__(self):
+        """Return the hash of the Cache."""
+        from edsl.utilities.utilities import dict_hash
+
+        return dict_hash(self.to_dict(add_edsl_version=False))
+
+    def to_dict(self, add_edsl_version=True) -> dict:
+        d = {k: v.to_dict() for k, v in self.data.items()}
+        if add_edsl_version:
+            from edsl import __version__
+
+            d["edsl_version"] = __version__
+            d["edsl_class_name"] = "Cache"
+
+        return d
+
+    def _summary(self):
+        return {"EDSL Class": "Cache", "Number of entries": len(self.data)}
 
-    def to_dict(self):
-        return {"schema": self.schema, "data": self.data}
+    def table(
+        self,
+        *fields,
+        tablefmt: Optional[str] = None,
+        pretty_labels: Optional[dict] = None,
+    ) -> str:
+        return self.to_dataset().table(
+            *fields, tablefmt=tablefmt, pretty_labels=pretty_labels
+        )
+
+    def select(self, *fields):
+        return self.to_dataset().select(*fields)
+
+    def tree(self, node_list: Optional[list[str]] = None):
+        return self.to_scenario_list().tree(node_list)
+
+    def to_dataset(self):
+        return self.to_scenario_list().to_dataset()
 
     @classmethod
-    def from_dict(cls, data):
-        raw_schema = data["schema"]
-        schema = [tuple(col) for col in raw_schema]  # Convert to tuples
-        cache = cls(data=data["data"], schema=schema)
-        return cache
+    @remove_edsl_version
+    def from_dict(cls, data) -> Cache:
+        """Construct a Cache from a dictionary."""
+        from edsl.data.CacheEntry import CacheEntry
+
+        newdata = {k: CacheEntry.from_dict(v) for k, v in data.items()}
+        return cls(data=newdata)
+
+    def __len__(self):
+        """Return the number of CacheEntry objects in the Cache."""
+        return len(self.data)
+
+    # TODO: Same inputs could give different results and this could be useful
+    # can't distinguish unless we do the ε trick or vary iterations
+    def __eq__(self, other_cache: "Cache") -> bool:
+        """
+        Check if two Cache objects are equal.
+        Does not verify their values are equal, only that they have the same keys.
+        """
+        if not isinstance(other_cache, Cache):
+            return False
+        return set(self.data.keys()) == set(other_cache.data.keys())
+
+    def __add__(self, other: "Cache"):
+        """
+        Combine two caches.
+        """
+        if not isinstance(other, Cache):
+            raise CacheError("Can only add two caches together")
+        self.data.update(other.data)
+        return self
+
+    def __repr__(self):
+        """
+        Return a string representation of the Cache object.
+        """
+        return (
+            f"Cache(data = {repr(self.data)}, immediate_write={self.immediate_write})"
+        )
+
+    ####################
+    # EXAMPLES
+    ####################
+    def fetch_input_example(self) -> dict:
+        """
+        Create an example input for a 'fetch' operation.
+        """
+        from edsl.data.CacheEntry import CacheEntry
+
+        return CacheEntry.fetch_input_example()
+
+    def to_html(self):
+        # json_str = json.dumps(self.data, indent=4)
+        d = {k: v.to_dict() for k, v in self.data.items()}
+        for key, value in d.items():
+            for k, v in value.items():
+                if isinstance(v, dict):
+                    d[key][k] = {kk: str(vv) for kk, vv in v.items()}
+                else:
+                    d[key][k] = str(v)
+
+        json_str = json.dumps(d, indent=4)
+
+        # HTML template with the JSON string embedded
+        html = f"""
+        <!DOCTYPE html>
+        <html>
+        <head>
+            <title>Display JSON</title>
+        </head>
+        <body>
+            <pre id="jsonData"></pre>
+            <script>
+                var json = {json_str};
+
+                // JSON.stringify with spacing to format
+                document.getElementById('jsonData').textContent = JSON.stringify(json, null, 4);
+            </script>
+        </body>
+        </html>
+        """
+        return html
+
+    def subset(self, keys: list[str]) -> Cache:
+        """
+        Return a subset of the Cache with the specified keys.
+        """
+        new_data = {k: v for k, v in self.data.items() if k in keys}
+        return Cache(data=new_data)
+
+    def view(self) -> None:
+        """View the Cache in a new browser tab."""
+        import tempfile
+        import webbrowser
+
+        html_content = self.to_html()
+        # Create a temporary file to hold the HTML
+        with tempfile.NamedTemporaryFile("w", delete=False, suffix=".html") as tmpfile:
+            tmpfile.write(html_content)
+            # Get the path to the temporary file
+            filepath = tmpfile.name
+
+        # Open the HTML file in a new browser tab
+        webbrowser.open("file://" + filepath)
+
+    @classmethod
+    def example(cls, randomize: bool = False) -> Cache:
+        """
+        Returns an example Cache instance.
+
+        :param randomize: If True, uses CacheEntry's randomize method.
+        """
+        from edsl.data.CacheEntry import CacheEntry
+
+        return cls(
+            data={
+                CacheEntry.example(randomize).key: CacheEntry.example(),
+                CacheEntry.example(randomize).key: CacheEntry.example(),
+            }
+        )
 
 
 if __name__ == "__main__":
-    cache = Cache()
-    cache.print()
-    # cache.load_data()
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)