edsl 0.1.36.dev5__py3-none-any.whl → 0.1.37__py3-none-any.whl

edsl/coop/coop.py

@@ -28,9 +28,18 @@ class Coop:
         - Provide a URL directly, or use the default one.
         """
         self.api_key = api_key or os.getenv("EXPECTED_PARROT_API_KEY")
+
         self.url = url or CONFIG.EXPECTED_PARROT_URL
         if self.url.endswith("/"):
             self.url = self.url[:-1]
+        if "chick.expectedparrot" in self.url:
+            self.api_url = "https://chickapi.expectedparrot.com"
+        elif "expectedparrot" in self.url:
+            self.api_url = "https://api.expectedparrot.com"
+        elif "localhost:1234" in self.url:
+            self.api_url = "http://localhost:8000"
+        else:
+            self.api_url = self.url
         self._edsl_version = edsl.__version__
 
     ################
@@ -59,7 +68,7 @@ class Coop:
         """
         Send a request to the server and return the response.
         """
-        url = f"{self.url}/{uri}"
+        url = f"{self.api_url}/{uri}"
         method = method.upper()
         if payload is None:
             timeout = 20
@@ -90,18 +99,83 @@ class Coop:
 
         return response
 
-    def _resolve_server_response(self, response: requests.Response) -> None:
+    def _resolve_server_response(
+        self, response: requests.Response, check_api_key: bool = True
+    ) -> None:
         """
         Check the response from the server and raise errors as appropriate.
         """
         if response.status_code >= 400:
             message = response.json().get("detail")
             # print(response.text)
-            if "Authorization" in message:
+            if "The API key you provided is invalid" in message and check_api_key:
+                import secrets
+                from edsl.utilities.utilities import write_api_key_to_env
+
+                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)
+                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.")
+                return
+
+            elif "Authorization" in message:
                 print(message)
                 message = "Please provide an Expected Parrot API key."
+
             raise CoopServerResponseError(message)
 
+    def _poll_for_api_key(
+        self, edsl_auth_token: str, timeout: int = 120
+    ) -> Union[str, None]:
+        """
+        Allows the user to retrieve their Expected Parrot API key by logging in with an EDSL auth token.
+
+        :param edsl_auth_token: The EDSL auth token to use for login
+        :param timeout: Maximum time to wait for login, in seconds (default: 120)
+        """
+        import time
+        from datetime import datetime
+
+        start_poll_time = time.time()
+        waiting_for_login = True
+        while waiting_for_login:
+            elapsed_time = time.time() - start_poll_time
+            if elapsed_time > timeout:
+                # Timed out waiting for the user to log in
+                print("\r" + " " * 80 + "\r", end="")
+                return None
+
+            api_key = self._get_api_key(edsl_auth_token)
+            if api_key is not None:
+                print("\r" + " " * 80 + "\r", end="")
+                return api_key
+            else:
+                duration = 5
+                time_checked = datetime.now().strftime("%Y-%m-%d %I:%M:%S %p")
+                frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
+                start_time = time.time()
+                i = 0
+                while time.time() - start_time < duration:
+                    print(
+                        f"\r{frames[i % len(frames)]} Waiting for login. Last checked: {time_checked}",
+                        end="",
+                        flush=True,
+                    )
+                    time.sleep(0.1)
+                    i += 1
+
     def _json_handle_none(self, value: Any) -> Any:
         """
         Handle None values during JSON serialization.
@@ -134,7 +208,7 @@ class Coop:
             response = self._send_server_request(
                 uri="api/v0/edsl-settings", method="GET", timeout=5
             )
-            self._resolve_server_response(response)
+            self._resolve_server_response(response, check_api_key=False)
             return response.json()
         except Timeout:
             return {}
@@ -489,6 +563,7 @@ class Coop:
         description: Optional[str] = None,
         status: RemoteJobStatus = "queued",
         visibility: Optional[VisibilityType] = "unlisted",
+        initial_results_visibility: Optional[VisibilityType] = "unlisted",
         iterations: Optional[int] = 1,
     ) -> dict:
         """
@@ -517,6 +592,7 @@ class Coop:
                 "iterations": iterations,
                 "visibility": visibility,
                 "version": self._edsl_version,
+                "initial_results_visibility": initial_results_visibility,
             },
         )
         self._resolve_server_response(response)
@@ -568,7 +644,9 @@ class Coop:
             "version": data.get("version"),
         }
 
-    def remote_inference_cost(self, input: Union[Jobs, Survey]) -> int:
+    def remote_inference_cost(
+        self, input: Union[Jobs, Survey], iterations: int = 1
+    ) -> int:
         """
         Get the cost of a remote inference job.
 
@@ -593,6 +671,7 @@ class Coop:
                     job.to_dict(),
                     default=self._json_handle_none,
                 ),
+                "iterations": iterations,
             },
         )
         self._resolve_server_response(response)
@@ -602,24 +681,6 @@ class Coop:
             "usd": response_json.get("cost_in_usd"),
         }
 
-    ################
-    # Remote Errors
-    ################
-    def error_create(self, error_data: str) -> dict:
-        """
-        Send an error message to the server.
-        """
-        response = self._send_server_request(
-            uri="api/v0/errors",
-            method="POST",
-            payload={
-                "json_string": json.dumps(error_data),
-                "version": self._edsl_version,
-            },
-        )
-        self._resolve_server_response(response)
-        return response.json()
-
     ################
     # DUNDER METHODS
     ################
@@ -633,7 +694,7 @@ class Coop:
     async def remote_async_execute_model_call(
         self, model_dict: dict, user_prompt: str, system_prompt: str
     ) -> dict:
-        url = self.url + "/inference/"
+        url = self.api_url + "/inference/"
         # print("Now using url: ", url)
         data = {
             "model_dict": model_dict,
@@ -654,7 +715,7 @@ class Coop:
         ] = "lime_survey",
         email=None,
     ):
-        url = f"{self.url}/api/v0/export_to_{platform}"
+        url = f"{self.api_url}/api/v0/export_to_{platform}"
         if email:
             data = {"json_string": json.dumps({"survey": survey, "email": email})}
         else:
@@ -679,6 +740,17 @@ class Coop:
         else:
             return {}
 
+    def fetch_models(self) -> dict:
+        """
+        Fetch a dict of available models from Coop.
+
+        Each key in the dict is an inference service, and each value is a list of models from that service.
+        """
+        response = self._send_server_request(uri="api/v0/models", method="GET")
+        self._resolve_server_response(response)
+        data = response.json()
+        return data
+
     def fetch_rate_limit_config_vars(self) -> dict:
         """
         Fetch a dict of rate limit config vars from Coop.
@@ -693,14 +765,58 @@ class Coop:
         data = response.json()
         return data
 
+    def _display_login_url(self, edsl_auth_token: str):
+        """
+        Uses rich.print to display a login URL.
+
+        - We need this function because URL detection with print() does not work alongside animations in VSCode.
+        """
+        from rich import print as rich_print
+
+        url = f"{CONFIG.EXPECTED_PARROT_URL}/login?edsl_auth_token={edsl_auth_token}"
+
+        rich_print(f"[#38bdf8][link={url}]{url}[/link][/#38bdf8]")
+
+    def _get_api_key(self, edsl_auth_token: str):
+        """
+        Given an EDSL auth token, find the corresponding user's API key.
+        """
+
+        response = self._send_server_request(
+            uri="api/v0/get-api-key",
+            method="POST",
+            payload={
+                "edsl_auth_token": edsl_auth_token,
+            },
+        )
+        data = response.json()
+        api_key = data.get("api_key")
+        return api_key
+
+    def login(self):
+        """
+        Starts the EDSL auth token login flow.
+        """
+        import secrets
+        from dotenv import load_dotenv
+        from edsl.utilities.utilities import write_api_key_to_env
+
+        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)
+        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.")
 
-if __name__ == "__main__":
-    sheet_data = fetch_sheet_data()
-    if sheet_data:
-        print(f"Successfully fetched {len(sheet_data)} rows of data.")
-        print("First row:", sheet_data[0])
-    else:
-        print("Failed to fetch sheet data.")
+        write_api_key_to_env(api_key)
+        print("\n✨ API key retrieved and written to .env file.")
+
+        # Add API key to environment
+        load_dotenv()
 
 
 def main():
@@ -840,10 +956,3 @@ def main():
     job_coop_object = coop.remote_inference_create(job)
     job_coop_results = coop.remote_inference_get(job_coop_object.get("uuid"))
     coop.get(uuid=job_coop_results.get("results_uuid"))
-
-    ##############
-    # E. Errors
-    ##############
-    coop.error_create({"something": "This is an error message"})
-    coop.api_key = None
-    coop.error_create({"something": "This is an error message"})

edsl/data/Cache.py

@@ -194,7 +194,7 @@ class Cache(Base):
         >>> c = Cache()
         >>> len(c)
         0
-        >>> results = Question.example("free_text").by(m).run(cache = c)
+        >>> results = Question.example("free_text").by(m).run(cache = c, disable_remote_cache = True, disable_remote_inference = True)
         >>> len(c)
         1
         """

edsl/data/RemoteCacheSync.py

@@ -1,5 +1,7 @@
 class RemoteCacheSync:
-    def __init__(self, coop, cache, output_func, remote_cache=True, remote_cache_description=""):
+    def __init__(
+        self, coop, cache, output_func, remote_cache=True, remote_cache_description=""
+    ):
         self.coop = coop
         self.cache = cache
         self._output = output_func
@@ -21,34 +23,42 @@ class RemoteCacheSync:
 
     def _sync_from_remote(self):
         cache_difference = self.coop.remote_cache_get_diff(self.cache.keys())
-        client_missing_cacheentries = cache_difference.get("client_missing_cacheentries", [])
+        client_missing_cacheentries = cache_difference.get(
+            "client_missing_cacheentries", []
+        )
         missing_entry_count = len(client_missing_cacheentries)
-        
+
         if missing_entry_count > 0:
             self._output(
                 f"Updating local cache with {missing_entry_count:,} new "
                 f"{'entry' if missing_entry_count == 1 else 'entries'} from remote..."
             )
-            self.cache.add_from_dict({entry.key: entry for entry in client_missing_cacheentries})
+            self.cache.add_from_dict(
+                {entry.key: entry for entry in client_missing_cacheentries}
+            )
             self._output("Local cache updated!")
         else:
             self._output("No new entries to add to local cache.")
 
     def _sync_to_remote(self):
         cache_difference = self.coop.remote_cache_get_diff(self.cache.keys())
-        server_missing_cacheentry_keys = cache_difference.get("server_missing_cacheentry_keys", [])
+        server_missing_cacheentry_keys = cache_difference.get(
+            "server_missing_cacheentry_keys", []
+        )
         server_missing_cacheentries = [
             entry
             for key in server_missing_cacheentry_keys
             if (entry := self.cache.data.get(key)) is not None
         ]
-        
+
         new_cache_entries = [
-            entry for entry in self.cache.values() if entry.key not in self.old_entry_keys
+            entry
+            for entry in self.cache.values()
+            if entry.key not in self.old_entry_keys
         ]
         server_missing_cacheentries.extend(new_cache_entries)
         new_entry_count = len(server_missing_cacheentries)
-        
+
         if new_entry_count > 0:
             self._output(
                 f"Updating remote cache with {new_entry_count:,} new "
@@ -62,8 +72,11 @@ class RemoteCacheSync:
             self._output("Remote cache updated!")
         else:
             self._output("No new entries to add to remote cache.")
-        
-        self._output(f"There are {len(self.cache.keys()):,} entries in the local cache.")
+
+        self._output(
+            f"There are {len(self.cache.keys()):,} entries in the local cache."
+        )
+
 
 # # Usage example
 # def run_job(self, n, progress_bar, cache, stop_on_exception, sidecar_model, print_exceptions, raise_validation_errors, use_remote_cache=True):
@@ -79,6 +92,6 @@ class RemoteCacheSync:
 #             raise_validation_errors=raise_validation_errors,
 #         )
 #         self._output("Job completed!")
-    
+
 #     results.cache = cache.new_entries_cache()
-#     return results
+#     return results

edsl/exceptions/BaseException.py

@@ -0,0 +1,21 @@
+class BaseException(Exception):
+    relevant_doc = "https://docs.expectedparrot.com/"
+
+    def __init__(self, message, *, show_docs=True):
+        # Format main error message
+        formatted_message = [message.strip()]
+
+        # Add documentation links if requested
+        if show_docs:
+            if hasattr(self, "relevant_doc"):
+                formatted_message.append(
+                    f"\nFor more information, see:\n{self.relevant_doc}"
+                )
+            if hasattr(self, "relevant_notebook"):
+                formatted_message.append(
+                    f"\nFor a usage example, see:\n{self.relevant_notebook}"
+                )
+
+        # Join with double newlines for clear separation
+        final_message = "\n\n".join(formatted_message)
+        super().__init__(final_message)

edsl/exceptions/__init__.py

@@ -1,8 +1,8 @@
 from .agents import (
-    AgentAttributeLookupCallbackError,
+    #    AgentAttributeLookupCallbackError,
     AgentCombinationError,
-    AgentLacksLLMError,
-    AgentRespondedWithBadJSONError,
+    #    AgentLacksLLMError,
+    #    AgentRespondedWithBadJSONError,
 )
 from .configuration import (
     InvalidEnvironmentVariableError,
@@ -14,6 +14,10 @@ from .data import (
     DatabaseIntegrityError,
 )
 
+from .scenarios import (
+    ScenarioError,
+)
+
 from .general import MissingAPIKeyError
 
 from .jobs import JobsRunError, InterviewErrorPriorTaskCanceled, InterviewTimeoutError

edsl/exceptions/agents.py

@@ -1,37 +1,35 @@
-class AgentErrors(Exception):
-    pass
+from edsl.exceptions.BaseException import BaseException
 
 
-class AgentDynamicTraitsFunctionError(AgentErrors):
-    pass
+class AgentErrors(BaseException):
+    relevant_doc = "https://docs.expectedparrot.com/en/latest/agents.html"
 
 
-class AgentDirectAnswerFunctionError(AgentErrors):
-    pass
+class AgentDynamicTraitsFunctionError(AgentErrors):
+    relevant_doc = (
+        "https://docs.expectedparrot.com/en/latest/agents.html#dynamic-traits-function"
+    )
+    relevant_notebook = "https://docs.expectedparrot.com/en/latest/notebooks/example_agent_dynamic_traits.html"
 
 
-class AgentAttributeLookupCallbackError(AgentErrors):
-    pass
+class AgentDirectAnswerFunctionError(AgentErrors):
+    relevant_doc = "https://docs.expectedparrot.com/en/latest/agents.html#agent-direct-answering-methods"
 
 
 class AgentCombinationError(AgentErrors):
-    pass
-
-
-class AgentLacksLLMError(AgentErrors):
-    pass
-
-
-class AgentRespondedWithBadJSONError(AgentErrors):
-    pass
+    relevant_doc = (
+        "https://docs.expectedparrot.com/en/latest/agents.html#combining-agents"
+    )
 
 
 class AgentNameError(AgentErrors):
-    pass
+    relevant_doc = "https://docs.expectedparrot.com/en/latest/agents.html#agent-names"
 
 
 class AgentTraitKeyError(AgentErrors):
-    pass
+    relevant_doc = (
+        "https://docs.expectedparrot.com/en/latest/agents.html#constructing-an-agent"
+    )
 
 
 class FailedTaskException(Exception):

edsl/exceptions/results.py

@@ -1,26 +1,29 @@
-class ResultsErrors(Exception):
-    pass
+from edsl.exceptions.BaseException import BaseException
+
+
+class ResultsError(BaseException):
+    relevant_docs = "https://docs.expectedparrot.com/en/latest/results.html"
 
 
-class ResultsDeserializationError(ResultsErrors):
+class ResultsDeserializationError(ResultsError):
     pass
 
 
-class ResultsBadMutationstringError(ResultsErrors):
+class ResultsBadMutationstringError(ResultsError):
     pass
 
 
-class ResultsColumnNotFoundError(ResultsErrors):
+class ResultsColumnNotFoundError(ResultsError):
     pass
 
 
-class ResultsInvalidNameError(ResultsErrors):
+class ResultsInvalidNameError(ResultsError):
     pass
 
 
-class ResultsMutateError(ResultsErrors):
+class ResultsMutateError(ResultsError):
     pass
 
 
-class ResultsFilterError(ResultsErrors):
+class ResultsFilterError(ResultsError):
     pass

edsl/exceptions/scenarios.py

@@ -0,0 +1,22 @@
+import re
+import textwrap
+
+
+class ScenarioError(Exception):
+    documentation = "https://docs.expectedparrot.com/en/latest/scenarios.html#module-edsl.scenarios.Scenario"
+
+    def __init__(self, message: str):
+        self.message = message + "\n" + "Documentation: " + self.documentation
+        super().__init__(self.message)
+
+    def __str__(self):
+        return self.make_urls_clickable(self.message)
+
+    @staticmethod
+    def make_urls_clickable(text):
+        url_pattern = r"https?://[^\s]+"
+        urls = re.findall(url_pattern, text)
+        for url in urls:
+            clickable_url = f"\033]8;;{url}\007{url}\033]8;;\007"
+            text = text.replace(url, clickable_url)
+        return text

edsl/exceptions/surveys.py

@@ -1,34 +1,37 @@
-class SurveyErrors(Exception):
-    pass
+from edsl.exceptions.BaseException import BaseException
+
+
+class SurveyError(BaseException):
+    relevant_doc = "https://docs.expectedparrot.com/en/latest/surveys.html"
 
 
-class SurveyCreationError(SurveyErrors):
+class SurveyCreationError(SurveyError):
     pass
 
 
-class SurveyHasNoRulesError(SurveyErrors):
+class SurveyHasNoRulesError(SurveyError):
     pass
 
 
-class SurveyRuleSendsYouBackwardsError(SurveyErrors):
+class SurveyRuleSendsYouBackwardsError(SurveyError):
     pass
 
 
-class SurveyRuleSkipLogicSyntaxError(SurveyErrors):
+class SurveyRuleSkipLogicSyntaxError(SurveyError):
     pass
 
 
-class SurveyRuleReferenceInRuleToUnknownQuestionError(SurveyErrors):
+class SurveyRuleReferenceInRuleToUnknownQuestionError(SurveyError):
     pass
 
 
-class SurveyRuleRefersToFutureStateError(SurveyErrors):
+class SurveyRuleRefersToFutureStateError(SurveyError):
     pass
 
 
-class SurveyRuleCollectionHasNoRulesAtNodeError(SurveyErrors):
+class SurveyRuleCollectionHasNoRulesAtNodeError(SurveyError):
     pass
 
 
-class SurveyRuleCannotEvaluateError(SurveyErrors):
+class SurveyRuleCannotEvaluateError(SurveyError):
     pass

edsl/inference_services/AwsBedrock.py

@@ -28,12 +28,16 @@ class AwsBedrockService(InferenceServiceABC):
         "ai21.j2-ultra",
         "ai21.j2-ultra-v1",
     ]
+    _models_list_cache: List[str] = []
 
     @classmethod
     def available(cls):
         """Fetch available models from AWS Bedrock."""
+
+        region = os.getenv("AWS_REGION", "us-east-1")
+
         if not cls._models_list_cache:
-            client = boto3.client("bedrock", region_name="us-west-2")
+            client = boto3.client("bedrock", region_name=region)
             all_models_ids = [
                 x["modelId"] for x in client.list_foundation_models()["modelSummaries"]
             ]
@@ -80,7 +84,8 @@ class AwsBedrockService(InferenceServiceABC):
                     self.api_token
                 )  # call to check the if env variables are set.
 
-                client = boto3.client("bedrock-runtime", region_name="us-west-2")
+                region = os.getenv("AWS_REGION", "us-east-1")
+                client = boto3.client("bedrock-runtime", region_name=region)
 
                 conversation = [
                     {