edsl 0.1.27.dev2__py3-none-any.whl → 0.1.29__py3-none-any.whl

edsl/coop/utils.py

@@ -1,45 +1,49 @@
-from edsl import Agent, AgentList, Cache, Jobs, Results, Scenario, ScenarioList, Survey
-from edsl.notebooks import Notebook
+from edsl import (
+    Agent,
+    AgentList,
+    Cache,
+    Notebook,
+    Results,
+    Scenario,
+    ScenarioList,
+    Survey,
+    Study,
+)
 from edsl.questions import QuestionBase
-from typing import Literal, Type, Union
+from typing import Literal, Optional, Type, Union
 
 EDSLObject = Union[
     Agent,
     AgentList,
     Cache,
-    Jobs,
     Notebook,
     Type[QuestionBase],
     Results,
     Scenario,
     ScenarioList,
     Survey,
+    Study,
 ]
 
 ObjectType = Literal[
     "agent",
     "agent_list",
     "cache",
-    "job",
-    "question",
     "notebook",
+    "question",
     "results",
     "scenario",
     "scenario_list",
     "survey",
+    "study",
 ]
 
-ObjectPage = Literal[
-    "agents",
-    "agentlists",
-    "caches",
-    "jobs",
-    "notebooks",
-    "questions",
-    "results",
-    "scenarios",
-    "scenariolists",
-    "surveys",
+
+RemoteJobStatus = Literal[
+    "queued",
+    "running",
+    "completed",
+    "failed",
 ]
 
 VisibilityType = Literal[
@@ -55,62 +59,21 @@ class ObjectRegistry:
     """
 
     objects = [
-        {
-            "object_type": "agent",
-            "edsl_class": Agent,
-            "object_page": "agents",
-        },
-        {
-            "object_type": "agent_list",
-            "edsl_class": AgentList,
-            "object_page": "agentlists",
-        },
-        {
-            "object_type": "cache",
-            "edsl_class": Cache,
-            "object_page": "caches",
-        },
-        {
-            "object_type": "job",
-            "edsl_class": Jobs,
-            "object_page": "jobs",
-        },
-        {
-            "object_type": "question",
-            "edsl_class": QuestionBase,
-            "object_page": "questions",
-        },
-        {
-            "object_type": "notebook",
-            "edsl_class": Notebook,
-            "object_page": "notebooks",
-        },
-        {
-            "object_type": "results",
-            "edsl_class": Results,
-            "object_page": "results",
-        },
-        {
-            "object_type": "scenario",
-            "edsl_class": Scenario,
-            "object_page": "scenarios",
-        },
-        {
-            "object_type": "scenario_list",
-            "edsl_class": ScenarioList,
-            "object_page": "scenariolists",
-        },
-        {
-            "object_type": "survey",
-            "edsl_class": Survey,
-            "object_page": "surveys",
-        },
+        {"object_type": "agent", "edsl_class": Agent},
+        {"object_type": "agent_list", "edsl_class": AgentList},
+        {"object_type": "cache", "edsl_class": Cache},
+        {"object_type": "question", "edsl_class": QuestionBase},
+        {"object_type": "notebook", "edsl_class": Notebook},
+        {"object_type": "results", "edsl_class": Results},
+        {"object_type": "scenario", "edsl_class": Scenario},
+        {"object_type": "scenario_list", "edsl_class": ScenarioList},
+        {"object_type": "survey", "edsl_class": Survey},
+        {"object_type": "study", "edsl_class": Study},
     ]
     object_type_to_edsl_class = {o["object_type"]: o["edsl_class"] for o in objects}
     edsl_class_to_object_type = {
         o["edsl_class"].__name__: o["object_type"] for o in objects
     }
-    object_type_to_object_page = {o["object_type"]: o["object_page"] for o in objects}
 
     @classmethod
     def get_object_type_by_edsl_class(cls, edsl_object: EDSLObject) -> ObjectType:
@@ -133,5 +96,28 @@ class ObjectRegistry:
         return EDSL_object
 
     @classmethod
-    def get_object_page_by_object_type(cls, object_type: ObjectType) -> ObjectPage:
-        return cls.object_type_to_object_page.get(object_type)
+    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,12 +7,13 @@ import json
 import os
 import warnings
 from typing import Optional, Union
-
+import time
 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.SQLiteDict import SQLiteDict
+from edsl.Base import Base
+from edsl.utilities.utilities import dict_hash
 from edsl.utilities.decorators import (
     add_edsl_version,
     remove_edsl_version,
@@ -24,7 +25,6 @@ class Cache(Base):
     A class that represents a cache of responses from a language model.
 
     :param data: The data to initialize the cache with.
-    :param remote: Whether to sync the Cache with the server.
     :param immediate_write: Whether to write to the cache immediately after storing a new entry.
 
     Deprecated:
@@ -37,24 +37,51 @@ class Cache(Base):
     def __init__(
         self,
         *,
-        data: Optional[Union[SQLiteDict, dict]] = None,
-        remote: bool = False,
+        filename: Optional[str] = None,
+        data: Optional[Union["SQLiteDict", dict]] = None,
         immediate_write: bool = True,
         method=None,
     ):
         """
         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 = data or {}
+
         # self.data_at_init = data or {}
         self.fetched_data = {}
-        self.remote = remote
         self.immediate_write = immediate_write
         self.method = method
         self.new_entries = {}
         self.new_entries_to_write_later = {}
         self.coop = None
+
+        self.filename = filename
+        if filename and data:
+            raise ValueError("Cannot provide both filename and data")
+        if filename is None and data is None:
+            data = {}
+        if data is not None:
+            self.data = data
+        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 ValueError("Invalid file extension. Must be .jsonl or .db")
+
         self._perform_checks()
 
     def rich_print(sefl):
@@ -77,14 +104,12 @@ 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:
             warnings.warn("Argument `method` is deprecated", DeprecationWarning)
-        if self.remote:
-            from edsl.coop import Coop
-
-            self.coop = Coop()
 
     ####################
     # READ/WRITE
@@ -115,6 +140,8 @@ class Cache(Base):
 
 
         """
+        from edsl.data.CacheEntry import CacheEntry
+
         key = CacheEntry.gen_key(
             model=model,
             parameters=parameters,
@@ -148,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,
@@ -165,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]:
@@ -208,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():
@@ -219,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
@@ -245,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")
 
@@ -263,10 +298,25 @@ 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
 
+    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 ValueError("Invalid file extension. Must be .jsonl or .db")
+
     def write_jsonl(self, filename: str) -> None:
         """
         Write the cache to a JSONL file.
@@ -295,11 +345,6 @@ class Cache(Base):
         """
         Run when a context is entered.
         """
-        if self.remote:
-            print("Syncing local and remote caches")
-            exclude_keys = list(self.data.keys())
-            cache_entries = self.coop.get_cache_entries(exclude_keys)
-            self.add_from_dict({c.key: c for c in cache_entries}, write_now=True)
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
@@ -308,16 +353,21 @@ class Cache(Base):
         """
         for key, entry in self.new_entries_to_write_later.items():
             self.data[key] = entry
-        if self.remote:
-            _ = self.coop.create_cache_entries(cache_dict=self.new_entries)
 
     ####################
     # DUNDER / USEFUL
     ####################
+    def __hash__(self):
+        """Return the hash of the Cache."""
+        return dict_hash(self._to_dict())
+
+    def _to_dict(self) -> dict:
+        return {k: v.to_dict() for k, v in self.data.items()}
+
     @add_edsl_version
     def to_dict(self) -> dict:
         """Return the Cache as a dictionary."""
-        return {k: v.to_dict() for k, v in self.data.items()}
+        return self._to_dict()
 
     def _repr_html_(self):
         from edsl.utilities.utilities import data_to_html
@@ -359,7 +409,9 @@ class Cache(Base):
         """
         Return a string representation of the Cache object.
         """
-        return f"Cache(data = {repr(self.data)}, immediate_write={self.immediate_write}, remote={self.remote})"
+        return (
+            f"Cache(data = {repr(self.data)}, immediate_write={self.immediate_write})"
+        )
 
     ####################
     # EXAMPLES

edsl/data/CacheHandler.py

@@ -9,22 +9,22 @@ 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.
     """
-    print("All calls to 'run' will now use this cache by default.")
-    global _CACHE
-    _CACHE = cache
+    CONFIG.EDSL_SESSION_CACHE = cache
 
 
 def unset_session_cache() -> None:
     """
     Unset the session cache.
     """
-    global _CACHE
-    _CACHE = None
+    if hasattr(CONFIG, "EDSL_SESSION_CACHE"):
+        del CONFIG.EDSL_SESSION_CACHE
 
 
 class CacheHandler:
@@ -49,7 +49,9 @@ class CacheHandler:
         dir_path = os.path.dirname(path)
         if dir_path and not os.path.exists(dir_path):
             os.makedirs(dir_path)
-            print(f"Created cache directory: {dir_path}")
+            import warnings
+
+            warnings.warn(f"Created cache directory: {dir_path}")
 
     def gen_cache(self) -> Cache:
         """
@@ -58,9 +60,8 @@ class CacheHandler:
         if self.test:
             return Cache(data={})
 
-        if "_CACHE" in globals() and _CACHE is not None:
-            # print("Using globally-set cache.")
-            return _CACHE
+        if hasattr(CONFIG, "EDSL_SESSION_CACHE"):
+            return CONFIG.EDSL_SESSION_CACHE
 
         cache = Cache(data=SQLiteDict(self.CACHE_PATH))
         return cache

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/inference_services/AnthropicService.py

@@ -16,6 +16,7 @@ class AnthropicService(InferenceServiceABC):
     def available(cls):
         # TODO - replace with an API call
         return [
+            "claude-3-5-sonnet-20240620",
             "claude-3-opus-20240229",
             "claude-3-sonnet-20240229",
             "claude-3-haiku-20240307",

edsl/inference_services/DeepInfraService.py

@@ -1,7 +1,7 @@
 import aiohttp
 import json
 import requests
-from typing import Any
+from typing import Any, List
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
 from edsl.language_models import LanguageModel
 
@@ -12,6 +12,8 @@ class DeepInfraService(InferenceServiceABC):
     _inference_service_ = "deep_infra"
     _env_key_name_ = "DEEP_INFRA_API_KEY"
 
+    _models_list_cache: List[str] = []
+
     @classmethod
     def available(cls):
         text_models = cls.full_details_available()
@@ -19,20 +21,25 @@ class DeepInfraService(InferenceServiceABC):
 
     @classmethod
     def full_details_available(cls, verbose=False):
-        url = "https://api.deepinfra.com/models/list"
-        response = requests.get(url)
-        if response.status_code == 200:
-            text_generation_models = [
-                r for r in response.json() if r["type"] == "text-generation"
-            ]
-            from rich import print_json
-            import json
+        if not cls._models_list_cache:
+            url = "https://api.deepinfra.com/models/list"
+            response = requests.get(url)
+            if response.status_code == 200:
+                text_generation_models = [
+                    r for r in response.json() if r["type"] == "text-generation"
+                ]
+                cls._models_list_cache = text_generation_models
+
+                from rich import print_json
+                import json
 
-            if verbose:
-                print_json(json.dumps(text_generation_models))
-            return text_generation_models
+                if verbose:
+                    print_json(json.dumps(text_generation_models))
+                return text_generation_models
+            else:
+                return f"Failed to fetch data: Status code {response.status_code}"
         else:
-            return f"Failed to fetch data: Status code {response.status_code}"
+            return cls._models_list_cache
 
     @classmethod
     def create_model(cls, model_name: str, model_class_name=None) -> LanguageModel:

edsl/inference_services/GoogleService.py

@@ -60,7 +60,13 @@ class GoogleService(InferenceServiceABC):
 
             def parse_response(self, raw_response: dict[str, Any]) -> str:
                 data = raw_response
-                return data["candidates"][0]["content"]["parts"][0]["text"]
+                try:
+                    return data["candidates"][0]["content"]["parts"][0]["text"]
+                except KeyError as e:
+                    print(
+                        f"The data return was {data}, which was missing the key 'candidates'"
+                    )
+                    raise e
 
         LLM.__name__ = model_name
 

edsl/inference_services/InferenceServicesCollection.py

@@ -1,21 +1,47 @@
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
+import warnings
 
 
 class InferenceServicesCollection:
+    added_models = {}
+
     def __init__(self, services: list[InferenceServiceABC] = None):
         self.services = services or []
 
+    @classmethod
+    def add_model(cls, service_name, model_name):
+        if service_name not in cls.added_models:
+            cls.added_models[service_name] = []
+        cls.added_models[service_name].append(model_name)
+
+    @staticmethod
+    def _get_service_available(service) -> list[str]:
+        from_api = True
+        try:
+            service_models = service.available()
+        except Exception as e:
+            warnings.warn(
+                f"Error getting models for {service._inference_service_}. Relying on cache.",
+                UserWarning,
+            )
+            from edsl.inference_services.models_available_cache import models_available
+
+            service_models = models_available.get(service._inference_service_, [])
+            # cache results
+            service._models_list_cache = service_models
+            from_api = False
+        return service_models  # , from_api
+
     def available(self):
         total_models = []
         for service in self.services:
-            try:
-                service_models = service.available()
-            except Exception as e:
-                print(f"Error getting models for {service._inference_service_}: {e}")
-                service_models = []
-                continue
+            service_models = self._get_service_available(service)
             for model in service_models:
                 total_models.append([model, service._inference_service_, -1])
+
+            for model in self.added_models.get(service._inference_service_, []):
+                total_models.append([model, service._inference_service_, -1])
+
         sorted_models = sorted(total_models)
         for i, model in enumerate(sorted_models):
             model[2] = i
@@ -27,7 +53,7 @@ class InferenceServicesCollection:
 
     def create_model_factory(self, model_name: str, service_name=None, index=None):
         for service in self.services:
-            if model_name in service.available():
+            if model_name in self._get_service_available(service):
                 if service_name is None or service_name == service._inference_service_:
                     return service.create_model(model_name)
 

edsl/inference_services/OpenAIService.py

@@ -4,6 +4,7 @@ from openai import AsyncOpenAI
 
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
 from edsl.language_models import LanguageModel
+from edsl.inference_services.rate_limits_cache import rate_limits
 
 
 class OpenAIService(InferenceServiceABC):
@@ -43,15 +44,16 @@ class OpenAIService(InferenceServiceABC):
                     if m.id not in cls.model_exclude_list
                 ]
             except Exception as e:
-                print(
-                    f"""Error retrieving models: {e}. 
-                    See instructions about storing your API keys: https://docs.expectedparrot.com/en/latest/api_keys.html"""
-                )
-                cls._models_list_cache = [
-                    "gpt-3.5-turbo",
-                    "gpt-4-1106-preview",
-                    "gpt-4",
-                ]  # Fallback list
+                raise
+                # print(
+                #     f"""Error retrieving models: {e}.
+                #     See instructions about storing your API keys: https://docs.expectedparrot.com/en/latest/api_keys.html"""
+                # )
+                # cls._models_list_cache = [
+                #     "gpt-3.5-turbo",
+                #     "gpt-4-1106-preview",
+                #     "gpt-4",
+                # ]  # Fallback list
         return cls._models_list_cache
 
     @classmethod
@@ -98,7 +100,12 @@ class OpenAIService(InferenceServiceABC):
 
             def get_rate_limits(self) -> dict[str, Any]:
                 try:
-                    headers = self.get_headers()
+                    if "openai" in rate_limits:
+                        headers = rate_limits["openai"]
+
+                    else:
+                        headers = self.get_headers()
+
                 except Exception as e:
                     return {
                         "rpm": 10_000,