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

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},

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/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):

edsl/prompts/Prompt.py

@@ -10,6 +10,48 @@ from edsl.Base import PersistenceMixin, RepresentationMixin
 
 MAX_NESTING = 100
 
+from jinja2 import Environment, meta, TemplateSyntaxError, Undefined
+from functools import lru_cache
+
+class PreserveUndefined(Undefined):
+    def __str__(self):
+        return "{{ " + str(self._undefined_name) + " }}"
+
+# Create environment once at module level
+_env = Environment(undefined=PreserveUndefined)
+
+@lru_cache(maxsize=1024)
+def _compile_template(text: str):
+    return _env.from_string(text)
+
+@lru_cache(maxsize=1024)
+def _find_template_variables(template: str) -> list[str]:
+    """Find and return the template variables."""
+    ast = _env.parse(template)
+    return list(meta.find_undeclared_variables(ast))
+
+def _make_hashable(value):
+    """Convert unhashable types to hashable ones."""
+    if isinstance(value, list):
+        return tuple(_make_hashable(item) for item in value)
+    if isinstance(value, dict):
+        return frozenset((k, _make_hashable(v)) for k, v in value.items())
+    return value
+
+@lru_cache(maxsize=1024)
+def _cached_render(text: str, frozen_replacements: frozenset) -> str:
+    """Cached version of template rendering with frozen replacements."""
+    # Print cache info on every call
+    cache_info = _cached_render.cache_info()
+    print(f"\t\t\t\t\t Cache status - hits: {cache_info.hits}, misses: {cache_info.misses}, current size: {cache_info.currsize}")
+    
+    # Convert back to dict with original types for rendering
+    replacements = {k: v for k, v in frozen_replacements}
+    
+    template = _compile_template(text)
+    result = template.render(replacements)
+    
+    return result
 
 class Prompt(PersistenceMixin, RepresentationMixin):
     """Class for creating a prompt to be used in a survey."""
@@ -145,33 +187,8 @@ class Prompt(PersistenceMixin, RepresentationMixin):
         return f'Prompt(text="""{self.text}""")'
 
     def template_variables(self) -> list[str]:
-        """Return the the variables in the template.
-
-        Example:
-
-        >>> p = Prompt("Hello, {{person}}")
-        >>> p.template_variables()
-        ['person']
-
-        """
-        return self._template_variables(self.text)
-
-    @staticmethod
-    def _template_variables(template: str) -> list[str]:
-        """Find and return the template variables.
-
-        :param template: The template to find the variables in.
-
-        """
-        from jinja2 import Environment, meta, Undefined
-
-        class PreserveUndefined(Undefined):
-            def __str__(self):
-                return "{{ " + str(self._undefined_name) + " }}"
-
-        env = Environment(undefined=PreserveUndefined)
-        ast = env.parse(template)
-        return list(meta.find_undeclared_variables(ast))
+        """Return the variables in the template."""
+        return _find_template_variables(self.text)
 
     def undefined_template_variables(self, replacement_dict: dict):
         """Return the variables in the template that are not in the replacement_dict.
@@ -239,45 +256,39 @@ class Prompt(PersistenceMixin, RepresentationMixin):
             return self
 
     @staticmethod
-    def _render(
-        text: str, primary_replacement, **additional_replacements
-    ) -> "PromptBase":
-        """Render the template text with variables replaced from the provided named dictionaries.
-
-        :param text: The text to render.
-        :param primary_replacement: The primary replacement dictionary.
-        :param additional_replacements: Additional replacement dictionaries.
-
-        Allows for nested variable resolution up to a specified maximum nesting depth.
-
-        Example:
-
-        >>> codebook = {"age": "Age"}
-        >>> p = Prompt("You are an agent named {{ name }}. {{ codebook['age']}}: {{ age }}")
-        >>> p.render({"name": "John", "age": 44}, codebook=codebook)
-        Prompt(text=\"""You are an agent named John. Age: 44\""")
-        """
-        from jinja2 import Environment, meta, TemplateSyntaxError, Undefined
-
-        class PreserveUndefined(Undefined):
-            def __str__(self):
-                return "{{ " + str(self._undefined_name) + " }}"
-
-        env = Environment(undefined=PreserveUndefined)
+    def _render(text: str, primary_replacement, **additional_replacements) -> "PromptBase":
+        """Render the template text with variables replaced."""
+        import time
+        
+        # if there are no replacements, return the text
+        if not primary_replacement and not additional_replacements:
+            return text
+     
         try:
+            variables = _find_template_variables(text)
+            
+            if not variables: # if there are no variables, return the text
+                return text
+            
+            # Combine all replacements
+            all_replacements = {**primary_replacement, **additional_replacements}
+            
             previous_text = None
+            current_text = text
+            iteration = 0
+            
             for _ in range(MAX_NESTING):
-                # breakpoint()
-                rendered_text = env.from_string(text).render(
-                    primary_replacement, **additional_replacements
-                )
-                if rendered_text == previous_text:
-                    # No more changes, so return the rendered text
+                iteration += 1
+                
+                template = _compile_template(current_text)
+                rendered_text = template.render(all_replacements)
+                
+                if rendered_text == current_text:
                     return rendered_text
-                previous_text = text
-                text = rendered_text
+                    
+                previous_text = current_text
+                current_text = rendered_text
 
-            # If the loop exits without returning, it indicates too much nesting
             raise TemplateRenderError(
                 "Too much nesting - you created an infinite loop here, pal"
             )
@@ -331,6 +342,58 @@ class Prompt(PersistenceMixin, RepresentationMixin):
         """Return an example of the prompt."""
         return cls(cls.default_instructions)
 
+    def get_prompts(self) -> Dict[str, Any]:
+        """Get the prompts for the question."""
+        start = time.time()
+        
+        # Build all the components
+        instr_start = time.time()
+        agent_instructions = self.agent_instructions_prompt
+        instr_end = time.time()
+        logger.debug(f"Time taken for agent instructions: {instr_end - instr_start:.4f}s")
+        
+        persona_start = time.time()
+        agent_persona = self.agent_persona_prompt
+        persona_end = time.time()
+        logger.debug(f"Time taken for agent persona: {persona_end - persona_start:.4f}s")
+        
+        q_instr_start = time.time()
+        question_instructions = self.question_instructions_prompt
+        q_instr_end = time.time()
+        logger.debug(f"Time taken for question instructions: {q_instr_end - q_instr_start:.4f}s")
+        
+        memory_start = time.time()
+        prior_question_memory = self.prior_question_memory_prompt
+        memory_end = time.time()
+        logger.debug(f"Time taken for prior question memory: {memory_end - memory_start:.4f}s")
+
+        # Get components dict
+        components = {
+            "agent_instructions": agent_instructions.text,
+            "agent_persona": agent_persona.text,
+            "question_instructions": question_instructions.text,
+            "prior_question_memory": prior_question_memory.text,
+        }
+
+        # Use PromptPlan's get_prompts method
+        plan_start = time.time()
+        prompts = self.prompt_plan.get_prompts(**components)
+        plan_end = time.time()
+        logger.debug(f"Time taken for prompt processing: {plan_end - plan_start:.4f}s")
+        
+        # Handle file keys if present
+        if hasattr(self, 'question_file_keys') and self.question_file_keys:
+            files_start = time.time()
+            files_list = []
+            for key in self.question_file_keys:
+                files_list.append(self.scenario[key])
+            prompts["files_list"] = files_list
+            files_end = time.time()
+            logger.debug(f"Time taken for file key processing: {files_end - files_start:.4f}s")
+        
+        end = time.time()
+        logger.debug(f"Total time in get_prompts: {end - start:.4f}s")
+        return prompts
 
 if __name__ == "__main__":
     print("Running doctests...")