edsl 0.1.14__py3-none-any.whl → 0.1.40__py3-none-any.whl

edsl/jobs/interviews/Interview.py

@@ -0,0 +1,392 @@
+"""This module contains the Interview class, which is responsible for conducting an interview asynchronously."""
+
+from __future__ import annotations
+import asyncio
+from typing import Any, Type, List, Generator, Optional, Union, TYPE_CHECKING
+import copy
+from dataclasses import dataclass
+
+# from edsl.jobs.Answers import Answers
+from edsl.jobs.data_structures import Answers
+from edsl.jobs.interviews.InterviewStatusLog import InterviewStatusLog
+from edsl.jobs.interviews.InterviewStatusDictionary import InterviewStatusDictionary
+from edsl.jobs.interviews.InterviewExceptionCollection import (
+    InterviewExceptionCollection,
+)
+from edsl.jobs.interviews.InterviewExceptionEntry import InterviewExceptionEntry
+from edsl.jobs.buckets.ModelBuckets import ModelBuckets
+from edsl.jobs.AnswerQuestionFunctionConstructor import (
+    AnswerQuestionFunctionConstructor,
+)
+from edsl.jobs.InterviewTaskManager import InterviewTaskManager
+from edsl.jobs.FetchInvigilator import FetchInvigilator
+from edsl.jobs.RequestTokenEstimator import RequestTokenEstimator
+
+
+if TYPE_CHECKING:
+    from edsl.agents.Agent import Agent
+    from edsl.surveys.Survey import Survey
+    from edsl.scenarios.Scenario import Scenario
+    from edsl.data.Cache import Cache
+    from edsl.language_models.LanguageModel import LanguageModel
+    from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
+    from edsl.agents.InvigilatorBase import InvigilatorBase
+    from edsl.language_models.key_management.KeyLookup import KeyLookup
+
+
+@dataclass
+class InterviewRunningConfig:
+    cache: Optional["Cache"] = (None,)
+    skip_retry: bool = (False,)  # COULD BE SET WITH CONFIG
+    raise_validation_errors: bool = (True,)
+    stop_on_exception: bool = (False,)
+
+
+class Interview:
+    """
+    An 'interview' is one agent answering one survey, with one language model, for a given scenario.
+
+    The main method is `async_conduct_interview`, which conducts the interview asynchronously.
+    Most of the class is dedicated to creating the tasks for each question in the survey, and then running them.
+    """
+
+    def __init__(
+        self,
+        agent: Agent,
+        survey: Survey,
+        scenario: Scenario,
+        model: Type["LanguageModel"],
+        iteration: int = 0,
+        indices: dict = None,  # explain?
+        cache: Optional["Cache"] = None,
+        skip_retry: bool = False,  # COULD BE SET WITH CONFIG
+        raise_validation_errors: bool = True,
+    ):
+        """Initialize the Interview instance.
+
+        :param agent: the agent being interviewed.
+        :param survey: the survey being administered to the agent.
+        :param scenario: the scenario that populates the survey questions.
+        :param model: the language model used to answer the questions.
+        # :param debug: if True, run without calls to the language model.
+        :param iteration: the iteration number of the interview.
+        :param cache: the cache used to store the answers.
+
+        >>> i = Interview.example()
+        >>> i.task_manager.task_creators
+        {}
+
+        >>> i.exceptions
+        {}
+
+        >>> _ = asyncio.run(i.async_conduct_interview())
+        >>> i.task_status_logs['q0']
+        [{'log_time': ..., 'value': <TaskStatus.NOT_STARTED: 1>}, {'log_time': ..., 'value': <TaskStatus.WAITING_FOR_DEPENDENCIES: 2>}, {'log_time': ..., 'value': <TaskStatus.API_CALL_IN_PROGRESS: 7>}, {'log_time': ..., 'value': <TaskStatus.SUCCESS: 8>}]
+
+        >>> i.to_index
+        {'q0': 0, 'q1': 1, 'q2': 2}
+
+        """
+        self.agent = agent
+        self.survey = copy.deepcopy(survey)  # why do we need to deepcopy the survey?
+        self.scenario = scenario
+        self.model = model
+        self.iteration = iteration
+
+        self.answers = Answers()  # will get filled in as interview progresses
+
+        self.task_manager = InterviewTaskManager(
+            survey=self.survey,
+            iteration=iteration,
+        )
+
+        self.exceptions = InterviewExceptionCollection()
+
+        self.running_config = InterviewRunningConfig(
+            cache=cache,
+            skip_retry=skip_retry,
+            raise_validation_errors=raise_validation_errors,
+        )
+
+        self.cache = cache
+        self.skip_retry = skip_retry
+        self.raise_validation_errors = raise_validation_errors
+
+        # dictionary mapping question names to their index in the survey.
+        self.to_index = {
+            question_name: index
+            for index, question_name in enumerate(self.survey.question_names)
+        }
+
+        self.failed_questions = []
+
+        self.indices = indices
+        self.initial_hash = hash(self)
+
+    @property
+    def has_exceptions(self) -> bool:
+        """Return True if there are exceptions."""
+        return len(self.exceptions) > 0
+
+    @property
+    def task_status_logs(self) -> InterviewStatusLog:
+        """Return the task status logs for the interview.
+
+        The keys are the question names; the values are the lists of status log changes for each task.
+        """
+        return self.task_manager.task_status_logs
+
+    @property
+    def token_usage(self) -> InterviewTokenUsage:
+        """Determine how many tokens were used for the interview."""
+        return self.task_manager.token_usage  # task_creators.token_usage
+
+    @property
+    def interview_status(self) -> InterviewStatusDictionary:
+        """Return a dictionary mapping task status codes to counts."""
+        # return self.task_creators.interview_status
+        return self.task_manager.interview_status
+
+    def to_dict(self, include_exceptions=True, add_edsl_version=True) -> dict[str, Any]:
+        """Return a dictionary representation of the Interview instance.
+        This is just for hashing purposes.
+
+        >>> i = Interview.example()
+        >>> hash(i)
+        193593189022259693
+        """
+        d = {
+            "agent": self.agent.to_dict(add_edsl_version=add_edsl_version),
+            "survey": self.survey.to_dict(add_edsl_version=add_edsl_version),
+            "scenario": self.scenario.to_dict(add_edsl_version=add_edsl_version),
+            "model": self.model.to_dict(add_edsl_version=add_edsl_version),
+            "iteration": self.iteration,
+            "exceptions": {},
+        }
+        if include_exceptions:
+            d["exceptions"] = self.exceptions.to_dict()
+        if hasattr(self, "indices"):
+            d["indices"] = self.indices
+        return d
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Interview":
+        """Return an Interview instance from a dictionary."""
+
+        from edsl.agents.Agent import Agent
+        from edsl.surveys.Survey import Survey
+        from edsl.scenarios.Scenario import Scenario
+        from edsl.language_models.LanguageModel import LanguageModel
+
+        agent = Agent.from_dict(d["agent"])
+        survey = Survey.from_dict(d["survey"])
+        scenario = Scenario.from_dict(d["scenario"])
+        model = LanguageModel.from_dict(d["model"])
+        iteration = d["iteration"]
+        params = {
+            "agent": agent,
+            "survey": survey,
+            "scenario": scenario,
+            "model": model,
+            "iteration": iteration,
+        }
+        if "indices" in d:
+            params["indices"] = d["indices"]
+        interview = cls(**params)
+        if "exceptions" in d:
+            exceptions = InterviewExceptionCollection.from_dict(d["exceptions"])
+            interview.exceptions = exceptions
+        return interview
+
+    def __hash__(self) -> int:
+        from edsl.utilities.utilities import dict_hash
+
+        return dict_hash(self.to_dict(include_exceptions=False, add_edsl_version=False))
+
+    def __eq__(self, other: "Interview") -> bool:
+        """
+        >>> from edsl.jobs.interviews.Interview import Interview; i = Interview.example(); d = i.to_dict(); i2 = Interview.from_dict(d); i == i2
+        True
+        """
+        return hash(self) == hash(other)
+
+    async def async_conduct_interview(
+        self,
+        run_config: Optional["RunConfig"] = None,
+    ) -> tuple["Answers", List[dict[str, Any]]]:
+        """
+        Conduct an Interview asynchronously.
+        It returns a tuple with the answers and a list of valid results.
+
+        :param model_buckets: a dictionary of token buckets for the model.
+        :param debug: run without calls to LLM.
+        :param stop_on_exception: if True, stops the interview if an exception is raised.
+
+        Example usage:
+
+        >>> i = Interview.example()
+        >>> result, _ = asyncio.run(i.async_conduct_interview())
+        >>> result['q0']
+        'yes'
+
+        >>> i = Interview.example(throw_exception = True)
+        >>> result, _ = asyncio.run(i.async_conduct_interview())
+        >>> i.exceptions
+        {'q0': ...
+        >>> i = Interview.example()
+        >>> from edsl.jobs.Jobs import RunConfig, RunParameters, RunEnvironment
+        >>> run_config = RunConfig(parameters = RunParameters(), environment = RunEnvironment())
+        >>> run_config.parameters.stop_on_exception = True
+        >>> result, _ = asyncio.run(i.async_conduct_interview(run_config))
+        Traceback (most recent call last):
+        ...
+        asyncio.exceptions.CancelledError
+        """
+        from edsl.jobs.Jobs import RunConfig, RunParameters, RunEnvironment
+
+        if run_config is None:
+            run_config = RunConfig(
+                parameters=RunParameters(),
+                environment=RunEnvironment(),
+            )
+        self.stop_on_exception = run_config.parameters.stop_on_exception
+
+        # if no model bucket is passed, create an 'infinity' bucket with no rate limits
+        bucket_collection = run_config.environment.bucket_collection
+
+        if bucket_collection:
+            model_buckets = bucket_collection.get(self.model)
+        else:
+            model_buckets = None
+
+        if model_buckets is None or hasattr(self.agent, "answer_question_directly"):
+            model_buckets = ModelBuckets.infinity_bucket()
+
+        # was "self.tasks" - is that necessary?
+        self.tasks = self.task_manager.build_question_tasks(
+            answer_func=AnswerQuestionFunctionConstructor(
+                self, key_lookup=run_config.environment.key_lookup
+            )(),
+            token_estimator=RequestTokenEstimator(self),
+            model_buckets=model_buckets,
+        )
+
+        ## This is the key part---it creates a task for each question,
+        ## with dependencies on the questions that must be answered before this one can be answered.
+
+        ## 'Invigilators' are used to administer the survey.
+        fetcher = FetchInvigilator(
+            interview=self,
+            current_answers=self.answers,
+            key_lookup=run_config.environment.key_lookup,
+        )
+        self.invigilators = [fetcher(question) for question in self.survey.questions]
+        await asyncio.gather(
+            *self.tasks, return_exceptions=not run_config.parameters.stop_on_exception
+        )
+        self.answers.replace_missing_answers_with_none(self.survey)
+        valid_results = list(
+            self._extract_valid_results(self.tasks, self.invigilators, self.exceptions)
+        )
+        return self.answers, valid_results
+
+    @staticmethod
+    def _extract_valid_results(
+        tasks: List["asyncio.Task"],
+        invigilators: List["InvigilatorBase"],
+        exceptions: InterviewExceptionCollection,
+    ) -> Generator["Answers", None, None]:
+        """Extract the valid results from the list of results.
+
+        It iterates through the tasks and invigilators, and yields the results of the tasks that are done.
+        If a task is not done, it raises a ValueError.
+        If an exception is raised in the task, it records the exception in the Interview instance except if the task was cancelled, which is expected behavior.
+
+        >>> i = Interview.example()
+        >>> result, _ = asyncio.run(i.async_conduct_interview())
+        """
+        assert len(tasks) == len(invigilators)
+
+        def handle_task(task, invigilator):
+            try:
+                result: Answers = task.result()
+            except asyncio.CancelledError as e:  # task was cancelled
+                result = invigilator.get_failed_task_result(
+                    failure_reason="Task was cancelled."
+                )
+            except Exception as e:  # any other kind of exception in the task
+                result = invigilator.get_failed_task_result(
+                    failure_reason=f"Task failed with exception: {str(e)}."
+                )
+                exception_entry = InterviewExceptionEntry(
+                    exception=e,
+                    invigilator=invigilator,
+                )
+                exceptions.add(task.get_name(), exception_entry)
+            return result
+
+        for task, invigilator in zip(tasks, invigilators):
+            if not task.done():
+                raise ValueError(f"Task {task.get_name()} is not done.")
+
+            yield handle_task(task, invigilator)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the Interview instance."""
+        return f"Interview(agent = {repr(self.agent)}, survey = {repr(self.survey)}, scenario = {repr(self.scenario)}, model = {repr(self.model)})"
+
+    def duplicate(
+        self, iteration: int, cache: "Cache", randomize_survey: Optional[bool] = True
+    ) -> Interview:
+        """Duplicate the interview, but with a new iteration number and cache.
+
+        >>> i = Interview.example()
+        >>> i2 = i.duplicate(1, None)
+        >>> i.iteration + 1 == i2.iteration
+        True
+
+        """
+        if randomize_survey:
+            new_survey = self.survey.draw()
+        else:
+            new_survey = self.survey
+
+        return Interview(
+            agent=self.agent,
+            survey=new_survey,
+            scenario=self.scenario,
+            model=self.model,
+            iteration=iteration,
+            cache=self.running_config.cache,
+            skip_retry=self.running_config.skip_retry,
+            indices=self.indices,
+        )
+
+    @classmethod
+    def example(self, throw_exception: bool = False) -> Interview:
+        """Return an example Interview instance."""
+        from edsl.agents import Agent
+        from edsl.surveys import Survey
+        from edsl.scenarios import Scenario
+        from edsl.language_models import LanguageModel
+
+        def f(self, question, scenario):
+            return "yes"
+
+        agent = Agent.example()
+        agent.add_direct_question_answering_method(f)
+        survey = Survey.example()
+        scenario = Scenario.example()
+        model = LanguageModel.example()
+        if throw_exception:
+            model = LanguageModel.example(test_model=True, throw_exception=True)
+            agent = Agent.example()
+            return Interview(agent=agent, survey=survey, scenario=scenario, model=model)
+        return Interview(agent=agent, survey=survey, scenario=scenario, model=model)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    # add ellipsis
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/jobs/interviews/InterviewExceptionCollection.py

@@ -0,0 +1,99 @@
+from collections import UserDict
+
+from edsl.jobs.interviews.InterviewExceptionEntry import InterviewExceptionEntry
+
+
+class InterviewExceptionCollection(UserDict):
+    """A collection of exceptions that occurred during the interview."""
+
+    def __init__(self):
+        super().__init__()
+        self.fixed = set()
+
+    def unfixed_exceptions(self) -> list:
+        """Return a list of unfixed exceptions."""
+        return {k: v for k, v in self.data.items() if k not in self.fixed}
+
+    def num_unfixed(self) -> list:
+        """Return a list of unfixed questions."""
+        return len([k for k in self.data.keys() if k not in self.fixed])
+
+    def record_fixed_question(self, question_name: str) -> None:
+        """Record that a question has been fixed."""
+        self.fixed.add(question_name)
+
+    def add(self, question_name: str, entry: InterviewExceptionEntry) -> None:
+        """Add an exception entry to the collection."""
+        question_name = question_name
+        if question_name not in self.data:
+            self.data[question_name] = []
+        self.data[question_name].append(entry)
+
+    def to_dict(self, include_traceback=True) -> dict:
+        """Return the collection of exceptions as a dictionary."""
+        newdata = {k: [e.to_dict() for e in v] for k, v in self.data.items()}
+        return newdata
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "InterviewExceptionCollection":
+        """Create an InterviewExceptionCollection from a dictionary."""
+        collection = cls()
+        for question_name, entries in data.items():
+            for entry in entries:
+                collection.add(question_name, InterviewExceptionEntry.from_dict(entry))
+        return collection
+
+    def _repr_html_(self) -> str:
+        from edsl.utilities.utilities import data_to_html
+
+        return data_to_html(self.to_dict(include_traceback=True))
+
+    def ascii_table(self, traceback: bool = False) -> None:
+        headers = ["Question name", "Exception", "Time", "Traceback"]
+        from tabulate import tabulate
+
+        data = []
+        for question, exceptions in self.data.items():
+            for exception in exceptions:
+                if traceback:
+                    row = [
+                        question,
+                        exception["exception"],
+                        exception["time"],
+                        exception["traceback"],
+                    ]
+                else:
+                    row = [question, exception["exception"], exception["time"]]
+                data.append(row)
+
+        print(tabulate(data, headers=headers, tablefmt="grid"))
+
+    def print(self, traceback=False):
+        """Print the collection of exceptions."""
+        console = Console()
+        table = Table(show_header=True, header_style="bold magenta")
+        table.add_column("Question name", style="dim", width=12)
+        table.add_column("Exception", width=32)
+        table.add_column("Time", justify="right")
+        table.add_column("Traceback", min_width=20)
+
+        for queue, exceptions in self.data.items():
+            for exception in exceptions:
+                if traceback:
+                    traceback_string = exception["traceback"].replace("\n", "\n\n")
+                else:
+                    traceback_string = ""
+                table.add_row(
+                    queue,
+                    exception["exception"],
+                    str(exception["time"]),
+                    traceback_string,  # Adding extra newlines for better readability
+                )
+
+        console.print(table)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/jobs/interviews/InterviewExceptionEntry.py

@@ -0,0 +1,186 @@
+import traceback
+import datetime
+
+
+class InterviewExceptionEntry:
+    """Class to record an exception that occurred during the interview."""
+
+    def __init__(
+        self,
+        *,
+        exception: Exception,
+        invigilator: "Invigilator",
+        traceback_format="text",
+        answers=None,
+    ):
+        self.time = datetime.datetime.now().isoformat()
+        self.exception = exception
+        # self.failed_question = failed_question
+        self.invigilator = invigilator
+        self.traceback_format = traceback_format
+        self.answers = answers
+
+    @property
+    def question_type(self):
+        # return self.failed_question.question.question_type
+        return self.invigilator.question.question_type
+
+    @property
+    def name(self):
+        return repr(self.exception)
+
+    @property
+    def rendered_prompts(self):
+        return self.invigilator.get_prompts()
+
+    @property
+    def key_sequence(self):
+        return self.invigilator.model.key_sequence
+
+    @property
+    def generated_token_string(self):
+        # return "POO"
+        if self.invigilator.raw_model_response is None:
+            return "No raw model response available."
+        else:
+            return self.invigilator.model.get_generated_token_string(
+                self.invigilator.raw_model_response
+            )
+
+    @property
+    def raw_model_response(self):
+        import json
+
+        if self.invigilator.raw_model_response is None:
+            return "No raw model response available."
+        return json.dumps(self.invigilator.raw_model_response, indent=2)
+
+    def __getitem__(self, key):
+        # Support dict-like access obj['a']
+        return str(getattr(self, key))
+
+    @classmethod
+    def example(cls):
+        from edsl import QuestionFreeText
+        from edsl.language_models import LanguageModel
+
+        m = LanguageModel.example(test_model=True)
+        q = QuestionFreeText.example(exception_to_throw=ValueError)
+        results = q.by(m).run(
+            skip_retry=True,
+            print_exceptions=False,
+            raise_validation_errors=True,
+            disable_remote_cache=True,
+            disable_remote_inference=True,
+        )
+        return results.task_history.exceptions[0]["how_are_you"][0]
+
+    @property
+    def code_to_reproduce(self):
+        return self.code(run=False)
+
+    def code(self, run=True):
+        lines = []
+        lines.append("from edsl import Question, Model, Scenario, Agent")
+
+        lines.append(f"q = {repr(self.invigilator.question)}")
+        lines.append(f"scenario = {repr(self.invigilator.scenario)}")
+        lines.append(f"agent = {repr(self.invigilator.agent)}")
+        lines.append(f"m = Model('{self.invigilator.model.model}')")
+        lines.append("results = q.by(m).by(agent).by(scenario).run()")
+        code_str = "\n".join(lines)
+
+        if run:
+            # Create a new namespace to avoid polluting the global namespace
+            namespace = {}
+            exec(code_str, namespace)
+            return namespace["results"]
+        return code_str
+
+    @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()
+
+    @staticmethod
+    def serialize_exception(exception: Exception) -> dict:
+        return {
+            "type": type(exception).__name__,
+            "message": str(exception),
+            "traceback": "".join(
+                traceback.format_exception(
+                    type(exception), exception, exception.__traceback__
+                )
+            ),
+        }
+
+    @staticmethod
+    def deserialize_exception(data: dict) -> Exception:
+        try:
+            exception_class = globals()[data["type"]]
+        except KeyError:
+            exception_class = Exception
+        return exception_class(data["message"])
+
+    def to_dict(self) -> dict:
+        """Return the exception as a dictionary.
+
+        >>> entry = InterviewExceptionEntry.example()
+        >>> _ = entry.to_dict()
+        """
+        return {
+            "exception": self.serialize_exception(self.exception),
+            "time": self.time,
+            "traceback": self.traceback,
+            "invigilator": self.invigilator.to_dict(),
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "InterviewExceptionEntry":
+        """Create an InterviewExceptionEntry from a dictionary."""
+        from edsl.agents.Invigilator import InvigilatorAI
+
+        exception = cls.deserialize_exception(data["exception"])
+        invigilator = InvigilatorAI.from_dict(data["invigilator"])
+        return cls(exception=exception, invigilator=invigilator)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)