edsl 0.1.39.dev1__py3-none-any.whl → 0.1.39.dev2__py3-none-any.whl

edsl/jobs/AnswerQuestionFunctionConstructor.py

@@ -0,0 +1,188 @@
+import copy
+import asyncio
+
+from typing import Union, Type, Callable, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.questions.QuestionBase import QuestionBase
+
+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 AnswerQuestionFunctionConstructor:
+    def __init__(self, interview):
+        self.interview = interview
+        self.had_language_model_no_response_error = False
+        self.question_index = self.interview.to_index
+
+        self.skip_function: Callable = (
+            self.interview.survey.rule_collection.skip_question_before_running
+        )
+
+    def _combined_answers(self) -> Answers:
+        return self.answers | self.interview.scenario | self.interview.agent["traits"]
+
+    @property
+    def answers(self) -> Answers:
+        return self.interview.answers
+
+    def _skip_this_question(self, current_question: "QuestionBase") -> bool:
+        current_question_index = self.question_index[current_question.question_name]
+        combined_answers = self._combined_answers()
+        return self.skip_function(current_question_index, combined_answers)
+
+    def _handle_exception(
+        self, e: Exception, invigilator: "InvigilatorBase", task=None
+    ):
+        from edsl.jobs.interviews.InterviewExceptionEntry import InterviewExceptionEntry
+
+        answers = copy.copy(self.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 _cancel_skipped_questions(self, current_question: "QuestionBase") -> None:
+        current_question_index: int = self.interview.to_index[
+            current_question.question_name
+        ]
+        answers = (
+            self.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):
+            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)
+
+    def __call__(self):
+        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"))
+
+        from tenacity import (
+            retry,
+            stop_after_attempt,
+            wait_exponential,
+            retry_if_exception_type,
+            RetryError,
+        )
+
+        async def answer_question_and_record_task(
+            *,
+            question: "QuestionBase",
+            task=None,
+        ) -> "AgentResponseDict":
+            @retry(
+                stop=stop_after_attempt(EDSL_MAX_ATTEMPTS),
+                wait=wait_exponential(
+                    multiplier=EDSL_BACKOFF_START_SEC, max=EDSL_BACKOFF_MAX_SEC
+                ),
+                retry=retry_if_exception_type(LanguageModelNoResponseError),
+                reraise=True,
+            )
+            async def attempt_answer():
+                invigilator = FetchInvigilator(self.interview)(question)
+
+                if self._skip_this_question(question):
+                    return invigilator.get_failed_task_result(
+                        failure_reason="Question skipped."
+                    )
+
+                try:
+                    response: EDSLResultObjectInput = (
+                        await invigilator.async_answer_question()
+                    )
+                    if response.validated:
+                        self.answers.add_answer(response=response, question=question)
+                        self._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, FetchInvigilator(self.interview)(question), task
+                )
+                raise original_error
+
+        return answer_question_and_record_task

edsl/jobs/Answers.py

@@ -1,7 +1,6 @@
 """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
 
 
@@ -22,7 +21,7 @@ class Answers(UserDict):
         if comment:
             self[question.question_name + "_comment"] = comment
 
-    def replace_missing_answers_with_none(self, survey) -> None:
+    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:
@@ -37,18 +36,6 @@ class Answers(UserDict):
         """Return an Answers object from a dictionary."""
         return cls(d)
 
-    def rich_print(self):
-        """Display 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))
-
-        return table
-
 
 if __name__ == "__main__":
     import doctest

edsl/jobs/FetchInvigilator.py

@@ -0,0 +1,40 @@
+from typing import List, Dict, Any, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.questions.QuestionBase import QuestionBase
+    from edsl.agents.InvigilatorBase import InvigilatorBase
+
+
+class FetchInvigilator:
+    def __init__(self, interview, current_answers: Optional[Dict[str, Any]] = None):
+        self.interview = interview
+        if current_answers is None:
+            self.current_answers = self.interview.answers
+        else:
+            self.current_answers = current_answers
+
+    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,
+            #            debug=False,
+            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,
+            #           sidecar_model=self.interview.sidecar_model,
+            raise_validation_errors=self.interview.raise_validation_errors,
+        )
+        """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,48 @@
+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):
+        self.jobs = jobs
+
+    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,
+                agent=agent,
+                scenario=scenario,
+                model=model,
+                skip_retry=self.jobs.skip_retry,
+                raise_validation_errors=self.jobs.raise_validation_errors,
+                indices={
+                    "agent": agent_index[hash(agent)],
+                    "model": model_index[hash(model)],
+                    "scenario": scenario_index[hash(scenario)],
+                },
+            )