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

edsl/data/RemoteCacheSync.py

@@ -100,7 +100,7 @@ class RemoteCacheSync(AbstractContextManager):
 
     def _get_cache_difference(self) -> CacheDifference:
         """Retrieves differences between local and remote caches."""
-        diff = self.coop.remote_cache_get_diff(self.cache.keys())
+        diff = self.coop.legacy_remote_cache_get_diff(self.cache.keys())
         return CacheDifference(
             client_missing_entries=diff.get("client_missing_cacheentries", []),
             server_missing_keys=diff.get("server_missing_cacheentry_keys", []),
@@ -112,7 +112,7 @@ class RemoteCacheSync(AbstractContextManager):
         missing_count = len(diff.client_missing_entries)
 
         if missing_count == 0:
-        #     self._output("No new entries to add to local cache.")
+            #     self._output("No new entries to add to local cache.")
             return
 
         # self._output(
@@ -154,22 +154,23 @@ class RemoteCacheSync(AbstractContextManager):
         upload_count = len(entries_to_upload)
 
         if upload_count > 0:
+            pass
             # self._output(
             #     f"Updating remote cache with {upload_count:,} new "
             #     f"{'entry' if upload_count == 1 else 'entries'}..."
             # )
 
-            self.coop.remote_cache_create_many(
-                entries_to_upload,
-                visibility="private",
-                description=self.remote_cache_description,
-            )
+            # self.coop.remote_cache_create_many(
+            #     entries_to_upload,
+            #     visibility="private",
+            #     description=self.remote_cache_description,
+            # )
             # self._output("Remote cache updated!")
         # else:
-            # self._output("No new entries to add to remote cache.")
+        # self._output("No new entries to add to remote cache.")
 
         # self._output(
-            # f"There are {len(self.cache.keys()):,} entries in the local cache."
+        # f"There are {len(self.cache.keys()):,} entries in the local cache."
         # )
 
 

edsl/enums.py

@@ -67,7 +67,7 @@ class InferenceServiceType(EnumWithChecks):
     TOGETHER = "together"
     PERPLEXITY = "perplexity"
     DEEPSEEK = "deepseek"
-    GROK = "grok"
+    XAI = "xai"
 
 
 # unavoidable violation of the DRY principle but it is necessary
@@ -87,7 +87,7 @@ InferenceServiceLiteral = Literal[
     "together",
     "perplexity",
     "deepseek",
-    "grok",
+    "xai",
 ]
 
 available_models_urls = {
@@ -111,7 +111,7 @@ service_to_api_keyname = {
     InferenceServiceType.TOGETHER.value: "TOGETHER_API_KEY",
     InferenceServiceType.PERPLEXITY.value: "PERPLEXITY_API_KEY",
     InferenceServiceType.DEEPSEEK.value: "DEEPSEEK_API_KEY",
-    InferenceServiceType.GROK.value: "XAI_API_KEY",
+    InferenceServiceType.XAI.value: "XAI_API_KEY",
 }
 
 

edsl/inference_services/AnthropicService.py

@@ -17,11 +17,10 @@ class AnthropicService(InferenceServiceABC):
     output_token_name = "output_tokens"
     model_exclude_list = []
 
-    available_models_url = 'https://docs.anthropic.com/en/docs/about-claude/models'
+    available_models_url = "https://docs.anthropic.com/en/docs/about-claude/models"
 
     @classmethod
     def get_model_list(cls, api_key: str = None):
-
         import requests
 
         if api_key is None:
@@ -94,13 +93,16 @@ class AnthropicService(InferenceServiceABC):
                 # breakpoint()
                 client = AsyncAnthropic(api_key=self.api_token)
 
-                response = await client.messages.create(
-                    model=model_name,
-                    max_tokens=self.max_tokens,
-                    temperature=self.temperature,
-                    system=system_prompt,  # note that the Anthropic API uses "system" parameter rather than put it in the message
-                    messages=messages,
-                )
+                try:
+                    response = await client.messages.create(
+                        model=model_name,
+                        max_tokens=self.max_tokens,
+                        temperature=self.temperature,
+                        system=system_prompt,  # note that the Anthropic API uses "system" parameter rather than put it in the message
+                        messages=messages,
+                    )
+                except Exception as e:
+                    return {"message": str(e)}
                 return response.model_dump()
 
         LLM.__name__ = model_class_name

edsl/inference_services/AvailableModelFetcher.py

@@ -69,6 +69,8 @@ class AvailableModelFetcher:
 
         Returns a list of [model, service_name, index] entries.
         """
+        if service == "azure" or service == "bedrock":
+            force_refresh = True  # Azure models are listed inside the .env AZURE_ENDPOINT_URL_AND_KEY variable
 
         if service:  # they passed a specific service
             matching_models, _ = self.get_available_models_by_service(

edsl/inference_services/AwsBedrock.py

@@ -110,8 +110,7 @@ class AwsBedrockService(InferenceServiceABC):
                     )
                     return response
                 except (ClientError, Exception) as e:
-                    print(e)
-                    return {"error": str(e)}
+                    return {"message": str(e)}
 
         LLM.__name__ = model_class_name
 

edsl/inference_services/AzureAI.py

@@ -179,15 +179,18 @@ class AzureAIService(InferenceServiceABC):
                         api_version=api_version,
                         api_key=api_key,
                     )
-                    response = await client.chat.completions.create(
-                        model=model_name,
-                        messages=[
-                            {
-                                "role": "user",
-                                "content": user_prompt,  # Your question can go here
-                            },
-                        ],
-                    )
+                    try:
+                        response = await client.chat.completions.create(
+                            model=model_name,
+                            messages=[
+                                {
+                                    "role": "user",
+                                    "content": user_prompt,  # Your question can go here
+                                },
+                            ],
+                        )
+                    except Exception as e:
+                        return {"message": str(e)}
                     return response.model_dump()
 
             # @staticmethod

edsl/inference_services/GoogleService.py

@@ -39,7 +39,9 @@ class GoogleService(InferenceServiceABC):
 
     model_exclude_list = []
 
-    available_models_url = 'https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models'
+    available_models_url = (
+        "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models"
+    )
 
     @classmethod
     def get_model_list(cls):
@@ -132,9 +134,12 @@ class GoogleService(InferenceServiceABC):
                     )
                     combined_prompt.append(gen_ai_file)
 
-                response = await self.generative_model.generate_content_async(
-                    combined_prompt, generation_config=generation_config
-                )
+                try:
+                    response = await self.generative_model.generate_content_async(
+                        combined_prompt, generation_config=generation_config
+                    )
+                except Exception as e:
+                    return {"message": str(e)}
                 return response.to_dict()
 
         LLM.__name__ = model_name

edsl/inference_services/InferenceServicesCollection.py

@@ -104,8 +104,9 @@ class InferenceServicesCollection:
     def available(
         self,
         service: Optional[str] = None,
+        force_refresh: bool = False,
     ) -> List[Tuple[str, str, int]]:
-        return self.availability_fetcher.available(service)
+        return self.availability_fetcher.available(service, force_refresh=force_refresh)
 
     def reset_cache(self) -> None:
         self.availability_fetcher.reset_cache()
@@ -120,7 +121,6 @@ class InferenceServicesCollection:
     def create_model_factory(
         self, model_name: str, service_name: Optional[InferenceServiceLiteral] = None
     ) -> "LanguageModel":
-
         if service_name is None:  # we try to find the right service
             service = self.resolver.resolve_model(model_name, service_name)
         else:  # if they passed a service, we'll use that

edsl/inference_services/MistralAIService.py

@@ -111,8 +111,7 @@ class MistralAIService(InferenceServiceABC):
                         ],
                     )
                 except Exception as e:
-                    raise LanguageModelBadResponseError(f"Error with Mistral API: {e}")
-
+                    return {"message": str(e)}
                 return res.model_dump()
 
         LLM.__name__ = model_class_name

edsl/inference_services/OpenAIService.py

@@ -207,8 +207,10 @@ class OpenAIService(InferenceServiceABC):
                     {"role": "user", "content": content},
                 ]
                 if (
-                    system_prompt == "" and self.omit_system_prompt_if_empty
-                ) or "o1" in self.model:
+                    (system_prompt == "" and self.omit_system_prompt_if_empty)
+                    or "o1" in self.model
+                    or "o3" in self.model
+                ):
                     messages = messages[1:]
 
                 params = {
@@ -222,14 +224,17 @@ class OpenAIService(InferenceServiceABC):
                     "logprobs": self.logprobs,
                     "top_logprobs": self.top_logprobs if self.logprobs else None,
                 }
-                if "o1" in self.model:
+                if "o1" in self.model or "o3" in self.model:
                     params.pop("max_tokens")
                     params["max_completion_tokens"] = self.max_tokens
                     params["temperature"] = 1
                 try:
                     response = await client.chat.completions.create(**params)
                 except Exception as e:
-                    print(e)
+                    #breakpoint()
+                    #print(e)
+                    #raise e
+                    return {'message': str(e)}
                 return response.model_dump()
 
         LLM.__name__ = "LanguageModel"

edsl/inference_services/PerplexityService.py

@@ -152,7 +152,8 @@ class PerplexityService(OpenAIService):
                 try:
                     response = await client.chat.completions.create(**params)
                 except Exception as e:
-                    print(e, flush=True)
+                    return {"message": str(e)}
+
                 return response.model_dump()
 
         LLM.__name__ = "LanguageModel"

edsl/inference_services/{GrokService.py → XAIService.py}

@@ -2,10 +2,10 @@ from typing import Any, List
 from edsl.inference_services.OpenAIService import OpenAIService
 
 
-class GrokService(OpenAIService):
+class XAIService(OpenAIService):
     """Openai service class."""
 
-    _inference_service_ = "grok"
+    _inference_service_ = "xai"
     _env_key_name_ = "XAI_API_KEY"
     _base_url_ = "https://api.x.ai/v1"
     _models_list_cache: List[str] = []

edsl/inference_services/registry.py

@@ -14,7 +14,7 @@ from edsl.inference_services.TestService import TestService
 from edsl.inference_services.TogetherAIService import TogetherAIService
 from edsl.inference_services.PerplexityService import PerplexityService
 from edsl.inference_services.DeepSeekService import DeepSeekService
-from edsl.inference_services.GrokService import GrokService
+from edsl.inference_services.XAIService import XAIService
 
 try:
     from edsl.inference_services.MistralAIService import MistralAIService
@@ -36,7 +36,7 @@ services = [
     TogetherAIService,
     PerplexityService,
     DeepSeekService,
-    GrokService,
+    XAIService,
 ]
 
 if mistral_available:

edsl/jobs/AnswerQuestionFunctionConstructor.py

@@ -66,10 +66,14 @@ class SkipHandler:
             )
         )
 
+
         def cancel_between(start, end):
             """Cancel the tasks for questions between the start and end indices."""
             for i in range(start, end):
-                self.interview.tasks[i].cancel()
+                #print(f"Cancelling task {i}")
+                #self.interview.tasks[i].cancel()
+                #self.interview.tasks[i].set_result("skipped")
+                self.interview.skip_flags[self.interview.survey.questions[i].question_name] = True
 
         if (next_question_index := next_question.next_q) == EndOfSurvey:
             cancel_between(
@@ -80,6 +84,8 @@ class SkipHandler:
         if next_question_index > (current_question_index + 1):
             cancel_between(current_question_index + 1, next_question_index)
 
+        
+
 
 class AnswerQuestionFunctionConstructor:
     """Constructs a function that answers a question and records the answer."""
@@ -161,6 +167,11 @@ class AnswerQuestionFunctionConstructor:
         async def attempt_answer():
             invigilator = self.invigilator_fetcher(question)
 
+            if self.interview.skip_flags.get(question.question_name, False):
+                return invigilator.get_failed_task_result(
+                    failure_reason="Question skipped."
+                )
+
             if self.skip_handler.should_skip(question):
                 return invigilator.get_failed_task_result(
                     failure_reason="Question skipped."

edsl/jobs/Jobs.py

@@ -277,7 +277,7 @@ class Jobs(Base):
 
         return JobsComponentConstructor(self).by(*args)
 
-    def prompts(self) -> "Dataset":
+    def prompts(self, iterations=1) -> "Dataset":
         """Return a Dataset of prompts that will be used.
 
 
@@ -285,7 +285,7 @@ class Jobs(Base):
         >>> Jobs.example().prompts()
         Dataset(...)
         """
-        return JobsPrompts(self).prompts()
+        return JobsPrompts(self).prompts(iterations=iterations)
 
     def show_prompts(self, all: bool = False) -> None:
         """Print the prompts."""
@@ -364,6 +364,15 @@ class Jobs(Base):
             self, cache=self.run_config.environment.cache
         ).create_interviews()
 
+    def show_flow(self, filename: Optional[str] = None) -> None:
+        """Show the flow of the survey."""
+        from edsl.surveys.SurveyFlowVisualization import SurveyFlowVisualization
+        if self.scenarios: 
+            scenario = self.scenarios[0]
+        else:
+            scenario = None
+        SurveyFlowVisualization(self.survey, scenario=scenario, agent=None).show_flow(filename=filename)
+
     def interviews(self) -> list[Interview]:
         """
         Return a list of :class:`edsl.jobs.interviews.Interview` objects.
@@ -409,11 +418,9 @@ class Jobs(Base):
         BucketCollection(...)
         """
         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
-            )
+            bc.update_from_key_lookup(self.run_config.environment.key_lookup)
         return bc
 
     def html(self):
@@ -475,25 +482,24 @@ class Jobs(Base):
     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,
+            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,
+            fresh=self.run_config.parameters.fresh,
         )
         return job_info
-    
-    def _create_remote_inference_handler(self) -> JobsRemoteInferenceHandler:
 
+    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,
@@ -507,7 +513,8 @@ class Jobs(Base):
         if jh.use_remote_inference(self.run_config.parameters.disable_remote_inference):
             job_info: RemoteJobInfo = self._start_remote_inference_job(jh)
             if background:
-                from edsl.results.Results import Results 
+                from edsl.results.Results import Results
+
                 results = Results.from_job_info(job_info)
                 return results
             else:
@@ -594,7 +601,7 @@ class Jobs(Base):
         # first try to run the job remotely
         if (results := self._remote_results(config)) is not None:
             return results
-        
+
         self._check_if_local_keys_ok()
 
         if config.environment.bucket_collection is None:

edsl/jobs/JobsChecks.py

@@ -24,7 +24,7 @@ class JobsChecks:
 
     def get_missing_api_keys(self) -> set:
         """
-        Returns a list of the api keys that a user needs to run this job, but does not currently have in their .env file.
+        Returns a list of the API keys that a user needs to run this job, but does not currently have in their .env file.
         """
         missing_api_keys = set()
 
@@ -134,22 +134,20 @@ class JobsChecks:
 
         edsl_auth_token = secrets.token_urlsafe(16)
 
-        print("API keys are required to run surveys with language models. The following keys are needed to run this survey: ")
+        print("\nThe following keys are needed to run this survey: \n")
         for api_key in missing_api_keys:
-            print(f"     🔑 {api_key}")
+            print(f"🔑 {api_key}")
         print(
-            "\nYou can provide your own keys or use an Expected Parrot key to access all available models."
+            """
+            \nYou can provide your own keys for language models or use an Expected Parrot key to access all available models.
+            \nClick the link below to create an account and run your survey with your Expected Parrot key:
+            """
         )
-        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➡️ Click the link below to create an account and get an Expected Parrot key:\n",
-        )
-
-        print(
-            "\nOnce you log in, your key will be stored on your computer and your survey will start running at the Expected Parrot server."
+            # link_description="",
         )
 
         api_key = coop._poll_for_api_key(edsl_auth_token)
@@ -159,8 +157,7 @@ class JobsChecks:
             return
 
         path_to_env = write_api_key_to_env(api_key)
-        print("\n✨ Your key has been stored at the following path: ")
-        print(f"    {path_to_env}")
+        print(f"\n✨ Your Expected Parrot key has been stored at the following path: {path_to_env}\n")
 
         # Retrieve API key so we can continue running the job
         load_dotenv()

edsl/jobs/JobsPrompts.py

@@ -18,6 +18,7 @@ from edsl.data.CacheEntry import CacheEntry
 
 logger = logging.getLogger(__name__)
 
+
 class JobsPrompts:
     def __init__(self, jobs: "Jobs"):
         self.interviews = jobs.interviews()
@@ -26,7 +27,9 @@ class JobsPrompts:
         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)}
+        self._scenario_lookup = {
+            scenario: idx for idx, scenario in enumerate(self.scenarios)
+        }
 
     @property
     def price_lookup(self):
@@ -37,7 +40,7 @@ class JobsPrompts:
             self._price_lookup = c.fetch_prices()
         return self._price_lookup
 
-    def prompts(self) -> "Dataset":
+    def prompts(self, iterations=1) -> "Dataset":
         """Return a Dataset of prompts that will be used.
 
         >>> from edsl.jobs import Jobs
@@ -54,11 +57,11 @@ 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 = [
@@ -66,8 +69,10 @@ class JobsPrompts:
                 for question in interview.survey.questions
             ]
             invig_end = time.time()
-            logger.debug(f"Time taken to fetch invigilators: {invig_end - invig_start:.4f}s")
-            
+            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):
@@ -75,13 +80,15 @@ class JobsPrompts:
                 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")
-                
+                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)
-                
+
                 # Index lookups timing
                 index_start = time.time()
                 agent_index = self._agent_lookup[invigilator.agent]
@@ -90,14 +97,18 @@ class JobsPrompts:
                 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")
-                
+                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")
+                logger.debug(
+                    f"Time taken for assignments: {assign_end - assign_start:.4f}s"
+                )
 
                 # Cost estimation timing
                 cost_start = time.time()
@@ -109,32 +120,44 @@ class JobsPrompts:
                     model=invigilator.model.model,
                 )
                 cost_end = time.time()
-                logger.debug(f"Time taken to estimate prompt cost: {cost_end - cost_start:.4f}s")
+                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,
-                )
+                for iteration in range(iterations):
+                    cache_key = CacheEntry.gen_key(
+                        model=invigilator.model.model,
+                        parameters=invigilator.model.parameters,
+                        system_prompt=system_prompt,
+                        user_prompt=user_prompt,
+                        iteration=iteration,
+                    )
+                    cache_keys.append(cache_key)
+
                 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(
+                    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")
-            
+            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(
+                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")
+            logger.info(
+                f"  Other overhead: {(interview_end - interview_start) - ((invig_end - invig_start) + (prompts_end - prompts_start)):.4f}s"
+            )
 
         d = Dataset(
             [