edsl 0.1.38.dev4__py3-none-any.whl → 0.1.39__py3-none-any.whl

edsl/coop/ExpectedParrotKeyHandler.py

@@ -0,0 +1,125 @@
+from pathlib import Path
+import os
+import platformdirs
+
+
+import sys
+import select
+
+
+def get_input_with_timeout(prompt, timeout=5, default="y"):
+    print(prompt, end="", flush=True)
+    ready, _, _ = select.select([sys.stdin], [], [], timeout)
+    if ready:
+        return sys.stdin.readline().strip()
+    print(f"\nNo input received within {timeout} seconds. Using default: {default}")
+    return default
+
+
+class ExpectedParrotKeyHandler:
+    asked_to_store_file_name = "asked_to_store.txt"
+    ep_key_file_name = "ep_api_key.txt"
+    application_name = "edsl"
+
+    @property
+    def config_dir(self):
+        return platformdirs.user_config_dir(self.application_name)
+
+    def _ep_key_file_exists(self) -> bool:
+        """Check if the Expected Parrot key file exists."""
+        return Path(self.config_dir).joinpath(self.ep_key_file_name).exists()
+
+    def ok_to_ask_to_store(self):
+        """Check if it's okay to ask the user to store the key."""
+        from edsl.config import CONFIG
+
+        if CONFIG.get("EDSL_RUN_MODE") != "production":
+            return False
+
+        return (
+            not Path(self.config_dir).joinpath(self.asked_to_store_file_name).exists()
+        )
+
+    def reset_asked_to_store(self):
+        """Reset the flag that indicates whether the user has been asked to store the key."""
+        asked_to_store_path = Path(self.config_dir).joinpath(
+            self.asked_to_store_file_name
+        )
+        if asked_to_store_path.exists():
+            os.remove(asked_to_store_path)
+            print(
+                "Deleted the file that indicates whether the user has been asked to store the key."
+            )
+
+    def ask_to_store(self, api_key) -> bool:
+        """Ask the user if they want to store the Expected Parrot key. If they say "yes", store it."""
+        if self.ok_to_ask_to_store():
+            # can_we_store = get_input_with_timeout(
+            #     "Would you like to store your Expected Parrot key for future use? (y/n): ",
+            #     timeout=5,
+            #     default="y",
+            # )
+            can_we_store = "y"
+            if can_we_store.lower() == "y":
+                Path(self.config_dir).mkdir(parents=True, exist_ok=True)
+                self.store_ep_api_key(api_key)
+                # print("Stored Expected Parrot API key at ", self.config_dir)
+                return True
+            else:
+                Path(self.config_dir).mkdir(parents=True, exist_ok=True)
+                with open(
+                    Path(self.config_dir).joinpath(self.asked_to_store_file_name), "w"
+                ) as f:
+                    f.write("Yes")
+        return False
+
+    def get_ep_api_key(self):
+        # check if the key is stored in the config_dir
+        api_key = None
+        api_key_from_cache = None
+        api_key_from_os = None
+
+        if self._ep_key_file_exists():
+            with open(Path(self.config_dir).joinpath(self.ep_key_file_name), "r") as f:
+                api_key_from_cache = f.read().strip()
+
+        api_key_from_os = os.getenv("EXPECTED_PARROT_API_KEY")
+
+        if api_key_from_os and api_key_from_cache:
+            if api_key_from_os != api_key_from_cache:
+                import warnings
+
+                warnings.warn(
+                    "WARNING: The Expected Parrot API key from the environment variable "
+                    "differs from the one stored in the config directory. Using the one "
+                    "from the environment variable."
+                )
+            api_key = api_key_from_os
+
+        if api_key_from_os and not api_key_from_cache:
+            api_key = api_key_from_os
+
+        if not api_key_from_os and api_key_from_cache:
+            api_key = api_key_from_cache
+
+        if api_key is not None:
+            _ = self.ask_to_store(api_key)
+        return api_key
+
+    def delete_ep_api_key(self):
+        key_path = Path(self.config_dir) / self.ep_key_file_name
+        if key_path.exists():
+            os.remove(key_path)
+            print("Deleted Expected Parrot API key at ", key_path)
+
+    def store_ep_api_key(self, api_key):
+        # Create the directory if it doesn't exist
+        os.makedirs(self.config_dir, exist_ok=True)
+
+        # Create the path for the key file
+        key_path = Path(self.config_dir) / self.ep_key_file_name
+
+        # Save the key
+        with open(key_path, "w") as f:
+            f.write(api_key)
+        # print("Stored Expected Parrot API key at ", key_path)

edsl/coop/PriceFetcher.py

@@ -18,7 +18,7 @@ class PriceFetcher:
 
         import os
         import requests
-        from edsl import CONFIG
+        from edsl.config import CONFIG
 
         try:
             # Fetch the pricing data

edsl/coop/coop.py

@@ -1,11 +1,19 @@
 import aiohttp
 import json
-import os
 import requests
-from typing import Any, Optional, Union, Literal
+
+from typing import Any, Optional, Union, Literal, TypedDict
 from uuid import UUID
+from collections import UserDict, defaultdict
+
 import edsl
-from edsl import CONFIG, CacheEntry, Jobs, Survey
+from pathlib import Path
+
+from edsl.config import CONFIG
+from edsl.data.CacheEntry import CacheEntry
+from edsl.jobs.Jobs import Jobs
+from edsl.surveys.Survey import Survey
+
 from edsl.exceptions.coop import CoopNoUUIDError, CoopServerResponseError
 from edsl.coop.utils import (
     EDSLObject,
@@ -15,19 +23,48 @@ from edsl.coop.utils import (
     VisibilityType,
 )
 
+from edsl.coop.CoopFunctionsMixin import CoopFunctionsMixin
+from edsl.coop.ExpectedParrotKeyHandler import ExpectedParrotKeyHandler
+
+from edsl.inference_services.data_structures import ServiceToModelsMapping
+
 
-class Coop:
+class RemoteInferenceResponse(TypedDict):
+    job_uuid: str
+    results_uuid: str
+    results_url: str
+    latest_error_report_uuid: str
+    latest_error_report_url: str
+    status: str
+    reason: str
+    credits_consumed: float
+    version: str
+
+
+class RemoteInferenceCreationInfo(TypedDict):
+    uuid: str
+    description: str
+    status: str
+    iterations: int
+    visibility: str
+    version: str
+
+
+class Coop(CoopFunctionsMixin):
     """
     Client for the Expected Parrot API.
     """
 
-    def __init__(self, api_key: str = None, url: str = None) -> None:
+    def __init__(
+        self, api_key: Optional[str] = None, url: Optional[str] = None
+    ) -> None:
         """
         Initialize the client.
         - Provide an API key directly, or through an env variable.
         - Provide a URL directly, or use the default one.
         """
-        self.api_key = api_key or os.getenv("EXPECTED_PARROT_API_KEY")
+        self.ep_key_handler = ExpectedParrotKeyHandler()
+        self.api_key = api_key or self.ep_key_handler.get_ep_api_key()
 
         self.url = url or CONFIG.EXPECTED_PARROT_URL
         if self.url.endswith("/"):
@@ -142,6 +179,7 @@ class Coop:
         Check the response from the server and raise errors as appropriate.
         """
         # Get EDSL version from header
+        # breakpoint()
         server_edsl_version = response.headers.get("X-EDSL-Version")
 
         if server_edsl_version:
@@ -150,11 +188,18 @@ class Coop:
                 server_version_str=server_edsl_version,
             ):
                 print(
-                    "Please upgrade your EDSL version to access our latest features. To upgrade, open your terminal and run `pip upgrade edsl`"
+                    "Please upgrade your EDSL version to access our latest features. To upgrade, open your terminal and run `pip install --upgrade edsl`"
                 )
 
         if response.status_code >= 400:
-            message = response.json().get("detail")
+            try:
+                message = response.json().get("detail")
+            except json.JSONDecodeError:
+                raise CoopServerResponseError(
+                    f"Server returned status code {response.status_code}."
+                    "JSON response could not be decoded.",
+                    "The server response was: " + response.text,
+                )
             # print(response.text)
             if "The API key you provided is invalid" in message and check_api_key:
                 import secrets
@@ -163,19 +208,27 @@ class Coop:
                 edsl_auth_token = secrets.token_urlsafe(16)
 
                 print("Your Expected Parrot API key is invalid.")
-                print(
-                    "\nUse the link below to log in to Expected Parrot so we can automatically update your API key."
+                self._display_login_url(
+                    edsl_auth_token=edsl_auth_token,
+                    link_description="\n🔗 Use the link below to log in to Expected Parrot so we can automatically update your API key.",
                 )
-                self._display_login_url(edsl_auth_token=edsl_auth_token)
                 api_key = self._poll_for_api_key(edsl_auth_token)
 
                 if api_key is None:
                     print("\nTimed out waiting for login. Please try again.")
                     return
 
-                write_api_key_to_env(api_key)
-                print("\n✨ API key retrieved and written to .env file.")
-                print("Rerun your code to try again with a valid API key.")
+                print("\n✨ API key retrieved.")
+
+                if stored_in_user_space := self.ep_key_handler.ask_to_store(api_key):
+                    pass
+                else:
+                    path_to_env = write_api_key_to_env(api_key)
+                    print(
+                        "\n✨ API key retrieved and written to .env file at the following path:"
+                    )
+                    print(f"    {path_to_env}")
+                    print("Rerun your code to try again with a valid API key.")
                 return
 
             elif "Authorization" in message:
@@ -268,6 +321,7 @@ class Coop:
         self,
         object: EDSLObject,
         description: Optional[str] = None,
+        alias: Optional[str] = None,
         visibility: Optional[VisibilityType] = "unlisted",
     ) -> dict:
         """
@@ -279,6 +333,7 @@ class Coop:
             method="POST",
             payload={
                 "description": description,
+                "alias": alias,
                 "json_string": json.dumps(
                     object.to_dict(),
                     default=self._json_handle_none,
@@ -373,6 +428,7 @@ class Coop:
         uuid: Union[str, UUID] = None,
         url: str = None,
         description: Optional[str] = None,
+        alias: Optional[str] = None,
         value: Optional[EDSLObject] = None,
         visibility: Optional[VisibilityType] = None,
     ) -> dict:
@@ -389,6 +445,7 @@ class Coop:
             params={"uuid": uuid},
             payload={
                 "description": description,
+                "alias": alias,
                 "json_string": (
                     json.dumps(
                         value.to_dict(),
@@ -602,9 +659,6 @@ class Coop:
         self._resolve_server_response(response)
         return response.json()
 
-    ################
-    # Remote Inference
-    ################
     def remote_inference_create(
         self,
         job: Jobs,
@@ -613,7 +667,7 @@ class Coop:
         visibility: Optional[VisibilityType] = "unlisted",
         initial_results_visibility: Optional[VisibilityType] = "unlisted",
         iterations: Optional[int] = 1,
-    ) -> dict:
+    ) -> RemoteInferenceCreationInfo:
         """
         Send a remote inference job to the server.
 
@@ -645,18 +699,21 @@ class Coop:
         )
         self._resolve_server_response(response)
         response_json = response.json()
-        return {
-            "uuid": response_json.get("job_uuid"),
-            "description": response_json.get("description"),
-            "status": response_json.get("status"),
-            "iterations": response_json.get("iterations"),
-            "visibility": response_json.get("visibility"),
-            "version": self._edsl_version,
-        }
+
+        return RemoteInferenceCreationInfo(
+            **{
+                "uuid": response_json.get("job_uuid"),
+                "description": response_json.get("description"),
+                "status": response_json.get("status"),
+                "iterations": response_json.get("iterations"),
+                "visibility": response_json.get("visibility"),
+                "version": self._edsl_version,
+            }
+        )
 
     def remote_inference_get(
         self, job_uuid: Optional[str] = None, results_uuid: Optional[str] = None
-    ) -> dict:
+    ) -> RemoteInferenceResponse:
         """
         Get the details of a remote inference job.
         You can pass either the job uuid or the results uuid as a parameter.
@@ -698,17 +755,30 @@ class Coop:
                 f"{self.url}/home/remote-inference/error/{latest_error_report_uuid}"
             )
 
-        return {
-            "job_uuid": data.get("job_uuid"),
-            "results_uuid": results_uuid,
-            "results_url": results_url,
-            "latest_error_report_uuid": latest_error_report_uuid,
-            "latest_error_report_url": latest_error_report_url,
-            "status": data.get("status"),
-            "reason": data.get("reason"),
-            "credits_consumed": data.get("price"),
-            "version": data.get("version"),
-        }
+        return RemoteInferenceResponse(
+            **{
+                "job_uuid": data.get("job_uuid"),
+                "results_uuid": results_uuid,
+                "results_url": results_url,
+                "latest_error_report_uuid": latest_error_report_uuid,
+                "latest_error_report_url": latest_error_report_url,
+                "status": data.get("status"),
+                "reason": data.get("reason"),
+                "credits_consumed": data.get("price"),
+                "version": data.get("version"),
+            }
+        )
+
+    def get_running_jobs(self) -> list[str]:
+        """
+        Get a list of currently running job IDs.
+
+        Returns:
+            list[str]: List of running job UUIDs
+        """
+        response = self._send_server_request(uri="jobs/status", method="GET")
+        self._resolve_server_response(response)
+        return response.json().get("running_jobs", [])
 
     def remote_inference_cost(
         self, input: Union[Jobs, Survey], iterations: int = 1
@@ -810,7 +880,7 @@ class Coop:
                 "Invalid EDSL_FETCH_TOKEN_PRICES value---should be 'True' or 'False'."
             )
 
-    def fetch_models(self) -> dict:
+    def fetch_models(self) -> ServiceToModelsMapping:
         """
         Fetch a dict of available models from Coop.
 
@@ -819,7 +889,7 @@ class Coop:
         response = self._send_server_request(uri="api/v0/models", method="GET")
         self._resolve_server_response(response)
         data = response.json()
-        return data
+        return ServiceToModelsMapping(data)
 
     def fetch_rate_limit_config_vars(self) -> dict:
         """
@@ -835,7 +905,9 @@ class Coop:
         data = response.json()
         return data
 
-    def _display_login_url(self, edsl_auth_token: str):
+    def _display_login_url(
+        self, edsl_auth_token: str, link_description: Optional[str] = None
+    ):
         """
         Uses rich.print to display a login URL.
 
@@ -845,7 +917,12 @@ class Coop:
 
         url = f"{CONFIG.EXPECTED_PARROT_URL}/login?edsl_auth_token={edsl_auth_token}"
 
-        rich_print(f"[#38bdf8][link={url}]{url}[/link][/#38bdf8]")
+        if link_description:
+            rich_print(
+                f"{link_description}\n    [#38bdf8][link={url}]{url}[/link][/#38bdf8]"
+            )
+        else:
+            rich_print(f"    [#38bdf8][link={url}]{url}[/link][/#38bdf8]")
 
     def _get_api_key(self, edsl_auth_token: str):
         """
@@ -873,17 +950,18 @@ class Coop:
 
         edsl_auth_token = secrets.token_urlsafe(16)
 
-        print(
-            "\nUse the link below to log in to Expected Parrot so we can automatically update your API key."
+        self._display_login_url(
+            edsl_auth_token=edsl_auth_token,
+            link_description="\n🔗 Use the link below to log in to Expected Parrot so we can automatically update your API key.",
         )
-        self._display_login_url(edsl_auth_token=edsl_auth_token)
         api_key = self._poll_for_api_key(edsl_auth_token)
 
         if api_key is None:
             raise Exception("Timed out waiting for login. Please try again.")
 
-        write_api_key_to_env(api_key)
-        print("\n✨ API key retrieved and written to .env file.")
+        path_to_env = write_api_key_to_env(api_key)
+        print("\n✨ API key retrieved and written to .env file at the following path:")
+        print(f"    {path_to_env}")
 
         # Add API key to environment
         load_dotenv()

edsl/coop/utils.py

@@ -1,19 +1,19 @@
-from edsl import (
-    Agent,
-    AgentList,
-    Cache,
-    ModelList,
-    Notebook,
-    Results,
-    Scenario,
-    ScenarioList,
-    Survey,
-    Study,
-)
-from edsl.language_models import LanguageModel
-from edsl.questions import QuestionBase
 from typing import Literal, Optional, Type, Union
 
+from edsl.agents.Agent import Agent
+from edsl.agents.AgentList import AgentList
+from edsl.data.Cache import Cache
+from edsl.language_models.ModelList import ModelList
+from edsl.notebooks.Notebook import Notebook
+from edsl.results.Results import Results
+from edsl.scenarios.Scenario import Scenario
+from edsl.scenarios.ScenarioList import ScenarioList
+from edsl.surveys.Survey import Survey
+from edsl.study.Study import Study
+
+from edsl.language_models.LanguageModel import LanguageModel
+from edsl.questions.QuestionBase import QuestionBase
+
 EDSLObject = Union[
     Agent,
     AgentList,

edsl/data/Cache.py

@@ -6,12 +6,10 @@ from __future__ import annotations
 import json
 import os
 import warnings
-import copy
-from typing import Optional, Union
+from typing import Optional, Union, TYPE_CHECKING
 from edsl.Base import Base
-from edsl.data.CacheEntry import CacheEntry
-from edsl.utilities.utilities import dict_hash
-from edsl.utilities.decorators import remove_edsl_version
+
+from edsl.utilities.remove_edsl_version import remove_edsl_version
 from edsl.exceptions.cache import CacheError
 
 
@@ -83,10 +81,6 @@ class Cache(Base):
 
         self._perform_checks()
 
-    def rich_print(sefl):
-        pass
-        # raise NotImplementedError("This method is not implemented yet.")
-
     def code(sefl):
         pass
         # raise NotImplementedError("This method is not implemented yet.")
@@ -201,6 +195,7 @@ class Cache(Base):
         >>> len(c)
         1
         """
+        from edsl.data.CacheEntry import CacheEntry
 
         entry = CacheEntry(
             model=model,
@@ -226,6 +221,7 @@ class Cache(Base):
 
         :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:
@@ -246,6 +242,8 @@ class Cache(Base):
 
         :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()
@@ -289,8 +287,8 @@ class Cache(Base):
 
         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(db_path=db_path)
+        # 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:
@@ -353,7 +351,8 @@ class Cache(Base):
                 f.write(json.dumps({key: value.to_dict()}) + "\n")
 
     def to_scenario_list(self):
-        from edsl import ScenarioList, Scenario
+        from edsl.scenarios.ScenarioList import ScenarioList
+        from edsl.scenarios.Scenario import Scenario
 
         scenarios = []
         for key, value in self.data.items():
@@ -363,12 +362,32 @@ class Cache(Base):
             scenarios.append(s)
         return ScenarioList(scenarios)
 
-    ####################
-    # REMOTE
-    ####################
-    # TODO: Make this work
-    # - Need to decide whether the cache belongs to a user and what can be shared
-    # - I.e., some cache entries? all or nothing?
+    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:
         """
@@ -394,11 +413,10 @@ class Cache(Base):
         if self.filename:
             self.write(self.filename)
 
-    ####################
-    # DUNDER / USEFUL
-    ####################
     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:
@@ -414,12 +432,6 @@ class Cache(Base):
     def _summary(self):
         return {"EDSL Class": "Cache", "Number of entries": len(self.data)}
 
-    def _repr_html_(self):
-        # from edsl.utilities.utilities import data_to_html
-        # return data_to_html(self.to_dict())
-        footer = f"<a href={self.__documentation__}>(docs)</a>"
-        return str(self.summary(format="html")) + footer
-
     def table(
         self,
         *fields,
@@ -443,6 +455,8 @@ class Cache(Base):
     @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)
 
@@ -485,6 +499,8 @@ class Cache(Base):
         """
         Create an example input for a 'fetch' operation.
         """
+        from edsl.data.CacheEntry import CacheEntry
+
         return CacheEntry.fetch_input_example()
 
     def to_html(self):
@@ -541,6 +557,8 @@ class Cache(Base):
 
         :param randomize: If True, uses CacheEntry's randomize method.
         """
+        from edsl.data.CacheEntry import CacheEntry
+
         return cls(
             data={
                 CacheEntry.example(randomize).key: CacheEntry.example(),