edsl 0.1.38.dev2__py3-none-any.whl → 0.1.38.dev4__py3-none-any.whl

edsl/jobs/JobsChecks.py

@@ -0,0 +1,147 @@
+import os
+from edsl.exceptions import MissingAPIKeyError
+
+
+class JobsChecks:
+    def __init__(self, jobs):
+        """ """
+        self.jobs = jobs
+
+    def check_api_keys(self) -> None:
+        from edsl import Model
+
+        for model in self.jobs.models + [Model()]:
+            if not model.has_valid_api_key():
+                raise MissingAPIKeyError(
+                    model_name=str(model.model),
+                    inference_service=model._inference_service_,
+                )
+
+    def get_missing_api_keys(self) -> set:
+        """
+        Returns a list of the api keys that a user needs to run this job, but does not currently have in their .env file.
+        """
+        missing_api_keys = set()
+
+        from edsl import Model
+        from edsl.enums import service_to_api_keyname
+
+        for model in self.jobs.models + [Model()]:
+            if not model.has_valid_api_key():
+                key_name = service_to_api_keyname.get(
+                    model._inference_service_, "NOT FOUND"
+                )
+                missing_api_keys.add(key_name)
+
+        return missing_api_keys
+
+    def user_has_ep_api_key(self) -> bool:
+        """
+        Returns True if the user has an EXPECTED_PARROT_API_KEY in their env.
+
+        Otherwise, returns False.
+        """
+
+        coop_api_key = os.getenv("EXPECTED_PARROT_API_KEY")
+
+        if coop_api_key is not None:
+            return True
+        else:
+            return False
+
+    def user_has_all_model_keys(self):
+        """
+        Returns True if the user has all model keys required to run their job.
+
+        Otherwise, returns False.
+        """
+
+        try:
+            self.check_api_keys()
+            return True
+        except MissingAPIKeyError:
+            return False
+        except Exception:
+            raise
+
+    def needs_external_llms(self) -> bool:
+        """
+        Returns True if the job needs external LLMs to run.
+
+        Otherwise, returns False.
+        """
+        # These cases are necessary to skip the API key check during doctests
+
+        # Accounts for Results.example()
+        all_agents_answer_questions_directly = len(self.jobs.agents) > 0 and all(
+            [hasattr(a, "answer_question_directly") for a in self.jobs.agents]
+        )
+
+        # Accounts for InterviewExceptionEntry.example()
+        only_model_is_test = set([m.model for m in self.jobs.models]) == set(["test"])
+
+        # Accounts for Survey.__call__
+        all_questions_are_functional = set(
+            [q.question_type for q in self.jobs.survey.questions]
+        ) == set(["functional"])
+
+        if (
+            all_agents_answer_questions_directly
+            or only_model_is_test
+            or all_questions_are_functional
+        ):
+            return False
+        else:
+            return True
+
+    def needs_key_process(self):
+        return (
+            not self.user_has_all_model_keys()
+            and not self.user_has_ep_api_key()
+            and self.needs_external_llms()
+        )
+
+    def key_process(self):
+        import secrets
+        from dotenv import load_dotenv
+        from edsl import CONFIG
+        from edsl.coop.coop import Coop
+        from edsl.utilities.utilities import write_api_key_to_env
+
+        missing_api_keys = self.get_missing_api_keys()
+
+        edsl_auth_token = secrets.token_urlsafe(16)
+
+        print("You're missing some of the API keys needed to run this job:")
+        for api_key in missing_api_keys:
+            print(f"     🔑 {api_key}")
+        print(
+            "\nYou can either add the missing keys to your .env file, or use remote inference."
+        )
+        print("Remote inference allows you to run jobs on our server.")
+        print("\n🚀 To use remote inference, sign up at the following link:")
+
+        coop = Coop()
+        coop._display_login_url(edsl_auth_token=edsl_auth_token)
+
+        print(
+            "\nOnce you log in, we will automatically retrieve your Expected Parrot API key and continue your job remotely."
+        )
+
+        api_key = coop._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("✨ API key retrieved and written to .env file.\n")
+
+        # Retrieve API key so we can continue running the job
+        load_dotenv()
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

edsl/jobs/JobsPrompts.py

@@ -0,0 +1,268 @@
+from typing import List, TYPE_CHECKING
+
+from edsl.results.Dataset import Dataset
+
+if TYPE_CHECKING:
+    from edsl.jobs import Jobs
+
+    # from edsl.jobs.interviews.Interview import Interview
+    # from edsl.results.Dataset import Dataset
+    # from edsl.agents.AgentList import AgentList
+    # from edsl.scenarios.ScenarioList import ScenarioList
+    # from edsl.surveys.Survey import Survey
+
+
+class JobsPrompts:
+    def __init__(self, jobs: "Jobs"):
+        self.interviews = jobs.interviews()
+        self.agents = jobs.agents
+        self.scenarios = jobs.scenarios
+        self.survey = jobs.survey
+        self._price_lookup = None
+
+    @property
+    def price_lookup(self):
+        if self._price_lookup is None:
+            from edsl import Coop
+
+            c = Coop()
+            self._price_lookup = c.fetch_prices()
+        return self._price_lookup
+
+    def prompts(self) -> "Dataset":
+        """Return a Dataset of prompts that will be used.
+
+        >>> from edsl.jobs import Jobs
+        >>> Jobs.example().prompts()
+        Dataset(...)
+        """
+        interviews = self.interviews
+        interview_indices = []
+        question_names = []
+        user_prompts = []
+        system_prompts = []
+        scenario_indices = []
+        agent_indices = []
+        models = []
+        costs = []
+
+        for interview_index, interview in enumerate(interviews):
+            invigilators = [
+                interview._get_invigilator(question)
+                for question in self.survey.questions
+            ]
+            for _, invigilator in enumerate(invigilators):
+                prompts = invigilator.get_prompts()
+                user_prompt = prompts["user_prompt"]
+                system_prompt = prompts["system_prompt"]
+                user_prompts.append(user_prompt)
+                system_prompts.append(system_prompt)
+                agent_index = self.agents.index(invigilator.agent)
+                agent_indices.append(agent_index)
+                interview_indices.append(interview_index)
+                scenario_index = self.scenarios.index(invigilator.scenario)
+                scenario_indices.append(scenario_index)
+                models.append(invigilator.model.model)
+                question_names.append(invigilator.question.question_name)
+
+                prompt_cost = self.estimate_prompt_cost(
+                    system_prompt=system_prompt,
+                    user_prompt=user_prompt,
+                    price_lookup=self.price_lookup,
+                    inference_service=invigilator.model._inference_service_,
+                    model=invigilator.model.model,
+                )
+                costs.append(prompt_cost["cost_usd"])
+
+        d = Dataset(
+            [
+                {"user_prompt": user_prompts},
+                {"system_prompt": system_prompts},
+                {"interview_index": interview_indices},
+                {"question_name": question_names},
+                {"scenario_index": scenario_indices},
+                {"agent_index": agent_indices},
+                {"model": models},
+                {"estimated_cost": costs},
+            ]
+        )
+        return d
+
+    @staticmethod
+    def estimate_prompt_cost(
+        system_prompt: str,
+        user_prompt: str,
+        price_lookup: dict,
+        inference_service: str,
+        model: str,
+    ) -> dict:
+        """Estimates the cost of a prompt. Takes piping into account."""
+        import math
+
+        def get_piping_multiplier(prompt: str):
+            """Returns 2 if a prompt includes Jinja braces, and 1 otherwise."""
+
+            if "{{" in prompt and "}}" in prompt:
+                return 2
+            return 1
+
+        # Look up prices per token
+        key = (inference_service, model)
+
+        try:
+            relevant_prices = price_lookup[key]
+
+            service_input_token_price = float(
+                relevant_prices["input"]["service_stated_token_price"]
+            )
+            service_input_token_qty = float(
+                relevant_prices["input"]["service_stated_token_qty"]
+            )
+            input_price_per_token = service_input_token_price / service_input_token_qty
+
+            service_output_token_price = float(
+                relevant_prices["output"]["service_stated_token_price"]
+            )
+            service_output_token_qty = float(
+                relevant_prices["output"]["service_stated_token_qty"]
+            )
+            output_price_per_token = (
+                service_output_token_price / service_output_token_qty
+            )
+
+        except KeyError:
+            # A KeyError is likely to occur if we cannot retrieve prices (the price_lookup dict is empty)
+            # Use a sensible default
+
+            import warnings
+
+            warnings.warn(
+                "Price data could not be retrieved. Using default estimates for input and output token prices. Input: $0.15 / 1M tokens; Output: $0.60 / 1M tokens"
+            )
+            input_price_per_token = 0.00000015  # $0.15 / 1M tokens
+            output_price_per_token = 0.00000060  # $0.60 / 1M tokens
+
+        # Compute the number of characters (double if the question involves piping)
+        user_prompt_chars = len(str(user_prompt)) * get_piping_multiplier(
+            str(user_prompt)
+        )
+        system_prompt_chars = len(str(system_prompt)) * get_piping_multiplier(
+            str(system_prompt)
+        )
+
+        # Convert into tokens (1 token approx. equals 4 characters)
+        input_tokens = (user_prompt_chars + system_prompt_chars) // 4
+
+        output_tokens = math.ceil(0.75 * input_tokens)
+
+        cost = (
+            input_tokens * input_price_per_token
+            + output_tokens * output_price_per_token
+        )
+
+        return {
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "cost_usd": cost,
+        }
+
+    def estimate_job_cost_from_external_prices(
+        self, price_lookup: dict, iterations: int = 1
+    ) -> dict:
+        """
+        Estimates the cost of a job according to the following assumptions:
+
+        - 1 token = 4 characters.
+        - For each prompt, output tokens = input tokens * 0.75, rounded up to the nearest integer.
+
+        price_lookup is an external pricing dictionary.
+        """
+
+        import pandas as pd
+
+        interviews = self.interviews
+        data = []
+        for interview in interviews:
+            invigilators = [
+                interview._get_invigilator(question)
+                for question in self.survey.questions
+            ]
+            for invigilator in invigilators:
+                prompts = invigilator.get_prompts()
+
+                # By this point, agent and scenario data has already been added to the prompts
+                user_prompt = prompts["user_prompt"]
+                system_prompt = prompts["system_prompt"]
+                inference_service = invigilator.model._inference_service_
+                model = invigilator.model.model
+
+                prompt_cost = self.estimate_prompt_cost(
+                    system_prompt=system_prompt,
+                    user_prompt=user_prompt,
+                    price_lookup=price_lookup,
+                    inference_service=inference_service,
+                    model=model,
+                )
+
+                data.append(
+                    {
+                        "user_prompt": user_prompt,
+                        "system_prompt": system_prompt,
+                        "estimated_input_tokens": prompt_cost["input_tokens"],
+                        "estimated_output_tokens": prompt_cost["output_tokens"],
+                        "estimated_cost_usd": prompt_cost["cost_usd"],
+                        "inference_service": inference_service,
+                        "model": model,
+                    }
+                )
+
+        df = pd.DataFrame.from_records(data)
+
+        df = (
+            df.groupby(["inference_service", "model"])
+            .agg(
+                {
+                    "estimated_cost_usd": "sum",
+                    "estimated_input_tokens": "sum",
+                    "estimated_output_tokens": "sum",
+                }
+            )
+            .reset_index()
+        )
+        df["estimated_cost_usd"] = df["estimated_cost_usd"] * iterations
+        df["estimated_input_tokens"] = df["estimated_input_tokens"] * iterations
+        df["estimated_output_tokens"] = df["estimated_output_tokens"] * iterations
+
+        estimated_costs_by_model = df.to_dict("records")
+
+        estimated_total_cost = sum(
+            model["estimated_cost_usd"] for model in estimated_costs_by_model
+        )
+        estimated_total_input_tokens = sum(
+            model["estimated_input_tokens"] for model in estimated_costs_by_model
+        )
+        estimated_total_output_tokens = sum(
+            model["estimated_output_tokens"] for model in estimated_costs_by_model
+        )
+
+        output = {
+            "estimated_total_cost_usd": estimated_total_cost,
+            "estimated_total_input_tokens": estimated_total_input_tokens,
+            "estimated_total_output_tokens": estimated_total_output_tokens,
+            "model_costs": estimated_costs_by_model,
+        }
+
+        return output
+
+    def estimate_job_cost(self, iterations: int = 1) -> dict:
+        """
+        Estimates the cost of a job according to the following assumptions:
+
+        - 1 token = 4 characters.
+        - For each prompt, output tokens = input tokens * 0.75, rounded up to the nearest integer.
+
+        Fetches prices from Coop.
+        """
+        return self.estimate_job_cost_from_external_prices(
+            price_lookup=self.price_lookup, iterations=iterations
+        )

edsl/jobs/JobsRemoteInferenceHandler.py

@@ -0,0 +1,239 @@
+from typing import Optional, Union, Literal
+import requests
+import sys
+from edsl.exceptions.coop import CoopServerResponseError
+
+# from edsl.enums import VisibilityType
+from edsl.results import Results
+
+
+class JobsRemoteInferenceHandler:
+    def __init__(self, jobs, verbose=False, poll_interval=3):
+        """
+        >>> from edsl.jobs import Jobs
+        >>> jh = JobsRemoteInferenceHandler(Jobs.example(), verbose=True)
+        >>> jh.use_remote_inference(True)
+        False
+        >>> jh._poll_remote_inference_job({'uuid':1234}, testing_simulated_response={"status": "failed"}) # doctest: +NORMALIZE_WHITESPACE
+        Job failed.
+        ...
+        >>> jh._poll_remote_inference_job({'uuid':1234}, testing_simulated_response={"status": "completed"}) # doctest: +NORMALIZE_WHITESPACE
+        Job completed and Results stored on Coop: None.
+        Results(...)
+        """
+        self.jobs = jobs
+        self.verbose = verbose
+        self.poll_interval = poll_interval
+
+        self._remote_job_creation_data = None
+        self._job_uuid = None
+
+    @property
+    def remote_job_creation_data(self):
+        return self._remote_job_creation_data
+
+    @property
+    def job_uuid(self):
+        return self._job_uuid
+
+    def use_remote_inference(self, disable_remote_inference: bool) -> bool:
+        if disable_remote_inference:
+            return False
+        if not disable_remote_inference:
+            try:
+                from edsl import Coop
+
+                user_edsl_settings = Coop().edsl_settings
+                return user_edsl_settings.get("remote_inference", False)
+            except requests.ConnectionError:
+                pass
+            except CoopServerResponseError as e:
+                pass
+
+        return False
+
+    def create_remote_inference_job(
+        self,
+        iterations: int = 1,
+        remote_inference_description: Optional[str] = None,
+        remote_inference_results_visibility: Optional["VisibilityType"] = "unlisted",
+        verbose=False,
+    ):
+        """ """
+        from edsl.config import CONFIG
+        from edsl.coop.coop import Coop
+        from rich import print as rich_print
+
+        coop = Coop()
+        print("Remote inference activated. Sending job to server...")
+        remote_job_creation_data = coop.remote_inference_create(
+            self.jobs,
+            description=remote_inference_description,
+            status="queued",
+            iterations=iterations,
+            initial_results_visibility=remote_inference_results_visibility,
+        )
+        job_uuid = remote_job_creation_data.get("uuid")
+        print(f"Job sent to server. (Job uuid={job_uuid}).")
+
+        expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
+        progress_bar_url = f"{expected_parrot_url}/home/remote-job-progress/{job_uuid}"
+
+        rich_print(
+            f"View job progress here: [#38bdf8][link={progress_bar_url}]{progress_bar_url}[/link][/#38bdf8]"
+        )
+
+        self._remote_job_creation_data = remote_job_creation_data
+        self._job_uuid = job_uuid
+        # return remote_job_creation_data
+
+    @staticmethod
+    def check_status(job_uuid):
+        from edsl.coop.coop import Coop
+
+        coop = Coop()
+        return coop.remote_inference_get(job_uuid)
+
+    def poll_remote_inference_job(self):
+        return self._poll_remote_inference_job(
+            self.remote_job_creation_data, verbose=self.verbose
+        )
+
+    def _poll_remote_inference_job(
+        self,
+        remote_job_creation_data: dict,
+        verbose=False,
+        poll_interval: Optional[float] = None,
+        testing_simulated_response: Optional[dict] = None,
+    ) -> Union[Results, None]:
+        import time
+        from datetime import datetime
+        from edsl.config import CONFIG
+        from edsl.coop.coop import Coop
+
+        if poll_interval is None:
+            poll_interval = self.poll_interval
+
+        expected_parrot_url = CONFIG.get("EXPECTED_PARROT_URL")
+
+        job_uuid = remote_job_creation_data.get("uuid")
+        coop = Coop()
+
+        if testing_simulated_response is not None:
+            remote_job_data_fetcher = lambda job_uuid: testing_simulated_response
+            object_fetcher = (
+                lambda results_uuid, expected_object_type: Results.example()
+            )
+        else:
+            remote_job_data_fetcher = coop.remote_inference_get
+            object_fetcher = coop.get
+
+        job_in_queue = True
+        while job_in_queue:
+            remote_job_data = remote_job_data_fetcher(job_uuid)
+            status = remote_job_data.get("status")
+            if status == "cancelled":
+                print("\r" + " " * 80 + "\r", end="")
+                print("Job cancelled by the user.")
+                print(
+                    f"See {expected_parrot_url}/home/remote-inference for more details."
+                )
+                return None
+            elif status == "failed":
+                print("\r" + " " * 80 + "\r", end="")
+                # write to stderr
+                latest_error_report_url = remote_job_data.get("latest_error_report_url")
+                if latest_error_report_url:
+                    print("Job failed.")
+                    print(
+                        f"Your job generated exceptions. Details on these exceptions can be found in the following report: {latest_error_report_url}"
+                    )
+                    print(
+                        f"Need support? Post a message at the Expected Parrot Discord channel (https://discord.com/invite/mxAYkjfy9m) or send an email to info@expectedparrot.com."
+                    )
+                else:
+                    print("Job failed.")
+                    print(
+                        f"See {expected_parrot_url}/home/remote-inference for more details."
+                    )
+                return None
+            elif status == "completed":
+                results_uuid = remote_job_data.get("results_uuid")
+                results_url = remote_job_data.get("results_url")
+                results = object_fetcher(results_uuid, expected_object_type="results")
+                print("\r" + " " * 80 + "\r", end="")
+                print(f"Job completed and Results stored on Coop: {results_url}.")
+                return results
+            else:
+                duration = poll_interval
+                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)]} Job status: {status} - last update: {time_checked}",
+                        end="",
+                        flush=True,
+                    )
+                    time.sleep(0.1)
+                    i += 1
+
+    def use_remote_inference(self, disable_remote_inference: bool) -> bool:
+        if disable_remote_inference:
+            return False
+        if not disable_remote_inference:
+            try:
+                from edsl import Coop
+
+                user_edsl_settings = Coop().edsl_settings
+                return user_edsl_settings.get("remote_inference", False)
+            except requests.ConnectionError:
+                pass
+            except CoopServerResponseError as e:
+                pass
+
+        return False
+
+    async def create_and_poll_remote_job(
+        self,
+        iterations: int = 1,
+        remote_inference_description: Optional[str] = None,
+        remote_inference_results_visibility: Optional[
+            Literal["private", "public", "unlisted"]
+        ] = "unlisted",
+    ) -> Union[Results, None]:
+        """
+        Creates and polls a remote inference job asynchronously.
+        Reuses existing synchronous methods but runs them in an async context.
+
+        :param iterations: Number of times to run each interview
+        :param remote_inference_description: Optional description for the remote job
+        :param remote_inference_results_visibility: Visibility setting for results
+        :return: Results object if successful, None if job fails or is cancelled
+        """
+        import asyncio
+        from functools import partial
+
+        # Create job using existing method
+        loop = asyncio.get_event_loop()
+        remote_job_creation_data = await loop.run_in_executor(
+            None,
+            partial(
+                self.create_remote_inference_job,
+                iterations=iterations,
+                remote_inference_description=remote_inference_description,
+                remote_inference_results_visibility=remote_inference_results_visibility,
+            ),
+        )
+
+        # Poll using existing method but with async sleep
+        return await loop.run_in_executor(
+            None, partial(self.poll_remote_inference_job, remote_job_creation_data)
+        )
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/jobs/interviews/Interview.py

@@ -110,9 +110,9 @@ class Interview:
         self.debug = debug
         self.iteration = iteration
         self.cache = cache
-        self.answers: dict[str, str] = (
-            Answers()
-        )  # will get filled in as interview progresses
+        self.answers: dict[
+            str, str
+        ] = Answers()  # will get filled in as interview progresses
         self.sidecar_model = sidecar_model
 
         # Trackers
@@ -143,9 +143,9 @@ class Interview:
         The keys are the question names; the values are the lists of status log changes for each task.
         """
         for task_creator in self.task_creators.values():
-            self._task_status_log_dict[task_creator.question.question_name] = (
-                task_creator.status_log
-            )
+            self._task_status_log_dict[
+                task_creator.question.question_name
+            ] = task_creator.status_log
         return self._task_status_log_dict
 
     @property
@@ -486,11 +486,11 @@ class Interview:
         """
         current_question_index: int = self.to_index[current_question.question_name]
 
-        next_question: Union[int, EndOfSurvey] = (
-            self.survey.rule_collection.next_question(
-                q_now=current_question_index,
-                answers=self.answers | self.scenario | self.agent["traits"],
-            )
+        next_question: Union[
+            int, EndOfSurvey
+        ] = self.survey.rule_collection.next_question(
+            q_now=current_question_index,
+            answers=self.answers | self.scenario | self.agent["traits"],
         )
 
         next_question_index = next_question.next_q