edsl 0.1.39.dev3__py3-none-any.whl → 0.1.39.dev4__py3-none-any.whl

edsl/jobs/Answers.py

@@ -1,56 +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: 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:
-        """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):
-        """Return a dictionary of the answers."""
-        return self.data
-
-    @classmethod
-    def from_dict(cls, d):
-        """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
-
-    doctest.testmod()
+"""This module contains the Answers class, which is a helper class to hold the answers to a survey."""
+
+from collections import UserDict
+from edsl.data_transfer_models import EDSLResultObjectInput
+
+
+class Answers(UserDict):
+    """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: "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):
+        """Return a dictionary of the answers."""
+        return self.data
+
+    @classmethod
+    def from_dict(cls, d):
+        """Return an Answers object from a dictionary."""
+        return cls(d)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    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)],
+                },
+            )