edsl 0.1.29.dev6__py3-none-any.whl → 0.1.30__py3-none-any.whl

edsl/base/Base.py

@@ -0,0 +1,289 @@
+"""Base class for all classes in the package. It provides rich printing and persistence of objects."""
+
+from abc import ABC, abstractmethod, ABCMeta
+import gzip
+import io
+import json
+from typing import Any, Optional, Union
+from uuid import UUID
+from IPython.display import display
+from rich.console import Console
+
+
+class RichPrintingMixin:
+    """Mixin for rich printing and persistence of objects."""
+
+    def _for_console(self):
+        """Return a string representation of the object for console printing."""
+        with io.StringIO() as buf:
+            console = Console(file=buf, record=True)
+            table = self.rich_print()
+            console.print(table)
+            return console.export_text()
+
+    def __str__(self):
+        """Return a string representation of the object for console printing."""
+        return self._for_console()
+
+    def print(self):
+        """Print the object to the console."""
+        from edsl.utilities.utilities import is_notebook
+
+        if is_notebook():
+            display(self.rich_print())
+        else:
+            from rich.console import Console
+
+            console = Console()
+            console.print(self.rich_print())
+
+
+class PersistenceMixin:
+    """Mixin for saving and loading objects to and from files."""
+
+    def push(
+        self,
+        description: Optional[str] = None,
+        visibility: Optional[str] = "unlisted",
+    ):
+        """Post the object to coop."""
+        from edsl.coop import Coop
+
+        c = Coop()
+        return c.create(self, description, visibility)
+
+    @classmethod
+    def pull(cls, id_or_url: Union[str, UUID], exec_profile=None):
+        """Pull the object from coop."""
+        from edsl.coop import Coop
+
+        if id_or_url.startswith("http"):
+            uuid_value = id_or_url.split("/")[-1]
+        else:
+            uuid_value = id_or_url
+
+        c = Coop()
+
+        return c._get_base(cls, uuid_value, exec_profile=exec_profile)
+
+    @classmethod
+    def delete(cls, id_or_url: Union[str, UUID]):
+        """Delete the object from coop."""
+        from edsl.coop import Coop
+
+        c = Coop()
+        return c._delete_base(cls, id_or_url)
+
+    @classmethod
+    def patch(
+        cls,
+        id_or_url: Union[str, UUID],
+        description: Optional[str] = None,
+        value: Optional[Any] = None,
+        visibility: Optional[str] = None,
+    ):
+        """
+        Patch an uploaded objects attributes.
+        - `description` changes the description of the object on Coop
+        - `value` changes the value of the object on Coop. **has to be an EDSL object**
+        - `visibility` changes the visibility of the object on Coop
+        """
+        from edsl.coop import Coop
+
+        c = Coop()
+        return c._patch_base(cls, id_or_url, description, value, visibility)
+
+    @classmethod
+    def search(cls, query):
+        """Search for objects on coop."""
+        from edsl.coop import Coop
+
+        c = Coop()
+        return c.search(cls, query)
+
+    def save(self, filename, compress=True):
+        """Save the object to a file as zippped JSON.
+
+        >>> obj.save("obj.json.gz")
+
+        """
+        if filename.endswith("json.gz"):
+            import warnings
+
+            warnings.warn(
+                "Do not apply the file extensions. The filename should not end with 'json.gz'."
+            )
+            filename = filename[:-7]
+        if filename.endswith("json"):
+            filename = filename[:-4]
+            warnings.warn(
+                "Do not apply the file extensions. The filename should not end with 'json'."
+            )
+
+        if compress:
+            with gzip.open(filename + ".json.gz", "wb") as f:
+                f.write(json.dumps(self.to_dict()).encode("utf-8"))
+        else:
+            with open(filename + ".json", "w") as f:
+                f.write(json.dumps(self.to_dict()))
+
+    @staticmethod
+    def open_compressed_file(filename):
+        with gzip.open(filename, "rb") as f:
+            file_contents = f.read()
+            file_contents_decoded = file_contents.decode("utf-8")
+            d = json.loads(file_contents_decoded)
+        return d
+
+    @staticmethod
+    def open_regular_file(filename):
+        with open(filename, "r") as f:
+            d = json.loads(f.read())
+        return d
+
+    @classmethod
+    def load(cls, filename):
+        """Load the object from a file.
+
+        >>> obj = cls.load("obj.json.gz")
+
+        """
+
+        if filename.endswith("json.gz"):
+            d = cls.open_compressed_file(filename)
+        elif filename.endswith("json"):
+            d = cls.open_regular_file(filename)
+        else:
+            try:
+                d = cls.open_compressed_file(filename)
+            except:
+                d = cls.open_regular_file(filename)
+            finally:
+                raise ValueError("File must be a json or json.gz file")
+
+        return cls.from_dict(d)
+
+
+class RegisterSubclassesMeta(ABCMeta):
+    """Metaclass for registering subclasses."""
+
+    _registry = {}
+
+    def __init__(cls, name, bases, nmspc):
+        """Register the class in the registry upon creation."""
+        super(RegisterSubclassesMeta, cls).__init__(name, bases, nmspc)
+        if cls.__name__ != "Base":
+            RegisterSubclassesMeta._registry[cls.__name__] = cls
+
+    @staticmethod
+    def get_registry():
+        """Return the registry of subclasses."""
+        return dict(RegisterSubclassesMeta._registry)
+
+
+class DiffMethodsMixin:
+    def __sub__(self, other):
+        """Return the difference between two objects."""
+        from edsl.BaseDiff import BaseDiff
+
+        return BaseDiff(self, other)
+
+
+class Base(
+    RichPrintingMixin,
+    PersistenceMixin,
+    DiffMethodsMixin,
+    ABC,
+    metaclass=RegisterSubclassesMeta,
+):
+    """Base class for all classes in the package."""
+
+    # def __getitem__(self, key):
+    #     return getattr(self, key)
+
+    # @abstractmethod
+    # def _repr_html_(self) -> str:
+    #     raise NotImplementedError("This method is not implemented yet.")
+
+    # @abstractmethod
+    # def _repr_(self) -> str:
+    #     raise NotImplementedError("This method is not implemented yet.")
+
+    def keys(self):
+        """Return the keys of the object."""
+        _keys = list(self.to_dict().keys())
+        if "edsl_version" in _keys:
+            _keys.remove("edsl_version")
+        if "edsl_class_name" in _keys:
+            _keys.remove("edsl_class_name")
+        return _keys
+
+    def values(self):
+        """Return the values of the object."""
+        data = self.to_dict()
+        keys = self.keys()
+        return {data[key] for key in keys}
+
+    def _repr_html_(self):
+        from edsl.utilities.utilities import data_to_html
+
+        return data_to_html(self.to_dict())
+
+    # def html(self):
+    #     html_string = self._repr_html_()
+    #     import tempfile
+    #     import webbrowser
+
+    #     with tempfile.NamedTemporaryFile("w", delete=False, suffix=".html") as f:
+    #         # print("Writing HTML to", f.name)
+    #         f.write(html_string)
+    #         webbrowser.open(f.name)
+
+    def __eq__(self, other):
+        """Return whether two objects are equal."""
+        import inspect
+
+        if not isinstance(other, self.__class__):
+            return False
+        if "sort" in inspect.signature(self._to_dict).parameters:
+            return self._to_dict(sort=True) == other._to_dict(sort=True)
+        else:
+            return self._to_dict() == other._to_dict()
+
+    @abstractmethod
+    def example():
+        """This method should be implemented by subclasses."""
+        raise NotImplementedError("This method is not implemented yet.")
+
+    @abstractmethod
+    def rich_print():
+        """This method should be implemented by subclasses."""
+        raise NotImplementedError("This method is not implemented yet.")
+
+    @abstractmethod
+    def to_dict():
+        """This method should be implemented by subclasses."""
+        raise NotImplementedError("This method is not implemented yet.")
+
+    @abstractmethod
+    def from_dict():
+        """This method should be implemented by subclasses."""
+        raise NotImplementedError("This method is not implemented yet.")
+
+    @abstractmethod
+    def code():
+        """This method should be implemented by subclasses."""
+        raise NotImplementedError("This method is not implemented yet.")
+
+    def show_methods(self, show_docstrings=True):
+        """Show the methods of the object."""
+        public_methods_with_docstrings = [
+            (method, getattr(self, method).__doc__)
+            for method in dir(self)
+            if callable(getattr(self, method)) and not method.startswith("_")
+        ]
+        if show_docstrings:
+            for method, documentation in public_methods_with_docstrings:
+                print(f"{method}: {documentation}")
+        else:
+            return [x[0] for x in public_methods_with_docstrings]

edsl/config.py

@@ -1,11 +1,11 @@
 """This module provides a Config class that loads environment variables from a .env file and sets them as class attributes."""
 
 import os
-from dotenv import load_dotenv, find_dotenv
 from edsl.exceptions import (
     InvalidEnvironmentVariableError,
     MissingEnvironmentVariableError,
 )
+from dotenv import load_dotenv, find_dotenv
 
 # valid values for EDSL_RUN_MODE
 EDSL_RUN_MODES = ["development", "development-testrun", "production"]
@@ -96,6 +96,7 @@ class Config:
         Loads the .env
         - Overrides existing env vars unless EDSL_RUN_MODE=="development-testrun"
         """
+
         override = True
         if self.EDSL_RUN_MODE == "development-testrun":
             override = False

edsl/conversation/car_buying.py

@@ -30,7 +30,7 @@ c1 = Conversation(agent_list=AgentList([a1, a3, a2]), max_turns=5, verbose=True)
 c2 = Conversation(agent_list=AgentList([a1, a2]), max_turns=5, verbose=True)
 
 c = Cache.load("car_talk.json.gz")
-breakpoint()
+# breakpoint()
 combo = ConversationList([c1, c2], cache=c)
 combo.run()
 results = combo.to_results()

edsl/coop/utils.py

@@ -10,7 +10,7 @@ from edsl import (
     Study,
 )
 from edsl.questions import QuestionBase
-from typing import Literal, Type, Union
+from typing import Literal, Optional, Type, Union
 
 EDSLObject = Union[
     Agent,
@@ -94,3 +94,30 @@ class ObjectRegistry:
         if EDSL_object is None:
             raise ValueError(f"EDSL class not found for {object_type=}")
         return EDSL_object
+
+    @classmethod
+    def get_registry(
+        cls,
+        subclass_registry: Optional[dict] = None,
+        exclude_classes: Optional[list] = None,
+    ) -> dict:
+        """
+        Return the registry of objects.
+
+        Exclude objects that are already registered in subclass_registry.
+        This allows the user to isolate Coop-only objects.
+
+        Also exclude objects if their class name is in the exclude_classes list.
+        """
+
+        if subclass_registry is None:
+            subclass_registry = {}
+        if exclude_classes is None:
+            exclude_classes = []
+
+        return {
+            class_name: o["edsl_class"]
+            for o in cls.objects
+            if (class_name := o["edsl_class"].__name__) not in subclass_registry
+            and class_name not in exclude_classes
+        }

edsl/data/Cache.py

@@ -7,17 +7,10 @@ import json
 import os
 import warnings
 from typing import Optional, Union
-
-from edsl.config import CONFIG
-from edsl.data.CacheEntry import CacheEntry
-from edsl.data.SQLiteDict import SQLiteDict
 from edsl.Base import Base
+from edsl.data.CacheEntry import CacheEntry
 from edsl.utilities.utilities import dict_hash
-
-from edsl.utilities.decorators import (
-    add_edsl_version,
-    remove_edsl_version,
-)
+from edsl.utilities.decorators import add_edsl_version, remove_edsl_version
 
 
 class Cache(Base):
@@ -38,9 +31,10 @@ class Cache(Base):
         self,
         *,
         filename: Optional[str] = None,
-        data: Optional[Union[SQLiteDict, dict]] = None,
+        data: Optional[Union["SQLiteDict", dict]] = None,
         immediate_write: bool = True,
         method=None,
+        verbose=False,
     ):
         """
         Create two dictionaries to store the cache data.
@@ -59,6 +53,7 @@ class Cache(Base):
         self.new_entries = {}
         self.new_entries_to_write_later = {}
         self.coop = None
+        self.verbose = verbose
 
         self.filename = filename
         if filename and data:
@@ -104,6 +99,8 @@ class Cache(Base):
 
     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 Exception("Not all values are CacheEntry instances")
         if self.method is not None:
@@ -120,7 +117,7 @@ class Cache(Base):
         system_prompt: str,
         user_prompt: str,
         iteration: int,
-    ) -> Union[None, str]:
+    ) -> tuple(Union[None, str], str):
         """
         Fetch a value (LLM output) from the cache.
 
@@ -133,11 +130,13 @@ class Cache(Base):
         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) is None
+        >>> 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
+
         key = CacheEntry.gen_key(
             model=model,
             parameters=parameters,
@@ -147,8 +146,13 @@ class Cache(Base):
         )
         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
-        return None if entry is None else entry.output
+        else:
+            if self.verbose:
+                print(f"Cache miss for key: {key}")
+        return None if entry is None else entry.output, key
 
     def store(
         self,
@@ -171,6 +175,7 @@ class Cache(Base):
         * 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`
         """
+
         entry = CacheEntry(
             model=model,
             parameters=parameters,
@@ -188,13 +193,14 @@ class Cache(Base):
         return key
 
     def add_from_dict(
-        self, new_data: dict[str, CacheEntry], write_now: Optional[bool] = True
+        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`).
         """
+
         for key, value in new_data.items():
             if key in self.data:
                 if value != self.data[key]:
@@ -231,6 +237,8 @@ class Cache(Base):
 
         :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():
@@ -242,6 +250,8 @@ class Cache(Base):
         """
         Construct a Cache from a SQLite database.
         """
+        from edsl.data.SQLiteDict import SQLiteDict
+
         return cls(data=SQLiteDict(db_path))
 
     @classmethod
@@ -268,6 +278,8 @@ class Cache(Base):
         * 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")
 
@@ -286,6 +298,8 @@ class Cache(Base):
         """
         ## 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
@@ -340,6 +354,9 @@ class Cache(Base):
         for key, entry in self.new_entries_to_write_later.items():
             self.data[key] = entry
 
+        if self.filename:
+            self.write(self.filename)
+
     ####################
     # DUNDER / USEFUL
     ####################
@@ -456,12 +473,18 @@ class Cache(Base):
         webbrowser.open("file://" + filepath)
 
     @classmethod
-    def example(cls) -> Cache:
+    def example(cls, randomize: bool = False) -> Cache:
         """
-        Return an example Cache.
-        The example Cache has one entry.
+        Returns an example Cache instance.
+
+        :param randomize: If True, uses CacheEntry's randomize method.
         """
-        return cls(data={CacheEntry.example().key: CacheEntry.example()})
+        return cls(
+            data={
+                CacheEntry.example(randomize).key: CacheEntry.example(),
+                CacheEntry.example(randomize).key: CacheEntry.example(),
+            }
+        )
 
 
 if __name__ == "__main__":

edsl/data/CacheEntry.py

@@ -2,11 +2,8 @@ from __future__ import annotations
 import json
 import datetime
 import hashlib
-import random
 from typing import Optional
-
-
-# TODO: Timestamp should probably be float?
+from uuid import uuid4
 
 
 class CacheEntry:
@@ -151,10 +148,12 @@ class CacheEntry:
     @classmethod
     def example(cls, randomize: bool = False) -> CacheEntry:
         """
-        Returns a CacheEntry example.
+        Returns an example CacheEntry instance.
+
+        :param randomize: If True, adds a random string to the system prompt.
         """
-        # if random, create a random number for 0-100
-        addition = "" if not randomize else str(random.randint(0, 1000))
+        # if random, create a uuid
+        addition = "" if not randomize else str(uuid4())
         return CacheEntry(
             model="gpt-3.5-turbo",
             parameters={"temperature": 0.5},

edsl/data/SQLiteDict.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 import json
-from sqlalchemy import create_engine
-from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.orm import sessionmaker
 from typing import Any, Generator, Optional, Union
+
 from edsl.config import CONFIG
 from edsl.data.CacheEntry import CacheEntry
 from edsl.data.orm import Base, Data
@@ -25,10 +23,16 @@ class SQLiteDict:
         >>> import os; os.unlink(temp_db_path)  # Clean up the temp file after the test
 
         """
+        from sqlalchemy.exc import SQLAlchemyError
+        from sqlalchemy.orm import sessionmaker
+        from sqlalchemy import create_engine
+
         self.db_path = db_path or CONFIG.get("EDSL_DATABASE_PATH")
         if not self.db_path.startswith("sqlite:///"):
             self.db_path = f"sqlite:///{self.db_path}"
         try:
+            from edsl.data.orm import Base, Data
+
             self.engine = create_engine(self.db_path, echo=False, future=True)
             Base.metadata.create_all(self.engine)
             self.Session = sessionmaker(bind=self.engine)
@@ -55,6 +59,8 @@ class SQLiteDict:
         if not isinstance(value, CacheEntry):
             raise ValueError(f"Value must be a CacheEntry object (got {type(value)}).")
         with self.Session() as db:
+            from edsl.data.orm import Base, Data
+
             db.merge(Data(key=key, value=json.dumps(value.to_dict())))
             db.commit()
 
@@ -69,6 +75,8 @@ class SQLiteDict:
         True
         """
         with self.Session() as db:
+            from edsl.data.orm import Base, Data
+
             value = db.query(Data).filter_by(key=key).first()
             if not value:
                 raise KeyError(f"Key '{key}' not found.")

edsl/data_transfer_models.py

@@ -17,6 +17,8 @@ class AgentResponseDict(UserDict):
         cached_response=None,
         raw_model_response=None,
         simple_model_raw_response=None,
+        cache_used=None,
+        cache_key=None,
     ):
         """Initialize the AgentResponseDict object."""
         usage = usage or {"prompt_tokens": 0, "completion_tokens": 0}
@@ -30,5 +32,7 @@ class AgentResponseDict(UserDict):
                 "cached_response": cached_response,
                 "raw_model_response": raw_model_response,
                 "simple_model_raw_response": simple_model_raw_response,
+                "cache_used": cache_used,
+                "cache_key": cache_key,
             }
         )

edsl/jobs/Answers.py

@@ -8,7 +8,15 @@ class Answers(UserDict):
     """Helper class to hold the answers to a survey."""
 
     def add_answer(self, response, question) -> None:
-        """Add a response to the answers dictionary."""
+        """Add a response to the answers dictionary.
+
+        >>> from edsl import QuestionFreeText
+        >>> q = QuestionFreeText.example()
+        >>> answers = Answers()
+        >>> answers.add_answer({"answer": "yes"}, q)
+        >>> answers[q.question_name]
+        'yes'
+        """
         answer = response.get("answer")
         comment = response.pop("comment", None)
         # record the answer
@@ -42,3 +50,9 @@ class Answers(UserDict):
             table.add_row(attr_name, repr(attr_value))
 
         return table
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()