edsl 0.1.39__py3-none-any.whl → 0.1.39.dev2__py3-none-any.whl

edsl/jobs/Jobs.py

@@ -1,17 +1,7 @@
 # """The Jobs class is a collection of agents, scenarios and models and one survey."""
 from __future__ import annotations
-import asyncio
-from inspect import signature
-from typing import (
-    Literal,
-    Optional,
-    Union,
-    Sequence,
-    Generator,
-    TYPE_CHECKING,
-    Callable,
-    Tuple,
-)
+import warnings
+from typing import Literal, Optional, Union, Sequence, Generator, TYPE_CHECKING
 
 from edsl.Base import Base
 
@@ -19,13 +9,10 @@ from edsl.jobs.buckets.BucketCollection import BucketCollection
 from edsl.jobs.JobsPrompts import JobsPrompts
 from edsl.jobs.interviews.Interview import Interview
 from edsl.utilities.remove_edsl_version import remove_edsl_version
-from edsl.jobs.runners.JobsRunnerAsyncio import JobsRunnerAsyncio
+
 from edsl.data.RemoteCacheSync import RemoteCacheSync
 from edsl.exceptions.coop import CoopServerResponseError
 
-from edsl.jobs.JobsChecks import JobsChecks
-from edsl.jobs.data_structures import RunEnvironment, RunParameters, RunConfig
-
 if TYPE_CHECKING:
     from edsl.agents.Agent import Agent
     from edsl.agents.AgentList import AgentList
@@ -36,66 +23,6 @@ if TYPE_CHECKING:
     from edsl.results.Results import Results
     from edsl.results.Dataset import Dataset
     from edsl.language_models.ModelList import ModelList
-    from edsl.data.Cache import Cache
-    from edsl.language_models.key_management.KeyLookup import KeyLookup
-
-VisibilityType = Literal["private", "public", "unlisted"]
-
-from dataclasses import dataclass
-from typing import Optional, Union, TypeVar, Callable, cast
-from functools import wraps
-
-try:
-    from typing import ParamSpec
-except ImportError:
-    from typing_extensions import ParamSpec
-
-
-P = ParamSpec("P")
-T = TypeVar("T")
-
-
-from edsl.jobs.check_survey_scenario_compatibility import (
-    CheckSurveyScenarioCompatibility,
-)
-
-
-def with_config(f: Callable[P, T]) -> Callable[P, T]:
-    "This decorator make it so that the run function parameters match the RunConfig dataclass."
-    parameter_fields = {
-        name: field.default
-        for name, field in RunParameters.__dataclass_fields__.items()
-    }
-    environment_fields = {
-        name: field.default
-        for name, field in RunEnvironment.__dataclass_fields__.items()
-    }
-    combined = {**parameter_fields, **environment_fields}
-
-    @wraps(f)
-    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
-        environment = RunEnvironment(
-            **{k: v for k, v in kwargs.items() if k in environment_fields}
-        )
-        parameters = RunParameters(
-            **{k: v for k, v in kwargs.items() if k in parameter_fields}
-        )
-        config = RunConfig(environment=environment, parameters=parameters)
-        return f(*args, config=config)
-
-    # Update the wrapper's signature to include all RunConfig parameters
-    # old_sig = signature(f)
-    # wrapper.__signature__ = old_sig.replace(
-    #     parameters=list(old_sig.parameters.values())[:-1]
-    #     + [
-    #         old_sig.parameters["config"].replace(
-    #             default=parameter_fields[name], name=name
-    #         )
-    #         for name in combined
-    #     ]
-    # )
-
-    return cast(Callable[P, T], wrapper)
 
 
 class Jobs(Base):
@@ -119,62 +46,15 @@ class Jobs(Base):
         :param models: a list of models
         :param scenarios: a list of scenarios
         """
-        self.run_config = RunConfig(
-            environment=RunEnvironment(), parameters=RunParameters()
-        )
-
         self.survey = survey
         self.agents: AgentList = agents
         self.scenarios: ScenarioList = scenarios
-        self.models: ModelList = models
-
-    def add_running_env(self, running_env: RunEnvironment):
-        self.run_config.add_environment(running_env)
-        return self
+        self.models = models
 
-    def using_cache(self, cache: "Cache") -> Jobs:
-        """
-        Add a Cache to the job.
-
-        :param cache: the cache to add
-        """
-        self.run_config.add_cache(cache)
-        return self
-
-    def using_bucket_collection(self, bucket_collection: BucketCollection) -> Jobs:
-        """
-        Add a BucketCollection to the job.
-
-        :param bucket_collection: the bucket collection to add
-        """
-        self.run_config.add_bucket_collection(bucket_collection)
-        return self
-
-    def using_key_lookup(self, key_lookup: KeyLookup) -> Jobs:
-        """
-        Add a KeyLookup to the job.
-
-        :param key_lookup: the key lookup to add
-        """
-        self.run_config.add_key_lookup(key_lookup)
-        return self
-
-    def using(self, obj: Union[Cache, BucketCollection, KeyLookup]) -> Jobs:
-        """
-        Add a Cache, BucketCollection, or KeyLookup to the job.
-
-        :param obj: the object to add
-        """
-        from edsl.data.Cache import Cache
-        from edsl.language_models.key_management.KeyLookup import KeyLookup
+        self.__bucket_collection = None
 
-        if isinstance(obj, Cache):
-            self.using_cache(obj)
-        elif isinstance(obj, BucketCollection):
-            self.using_bucket_collection(obj)
-        elif isinstance(obj, KeyLookup):
-            self.using_key_lookup(obj)
-        return self
+    # these setters and getters are used to ensure that the agents, models, and scenarios
+    # are stored as AgentList, ModelList, and ScenarioList objects.
 
     @property
     def models(self):
@@ -192,12 +72,6 @@ class Jobs(Base):
         else:
             self._models = ModelList([])
 
-        # update the bucket collection if it exists
-        if self.run_config.environment.bucket_collection is None:
-            self.run_config.environment.bucket_collection = (
-                self.create_bucket_collection()
-            )
-
     @property
     def agents(self):
         return self._agents
@@ -340,29 +214,13 @@ class Jobs(Base):
 
     def replace_missing_objects(self) -> None:
         from edsl.agents.Agent import Agent
-        from edsl.language_models.model import Model
+        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()]
         self.scenarios = self.scenarios or [Scenario()]
 
-    def generate_interviews(self) -> Generator[Interview, None, None]:
-        """
-        Generate interviews.
-
-        Note that this sets the agents, model and scenarios if they have not been set. This is a side effect of the method.
-        This is useful because a user can create a job without setting the agents, models, or scenarios, and the job will still run,
-        with us filling in defaults.
-
-        """
-        from edsl.jobs.InterviewsConstructor import InterviewsConstructor
-
-        self.replace_missing_objects()
-        yield from InterviewsConstructor(
-            self, cache=self.run_config.environment.cache
-        ).create_interviews()
-
     def interviews(self) -> list[Interview]:
         """
         Return a list of :class:`edsl.jobs.interviews.Interview` objects.
@@ -377,10 +235,18 @@ class Jobs(Base):
         >>> j.interviews()[0]
         Interview(agent = Agent(traits = {'status': 'Joyful'}), survey = Survey(...), scenario = Scenario({'period': 'morning'}), model = Model(...))
         """
-        return list(self.generate_interviews())
+        if hasattr(self, "_interviews"):
+            return self._interviews
+        else:
+            self.replace_missing_objects()
+            from edsl.jobs.InterviewsConstructor import InterviewsConstructor
+
+            self._interviews = list(InterviewsConstructor(self).create_interviews())
+
+        return self._interviews
 
     @classmethod
-    def from_interviews(cls, interview_list) -> "Jobs":
+    def from_interviews(cls, interview_list):
         """Return a Jobs instance from a list of interviews.
 
         This is useful when you have, say, a list of failed interviews and you want to create
@@ -407,8 +273,16 @@ class Jobs(Base):
         >>> bc
         BucketCollection(...)
         """
+        self.replace_missing_objects()  # ensure that all objects are present
         return BucketCollection.from_models(self.models)
 
+    @property
+    def bucket_collection(self) -> BucketCollection:
+        """Return the bucket collection. If it does not exist, create it."""
+        if self.__bucket_collection is None:
+            self.__bucket_collection = self.create_bucket_collection()
+        return self.__bucket_collection
+
     def html(self):
         """Return the HTML representations for each scenario"""
         links = []
@@ -434,12 +308,10 @@ class Jobs(Base):
 
     def _output(self, message) -> None:
         """Check if a Job is verbose. If so, print the message."""
-        if self.run_config.parameters.verbose:
+        if hasattr(self, "verbose") and self.verbose:
             print(message)
-        # if hasattr(self, "verbose") and self.verbose:
-        #    print(message)
 
-    def all_question_parameters(self) -> set:
+    def all_question_parameters(self):
         """Return all the fields in the questions in the survey.
         >>> from edsl.jobs import Jobs
         >>> Jobs.example().all_question_parameters()
@@ -447,12 +319,86 @@ class Jobs(Base):
         """
         return set.union(*[question.parameters for question in self.survey.questions])
 
-    def use_remote_cache(self) -> bool:
+    def _check_parameters(self, strict=False, warn=False) -> None:
+        """Check if the parameters in the survey and scenarios are consistent.
+
+        >>> from edsl.questions.QuestionFreeText import QuestionFreeText
+        >>> from edsl.surveys.Survey import Survey
+        >>> from edsl.scenarios.Scenario 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'}
+
+        >>> q = QuestionFreeText(question_text = "Hello", question_name = "ugly_question")
+        >>> s = Scenario({'ugly_question': "B"})
+        >>> j = Jobs(survey = Survey(questions=[q])).by(s)
+        >>> j._check_parameters()
+        Traceback (most recent call last):
+        ...
+        ValueError: The following names are in both the survey question_names and the scenario keys: {'ugly_question'}. This will create issues.
+        """
+        survey_parameters: set = self.survey.parameters
+        scenario_parameters: set = self.scenarios.parameters
+
+        msg0, msg1, msg2 = None, None, None
+
+        # look for key issues
+        if intersection := set(self.scenarios.parameters) & set(
+            self.survey.question_names
+        ):
+            msg0 = f"The following names are in both the survey question_names and the scenario keys: {intersection}. This will create issues."
+
+            raise ValueError(msg0)
+
+        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)
+
+        if self.scenarios.has_jinja_braces:
+            warnings.warn(
+                "The scenarios have Jinja braces ({{ and }}). Converting to '<<' and '>>'. If you want a different conversion, use the convert_jinja_braces method first to modify the scenario."
+            )
+            self.scenarios = self.scenarios._convert_jinja_braces()
+
+    @property
+    def skip_retry(self):
+        if not hasattr(self, "_skip_retry"):
+            return False
+        return self._skip_retry
+
+    @property
+    def raise_validation_errors(self):
+        if not hasattr(self, "_raise_validation_errors"):
+            return False
+        return self._raise_validation_errors
+
+    def use_remote_cache(self, disable_remote_cache: bool) -> bool:
         import requests
 
-        if self.run_config.parameters.disable_remote_cache:
+        if disable_remote_cache:
             return False
-        if not self.run_config.parameters.disable_remote_cache:
+        if not disable_remote_cache:
             try:
                 from edsl.coop.coop import Coop
 
@@ -465,173 +411,154 @@ class Jobs(Base):
 
         return False
 
-    def _remote_results(
+    def run(
         self,
-    ) -> Union["Results", None]:
-        from edsl.jobs.JobsRemoteInferenceHandler import JobsRemoteInferenceHandler
+        n: int = 1,
+        progress_bar: bool = False,
+        stop_on_exception: bool = False,
+        cache: Union["Cache", bool] = None,
+        check_api_keys: bool = False,
+        sidecar_model: Optional[LanguageModel] = None,
+        verbose: bool = True,
+        print_exceptions=True,
+        remote_cache_description: Optional[str] = None,
+        remote_inference_description: Optional[str] = None,
+        remote_inference_results_visibility: Optional[
+            Literal["private", "public", "unlisted"]
+        ] = "unlisted",
+        skip_retry: bool = False,
+        raise_validation_errors: bool = False,
+        disable_remote_cache: bool = False,
+        disable_remote_inference: bool = False,
+        bucket_collection: Optional[BucketCollection] = None,
+    ) -> Results:
+        """
+        Runs the Job: conducts Interviews and returns their results.
 
-        jh = JobsRemoteInferenceHandler(
-            self, verbose=self.run_config.parameters.verbose
-        )
-        if jh.use_remote_inference(self.run_config.parameters.disable_remote_inference):
-            job_info = jh.create_remote_inference_job(
-                iterations=self.run_config.parameters.n,
-                remote_inference_description=self.run_config.parameters.remote_inference_description,
-                remote_inference_results_visibility=self.run_config.parameters.remote_inference_results_visibility,
-            )
-            results = jh.poll_remote_inference_job(job_info)
-            return results
-        else:
-            return None
+        :param n: How many times to run each interview
+        :param progress_bar: Whether to show a progress bar
+        :param stop_on_exception: Stops the job if an exception is raised
+        :param cache: A Cache object to store results
+        :param check_api_keys: Raises an error if API keys are invalid
+        :param verbose: Prints extra 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
+        :param remote_inference_results_visibility: The initial visibility of the Results object on Coop. This will only be used for remote jobs!
+        :param disable_remote_cache: If True, the job will not use remote cache. This only works for local jobs!
+        :param disable_remote_inference: If True, the job will not use remote inference
+        """
+        from edsl.coop.coop import Coop
+        from edsl.jobs.JobsChecks import JobsChecks
+        from edsl.jobs.JobsRemoteInferenceHandler import JobsRemoteInferenceHandler
 
-    def _prepare_to_run(self) -> None:
-        "This makes sure that the job is ready to run and that keys are in place for a remote job."
-        CheckSurveyScenarioCompatibility(self.survey, self.scenarios).check()
+        self._check_parameters()
+        self._skip_retry = skip_retry
+        self._raise_validation_errors = raise_validation_errors
+        self.verbose = verbose
 
-    def _check_if_remote_keys_ok(self):
         jc = JobsChecks(self)
+
+        # check if the user has all the keys they need
         if jc.needs_key_process():
             jc.key_process()
 
-    def _check_if_local_keys_ok(self):
-        jc = JobsChecks(self)
-        if self.run_config.parameters.check_api_keys:
+        jh = JobsRemoteInferenceHandler(self, verbose=verbose)
+        if jh.use_remote_inference(disable_remote_inference):
+            jh.create_remote_inference_job(
+                iterations=n,
+                remote_inference_description=remote_inference_description,
+                remote_inference_results_visibility=remote_inference_results_visibility,
+            )
+            results = jh.poll_remote_inference_job()
+            return results
+
+        if check_api_keys:
             jc.check_api_keys()
 
-    async def _execute_with_remote_cache(self, run_job_async: bool) -> Results:
+        # handle cache
+        if cache is None or cache is True:
+            from edsl.data.CacheHandler import CacheHandler
 
-        use_remote_cache = self.use_remote_cache()
+            cache = CacheHandler().get_cache()
+        if cache is False:
+            from edsl.data.Cache import Cache
 
-        from edsl.coop.coop import Coop
-        from edsl.jobs.runners.JobsRunnerAsyncio import JobsRunnerAsyncio
-        from edsl.data.Cache import Cache
+            cache = Cache()
 
-        assert isinstance(self.run_config.environment.cache, Cache)
+        if bucket_collection is None:
+            bucket_collection = self.create_bucket_collection()
 
+        remote_cache = self.use_remote_cache(disable_remote_cache)
         with RemoteCacheSync(
             coop=Coop(),
-            cache=self.run_config.environment.cache,
+            cache=cache,
             output_func=self._output,
-            remote_cache=use_remote_cache,
-            remote_cache_description=self.run_config.parameters.remote_cache_description,
-        ):
-            runner = JobsRunnerAsyncio(self, environment=self.run_config.environment)
-            if run_job_async:
-                results = await runner.run_async(self.run_config.parameters)
-            else:
-                results = runner.run(self.run_config.parameters)
+            remote_cache=remote_cache,
+            remote_cache_description=remote_cache_description,
+        ) as r:
+            results = self._run_local(
+                n=n,
+                progress_bar=progress_bar,
+                cache=cache,
+                stop_on_exception=stop_on_exception,
+                sidecar_model=sidecar_model,
+                print_exceptions=print_exceptions,
+                raise_validation_errors=raise_validation_errors,
+                bucket_collection=bucket_collection,
+            )
         return results
 
-    def _setup_and_check(self) -> Tuple[RunConfig, Optional[Results]]:
-
-        self._prepare_to_run()
-        self._check_if_remote_keys_ok()
-
-        # first try to run the job remotely
-        if results := self._remote_results():
-            return results
-
-        self._check_if_local_keys_ok()
-        return None
-
-    @property
-    def num_interviews(self):
-        if self.run_config.parameters.n is None:
-            return len(self)
-        else:
-            len(self) * self.run_config.parameters.n
+    async def run_async(
+        self,
+        cache=None,
+        n=1,
+        disable_remote_inference: bool = False,
+        remote_inference_description: Optional[str] = None,
+        remote_inference_results_visibility: Optional[
+            Literal["private", "public", "unlisted"]
+        ] = "unlisted",
+        bucket_collection: Optional[BucketCollection] = None,
+        **kwargs,
+    ):
+        """Run the job asynchronously, either locally or remotely.
 
-    def _run(self, config: RunConfig):
-        "Shared code for run and run_async"
-        if config.environment.cache is not None:
-            self.run_config.environment.cache = config.environment.cache
+        :param cache: Cache object or boolean
+        :param n: Number of iterations
+        :param disable_remote_inference: If True, forces local execution
+        :param remote_inference_description: Description for remote jobs
+        :param remote_inference_results_visibility: Visibility setting for remote results
+        :param kwargs: Additional arguments passed to local execution
+        :return: Results object
+        """
+        # Check if we should use remote inference
+        from edsl.jobs.JobsRemoteInferenceHandler import JobsRemoteInferenceHandler
+        from edsl.jobs.runners.JobsRunnerAsyncio import JobsRunnerAsyncio
 
-        if config.environment.bucket_collection is not None:
-            self.run_config.environment.bucket_collection = (
-                config.environment.bucket_collection
+        jh = JobsRemoteInferenceHandler(self, verbose=False)
+        if jh.use_remote_inference(disable_remote_inference):
+            results = await jh.create_and_poll_remote_job(
+                iterations=n,
+                remote_inference_description=remote_inference_description,
+                remote_inference_results_visibility=remote_inference_results_visibility,
             )
-
-        if config.environment.key_lookup is not None:
-            self.run_config.environment.key_lookup = config.environment.key_lookup
-
-        # replace the parameters with the ones from the config
-        self.run_config.parameters = config.parameters
-
-        self.replace_missing_objects()
-
-        # try to run remotely first
-        self._prepare_to_run()
-        self._check_if_remote_keys_ok()
-
-        if (
-            self.run_config.environment.cache is None
-            or self.run_config.environment.cache is True
-        ):
-            from edsl.data.CacheHandler import CacheHandler
-
-            self.run_config.environment.cache = CacheHandler().get_cache()
-
-        if self.run_config.environment.cache is False:
-            from edsl.data.Cache import Cache
-
-            self.run_config.environment.cache = Cache(immediate_write=False)
-
-        # first try to run the job remotely
-        if results := self._remote_results():
             return results
 
-        self._check_if_local_keys_ok()
-
-        if config.environment.bucket_collection is None:
-            self.run_config.environment.bucket_collection = (
-                self.create_bucket_collection()
-            )
-
-    @with_config
-    def run(self, *, config: RunConfig) -> "Results":
-        """
-        Runs the Job: conducts Interviews and returns their results.
-
-        :param n: How many times to run each interview
-        :param progress_bar: Whether to show a progress bar
-        :param stop_on_exception: Stops the job if an exception is raised
-        :param check_api_keys: Raises an error if API keys are invalid
-        :param verbose: Prints extra 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
-        :param remote_inference_results_visibility: The initial visibility of the Results object on Coop. This will only be used for remote jobs!
-        :param disable_remote_cache: If True, the job will not use remote cache. This only works for local jobs!
-        :param disable_remote_inference: If True, the job will not use remote inference
-        :param cache: A Cache object to store results
-        :param bucket_collection: A BucketCollection object to track API calls
-        :param key_lookup: A KeyLookup object to manage API keys
-        """
-        self._run(config)
-
-        return asyncio.run(self._execute_with_remote_cache(run_job_async=False))
+        if bucket_collection is None:
+            bucket_collection = self.create_bucket_collection()
 
-    @with_config
-    async def run_async(self, *, config: RunConfig) -> "Results":
-        """
-        Runs the Job: conducts Interviews and returns their results.
+        # If not using remote inference, run locally with async
+        return await JobsRunnerAsyncio(
+            self, bucket_collection=bucket_collection
+        ).run_async(cache=cache, n=n, **kwargs)
 
-        :param n: How many times to run each interview
-        :param progress_bar: Whether to show a progress bar
-        :param stop_on_exception: Stops the job if an exception is raised
-        :param check_api_keys: Raises an error if API keys are invalid
-        :param verbose: Prints extra 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
-        :param remote_inference_results_visibility: The initial visibility of the Results object on Coop. This will only be used for remote jobs!
-        :param disable_remote_cache: If True, the job will not use remote cache. This only works for local jobs!
-        :param disable_remote_inference: If True, the job will not use remote inference
-        :param cache: A Cache object to store results
-        :param bucket_collection: A BucketCollection object to track API calls
-        :param key_lookup: A KeyLookup object to manage API keys
-        """
-        self._run(config)
+    def _run_local(self, bucket_collection, *args, **kwargs):
+        """Run the job locally."""
+        from edsl.jobs.runners.JobsRunnerAsyncio import JobsRunnerAsyncio
 
-        return await self._execute_with_remote_cache(run_job_async=True)
+        results = JobsRunnerAsyncio(self, bucket_collection=bucket_collection).run(
+            *args, **kwargs
+        )
+        return results
 
     def __repr__(self) -> str:
         """Return an eval-able string representation of the Jobs instance."""
@@ -661,6 +588,10 @@ class Jobs(Base):
         )
         return number_of_questions
 
+    #######################
+    # Serialization methods
+    #######################
+
     def to_dict(self, add_edsl_version=True):
         d = {
             "survey": self.survey.to_dict(add_edsl_version=add_edsl_version),
@@ -714,6 +645,9 @@ class Jobs(Base):
         """
         return hash(self) == hash(other)
 
+    #######################
+    # Example methods
+    #######################
     @classmethod
     def example(
         cls,