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

edsl/jobs/Jobs.py

@@ -1,24 +1,15 @@
 # """The Jobs class is a collection of agents, scenarios and models and one survey."""
 from __future__ import annotations
-import os
+import warnings
 from itertools import product
 from typing import Optional, Union, Sequence, Generator
-from edsl import Model
-from edsl.agents import Agent
+
 from edsl.Base import Base
-from edsl.data.Cache import Cache
-from edsl.data.CacheHandler import CacheHandler
-from edsl.results.Dataset import Dataset
 
-from edsl.exceptions.jobs import MissingRemoteInferenceError
 from edsl.exceptions import MissingAPIKeyError
 from edsl.jobs.buckets.BucketCollection import BucketCollection
 from edsl.jobs.interviews.Interview import Interview
-from edsl.language_models import LanguageModel
-from edsl.results import Results
-from edsl.scenarios import Scenario
-from edsl.surveys import Survey
-
+from edsl.jobs.runners.JobsRunnerAsyncio import JobsRunnerAsyncio
 from edsl.utilities.decorators import add_edsl_version, remove_edsl_version
 
 
@@ -31,10 +22,10 @@ class Jobs(Base):
 
     def __init__(
         self,
-        survey: Survey,
-        agents: Optional[list[Agent]] = None,
-        models: Optional[list[LanguageModel]] = None,
-        scenarios: Optional[list[Scenario]] = None,
+        survey: "Survey",
+        agents: Optional[list["Agent"]] = None,
+        models: Optional[list["LanguageModel"]] = None,
+        scenarios: Optional[list["Scenario"]] = None,
     ):
         """Initialize a Jobs instance.
 
@@ -44,18 +35,67 @@ class Jobs(Base):
         :param scenarios: a list of scenarios
         """
         self.survey = survey
-        self.agents = agents or []
-        self.models = models or []
-        self.scenarios = scenarios or []
+        self.agents: "AgentList" = agents
+        self.scenarios: "ScenarioList" = scenarios
+        self.models = models
+
         self.__bucket_collection = None
 
+    @property
+    def models(self):
+        return self._models
+
+    @models.setter
+    def models(self, value):
+        from edsl import ModelList
+
+        if value:
+            if not isinstance(value, ModelList):
+                self._models = ModelList(value)
+            else:
+                self._models = value
+        else:
+            self._models = ModelList([])
+
+    @property
+    def agents(self):
+        return self._agents
+
+    @agents.setter
+    def agents(self, value):
+        from edsl import AgentList
+
+        if value:
+            if not isinstance(value, AgentList):
+                self._agents = AgentList(value)
+            else:
+                self._agents = value
+        else:
+            self._agents = AgentList([])
+
+    @property
+    def scenarios(self):
+        return self._scenarios
+
+    @scenarios.setter
+    def scenarios(self, value):
+        from edsl import ScenarioList
+
+        if value:
+            if not isinstance(value, ScenarioList):
+                self._scenarios = ScenarioList(value)
+            else:
+                self._scenarios = value
+        else:
+            self._scenarios = ScenarioList([])
+
     def by(
         self,
         *args: Union[
-            Agent,
-            Scenario,
-            LanguageModel,
-            Sequence[Union[Agent, Scenario, LanguageModel]],
+            "Agent",
+            "Scenario",
+            "LanguageModel",
+            Sequence[Union["Agent", "Scenario", "LanguageModel"]],
         ],
     ) -> Jobs:
         """
@@ -68,10 +108,10 @@ class Jobs(Base):
         >>> q = QuestionFreeText(question_name="name", question_text="What is your name?")
         >>> j = Jobs(survey = Survey(questions=[q]))
         >>> j
-        Jobs(survey=Survey(...), agents=[], models=[], scenarios=[])
+        Jobs(survey=Survey(...), agents=AgentList([]), models=ModelList([]), scenarios=ScenarioList([]))
         >>> from edsl import Agent; a = Agent(traits = {"status": "Sad"})
         >>> j.by(a).agents
-        [Agent(traits = {'status': 'Sad'})]
+        AgentList([Agent(traits = {'status': 'Sad'})])
 
         :param args: objects or a sequence (list, tuple, ...) of objects of the same type
 
@@ -95,13 +135,13 @@ class Jobs(Base):
         setattr(self, objects_key, new_objects)  # update the job
         return self
 
-    def prompts(self) -> Dataset:
+    def prompts(self) -> "Dataset":
         """Return a Dataset of prompts that will be used.
 
 
         >>> from edsl.jobs import Jobs
         >>> Jobs.example().prompts()
-        Dataset([{'interview_index': [0, 0, 1, 1, 2, 2, 3, 3]}, {'question_index': ['how_feeling', 'how_feeling_yesterday', 'how_feeling', 'how_feeling_yesterday', 'how_feeling', 'how_feeling_yesterday', 'how_feeling', 'how_feeling_yesterday']}, {'user_prompt': [Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA')]}, {'scenario_index': [Scenario({'period': 'morning'}), Scenario({'period': 'morning'}), Scenario({'period': 'afternoon'}), Scenario({'period': 'afternoon'}), Scenario({'period': 'morning'}), Scenario({'period': 'morning'}), Scenario({'period': 'afternoon'}), Scenario({'period': 'afternoon'})]}, {'system_prompt': [Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA'), Prompt(text='NA')]}])
+        Dataset(...)
         """
 
         interviews = self.interviews()
@@ -111,6 +151,7 @@ class Jobs(Base):
         user_prompts = []
         system_prompts = []
         scenario_indices = []
+        from edsl.results.Dataset import Dataset
 
         for interview_index, interview in enumerate(interviews):
             invigilators = list(interview._build_invigilators(debug=False))
@@ -131,6 +172,20 @@ class Jobs(Base):
             ]
         )
 
+    @staticmethod
+    def _get_container_class(object):
+        from edsl.agents.AgentList import AgentList
+        from edsl.agents.Agent import Agent
+        from edsl.scenarios.Scenario import Scenario
+        from edsl.scenarios.ScenarioList import ScenarioList
+
+        if isinstance(object, Agent):
+            return AgentList
+        elif isinstance(object, Scenario):
+            return ScenarioList
+        else:
+            return list
+
     @staticmethod
     def _turn_args_to_list(args):
         """Return a list of the first argument if it is a sequence, otherwise returns a list of all the arguments."""
@@ -149,19 +204,25 @@ class Jobs(Base):
             return len(args) == 1 and isinstance(args[0], Sequence)
 
         if did_user_pass_a_sequence(args):
-            return list(args[0])
+            container_class = Jobs._get_container_class(args[0][0])
+            return container_class(args[0])
         else:
-            return list(args)
+            container_class = Jobs._get_container_class(args[0])
+            return container_class(args)
 
     def _get_current_objects_of_this_type(
         self, object: Union[Agent, Scenario, LanguageModel]
     ) -> tuple[list, str]:
+        from edsl.agents.Agent import Agent
+        from edsl.scenarios.Scenario import Scenario
+        from edsl.language_models.LanguageModel import LanguageModel
+
         """Return the current objects of the same type as the first argument.
 
         >>> from edsl.jobs import Jobs
         >>> j = Jobs.example()
         >>> j._get_current_objects_of_this_type(j.agents[0])
-        ([Agent(traits = {'status': 'Joyful'}), Agent(traits = {'status': 'Sad'})], 'agents')
+        (AgentList([Agent(traits = {'status': 'Joyful'}), Agent(traits = {'status': 'Sad'})]), 'agents')
         """
         class_to_key = {
             Agent: "agents",
@@ -181,6 +242,20 @@ class Jobs(Base):
         current_objects = getattr(self, key, None)
         return current_objects, key
 
+    @staticmethod
+    def _get_empty_container_object(object):
+        from edsl import AgentList
+        from edsl import Agent
+        from edsl import Scenario
+        from edsl import ScenarioList
+
+        if isinstance(object, Agent):
+            return AgentList([])
+        elif isinstance(object, Scenario):
+            return ScenarioList([])
+        else:
+            return []
+
     @staticmethod
     def _merge_objects(passed_objects, current_objects) -> list:
         """
@@ -192,7 +267,7 @@ class Jobs(Base):
         >>> Jobs(survey = [])._merge_objects([1,2,3], [4,5,6])
         [5, 6, 7, 6, 7, 8, 7, 8, 9]
         """
-        new_objects = []
+        new_objects = Jobs._get_empty_container_object(passed_objects[0])
         for current_object in current_objects:
             for new_object in passed_objects:
                 new_objects.append(current_object + new_object)
@@ -237,12 +312,12 @@ class Jobs(Base):
         with us filling in defaults.
         """
         # if no agents, models, or scenarios are set, set them to defaults
+        from edsl.agents.Agent import Agent
+        from edsl.language_models.registry import Model
+        from edsl.scenarios.Scenario import Scenario
+
         self.agents = self.agents or [Agent()]
         self.models = self.models or [Model()]
-        # if remote, set all the models to remote
-        if hasattr(self, "remote") and self.remote:
-            for model in self.models:
-                model.remote = True
         self.scenarios = self.scenarios or [Scenario()]
         for agent, scenario, model in product(self.agents, self.scenarios, self.models):
             yield Interview(
@@ -256,6 +331,7 @@ class Jobs(Base):
         These buckets are used to track API calls and token usage.
 
         >>> from edsl.jobs import Jobs
+        >>> from edsl import Model
         >>> j = Jobs.example().by(Model(temperature = 1), Model(temperature = 0.5))
         >>> bc = j.create_bucket_collection()
         >>> bc
@@ -284,6 +360,57 @@ class Jobs(Base):
             )
         return links
 
+    def __hash__(self):
+        """Allow the model to be used as a key in a dictionary."""
+        from edsl.utilities.utilities import dict_hash
+
+        return dict_hash(self.to_dict())
+
+    def _output(self, message) -> None:
+        """Check if a Job is verbose. If so, print the message."""
+        if self.verbose:
+            print(message)
+
+    def _check_parameters(self, strict=False, warn=False) -> None:
+        """Check if the parameters in the survey and scenarios are consistent.
+
+        >>> from edsl import QuestionFreeText
+        >>> from edsl import Survey
+        >>> from edsl import Scenario
+        >>> q = QuestionFreeText(question_text = "{{poo}}", question_name = "ugly_question")
+        >>> j = Jobs(survey = Survey(questions=[q]))
+        >>> with warnings.catch_warnings(record=True) as w:
+        ...     j._check_parameters(warn = True)
+        ...     assert len(w) == 1
+        ...     assert issubclass(w[-1].category, UserWarning)
+        ...     assert "The following parameters are in the survey but not in the scenarios" in str(w[-1].message)
+
+        >>> q = QuestionFreeText(question_text = "{{poo}}", question_name = "ugly_question")
+        >>> s = Scenario({'plop': "A", 'poo': "B"})
+        >>> j = Jobs(survey = Survey(questions=[q])).by(s)
+        >>> j._check_parameters(strict = True)
+        Traceback (most recent call last):
+        ...
+        ValueError: The following parameters are in the scenarios but not in the survey: {'plop'}
+        """
+        survey_parameters: set = self.survey.parameters
+        scenario_parameters: set = self.scenarios.parameters
+
+        msg1, msg2 = None, None
+
+        if in_survey_but_not_in_scenarios := survey_parameters - scenario_parameters:
+            msg1 = f"The following parameters are in the survey but not in the scenarios: {in_survey_but_not_in_scenarios}"
+        if in_scenarios_but_not_in_survey := scenario_parameters - survey_parameters:
+            msg2 = f"The following parameters are in the scenarios but not in the survey: {in_scenarios_but_not_in_survey}"
+
+        if msg1 or msg2:
+            message = "\n".join(filter(None, [msg1, msg2]))
+            if strict:
+                raise ValueError(message)
+            else:
+                if warn:
+                    warnings.warn(message)
+
     def run(
         self,
         n: int = 1,
@@ -291,41 +418,81 @@ class Jobs(Base):
         progress_bar: bool = False,
         stop_on_exception: bool = False,
         cache: Union[Cache, bool] = None,
-        remote: bool = (
-            False if os.getenv("DEFAULT_RUN_MODE", "local") == "local" else True
-        ),
         check_api_keys: bool = False,
         sidecar_model: Optional[LanguageModel] = None,
         batch_mode: Optional[bool] = None,
-        print_exceptions=False,
+        verbose: bool = False,
+        print_exceptions=True,
+        remote_cache_description: Optional[str] = None,
+        remote_inference_description: Optional[str] = None,
     ) -> Results:
         """
         Runs the Job: conducts Interviews and returns their results.
 
         :param n: how many times to run each interview
         :param debug: prints debug messages
-        :param verbose: prints messages
         :param progress_bar: shows a progress bar
         :param stop_on_exception: stops the job if an exception is raised
         :param cache: a cache object to store results
-        :param remote: run the job remotely
         :param check_api_keys: check if the API keys are valid
-        :batch_mode: run the job in batch mode i.e., no expecation of interaction with the user
-
+        :param batch_mode: run the job in batch mode i.e., no expecation of interaction with the user
+        :param verbose: prints messages
+        :param remote_cache_description: specifies a description for this group of entries in the remote cache
+        :param remote_inference_description: specifies a description for the remote inference job
         """
+        from edsl.coop.coop import Coop
+
+        self._check_parameters()
+
         if batch_mode is not None:
             raise NotImplementedError(
                 "Batch mode is deprecated. Please update your code to not include 'batch_mode' in the 'run' method."
             )
 
-        self.remote = remote
-
-        if self.remote:
-            ## TODO: This should be a coop check
-            if os.getenv("EXPECTED_PARROT_API_KEY", None) is None:
-                raise MissingRemoteInferenceError()
+        self.verbose = verbose
+
+        try:
+            coop = Coop()
+            user_edsl_settings = coop.edsl_settings
+            remote_cache = user_edsl_settings["remote_caching"]
+            remote_inference = user_edsl_settings["remote_inference"]
+        except Exception:
+            remote_cache = False
+            remote_inference = False
+
+        if remote_inference:
+            self._output("Remote inference activated. Sending job to server...")
+            if remote_cache:
+                self._output(
+                    "Remote caching activated. The remote cache will be used for this job."
+                )
 
-        if not self.remote:
+            remote_job_data = coop.remote_inference_create(
+                self,
+                description=remote_inference_description,
+                status="queued",
+            )
+            self._output("Job sent!")
+            # Create mock results object to store job data
+            results = Results(
+                survey=Survey(),
+                data=[
+                    Result(
+                        agent=Agent.example(),
+                        scenario=Scenario.example(),
+                        model=Model(),
+                        iteration=1,
+                        answer={"info": "Remote job details"},
+                    )
+                ],
+            )
+            results.add_columns_from_dict([remote_job_data])
+            if self.verbose:
+                results.select(["info", "uuid", "status", "version"]).print(
+                    format="rich"
+                )
+            return results
+        else:
             if check_api_keys:
                 for model in self.models + [Model()]:
                     if not model.has_valid_api_key():
@@ -336,30 +503,111 @@ class Jobs(Base):
 
         # handle cache
         if cache is None:
+            from edsl.data.CacheHandler import CacheHandler
+
             cache = CacheHandler().get_cache()
         if cache is False:
+            from edsl.data.Cache import Cache
+
             cache = Cache()
 
-        results = self._run_local(
-            n=n,
-            debug=debug,
-            progress_bar=progress_bar,
-            cache=cache,
-            stop_on_exception=stop_on_exception,
-            sidecar_model=sidecar_model,
-            print_exceptions=print_exceptions,
-        )
-        results.cache = cache.new_entries_cache()
+        if not remote_cache:
+            results = self._run_local(
+                n=n,
+                debug=debug,
+                progress_bar=progress_bar,
+                cache=cache,
+                stop_on_exception=stop_on_exception,
+                sidecar_model=sidecar_model,
+                print_exceptions=print_exceptions,
+            )
+
+            results.cache = cache.new_entries_cache()
+
+            self._output(f"There are {len(cache.keys()):,} entries in the local cache.")
+        else:
+            cache_difference = coop.remote_cache_get_diff(cache.keys())
+
+            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..."
+                )
+                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.")
+
+            server_missing_cacheentry_keys = cache_difference.get(
+                "server_missing_cacheentry_keys", []
+            )
+            server_missing_cacheentries = [
+                entry
+                for key in server_missing_cacheentry_keys
+                if (entry := cache.data.get(key)) is not None
+            ]
+            old_entry_keys = [key for key in cache.keys()]
+
+            self._output("Running job...")
+            results = self._run_local(
+                n=n,
+                debug=debug,
+                progress_bar=progress_bar,
+                cache=cache,
+                stop_on_exception=stop_on_exception,
+                sidecar_model=sidecar_model,
+                print_exceptions=print_exceptions,
+            )
+            self._output("Job completed!")
+
+            new_cache_entries = list(
+                [entry for entry in cache.values() if entry.key not in 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 "
+                    f"{'entry' if new_entry_count == 1 else 'entries'}..."
+                )
+                coop.remote_cache_create_many(
+                    server_missing_cacheentries,
+                    visibility="private",
+                    description=remote_cache_description,
+                )
+                self._output("Remote cache updated!")
+            else:
+                self._output("No new entries to add to remote cache.")
+
+            results.cache = cache.new_entries_cache()
+
+            self._output(f"There are {len(cache.keys()):,} entries in the local cache.")
 
         return results
 
     def _run_local(self, *args, **kwargs):
         """Run the job locally."""
-        from edsl.jobs.runners.JobsRunnerAsyncio import JobsRunnerAsyncio
 
         results = JobsRunnerAsyncio(self).run(*args, **kwargs)
         return results
 
+    async def run_async(self, cache=None, **kwargs):
+        """Run the job asynchronously."""
+        results = await JobsRunnerAsyncio(self).run_async(cache=cache, **kwargs)
+        return results
+
+    def all_question_parameters(self):
+        """Return all the fields in the questions in the survey."""
+        return set.union(*[question.parameters for question in self.survey.questions])
+
     #######################
     # Dunder methods
     #######################
@@ -412,6 +660,11 @@ class Jobs(Base):
     @remove_edsl_version
     def from_dict(cls, data: dict) -> Jobs:
         """Creates a Jobs instance from a dictionary."""
+        from edsl import Survey
+        from edsl.agents.Agent import Agent
+        from edsl.language_models.LanguageModel import LanguageModel
+        from edsl.scenarios.Scenario import Scenario
+
         return cls(
             survey=Survey.from_dict(data["survey"]),
             agents=[Agent.from_dict(agent) for agent in data["agents"]],
@@ -438,7 +691,8 @@ class Jobs(Base):
         """
         import random
         from edsl.questions import QuestionMultipleChoice
-        from edsl import Agent
+        from edsl.agents.Agent import Agent
+        from edsl.scenarios.Scenario import Scenario
 
         # (status, question, period)
         agent_answers = {
@@ -477,11 +731,14 @@ class Jobs(Base):
             question_options=["Good", "Great", "OK", "Terrible"],
             question_name="how_feeling_yesterday",
         )
+        from edsl import Survey, ScenarioList
+
         base_survey = Survey(questions=[q1, q2])
 
-        job = base_survey.by(
-            Scenario({"period": "morning"}), Scenario({"period": "afternoon"})
-        ).by(joy_agent, sad_agent)
+        scenario_list = ScenarioList(
+            [Scenario({"period": "morning"}), Scenario({"period": "afternoon"})]
+        )
+        job = base_survey.by(scenario_list).by(joy_agent, sad_agent)
 
         return job
 
@@ -516,11 +773,3 @@ if __name__ == "__main__":
     import doctest
 
     doctest.testmod(optionflags=doctest.ELLIPSIS)
-
-    # from edsl.jobs import Jobs
-
-    # job = Jobs.example()
-    # len(job) == 8
-    # results, info = job.run(debug=True)
-    # len(results) == 8
-    # results

edsl/jobs/buckets/BucketCollection.py

@@ -10,8 +10,9 @@ class BucketCollection(UserDict):
     Models themselves are hashable, so this works.
     """
 
-    def __init__(self):
+    def __init__(self, infinity_buckets=False):
         super().__init__()
+        self.infinity_buckets = infinity_buckets
 
     def __repr__(self):
         return f"BucketCollection({self.data})"
@@ -21,8 +22,13 @@ class BucketCollection(UserDict):
 
         This will create the token and request buckets for the model."""
         # compute the TPS and RPS from the model
-        TPS = model.TPM / 60.0
-        RPS = model.RPM / 60.0
+        if not self.infinity_buckets:
+            TPS = model.TPM / 60.0
+            RPS = model.RPM / 60.0
+        else:
+            TPS = float("inf")
+            RPS = float("inf")
+
         # create the buckets
         requests_bucket = TokenBucket(
             bucket_name=model.model,

edsl/jobs/buckets/ModelBuckets.py

@@ -1,4 +1,4 @@
-from edsl.jobs.buckets.TokenBucket import TokenBucket
+# from edsl.jobs.buckets.TokenBucket import TokenBucket
 
 
 class ModelBuckets:
@@ -8,7 +8,7 @@ class ModelBuckets:
     A request is one call to the service. The number of tokens required for a request depends on parameters.
     """
 
-    def __init__(self, requests_bucket: TokenBucket, tokens_bucket: TokenBucket):
+    def __init__(self, requests_bucket: "TokenBucket", tokens_bucket: "TokenBucket"):
         """Initialize the model buckets.
 
         The requests bucket captures requests per unit of time.
@@ -28,6 +28,8 @@ class ModelBuckets:
     @classmethod
     def infinity_bucket(cls, model_name: str = "not_specified") -> "ModelBuckets":
         """Create a bucket with infinite capacity and refill rate."""
+        from edsl.jobs.buckets.TokenBucket import TokenBucket
+
         return cls(
             requests_bucket=TokenBucket(
                 bucket_name=model_name,

edsl/jobs/buckets/TokenBucket.py

@@ -1,8 +1,6 @@
 from typing import Union, List, Any
 import asyncio
 import time
-from collections import UserDict
-from matplotlib import pyplot as plt
 
 
 class TokenBucket:
@@ -114,6 +112,7 @@ class TokenBucket:
         times, tokens = zip(*self.get_log())
         start_time = times[0]
         times = [t - start_time for t in times]  # Normalize time to start from 0
+        from matplotlib import pyplot as plt
 
         plt.figure(figsize=(10, 6))
         plt.plot(times, tokens, label="Tokens Available")