edsl 0.1.42__py3-none-any.whl → 0.1.44__py3-none-any.whl

edsl/jobs/Jobs.py

@@ -38,6 +38,7 @@ if TYPE_CHECKING:
     from edsl.language_models.ModelList import ModelList
     from edsl.data.Cache import Cache
     from edsl.language_models.key_management.KeyLookup import KeyLookup
+    from edsl.jobs.JobsRemoteInferenceHandler import JobsRemoteInferenceHandler
 
 VisibilityType = Literal["private", "public", "unlisted"]
 
@@ -407,7 +408,13 @@ class Jobs(Base):
         >>> bc
         BucketCollection(...)
         """
-        return BucketCollection.from_models(self.models)
+        bc = BucketCollection.from_models(self.models)
+        
+        if self.run_config.environment.key_lookup is not None:
+            bc.update_from_key_lookup(
+                self.run_config.environment.key_lookup
+            )
+        return bc
 
     def html(self):
         """Return the HTML representations for each scenario"""
@@ -465,22 +472,47 @@ class Jobs(Base):
 
         return False
 
+    def _start_remote_inference_job(
+        self, job_handler: Optional[JobsRemoteInferenceHandler] = None
+    ) -> Union["Results", None]:
+
+        if job_handler is None:
+            job_handler = self._create_remote_inference_handler()
+            
+        job_info = job_handler.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,
+        )
+        return job_info
+    
+    def _create_remote_inference_handler(self) -> JobsRemoteInferenceHandler:
+
+        from edsl.jobs.JobsRemoteInferenceHandler import JobsRemoteInferenceHandler
+        
+        return JobsRemoteInferenceHandler(
+            self, verbose=self.run_config.parameters.verbose
+        )
+        
     def _remote_results(
         self,
+        config: RunConfig,
     ) -> Union["Results", None]:
         from edsl.jobs.JobsRemoteInferenceHandler import JobsRemoteInferenceHandler
+        from edsl.jobs.JobsRemoteInferenceHandler import RemoteJobInfo
 
-        jh = JobsRemoteInferenceHandler(
-            self, verbose=self.run_config.parameters.verbose
-        )
+        background = config.parameters.background
+
+        jh = self._create_remote_inference_handler()
         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
+            job_info: RemoteJobInfo = self._start_remote_inference_job(jh)
+            if background:
+                from edsl.results.Results import Results 
+                results = Results.from_job_info(job_info)
+                return results
+            else:
+                results = jh.poll_remote_inference_job(job_info)
+                return results
         else:
             return None
 
@@ -507,13 +539,6 @@ class Jobs(Base):
 
         assert isinstance(self.run_config.environment.cache, Cache)
 
-        # with RemoteCacheSync(
-        #     coop=Coop(),
-        #     cache=self.run_config.environment.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)
@@ -521,19 +546,6 @@ class Jobs(Base):
             results = runner.run(self.run_config.parameters)
         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
-    #     results = self._remote_results()
-    #     #breakpoint()
-    #     if results is not None:
-    #         return results
-
-    #     self._check_if_local_keys_ok()
-    #     return None
-
     @property
     def num_interviews(self):
         if self.run_config.parameters.n is None:
@@ -563,7 +575,6 @@ class Jobs(Base):
 
         self.replace_missing_objects()
 
-        # try to run remotely first
         self._prepare_to_run()
         self._check_if_remote_keys_ok()
 
@@ -581,9 +592,9 @@ class Jobs(Base):
             self.run_config.environment.cache = Cache(immediate_write=False)
 
         # first try to run the job remotely
-        if results := self._remote_results():
+        if (results := self._remote_results(config)) is not None:
             return results
-                
+        
         self._check_if_local_keys_ok()
 
         if config.environment.bucket_collection is None:
@@ -591,6 +602,14 @@ class Jobs(Base):
                 self.create_bucket_collection()
             )
 
+        if (
+            self.run_config.environment.key_lookup is not None
+            and self.run_config.environment.bucket_collection is not None
+        ):
+            self.run_config.environment.bucket_collection.update_from_key_lookup(
+                self.run_config.environment.key_lookup
+            )
+
         return None
 
     @with_config
@@ -613,7 +632,7 @@ class Jobs(Base):
         :param key_lookup: A KeyLookup object to manage API keys
         """
         potentially_completed_results = self._run(config)
-        
+
         if potentially_completed_results is not None:
             return potentially_completed_results
 

edsl/jobs/JobsChecks.py

@@ -31,7 +31,7 @@ class JobsChecks:
         from edsl.language_models.model import Model
         from edsl.enums import service_to_api_keyname
 
-        for model in self.jobs.models + [Model()]:
+        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"
@@ -134,22 +134,22 @@ class JobsChecks:
 
         edsl_auth_token = secrets.token_urlsafe(16)
 
-        print("You're missing some of the API keys needed to run this job:")
+        print("API keys are required to run surveys with language models. The following keys are needed to run this survey: ")
         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."
+            "\nYou can provide your own keys or use an Expected Parrot key to access all available models."
         )
-        print("Remote inference allows you to run jobs on our server.")
+        print("Please see the documentation page to learn about options for managing keys: https://docs.expectedparrot.com/en/latest/api_keys.html")
 
         coop = Coop()
         coop._display_login_url(
             edsl_auth_token=edsl_auth_token,
-            link_description="\n🚀 To use remote inference, sign up at the following link:",
+            link_description="\n➡️ Click the link below to create an account and get an Expected Parrot key:\n",
         )
 
         print(
-            "\nOnce you log in, we will automatically retrieve your Expected Parrot API key and continue your job remotely."
+            "\nOnce you log in, your key will be stored on your computer and your survey will start running at the Expected Parrot server."
         )
 
         api_key = coop._poll_for_api_key(edsl_auth_token)
@@ -159,7 +159,7 @@ class JobsChecks:
             return
 
         path_to_env = write_api_key_to_env(api_key)
-        print("\n✨ API key retrieved and written to .env file at the following path:")
+        print("\n✨ Your key has been stored at the following path: ")
         print(f"    {path_to_env}")
 
         # Retrieve API key so we can continue running the job

edsl/jobs/JobsPrompts.py

@@ -1,3 +1,5 @@
+import time
+import logging
 from typing import List, TYPE_CHECKING
 
 from edsl.results.Dataset import Dataset
@@ -14,6 +16,7 @@ if TYPE_CHECKING:
 from edsl.jobs.FetchInvigilator import FetchInvigilator
 from edsl.data.CacheEntry import CacheEntry
 
+logger = logging.getLogger(__name__)
 
 class JobsPrompts:
     def __init__(self, jobs: "Jobs"):
@@ -22,6 +25,8 @@ class JobsPrompts:
         self.scenarios = jobs.scenarios
         self.survey = jobs.survey
         self._price_lookup = None
+        self._agent_lookup = {agent: idx for idx, agent in enumerate(self.agents)}
+        self._scenario_lookup = {scenario: idx for idx, scenario in enumerate(self.scenarios)}
 
     @property
     def price_lookup(self):
@@ -49,25 +54,53 @@ class JobsPrompts:
         models = []
         costs = []
         cache_keys = []
+        
         for interview_index, interview in enumerate(interviews):
+            logger.info(f"Processing interview {interview_index} of {len(interviews)}")
+            interview_start = time.time()
+            
+            # Fetch invigilators timing
+            invig_start = time.time()
             invigilators = [
                 FetchInvigilator(interview)(question)
                 for question in interview.survey.questions
             ]
+            invig_end = time.time()
+            logger.debug(f"Time taken to fetch invigilators: {invig_end - invig_start:.4f}s")
+            
+            # Process prompts timing
+            prompts_start = time.time()
             for _, invigilator in enumerate(invigilators):
+                # Get prompts timing
+                get_prompts_start = time.time()
                 prompts = invigilator.get_prompts()
+                get_prompts_end = time.time()
+                logger.debug(f"Time taken to get prompts: {get_prompts_end - get_prompts_start:.4f}s")
+                
                 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)
+                
+                # Index lookups timing
+                index_start = time.time()
+                agent_index = self._agent_lookup[invigilator.agent]
                 agent_indices.append(agent_index)
                 interview_indices.append(interview_index)
-                scenario_index = self.scenarios.index(invigilator.scenario)
+                scenario_index = self._scenario_lookup[invigilator.scenario]
                 scenario_indices.append(scenario_index)
+                index_end = time.time()
+                logger.debug(f"Time taken for index lookups: {index_end - index_start:.4f}s")
+                
+                # Model and question name assignment timing
+                assign_start = time.time()
                 models.append(invigilator.model.model)
                 question_names.append(invigilator.question.question_name)
+                assign_end = time.time()
+                logger.debug(f"Time taken for assignments: {assign_end - assign_start:.4f}s")
 
+                # Cost estimation timing
+                cost_start = time.time()
                 prompt_cost = self.estimate_prompt_cost(
                     system_prompt=system_prompt,
                     user_prompt=user_prompt,
@@ -75,16 +108,34 @@ class JobsPrompts:
                     inference_service=invigilator.model._inference_service_,
                     model=invigilator.model.model,
                 )
+                cost_end = time.time()
+                logger.debug(f"Time taken to estimate prompt cost: {cost_end - cost_start:.4f}s")
                 costs.append(prompt_cost["cost_usd"])
 
+                # Cache key generation timing
+                cache_key_gen_start = time.time()
                 cache_key = CacheEntry.gen_key(
                     model=invigilator.model.model,
                     parameters=invigilator.model.parameters,
                     system_prompt=system_prompt,
                     user_prompt=user_prompt,
-                    iteration=0,  # TODO how to handle when there are multiple iterations?
+                    iteration=0,
                 )
+                cache_key_gen_end = time.time()
                 cache_keys.append(cache_key)
+                logger.debug(f"Time taken to generate cache key: {cache_key_gen_end - cache_key_gen_start:.4f}s")
+                logger.debug("-" * 50)  # Separator between iterations
+
+            prompts_end = time.time()
+            logger.info(f"Time taken to process prompts: {prompts_end - prompts_start:.4f}s")
+            
+            interview_end = time.time()
+            logger.info(f"Overall time taken for interview: {interview_end - interview_start:.4f}s")
+            logger.info("Time breakdown:")
+            logger.info(f"  Invigilators: {invig_end - invig_start:.4f}s")
+            logger.info(f"  Prompts processing: {prompts_end - prompts_start:.4f}s")
+            logger.info(f"  Other overhead: {(interview_end - interview_start) - ((invig_end - invig_start) + (prompts_end - prompts_start)):.4f}s")
+
         d = Dataset(
             [
                 {"user_prompt": user_prompts},
@@ -149,10 +200,10 @@ class JobsPrompts:
             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"
+                "Price data could not be retrieved. Using default estimates for input and output token prices. Input: $1.00 / 1M tokens; Output: $1.00 / 1M tokens"
             )
-            input_price_per_token = 0.00000015  # $0.15 / 1M tokens
-            output_price_per_token = 0.00000060  # $0.60 / 1M tokens
+            input_price_per_token = 0.000001  # $1.00 / 1M tokens
+            output_price_per_token = 0.000001  # $1.00 / 1M tokens
 
         # Compute the number of characters (double if the question involves piping)
         user_prompt_chars = len(str(user_prompt)) * get_piping_multiplier(

edsl/jobs/JobsRemoteInferenceHandler.py

@@ -228,6 +228,40 @@ class JobsRemoteInferenceHandler:
         results.results_uuid = results_uuid
         return results
 
+    def _attempt_fetch_job(
+        self,
+        job_info: RemoteJobInfo,
+        remote_job_data_fetcher: Callable,
+        object_fetcher: Callable,
+    ) -> Union[None, "Results", Literal["continue"]]:
+        """Makes one attempt to fetch and process a remote job's status and results."""
+        remote_job_data = remote_job_data_fetcher(job_info.job_uuid)
+        status = remote_job_data.get("status")
+
+        if status == "cancelled":
+            self._handle_cancelled_job(job_info)
+            return None
+
+        elif status == "failed" or status == "completed":
+            if status == "failed":
+                self._handle_failed_job(job_info, remote_job_data)
+
+            results_uuid = remote_job_data.get("results_uuid")
+            if results_uuid:
+                results = self._fetch_results_and_log(
+                    job_info=job_info,
+                    results_uuid=results_uuid,
+                    remote_job_data=remote_job_data,
+                    object_fetcher=object_fetcher,
+                )
+                return results
+            else:
+                return None
+
+        else:
+            self._sleep_for_a_bit(job_info, status)
+            return "continue"
+
     def poll_remote_inference_job(
         self,
         job_info: RemoteJobInfo,
@@ -242,31 +276,13 @@ class JobsRemoteInferenceHandler:
 
         job_in_queue = True
         while job_in_queue:
-            remote_job_data = remote_job_data_fetcher(job_info.job_uuid)
-            status = remote_job_data.get("status")
-
-            if status == "cancelled":
-                self._handle_cancelled_job(job_info)
-                return None
-
-            elif status == "failed" or status == "completed":
-                if status == "failed":
-                    self._handle_failed_job(job_info, remote_job_data)
-
-                results_uuid = remote_job_data.get("results_uuid")
-                if results_uuid:
-                    results = self._fetch_results_and_log(
-                        job_info=job_info,
-                        results_uuid=results_uuid,
-                        remote_job_data=remote_job_data,
-                        object_fetcher=object_fetcher,
-                    )
-                    return results
-                else:
-                    return None
-
-            else:
-                self._sleep_for_a_bit(job_info, status)
+            result = self._attempt_fetch_job(
+                job_info,
+                remote_job_data_fetcher,
+                object_fetcher
+            )
+            if result != "continue":
+                return result
 
     async def create_and_poll_remote_job(
         self,

edsl/jobs/buckets/BucketCollection.py

@@ -96,6 +96,36 @@ class BucketCollection(UserDict):
         else:
             self[model] = self.services_to_buckets[self.models_to_services[model.model]]
 
+    def update_from_key_lookup(self, key_lookup: "KeyLookup") -> None:
+        """Updates the bucket collection rates based on model RPM/TPM from KeyLookup"""
+
+        for model_name, service in self.models_to_services.items():
+            if service in key_lookup and not self.infinity_buckets:
+
+                if key_lookup[service].rpm is not None:
+                    new_rps = key_lookup[service].rpm / 60.0
+                    new_requests_bucket = TokenBucket(
+                        bucket_name=service,
+                        bucket_type="requests",
+                        capacity=new_rps,
+                        refill_rate=new_rps,
+                        remote_url=self.remote_url,
+                    )
+                    self.services_to_buckets[service].requests_bucket = (
+                        new_requests_bucket
+                    )
+
+                if key_lookup[service].tpm is not None:
+                    new_tps = key_lookup[service].tpm / 60.0
+                    new_tokens_bucket = TokenBucket(
+                        bucket_name=service,
+                        bucket_type="tokens",
+                        capacity=new_tps,
+                        refill_rate=new_tps,
+                        remote_url=self.remote_url,
+                    )
+                    self.services_to_buckets[service].tokens_bucket = new_tokens_bucket
+
     def visualize(self) -> dict:
         """Visualize the token and request buckets for each model."""
         plots = {}

edsl/jobs/data_structures.py

@@ -32,6 +32,7 @@ class RunParameters(Base):
     remote_inference_results_visibility: Optional[VisibilityType] = "unlisted"
     skip_retry: bool = False
     raise_validation_errors: bool = False
+    background: bool = False
     disable_remote_cache: bool = False
     disable_remote_inference: bool = False
     job_uuid: Optional[str] = None

edsl/language_models/LanguageModel.py

@@ -518,7 +518,11 @@ class LanguageModel(
         """
         from edsl.language_models.model import get_model_class
 
-        model_class = get_model_class(data["model"])
+        # breakpoint()
+
+        model_class = get_model_class(
+            data["model"], service_name=data.get("inference_service", None)
+        )
         return model_class(**data)
 
     def __repr__(self) -> str:
@@ -574,7 +578,6 @@ class LanguageModel(
             return Model(skip_api_key_check=True)
 
     def from_cache(self, cache: "Cache") -> LanguageModel:
-
         from copy import deepcopy
         from types import MethodType
         from edsl import Cache

edsl/language_models/key_management/KeyLookupBuilder.py

@@ -61,7 +61,14 @@ class KeyLookupBuilder:
     DEFAULT_RPM = int(CONFIG.get("EDSL_SERVICE_RPM_BASELINE"))
     DEFAULT_TPM = int(CONFIG.get("EDSL_SERVICE_TPM_BASELINE"))
 
-    def __init__(self, fetch_order: Optional[tuple[str]] = None):
+    def __init__(
+        self,
+        fetch_order: Optional[tuple[str]] = None,
+        coop: Optional["Coop"] = None,
+    ):
+        from edsl.coop import Coop
+
+        # Fetch order goes from lowest priority to highest priority
         if fetch_order is None:
             self.fetch_order = ("config", "env")
         else:
@@ -70,6 +77,11 @@ class KeyLookupBuilder:
         if not isinstance(self.fetch_order, tuple):
             raise ValueError("fetch_order must be a tuple")
 
+        if coop is None:
+            self.coop = Coop()
+        else:
+            self.coop = coop
+
         self.limit_data = {}
         self.key_data = {}
         self.id_data = {}
@@ -131,7 +143,8 @@ class KeyLookupBuilder:
                 service=service,
                 rpm=self.DEFAULT_RPM,
                 tpm=self.DEFAULT_TPM,
-                source="default",
+                rpm_source="default",
+                tpm_source="default",
             )
 
         if limit_entry.rpm is None:
@@ -145,7 +158,8 @@ class KeyLookupBuilder:
             tpm=int(limit_entry.tpm),
             api_id=api_id,
             token_source=api_key_entry.source,
-            limit_source=limit_entry.source,
+            rpm_source=limit_entry.rpm_source,
+            tpm_source=limit_entry.tpm_source,
             id_source=id_source,
         )
 
@@ -156,10 +170,7 @@ class KeyLookupBuilder:
         return dict(list(os.environ.items()))
 
     def _coop_key_value_pairs(self):
-        from edsl.coop import Coop
-
-        c = Coop()
-        return dict(list(c.fetch_rate_limit_config_vars().items()))
+        return dict(list(self.coop.fetch_rate_limit_config_vars().items()))
 
     def _config_key_value_pairs(self):
         from edsl.config import CONFIG
@@ -169,7 +180,7 @@ class KeyLookupBuilder:
     @staticmethod
     def extract_service(key: str) -> str:
         """Extract the service and limit type from the key"""
-        limit_type, service_raw = key.replace("EDSL_SERVICE_", "").split("_")
+        limit_type, service_raw = key.replace("EDSL_SERVICE_", "").split("_", 1)
         return service_raw.lower(), limit_type.lower()
 
     def get_key_value_pairs(self) -> dict:
@@ -187,17 +198,17 @@ class KeyLookupBuilder:
                 d[k] = (v, source)
         return d
 
-    def _entry_type(self, key, value) -> str:
+    def _entry_type(self, key: str) -> str:
         """Determine the type of entry from a key.
 
         >>> builder = KeyLookupBuilder()
-        >>> builder._entry_type("EDSL_SERVICE_RPM_OPENAI", "60")
+        >>> builder._entry_type("EDSL_SERVICE_RPM_OPENAI")
         'limit'
-        >>> builder._entry_type("OPENAI_API_KEY", "sk-1234")
+        >>> builder._entry_type("OPENAI_API_KEY")
         'api_key'
-        >>> builder._entry_type("AWS_ACCESS_KEY_ID", "AKIA1234")
+        >>> builder._entry_type("AWS_ACCESS_KEY_ID")
         'api_id'
-        >>> builder._entry_type("UNKNOWN_KEY", "value")
+        >>> builder._entry_type("UNKNOWN_KEY")
         'unknown'
         """
         if key.startswith("EDSL_SERVICE_"):
@@ -243,11 +254,13 @@ class KeyLookupBuilder:
         service, limit_type = self.extract_service(key)
         if service in self.limit_data:
             setattr(self.limit_data[service], limit_type.lower(), value)
+            setattr(self.limit_data[service], f"{limit_type}_source", source)
         else:
             new_limit_entry = LimitEntry(
-                service=service, rpm=None, tpm=None, source=source
+                service=service, rpm=None, tpm=None, rpm_source=None, tpm_source=None
             )
             setattr(new_limit_entry, limit_type.lower(), value)
+            setattr(new_limit_entry, f"{limit_type}_source", source)
             self.limit_data[service] = new_limit_entry
 
     def _add_api_key(self, key: str, value: str, source: str) -> None:
@@ -265,13 +278,27 @@ class KeyLookupBuilder:
         else:
             self.key_data[service].append(new_entry)
 
-    def process_key_value_pairs(self) -> None:
-        """Process all key-value pairs from the configured sources."""
-        for key, value_pair in self.get_key_value_pairs().items():
+    def update_from_dict(self, d: dict) -> None:
+        """
+        Update data from a dictionary of key-value pairs.
+        Each key is a key name, and each value is a tuple of (value, source).
+
+        >>> builder = KeyLookupBuilder()
+        >>> builder.update_from_dict({"OPENAI_API_KEY": ("sk-1234", "custodial_keys")})
+        >>> 'sk-1234' == builder.key_data["openai"][-1].value
+        True
+        >>> 'custodial_keys' == builder.key_data["openai"][-1].source
+        True
+        """
+        for key, value_pair in d.items():
             value, source = value_pair
-            if (entry_type := self._entry_type(key, value)) == "limit":
+            if self._entry_type(key) == "limit":
                 self._add_limit(key, value, source)
-            elif entry_type == "api_key":
+            elif self._entry_type(key) == "api_key":
                 self._add_api_key(key, value, source)
-            elif entry_type == "api_id":
+            elif self._entry_type(key) == "api_id":
                 self._add_id(key, value, source)
+
+    def process_key_value_pairs(self) -> None:
+        """Process all key-value pairs from the configured sources."""
+        self.update_from_dict(self.get_key_value_pairs())

edsl/language_models/key_management/models.py

@@ -40,18 +40,23 @@ class LimitEntry:
     60
     >>> limit.tpm
     100000
-    >>> limit.source
+    >>> limit.rpm_source
     'config'
+    >>> limit.tpm_source
+    'env'
     """
 
     service: str
     rpm: int
     tpm: int
-    source: Optional[str] = None
+    rpm_source: Optional[str] = None
+    tpm_source: Optional[str] = None
 
     @classmethod
     def example(cls):
-        return LimitEntry(service="openai", rpm=60, tpm=100000, source="config")
+        return LimitEntry(
+            service="openai", rpm=60, tpm=100000, rpm_source="config", tpm_source="env"
+        )
 
 
 @dataclass
@@ -108,7 +113,8 @@ class LanguageModelInput:
     tpm: int
     api_id: Optional[str] = None
     token_source: Optional[str] = None
-    limit_source: Optional[str] = None
+    rpm_source: Optional[str] = None
+    tpm_source: Optional[str] = None
     id_source: Optional[str] = None
 
     def to_dict(self):