edsl 0.1.31.dev4__py3-none-any.whl → 0.1.32__py3-none-any.whl

edsl/inference_services/OpenAIService.py

@@ -1,12 +1,14 @@
 from typing import Any, List
 import re
 import os
-#from openai import AsyncOpenAI
+
+# from openai import AsyncOpenAI
 import openai
 
 from edsl.inference_services.InferenceServiceABC import InferenceServiceABC
 from edsl.language_models import LanguageModel
 from edsl.inference_services.rate_limits_cache import rate_limits
+from edsl.utilities.utilities import fix_partial_correct_response
 
 
 class OpenAIService(InferenceServiceABC):
@@ -18,18 +20,18 @@ class OpenAIService(InferenceServiceABC):
 
     _sync_client_ = openai.OpenAI
     _async_client_ = openai.AsyncOpenAI
-    
+
     @classmethod
     def sync_client(cls):
         return cls._sync_client_(
-            api_key = os.getenv(cls._env_key_name_), 
-            base_url = cls._base_url_)
-    
+            api_key=os.getenv(cls._env_key_name_), base_url=cls._base_url_
+        )
+
     @classmethod
     def async_client(cls):
         return cls._async_client_(
-            api_key = os.getenv(cls._env_key_name_), 
-            base_url = cls._base_url_)
+            api_key=os.getenv(cls._env_key_name_), base_url=cls._base_url_
+        )
 
     # TODO: Make this a coop call
     model_exclude_list = [
@@ -59,14 +61,14 @@ class OpenAIService(InferenceServiceABC):
 
     @classmethod
     def available(cls) -> List[str]:
-        #from openai import OpenAI
+        # from openai import OpenAI
 
         if not cls._models_list_cache:
             try:
-                #client = OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
+                # client = OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
                 cls._models_list_cache = [
                     m.id
-                    for m in cls.get_model_list() 
+                    for m in cls.get_model_list()
                     if m.id not in cls.model_exclude_list
                 ]
             except Exception as e:
@@ -106,21 +108,21 @@ class OpenAIService(InferenceServiceABC):
 
             def sync_client(self):
                 return cls.sync_client()
-            
+
             def async_client(self):
                 return cls.async_client()
 
             @classmethod
             def available(cls) -> list[str]:
-                #import openai
-                #client = openai.OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
-                #return client.models.list()
+                # import openai
+                # client = openai.OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
+                # return client.models.list()
                 return cls.sync_client().models.list()
-            
+
             def get_headers(self) -> dict[str, Any]:
-                #from openai import OpenAI
+                # from openai import OpenAI
 
-                #client = OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
+                # client = OpenAI(api_key = os.getenv(cls._env_key_name_), base_url = cls._base_url_)
                 client = self.sync_client()
                 response = client.chat.completions.with_raw_response.create(
                     messages=[
@@ -172,7 +174,7 @@ class OpenAIService(InferenceServiceABC):
                 else:
                     content = user_prompt
                 # self.client = AsyncOpenAI(
-                #     api_key = os.getenv(cls._env_key_name_), 
+                #     api_key = os.getenv(cls._env_key_name_),
                 #     base_url = cls._base_url_
                 #     )
                 client = self.async_client()
@@ -206,6 +208,9 @@ class OpenAIService(InferenceServiceABC):
                 if match:
                     return match.group(1)
                 else:
+                    out = fix_partial_correct_response(response)
+                    if "error" not in out:
+                        response = out["extracted_json"]
                     return response
 
         LLM.__name__ = "LanguageModel"

edsl/inference_services/models_available_cache.py

@@ -66,4 +66,35 @@ models_available = {
         "openchat/openchat_3.5",
     ],
     "google": ["gemini-pro"],
+    "bedrock": [
+        "amazon.titan-tg1-large",
+        "amazon.titan-text-lite-v1",
+        "amazon.titan-text-express-v1",
+        "ai21.j2-grande-instruct",
+        "ai21.j2-jumbo-instruct",
+        "ai21.j2-mid",
+        "ai21.j2-mid-v1",
+        "ai21.j2-ultra",
+        "ai21.j2-ultra-v1",
+        "anthropic.claude-instant-v1",
+        "anthropic.claude-v2:1",
+        "anthropic.claude-v2",
+        "anthropic.claude-3-sonnet-20240229-v1:0",
+        "anthropic.claude-3-haiku-20240307-v1:0",
+        "anthropic.claude-3-opus-20240229-v1:0",
+        "anthropic.claude-3-5-sonnet-20240620-v1:0",
+        "cohere.command-text-v14",
+        "cohere.command-r-v1:0",
+        "cohere.command-r-plus-v1:0",
+        "cohere.command-light-text-v14",
+        "meta.llama3-8b-instruct-v1:0",
+        "meta.llama3-70b-instruct-v1:0",
+        "meta.llama3-1-8b-instruct-v1:0",
+        "meta.llama3-1-70b-instruct-v1:0",
+        "meta.llama3-1-405b-instruct-v1:0",
+        "mistral.mistral-7b-instruct-v0:2",
+        "mistral.mixtral-8x7b-instruct-v0:1",
+        "mistral.mistral-large-2402-v1:0",
+        "mistral.mistral-large-2407-v1:0",
+    ],
 }

edsl/inference_services/registry.py

@@ -7,7 +7,19 @@ from edsl.inference_services.AnthropicService import AnthropicService
 from edsl.inference_services.DeepInfraService import DeepInfraService
 from edsl.inference_services.GoogleService import GoogleService
 from edsl.inference_services.GroqService import GroqService
+from edsl.inference_services.AwsBedrock import AwsBedrockService
+from edsl.inference_services.AzureAI import AzureAIService
+from edsl.inference_services.OllamaService import OllamaService
 
 default = InferenceServicesCollection(
-    [OpenAIService, AnthropicService, DeepInfraService, GoogleService, GroqService]
+    [
+        OpenAIService,
+        AnthropicService,
+        DeepInfraService,
+        GoogleService,
+        GroqService,
+        AwsBedrockService,
+        AzureAIService,
+        OllamaService,
+    ]
 )

edsl/jobs/Jobs.py

@@ -39,6 +39,8 @@ class Jobs(Base):
 
         self.__bucket_collection = None
 
+    # 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):
         return self._models
@@ -119,7 +121,9 @@ class Jobs(Base):
         - scenarios: traits of new scenarios are combined with traits of old existing. New scenarios will overwrite overlapping traits, and do not increase the number of scenarios in the instance
         - models: new models overwrite old models.
         """
-        passed_objects = self._turn_args_to_list(args)
+        passed_objects = self._turn_args_to_list(
+            args
+        )  # objects can also be passed comma-separated
 
         current_objects, objects_key = self._get_current_objects_of_this_type(
             passed_objects[0]
@@ -176,17 +180,27 @@ class Jobs(Base):
         from edsl.agents.Agent import Agent
         from edsl.scenarios.Scenario import Scenario
         from edsl.scenarios.ScenarioList import ScenarioList
+        from edsl.language_models.ModelList import ModelList
 
         if isinstance(object, Agent):
             return AgentList
         elif isinstance(object, Scenario):
             return ScenarioList
+        elif isinstance(object, ModelList):
+            return ModelList
         else:
             return list
 
     @staticmethod
     def _turn_args_to_list(args):
-        """Return a list of the first argument if it is a sequence, otherwise returns a list of all the arguments."""
+        """Return a list of the first argument if it is a sequence, otherwise returns a list of all the arguments.
+
+        Example:
+
+        >>> Jobs._turn_args_to_list([1,2,3])
+        [1, 2, 3]
+
+        """
 
         def did_user_pass_a_sequence(args):
             """Return True if the user passed a sequence, False otherwise.
@@ -209,7 +223,7 @@ class Jobs(Base):
             return container_class(args)
 
     def _get_current_objects_of_this_type(
-        self, object: Union[Agent, Scenario, LanguageModel]
+        self, object: Union["Agent", "Scenario", "LanguageModel"]
     ) -> tuple[list, str]:
         from edsl.agents.Agent import Agent
         from edsl.scenarios.Scenario import Scenario
@@ -292,7 +306,11 @@ class Jobs(Base):
 
     @classmethod
     def from_interviews(cls, interview_list):
-        """Return a Jobs instance from a list of interviews."""
+        """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
+        a new job with only those interviews.
+        """
         survey = interview_list[0].survey
         # get all the models
         models = list(set([interview.model for interview in interview_list]))
@@ -308,6 +326,8 @@ class Jobs(Base):
         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.
+
+
         """
         # if no agents, models, or scenarios are set, set them to defaults
         from edsl.agents.Agent import Agent
@@ -319,7 +339,11 @@ class Jobs(Base):
         self.scenarios = self.scenarios or [Scenario()]
         for agent, scenario, model in product(self.agents, self.scenarios, self.models):
             yield Interview(
-                survey=self.survey, agent=agent, scenario=scenario, model=model
+                survey=self.survey,
+                agent=agent,
+                scenario=scenario,
+                model=model,
+                skip_retry=self.skip_retry,
             )
 
     def create_bucket_collection(self) -> BucketCollection:
@@ -359,10 +383,16 @@ class Jobs(Base):
         return links
 
     def __hash__(self):
-        """Allow the model to be used as a key in a dictionary."""
+        """Allow the model to be used as a key in a dictionary.
+
+        >>> from edsl.jobs import Jobs
+        >>> hash(Jobs.example())
+        846655441787442972
+
+        """
         from edsl.utilities.utilities import dict_hash
 
-        return dict_hash(self.to_dict())
+        return dict_hash(self._to_dict())
 
     def _output(self, message) -> None:
         """Check if a Job is verbose. If so, print the message."""
@@ -390,11 +420,27 @@ class Jobs(Base):
         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
 
-        msg1, msg2 = None, None
+        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}"
@@ -409,6 +455,12 @@ class Jobs(Base):
                 if warn:
                     warnings.warn(message)
 
+    @property
+    def skip_retry(self):
+        if not hasattr(self, "_skip_retry"):
+            return False
+        return self._skip_retry
+
     def run(
         self,
         n: int = 1,
@@ -423,6 +475,7 @@ class Jobs(Base):
         print_exceptions=True,
         remote_cache_description: Optional[str] = None,
         remote_inference_description: Optional[str] = None,
+        skip_retry: bool = False,
     ) -> Results:
         """
         Runs the Job: conducts Interviews and returns their results.
@@ -441,6 +494,7 @@ class Jobs(Base):
         from edsl.coop.coop import Coop
 
         self._check_parameters()
+        self._skip_retry = skip_retry
 
         if batch_mode is not None:
             raise NotImplementedError(
@@ -631,12 +685,16 @@ class Jobs(Base):
         return results
 
     async def run_async(self, cache=None, n=1, **kwargs):
-        """Run the job asynchronously."""
+        """Run asynchronously."""
         results = await JobsRunnerAsyncio(self).run_async(cache=cache, n=n, **kwargs)
         return results
 
     def all_question_parameters(self):
-        """Return all the fields in the questions in the survey."""
+        """Return all the fields in the questions in the survey.
+        >>> from edsl.jobs import Jobs
+        >>> Jobs.example().all_question_parameters()
+        {'period'}
+        """
         return set.union(*[question.parameters for question in self.survey.questions])
 
     #######################
@@ -677,15 +735,19 @@ class Jobs(Base):
     #######################
     # Serialization methods
     #######################
+
+    def _to_dict(self):
+        return {
+            "survey": self.survey._to_dict(),
+            "agents": [agent._to_dict() for agent in self.agents],
+            "models": [model._to_dict() for model in self.models],
+            "scenarios": [scenario._to_dict() for scenario in self.scenarios],
+        }
+
     @add_edsl_version
     def to_dict(self) -> dict:
         """Convert the Jobs instance to a dictionary."""
-        return {
-            "survey": self.survey.to_dict(),
-            "agents": [agent.to_dict() for agent in self.agents],
-            "models": [model.to_dict() for model in self.models],
-            "scenarios": [scenario.to_dict() for scenario in self.scenarios],
-        }
+        return self._to_dict()
 
     @classmethod
     @remove_edsl_version
@@ -704,7 +766,13 @@ class Jobs(Base):
         )
 
     def __eq__(self, other: Jobs) -> bool:
-        """Return True if the Jobs instance is equal to another Jobs instance."""
+        """Return True if the Jobs instance is equal to another Jobs instance.
+
+        >>> from edsl.jobs import Jobs
+        >>> Jobs.example() == Jobs.example()
+        True
+
+        """
         return self.to_dict() == other.to_dict()
 
     #######################
@@ -712,11 +780,16 @@ class Jobs(Base):
     #######################
     @classmethod
     def example(
-        cls, throw_exception_probability: int = 0, randomize: bool = False
+        cls,
+        throw_exception_probability: float = 0.0,
+        randomize: bool = False,
+        test_model=False,
     ) -> Jobs:
         """Return an example Jobs instance.
 
         :param throw_exception_probability: the probability that an exception will be thrown when answering a question. This is useful for testing error handling.
+        :param randomize: whether to randomize the job by adding a random string to the period
+        :param test_model: whether to use a test model
 
         >>> Jobs.example()
         Jobs(...)
@@ -730,6 +803,11 @@ class Jobs(Base):
 
         addition = "" if not randomize else str(uuid4())
 
+        if test_model:
+            from edsl.language_models import LanguageModel
+
+            m = LanguageModel.example(test_model=True)
+
         # (status, question, period)
         agent_answers = {
             ("Joyful", "how_feeling", "morning"): "OK",
@@ -777,7 +855,10 @@ class Jobs(Base):
                 Scenario({"period": "afternoon"}),
             ]
         )
-        job = base_survey.by(scenario_list).by(joy_agent, sad_agent)
+        if test_model:
+            job = base_survey.by(m).by(scenario_list).by(joy_agent, sad_agent)
+        else:
+            job = base_survey.by(scenario_list).by(joy_agent, sad_agent)
 
         return job
 

edsl/jobs/buckets/TokenBucket.py

@@ -100,7 +100,9 @@ class TokenBucket:
         available_tokens = min(self.capacity, self.tokens + refill_amount)
         return max(0, requested_tokens - available_tokens) / self.refill_rate
 
-    async def get_tokens(self, amount: Union[int, float] = 1) -> None:
+    async def get_tokens(
+        self, amount: Union[int, float] = 1, cheat_bucket_capacity=True
+    ) -> None:
         """Wait for the specified number of tokens to become available.
 
 
@@ -116,14 +118,20 @@ class TokenBucket:
         True
 
         >>> bucket = TokenBucket(bucket_name="test", bucket_type="test", capacity=10, refill_rate=1)
-        >>> asyncio.run(bucket.get_tokens(11))
+        >>> asyncio.run(bucket.get_tokens(11, cheat_bucket_capacity=False))
         Traceback (most recent call last):
         ...
         ValueError: Requested amount exceeds bucket capacity. Bucket capacity: 10, requested amount: 11. As the bucket never overflows, the requested amount will never be available.
+        >>> asyncio.run(bucket.get_tokens(11, cheat_bucket_capacity=True))
         """
         if amount > self.capacity:
-            msg = f"Requested amount exceeds bucket capacity. Bucket capacity: {self.capacity}, requested amount: {amount}. As the bucket never overflows, the requested amount will never be available."
-            raise ValueError(msg)
+            if not cheat_bucket_capacity:
+                msg = f"Requested amount exceeds bucket capacity. Bucket capacity: {self.capacity}, requested amount: {amount}. As the bucket never overflows, the requested amount will never be available."
+                raise ValueError(msg)
+            else:
+                self.tokens = 0  # clear the bucket but let it go through
+                return
+
         while self.tokens < amount:
             self.refill()
             await asyncio.sleep(0.01)  # Sleep briefly to prevent busy waiting

edsl/jobs/interviews/Interview.py

@@ -14,8 +14,8 @@ from edsl.jobs.tasks.TaskCreators import TaskCreators
 from edsl.jobs.interviews.InterviewStatusLog import InterviewStatusLog
 from edsl.jobs.interviews.interview_exception_tracking import (
     InterviewExceptionCollection,
-    InterviewExceptionEntry,
 )
+from edsl.jobs.interviews.InterviewExceptionEntry import InterviewExceptionEntry
 from edsl.jobs.interviews.retry_management import retry_strategy
 from edsl.jobs.interviews.InterviewTaskBuildingMixin import InterviewTaskBuildingMixin
 from edsl.jobs.interviews.InterviewStatusMixin import InterviewStatusMixin
@@ -44,6 +44,7 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         iteration: int = 0,
         cache: Optional["Cache"] = None,
         sidecar_model: Optional["LanguageModel"] = None,
+        skip_retry=False,
     ):
         """Initialize the Interview instance.
 
@@ -87,6 +88,7 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         self.task_creators = TaskCreators()  # tracks the task creators
         self.exceptions = InterviewExceptionCollection()
         self._task_status_log_dict = InterviewStatusLog()
+        self.skip_retry = skip_retry
 
         # dictionary mapping question names to their index in the survey.
         self.to_index = {
@@ -94,6 +96,30 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
             for index, question_name in enumerate(self.survey.question_names)
         }
 
+    def _to_dict(self, include_exceptions=False) -> dict[str, Any]:
+        """Return a dictionary representation of the Interview instance.
+        This is just for hashing purposes.
+
+        >>> i = Interview.example()
+        >>> hash(i)
+        1646262796627658719
+        """
+        d = {
+            "agent": self.agent._to_dict(),
+            "survey": self.survey._to_dict(),
+            "scenario": self.scenario._to_dict(),
+            "model": self.model._to_dict(),
+            "iteration": self.iteration,
+            "exceptions": {},
+        }
+        if include_exceptions:
+            d["exceptions"] = self.exceptions.to_dict()
+
+    def __hash__(self) -> int:
+        from edsl.utilities.utilities import dict_hash
+
+        return dict_hash(self._to_dict())
+
     async def async_conduct_interview(
         self,
         *,
@@ -134,8 +160,7 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         <BLANKLINE>
 
         >>> i.exceptions
-        {'q0': [{'exception': "Exception('This is a test error')", 'time': ..., 'traceback': ...
-
+        {'q0': ...
         >>> i = Interview.example()
         >>> result, _ = asyncio.run(i.async_conduct_interview(stop_on_exception = True))
         Traceback (most recent call last):
@@ -204,13 +229,9 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
         {}
         >>> i._record_exception(i.tasks[0], Exception("An exception occurred."))
         >>> i.exceptions
-        {'q0': [{'exception': "Exception('An exception occurred.')", 'time': ..., 'traceback': 'NoneType: None\\n'}]}
+        {'q0': ...
         """
-        exception_entry = InterviewExceptionEntry(
-            exception=repr(exception),
-            time=time.time(),
-            traceback=traceback.format_exc(),
-        )
+        exception_entry = InterviewExceptionEntry(exception)
         self.exceptions.add(task.get_name(), exception_entry)
 
     @property
@@ -251,6 +272,7 @@ class Interview(InterviewStatusMixin, InterviewTaskBuildingMixin):
             model=self.model,
             iteration=iteration,
             cache=cache,
+            skip_retry=self.skip_retry,
         )
 
     @classmethod

edsl/jobs/interviews/InterviewExceptionEntry.py

@@ -0,0 +1,101 @@
+import traceback
+import datetime
+import time
+from collections import UserDict
+
+# traceback=traceback.format_exc(),
+# traceback = frame_summary_to_dict(traceback.extract_tb(e.__traceback__))
+# traceback = [frame_summary_to_dict(f) for f in traceback.extract_tb(e.__traceback__)]
+
+
+class InterviewExceptionEntry:
+    """Class to record an exception that occurred during the interview.
+
+    >>> entry = InterviewExceptionEntry.example()
+    >>> entry.to_dict()['exception']
+    "ValueError('An error occurred.')"
+    """
+
+    def __init__(self, exception: Exception, traceback_format="html"):
+        self.time = datetime.datetime.now().isoformat()
+        self.exception = exception
+        self.traceback_format = traceback_format
+
+    def __getitem__(self, key):
+        # Support dict-like access obj['a']
+        return str(getattr(self, key))
+
+    @classmethod
+    def example(cls):
+        try:
+            raise ValueError("An error occurred.")
+        except Exception as e:
+            entry = InterviewExceptionEntry(e)
+        return entry
+
+    @property
+    def traceback(self):
+        """Return the exception as HTML."""
+        if self.traceback_format == "html":
+            return self.html_traceback
+        else:
+            return self.text_traceback
+
+    @property
+    def text_traceback(self):
+        """
+        >>> entry = InterviewExceptionEntry.example()
+        >>> entry.text_traceback
+        'Traceback (most recent call last):...'
+        """
+        e = self.exception
+        tb_str = "".join(traceback.format_exception(type(e), e, e.__traceback__))
+        return tb_str
+
+    @property
+    def html_traceback(self):
+        from rich.console import Console
+        from rich.table import Table
+        from rich.traceback import Traceback
+
+        from io import StringIO
+
+        html_output = StringIO()
+
+        console = Console(file=html_output, record=True)
+
+        tb = Traceback.from_exception(
+            type(self.exception),
+            self.exception,
+            self.exception.__traceback__,
+            show_locals=True,
+        )
+        console.print(tb)
+        return html_output.getvalue()
+
+    def to_dict(self) -> dict:
+        """Return the exception as a dictionary.
+
+        >>> entry = InterviewExceptionEntry.example()
+        >>> entry.to_dict()['exception']
+        "ValueError('An error occurred.')"
+
+        """
+        return {
+            "exception": repr(self.exception),
+            "time": self.time,
+            "traceback": self.traceback,
+        }
+
+    def push(self):
+        from edsl import Coop
+
+        coop = Coop()
+        results = coop.error_create(self.to_dict())
+        return results
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)