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

edsl/jobs/AnswerQuestionFunctionConstructor.py

@@ -0,0 +1,223 @@
+import copy
+import asyncio
+
+from typing import Union, Type, Callable, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.questions.QuestionBase import QuestionBase
+    from edsl.jobs.interviews.Interview import Interview
+    from edsl.language_models.key_management.KeyLookup import KeyLookup
+
+from edsl.surveys.base import EndOfSurvey
+from edsl.jobs.tasks.task_status_enum import TaskStatus
+
+from edsl.jobs.FetchInvigilator import FetchInvigilator
+from edsl.exceptions.language_models import LanguageModelNoResponseError
+from edsl.exceptions.questions import QuestionAnswerValidationError
+from edsl.data_transfer_models import AgentResponseDict, EDSLResultObjectInput
+
+from edsl.jobs.Answers import Answers
+
+
+class RetryConfig:
+    from edsl.config import CONFIG
+
+    EDSL_BACKOFF_START_SEC = float(CONFIG.get("EDSL_BACKOFF_START_SEC"))
+    EDSL_BACKOFF_MAX_SEC = float(CONFIG.get("EDSL_BACKOFF_MAX_SEC"))
+    EDSL_MAX_ATTEMPTS = int(CONFIG.get("EDSL_MAX_ATTEMPTS"))
+
+
+class SkipHandler:
+
+    def __init__(self, interview: "Interview"):
+        self.interview = interview
+        self.question_index = self.interview.to_index
+
+        self.skip_function: Callable = (
+            self.interview.survey.rule_collection.skip_question_before_running
+        )
+
+    def should_skip(self, current_question: "QuestionBase") -> bool:
+        """Determine if the current question should be skipped."""
+        current_question_index = self.question_index[current_question.question_name]
+        combined_answers = (
+            self.interview.answers
+            | self.interview.scenario
+            | self.interview.agent["traits"]
+        )
+        return self.skip_function(current_question_index, combined_answers)
+
+    def cancel_skipped_questions(self, current_question: "QuestionBase") -> None:
+        """Cancel the tasks for questions that should be skipped."""
+        current_question_index: int = self.interview.to_index[
+            current_question.question_name
+        ]
+        answers = (
+            self.interview.answers
+            | self.interview.scenario
+            | self.interview.agent["traits"]
+        )
+
+        # Get the index of the next question, which could also be the end of the survey
+        next_question: Union[int, EndOfSurvey] = (
+            self.interview.survey.rule_collection.next_question(
+                q_now=current_question_index,
+                answers=answers,
+            )
+        )
+
+        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()
+
+        if (next_question_index := next_question.next_q) == EndOfSurvey:
+            cancel_between(
+                current_question_index + 1, len(self.interview.survey.questions)
+            )
+            return
+
+        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."""
+
+    def __init__(self, interview: "Interview", key_lookup: "KeyLookup"):
+        self.interview = interview
+        self.key_lookup = key_lookup
+
+        self.had_language_model_no_response_error: bool = False
+        self.question_index = self.interview.to_index
+
+        self.skip_function: Callable = (
+            self.interview.survey.rule_collection.skip_question_before_running
+        )
+
+        self.invigilator_fetcher = FetchInvigilator(
+            self.interview, key_lookup=self.key_lookup
+        )
+        self.skip_handler = SkipHandler(self.interview)
+
+    def _handle_exception(
+        self, e: Exception, invigilator: "InvigilatorBase", task=None
+    ):
+        """Handle an exception that occurred while answering a question."""
+
+        from edsl.jobs.interviews.InterviewExceptionEntry import InterviewExceptionEntry
+
+        answers = copy.copy(
+            self.interview.answers
+        )  # copy to freeze the answers here for logging
+        exception_entry = InterviewExceptionEntry(
+            exception=e,
+            invigilator=invigilator,
+            answers=answers,
+        )
+        if task:
+            task.task_status = TaskStatus.FAILED
+
+        self.interview.exceptions.add(
+            invigilator.question.question_name, exception_entry
+        )
+
+        if self.interview.raise_validation_errors and isinstance(
+            e, QuestionAnswerValidationError
+        ):
+            raise e
+
+        stop_on_exception = getattr(self.interview, "stop_on_exception", False)
+        if stop_on_exception:
+            raise e
+
+    def __call__(self):
+        return self.answer_question_and_record_task
+
+    async def answer_question_and_record_task(
+        self,
+        *,
+        question: "QuestionBase",
+        task=None,
+    ) -> "EDSLResultObjectInput":
+
+        from tenacity import (
+            retry,
+            stop_after_attempt,
+            wait_exponential,
+            retry_if_exception_type,
+            RetryError,
+        )
+
+        @retry(
+            stop=stop_after_attempt(RetryConfig.EDSL_MAX_ATTEMPTS),
+            wait=wait_exponential(
+                multiplier=RetryConfig.EDSL_BACKOFF_START_SEC,
+                max=RetryConfig.EDSL_BACKOFF_MAX_SEC,
+            ),
+            retry=retry_if_exception_type(LanguageModelNoResponseError),
+            reraise=True,
+        )
+        async def attempt_answer():
+            invigilator = self.invigilator_fetcher(question)
+
+            if self.skip_handler.should_skip(question):
+                return invigilator.get_failed_task_result(
+                    failure_reason="Question skipped."
+                )
+
+            try:
+                response: EDSLResultObjectInput = (
+                    await invigilator.async_answer_question()
+                )
+                if response.validated:
+                    self.interview.answers.add_answer(
+                        response=response, question=question
+                    )
+
+                    self.skip_handler.cancel_skipped_questions(question)
+                else:
+                    if (
+                        hasattr(response, "exception_occurred")
+                        and response.exception_occurred
+                    ):
+                        raise response.exception_occurred
+
+            except QuestionAnswerValidationError as e:
+                self._handle_exception(e, invigilator, task)
+                return invigilator.get_failed_task_result(
+                    failure_reason="Question answer validation failed."
+                )
+
+            except asyncio.TimeoutError as e:
+                self._handle_exception(e, invigilator, task)
+                had_language_model_no_response_error = True
+                raise LanguageModelNoResponseError(
+                    f"Language model timed out for question '{question.question_name}.'"
+                )
+
+            except Exception as e:
+                self._handle_exception(e, invigilator, task)
+
+            if "response" not in locals():
+                had_language_model_no_response_error = True
+                raise LanguageModelNoResponseError(
+                    f"Language model did not return a response for question '{question.question_name}.'"
+                )
+
+            if (
+                question.question_name in self.interview.exceptions
+                and had_language_model_no_response_error
+            ):
+                self.interview.exceptions.record_fixed_question(question.question_name)
+
+            return response
+
+        try:
+            return await attempt_answer()
+        except RetryError as retry_error:
+            original_error = retry_error.last_attempt.exception()
+            self._handle_exception(
+                original_error, self.invigilator_fetcher(question), task
+            )
+            raise original_error

edsl/jobs/Answers.py

@@ -1,42 +1,43 @@
+"""This module contains the Answers class, which is a helper class to hold the answers to a survey."""
+
 from collections import UserDict
-from rich.table import Table
+from edsl.data_transfer_models import EDSLResultObjectInput
 
 
 class Answers(UserDict):
-    "Helper class to hold the answers to a survey"
-
-    def add_answer(self, response, question) -> None:
-        "Adds a response to the answers dictionary"
-        answer = response.get("answer")
-        comment = response.pop("comment", None)
+    """Helper class to hold the answers to a survey."""
+
+    def add_answer(
+        self, response: EDSLResultObjectInput, question: "QuestionBase"
+    ) -> None:
+        """Add a response to the answers dictionary."""
+        answer = response.answer
+        comment = response.comment
+        generated_tokens = response.generated_tokens
         # record the answer
+        if generated_tokens:
+            self[question.question_name + "_generated_tokens"] = generated_tokens
         self[question.question_name] = answer
         if comment:
             self[question.question_name + "_comment"] = comment
 
-    def replace_missing_answers_with_none(self, survey) -> None:
-        "Replaces missing answers with None. Answers can be missing if the agent skips a question."
+    def replace_missing_answers_with_none(self, survey: "Survey") -> None:
+        """Replace missing answers with None. Answers can be missing if the agent skips a question."""
         for question_name in survey.question_names:
             if question_name not in self:
                 self[question_name] = None
 
     def to_dict(self):
-        "Returns a dictionary of the answers"
+        """Return a dictionary of the answers."""
         return self.data
 
     @classmethod
     def from_dict(cls, d):
-        "Returns an Answers object from a dictionary"
+        """Return an Answers object from a dictionary."""
         return cls(d)
 
-    def rich_print(self):
-        """Displays an object as a table."""
-        table = Table(title="Answers")
-        table.add_column("Attribute", style="bold")
-        table.add_column("Value")
 
-        to_display = self
-        for attr_name, attr_value in to_display.items():
-            table.add_row(attr_name, repr(attr_value))
+if __name__ == "__main__":
+    import doctest
 
-        return table
+    doctest.testmod()

edsl/jobs/FetchInvigilator.py

@@ -0,0 +1,47 @@
+from typing import List, Dict, Any, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.questions.QuestionBase import QuestionBase
+    from edsl.agents.InvigilatorBase import InvigilatorBase
+    from edsl.language_models.key_management.KeyLookup import KeyLookup
+    from edsl.jobs.interviews.Interview import Interview
+
+
+class FetchInvigilator:
+    def __init__(
+        self,
+        interview: "Interview",
+        current_answers: Optional[Dict[str, Any]] = None,
+        key_lookup: Optional["KeyLookup"] = None,
+    ):
+        self.interview = interview
+        if current_answers is None:
+            self.current_answers = self.interview.answers
+        else:
+            self.current_answers = current_answers
+        self.key_lookup = key_lookup
+
+    def get_invigilator(self, question: "QuestionBase") -> "InvigilatorBase":
+        """Return an invigilator for the given question.
+
+        :param question: the question to be answered
+        :param debug: whether to use debug mode, in which case `InvigilatorDebug` is used.
+        """
+
+        invigilator = self.interview.agent.create_invigilator(
+            question=question,
+            scenario=self.interview.scenario,
+            model=self.interview.model,
+            survey=self.interview.survey,
+            memory_plan=self.interview.survey.memory_plan,
+            current_answers=self.current_answers,  # not yet known
+            iteration=self.interview.iteration,
+            cache=self.interview.cache,
+            raise_validation_errors=self.interview.raise_validation_errors,
+            key_lookup=self.key_lookup,
+        )
+        """Return an invigilator for the given question."""
+        return invigilator
+
+    def __call__(self, question):
+        return self.get_invigilator(question)

edsl/jobs/InterviewTaskManager.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+import asyncio
+from typing import Any, Type, List, Generator, Optional, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.questions import QuestionBase
+    from edsl.jobs.tokens.InterviewTokenUsage import InterviewTokenUsage
+    from edsl.jobs.interviews.InterviewStatusDictionary import InterviewStatusDictionary
+    from edsl.jobs.interviews.InterviewStatusLog import InterviewStatusLog
+
+
+class InterviewTaskManager:
+    """Handles creation and management of interview tasks."""
+
+    def __init__(self, survey, iteration=0):
+        from edsl.jobs.tasks.TaskCreators import TaskCreators
+        from edsl.jobs.interviews.InterviewStatusLog import InterviewStatusLog
+
+        self.survey = survey
+        self.iteration = iteration
+        self.task_creators = TaskCreators()
+        self.to_index = {
+            question_name: index
+            for index, question_name in enumerate(self.survey.question_names)
+        }
+        self._task_status_log_dict = InterviewStatusLog()
+
+    def build_question_tasks(
+        self, answer_func, token_estimator, model_buckets
+    ) -> list[asyncio.Task]:
+        """Create tasks for all questions with proper dependencies."""
+        tasks = []
+        for question in self.survey.questions:
+            dependencies = self._get_task_dependencies(tasks, question)
+            task = self._create_single_task(
+                question=question,
+                dependencies=dependencies,
+                answer_func=answer_func,
+                token_estimator=token_estimator,
+                model_buckets=model_buckets,
+            )
+            tasks.append(task)
+        return tuple(tasks)
+
+    def _get_task_dependencies(
+        self, existing_tasks: list[asyncio.Task], question: "QuestionBase"
+    ) -> list[asyncio.Task]:
+        """Get tasks that must be completed before the given question."""
+        dag = self.survey.dag(textify=True)
+        parents = dag.get(question.question_name, [])
+        return [existing_tasks[self.to_index[parent_name]] for parent_name in parents]
+
+    def _create_single_task(
+        self,
+        question: "QuestionBase",
+        dependencies: list[asyncio.Task],
+        answer_func,
+        token_estimator,
+        model_buckets,
+    ) -> asyncio.Task:
+        """Create a single question task with its dependencies."""
+        from edsl.jobs.tasks.QuestionTaskCreator import QuestionTaskCreator
+
+        task_creator = QuestionTaskCreator(
+            question=question,
+            answer_question_func=answer_func,
+            token_estimator=token_estimator,
+            model_buckets=model_buckets,
+            iteration=self.iteration,
+        )
+
+        for dependency in dependencies:
+            task_creator.add_dependency(dependency)
+
+        self.task_creators[question.question_name] = task_creator
+        return task_creator.generate_task()
+
+    @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.
+        """
+        for task_creator in self.task_creators.values():
+            self._task_status_log_dict[
+                task_creator.question.question_name
+            ] = task_creator.status_log
+        return self._task_status_log_dict
+
+    @property
+    def token_usage(self) -> "InterviewTokenUsage":
+        """Determine how many tokens were used for the interview."""
+        return self.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

edsl/jobs/InterviewsConstructor.py

@@ -0,0 +1,50 @@
+from typing import Generator, TYPE_CHECKING
+from itertools import product
+
+if TYPE_CHECKING:
+    from edsl.jobs.interviews.Interview import Interview
+
+
+class InterviewsConstructor:
+    def __init__(self, jobs: "Jobs", cache: "Cache"):
+        self.jobs = jobs
+        self.cache = cache
+
+    def create_interviews(self) -> Generator["Interview", None, None]:
+        """
+        Generate interviews.
+
+        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.
+
+        """
+        from edsl.jobs.interviews.Interview import Interview
+
+        agent_index = {
+            hash(agent): index for index, agent in enumerate(self.jobs.agents)
+        }
+        model_index = {
+            hash(model): index for index, model in enumerate(self.jobs.models)
+        }
+        scenario_index = {
+            hash(scenario): index for index, scenario in enumerate(self.jobs.scenarios)
+        }
+
+        for agent, scenario, model in product(
+            self.jobs.agents, self.jobs.scenarios, self.jobs.models
+        ):
+            yield Interview(
+                survey=self.jobs.survey.draw(),
+                agent=agent,
+                scenario=scenario,
+                model=model,
+                cache=self.cache,
+                skip_retry=self.jobs.run_config.parameters.skip_retry,
+                raise_validation_errors=self.jobs.run_config.parameters.raise_validation_errors,
+                indices={
+                    "agent": agent_index[hash(agent)],
+                    "model": model_index[hash(model)],
+                    "scenario": scenario_index[hash(scenario)],
+                },
+            )