edsl 0.1.46__py3-none-any.whl → 0.1.48__py3-none-any.whl

edsl/surveys/{SurveyExportMixin.py → survey_export.py}

@@ -1,4 +1,4 @@
-"""A mixin class for exporting surveys to different formats."""
+"""A class for exporting surveys to different formats."""
 
 from typing import Union, Optional
 
@@ -24,11 +24,15 @@ def open_docx(file_path):
         subprocess.call(("xdg-open", file_path))
 
 
-class SurveyExportMixin:
-    """A mixin class for exporting surveys to different formats."""
+class SurveyExport:
+    """A class for exporting surveys to different formats."""
+
+    def __init__(self, survey):
+        """Initialize with a Survey object."""
+        self.survey = survey
 
     def css(self):
-        from edsl.surveys.SurveyCSS import SurveyCSS
+        from .survey_css import SurveyCSS
 
         return SurveyCSS.default_style().generate_css()
 
@@ -36,7 +40,7 @@ class SurveyExportMixin:
         """Return the description of the survey."""
         from edsl import QuestionFreeText
 
-        question_texts = "\n".join([q.question_text for q in self._questions])
+        question_texts = "\n".join([q.question_text for q in self.survey._questions])
         q = QuestionFreeText(
             question_name="description",
             question_text=f"""A survey was conducted with the following questions: 
@@ -58,7 +62,7 @@ class SurveyExportMixin:
         doc = Document()
         doc.add_heading("EDSL Survey")
         doc.add_paragraph(f"\n")
-        for index, question in enumerate(self._questions):
+        for index, question in enumerate(self.survey._questions):
             h = doc.add_paragraph()  # Add question as a paragraph
             h.add_run(f"Question {index + 1} ({question.question_name})").bold = True
             h.add_run(f"; {question.question_type}").italic = True
@@ -72,20 +76,13 @@ class SurveyExportMixin:
                     for option in getattr(question, "question_options", []):
                         doc.add_paragraph(str(option), style="ListBullet")
 
-        if return_document_object and filename is None:
+        if return_document_object:
             return doc
-
-        if filename is None:
-            with tempfile.NamedTemporaryFile(
-                "w", delete=False, suffix=".docx", dir=os.getcwd()
-            ) as f:
-                filename = f.name
-
-        doc.save(filename)
-        print("The survey has been saved to", filename)
-        if open_file:
-            open_docx(filename)
-        return
+        else:
+            doc.save(filename)
+            if open_file:
+                os.system(f"open {filename}")
+            return None
 
     def show(self):
         self.to_scenario_list(questions_only=False, rename=True).print(format="rich")
@@ -93,14 +90,12 @@ class SurveyExportMixin:
     def to_scenario_list(
         self, questions_only: bool = True, rename=False
     ) -> "ScenarioList":
-        from edsl import ScenarioList, Scenario
-
-        # from edsl.questions import QuestionBase
+        from ..scenarios import ScenarioList, Scenario
 
         if questions_only:
-            to_iterate_over = self._questions
+            to_iterate_over = self.survey._questions
         else:
-            to_iterate_over = self.recombined_questions_and_instructions()
+            to_iterate_over = self.survey.recombined_questions_and_instructions()
 
         if rename:
             renaming_dict = {
@@ -150,13 +145,13 @@ class SurveyExportMixin:
         header_lines = ["from edsl.surveys.Survey import Survey"]
         header_lines.append("from edsl import Question")
         lines = ["\n".join(header_lines)]
-        for question in self._questions:
+        for question in self.survey._questions:
             question.question_text = question["question_text"].replace("\n", " ")
             # remove dublicate spaces
             question.question_text = " ".join(question.question_text.split())
             lines.append(f"{question.question_name} = " + repr(question))
         lines.append(
-            f"{survey_var_name} = Survey(questions = [{', '.join(self.question_names)}])"
+            f"{survey_var_name} = Survey(questions = [{', '.join(self.survey.question_names)}])"
         )
         # return lines
         code_string = "\n".join(lines)
@@ -216,7 +211,7 @@ class SurveyExportMixin:
 
         with open(filename, "w") as f:
             f.write(html_header)
-            for question in self._questions:
+            for question in self.survey._questions:
                 f.write(
                     question.html(
                         scenario=scenario, include_question_name=include_question_name

edsl/surveys/{SurveyFlowVisualization.py → survey_flow_visualization.py}

@@ -47,7 +47,12 @@ class SurveyFlowVisualization:
                 for param in question.detailed_parameters:
                     if "agent." in param:
                         # Handle agent trait references
-                        trait_name = param.replace("agent.", "")
+                        #trait_name = param.replace("agent.", "")
+                        params_and_refs.add(param)
+                        if param not in param_to_questions:
+                            param_to_questions[param] = []
+                        param_to_questions[param].append(index)
+                    if "scenario." in param:
                         params_and_refs.add(param)
                         if param not in param_to_questions:
                             param_to_questions[param] = []
@@ -108,7 +113,7 @@ class SurveyFlowVisualization:
                     "label": f"Agent Trait\n{{{{ {param} }}}}"
                 })
             # Check if parameter exists in scenario
-            elif self.scenario and param in self.scenario:
+            elif self.scenario and param.startswith("scenario."):
                 node_attrs.update({
                     "fillcolor": "lightgreen",
                     "label": f"Scenario\n{{{{ {param} }}}}"
@@ -207,6 +212,10 @@ class SurveyFlowVisualization:
                     """File not found. Most likely it's because you don't have graphviz installed. Please install it and try again.
                         On Ubuntu, you can install it by running:
                         $ sudo apt-get install graphviz 
+                        On Mac, you can install it by running:
+                        $ brew install graphviz
+                        On Windows, you can install it by running:
+                        $ choco install graphviz
                     """
                 )
             from edsl.utilities.is_notebook import is_notebook

edsl/surveys/{Simulator.py → survey_simulator.py}

@@ -1,8 +1,15 @@
+from typing import TYPE_CHECKING
 from typing import Callable
+from edsl.agents import Agent
+#from edsl.surveys import Survey
 
+if TYPE_CHECKING:
+    from edsl.surveys import Survey
+    from edsl.results import Results
+    from edsl.questions import QuestionBase
 
 class Simulator:
-    def __init__(self, survey):
+    def __init__(self, survey: "Survey"):
         self.survey = survey
 
     @classmethod
@@ -10,7 +17,7 @@ class Simulator:
         """Create a random survey."""
         from edsl.questions import QuestionMultipleChoice, QuestionFreeText
         from random import choice
-        from edsl.surveys.Survey import Survey
+        from edsl.surveys import Survey
 
         num_questions = 10
         questions = []
@@ -57,7 +64,7 @@ class Simulator:
     def create_agent(self) -> "Agent":
         """Create an agent from the simulated answers."""
         answers_dict = self.survey.simulate()
-        from edsl.agents.Agent import Agent
+        from edsl.agents import Agent
 
         def construct_answer_dict_function(traits: dict) -> Callable:
             def func(self, question: "QuestionBase", scenario=None):

edsl/tasks/__init__.py

@@ -0,0 +1,32 @@
+"""
+The tasks module provides EDSL's task management system for tracking and controlling 
+interview execution.
+
+This module implements a comprehensive system for creating, scheduling, executing, and 
+monitoring tasks in EDSL. Tasks represent individual units of work, typically answering 
+a question with an LLM, with features for dependency management, error handling, and 
+execution status tracking.
+
+Key components:
+
+1. TaskHistory - Records and analyzes the execution history of tasks with error reporting
+2. QuestionTaskCreator - Creates and manages tasks for individual questions
+3. TaskCreators - Manages collections of tasks for an entire interview
+4. TaskStatus - Enumeration of possible task states (running, waiting, success, etc.)
+5. TaskStatusLog - Records the status changes of tasks over time
+
+The tasks system helps EDSL manage complex interview workflows by:
+- Handling dependencies between questions
+- Managing API rate limits and token usage
+- Providing detailed execution metrics
+- Generating error reports and visualizations
+- Supporting both synchronous and asynchronous execution
+
+For most users, this module works behind the scenes, but understanding it can 
+be helpful when debugging or optimizing complex EDSL workflows.
+"""
+
+from .task_history import TaskHistory
+from .question_task_creator import QuestionTaskCreator
+from .task_creators import TaskCreators
+from .task_status_enum import TaskStatus, TaskStatusDescriptor

edsl/{jobs/tasks/QuestionTaskCreator.py → tasks/question_task_creator.py}

@@ -1,34 +1,77 @@
+"""
+This module provides the QuestionTaskCreator class for executing individual questions as tasks.
+
+The QuestionTaskCreator is responsible for executing a single question within the EDSL system. 
+It manages the entire lifecycle of a question task, including dependency resolution, rate 
+limiting, token management, execution, and status tracking. It serves as the fundamental 
+execution unit in EDSL's task system.
+"""
+
 import asyncio
-from typing import Callable, Union, List, TYPE_CHECKING
+from typing import Callable, Union, List, Dict, Any, Optional, TYPE_CHECKING
 from collections import UserList, UserDict
 
-from edsl.exceptions.jobs import InterviewErrorPriorTaskCanceled
+from ..jobs.exceptions import InterviewErrorPriorTaskCanceled
+from ..tokens import TokenUsage
+from ..data_transfer_models import Answers
 
-from edsl.jobs.tasks.task_status_enum import TaskStatus, TaskStatusDescriptor
-from edsl.jobs.tasks.TaskStatusLog import TaskStatusLog
-from edsl.jobs.tokens.TokenUsage import TokenUsage
-from edsl.jobs.Answers import Answers
+from .task_status_enum import TaskStatus, TaskStatusDescriptor
+from .task_status_log import TaskStatusLog
 
 if TYPE_CHECKING:
-    from edsl.questions.QuestionBase import QuestionBase
-    from edsl.jobs.buckets import ModelBuckets
+    from ..questions import QuestionBase
+    from ..buckets import ModelBuckets
 
 
 class TokensUsed(UserDict):
-    """ "Container for tokens used by a task."""
+    """
+    Container for tracking token usage for a task, separating cached and new tokens.
+    
+    This class provides a structured way to track token usage for a single task,
+    distinguishing between tokens reused from cache and tokens freshly generated.
+    It uses a UserDict interface for convenient access to the underlying data.
+    
+    Attributes:
+        cached_tokens: TokenUsage object tracking reused tokens from cache
+        new_tokens: TokenUsage object tracking freshly generated tokens
+    """
 
-    def __init__(self, cached_tokens, new_tokens):
+    def __init__(self, cached_tokens: TokenUsage, new_tokens: TokenUsage):
+        """
+        Initialize a TokensUsed container.
+        
+        Parameters:
+            cached_tokens: TokenUsage object for tokens reused from cache
+            new_tokens: TokenUsage object for newly generated tokens
+        """
         d = {"cached_tokens": cached_tokens, "new_tokens": new_tokens}
         super().__init__(d)
 
 
 class QuestionTaskCreator(UserList):
-    """Class to create and manage a single question and its dependencies.
-
-    It is a UserList with all the tasks that must be completed before the focal task can be run.
-    The focal task is the question that we are interested in answering.
     """
-
+    Creates and manages the execution of a single question as an asyncio task.
+    
+    The QuestionTaskCreator is a fundamental component of EDSL's task system,
+    responsible for executing a single question with its dependencies. It extends
+    UserList to maintain a list of dependent tasks that must complete before this
+    task can execute.
+    
+    Key responsibilities:
+    1. Task Dependency Management - Tracks prerequisite tasks that must complete first
+    2. Resource Management - Handles rate limiting and token quota management
+    3. Task Status Tracking - Monitors and logs task state transitions
+    4. Token Usage Tracking - Records token consumption for both cached and new tokens
+    5. Task Execution - Runs the question answering function when dependencies are met
+    
+    The class follows the state machine pattern, with task_status transitioning through
+    various TaskStatus states (NOT_STARTED, WAITING_FOR_DEPENDENCIES, etc.) as execution
+    progresses. All status changes are automatically logged to enable detailed analysis
+    and visualization.
+    
+    This class is designed to work with asyncio for concurrent task execution, enabling
+    efficient processing of interviews with multiple questions and dependencies.
+    """
     task_status = TaskStatusDescriptor()
 
     def __init__(
@@ -37,17 +80,24 @@ class QuestionTaskCreator(UserList):
         question: "QuestionBase",
         answer_question_func: Callable,
         model_buckets: "ModelBuckets",
-        token_estimator: Union[Callable, None] = None,
+        token_estimator: Optional[Callable] = None,
         iteration: int = 0,
     ):
-        """Initialize the QuestionTaskCreator instance.
-
-        :param question: the question that we are interested in answering.
-        :param answer_question_func: the function that will answer the question.
-        :param model_buckets: the bucket collection that contains the requests and tokens buckets which control the rate of API calls and token usage.
-        :param token_estimator: a function that estimates the number of tokens required to answer the question.
-        :param iteration: the iteration number of the question.
-
+        """
+        Initialize a QuestionTaskCreator for a specific question.
+        
+        Parameters:
+            question: The Question object to be answered
+            answer_question_func: Function that will execute the LLM call to answer the question
+            model_buckets: Container for rate limiting buckets (requests and tokens)
+            token_estimator: Function to estimate token usage for the question (for quota management)
+            iteration: The iteration number of this question (for repeated questions)
+            
+        Notes:
+            - The QuestionTaskCreator starts in the NOT_STARTED state
+            - Dependencies can be added after initialization with add_dependency()
+            - Token usage is tracked separately for cached vs. new tokens
+            - This class works with asyncio for concurrent execution
         """
         super().__init__([])
         self.answer_question_func = answer_question_func
@@ -159,8 +209,8 @@ class QuestionTaskCreator(UserList):
     @classmethod
     def example(cls):
         """Return an example instance of the class."""
-        from edsl.questions.QuestionFreeText import QuestionFreeText
-        from edsl.jobs.buckets.ModelBuckets import ModelBuckets
+        from ..questions import QuestionFreeText
+        from ..buckets import ModelBuckets
 
         m = ModelBuckets.infinity_bucket()
 
@@ -180,37 +230,45 @@ class QuestionTaskCreator(UserList):
             iteration=0,
         )
 
-    async def _run_task_async(self) -> None:
-        """Run the task asynchronously, awaiting the tasks that must be completed before this one can be run.
-
-        >>> qt1 = QuestionTaskCreator.example()
-        >>> qt2 = QuestionTaskCreator.example()
-        >>> qt2.add_dependency(qt1)
-
-        The method follows these steps:
-            1. Set the task_status to TaskStatus.WAITING_FOR_DEPENDENCIES, indicating that the task is waiting for its dependencies to complete.
-            2. Await asyncio.gather(*self, return_exceptions=True) to run all the dependent tasks concurrently.
-
-            - the return_exceptions=True flag ensures that the task does not raise an exception if any of the dependencies fail.
-
-            3. If any of the dependencies raise an exception:
-            - If it is a CancelledError, set the current task's task_status to TaskStatus.CANCELLED, and re-raise the CancelledError,
-                terminating the execution of the current task.
-            - If it is any other exception, set the task_status to TaskStatus.PARENT_FAILED, and raise a custom exception
-                InterviewErrorPriorTaskCanceled with the original exception as the cause, terminating the execution of the current task.
-            4. If all the dependencies complete successfully without raising any exceptions, the code reaches the else block.
-            5. In the else block, run the focal task (self._run_focal_task(debug)).
-
-            If any of the dependencies fail (raise an exception), the focal task will not run. The execution will be terminated,
-            and an exception will be raised to indicate the failure of the dependencies.
-
-            The focal task (self._run_focal_task(debug)) is only executed if all the dependencies complete successfully.
-
-            Args:
-                debug: A boolean value indicating whether to run the task in debug mode.
-
-            Returns:
-                None
+    async def _run_task_async(self) -> Answers:
+        """
+        Execute the task with its dependencies in an async workflow.
+        
+        This method implements the core task execution logic with dependency handling.
+        It manages the complete lifecycle of a task:
+        
+        1. Waiting for dependencies to complete
+        2. Handling dependency failures appropriately
+        3. Executing the task itself when dependencies are satisfied
+        4. Tracking status transitions throughout execution
+        
+        The method maintains the state machine pattern by updating task_status
+        at each stage of execution, allowing for detailed monitoring and visualization
+        of task progress.
+        
+        Returns:
+            Answers object containing the question's answer and metadata
+            
+        Raises:
+            asyncio.CancelledError: If the task is cancelled
+            InterviewErrorPriorTaskCanceled: If any dependency task fails
+            
+        Example:
+            >>> qt1 = QuestionTaskCreator.example()
+            >>> qt2 = QuestionTaskCreator.example()
+            >>> qt2.add_dependency(qt1)
+        
+        Implementation details:
+        
+        1. Set status to WAITING_FOR_DEPENDENCIES and await all dependencies
+           - Using gather with return_exceptions=True allows collecting all results
+        
+        2. Check dependency results for exceptions:
+           - If CancelledError: Set status to CANCELLED and propagate the cancellation
+           - If other exception: Set status to PARENT_FAILED and wrap in InterviewErrorPriorTaskCanceled
+        
+        3. If all dependencies succeed, execute the focal task (_run_focal_task)
+           - The focal task handles its own status transitions during execution
         """
         try:
             self.task_status = TaskStatus.WAITING_FOR_DEPENDENCIES

edsl/tasks/task_creators.py

@@ -0,0 +1,135 @@
+"""
+This module provides the TaskCreators class, which manages all task creators for an interview.
+
+The TaskCreators class maintains a dictionary of QuestionTaskCreator objects, each responsible
+for executing a single question within an interview. It aggregates status and token usage
+information across all tasks, providing a complete view of interview execution.
+"""
+
+from typing import Callable, Union, List, Dict, TYPE_CHECKING
+from collections import UserDict
+
+if TYPE_CHECKING:
+    from ..tokens import TokenUsage
+    from ..tokens import InterviewTokenUsage
+    from ..interviews import InterviewStatusDictionary
+    from .task_status_log import TaskStatusLog
+    from .question_task_creator import QuestionTaskCreator
+
+class TaskCreators(UserDict):
+    """
+    A collection manager for all question tasks within an interview.
+    
+    The TaskCreators class maintains a dictionary of QuestionTaskCreator objects,
+    where each key is a question name and each value is the corresponding task creator.
+    This class provides methods to aggregate information across all tasks, such as
+    token usage and status counts, enabling a holistic view of interview execution.
+    
+    In the EDSL architecture, an interview consists of multiple questions, each executed
+    as a separate task. The TaskCreators class helps track and manage these tasks,
+    maintaining their execution status and resource usage.
+    
+    Key features:
+    - Maintains a mapping of question names to their task creators
+    - Aggregates token usage across all tasks
+    - Tracks the status of all tasks in the interview
+    - Provides access to status logs for visualization and analysis
+    
+    This class is typically used by the Interview class to manage task execution
+    and track the overall status of the interview.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @property
+    def token_usage(self) -> 'InterviewTokenUsage':
+        """
+        Calculate the total token usage across all tasks in the interview.
+        
+        This property aggregates token usage statistics from all task creators,
+        separating cached tokens (reused from cache) from new tokens (freshly generated).
+        The resulting InterviewTokenUsage object provides a complete picture of token
+        consumption for the entire interview.
+        
+        Returns:
+            An InterviewTokenUsage object containing:
+            - Cached token usage (tokens reused from cache)
+            - New token usage (tokens freshly generated)
+            
+        Notes:
+            - This is useful for cost estimation and quota management
+            - The separation of cached vs. new tokens helps analyze caching effectiveness
+            - Token usage is tracked separately for prompts and completions
+        """
+        from ..tokens import TokenUsage
+        from ..tokens import InterviewTokenUsage
+
+        cached_tokens = TokenUsage(from_cache=True)
+        new_tokens = TokenUsage(from_cache=False)
+        for task_creator in self.values():
+            token_usage = task_creator.token_usage()
+            cached_tokens += token_usage["cached_tokens"]
+            new_tokens += token_usage["new_tokens"]
+        return InterviewTokenUsage(
+            new_token_usage=new_tokens, cached_token_usage=cached_tokens
+        )
+
+    @property
+    def interview_status(self) -> 'InterviewStatusDictionary':
+        """
+        Get a summary of task statuses across the entire interview.
+        
+        This property counts how many tasks are in each possible status state,
+        providing a snapshot of the interview's current execution state. The
+        resulting InterviewStatusDictionary maps each TaskStatus to a count
+        of tasks in that state, plus a special 'number_from_cache' counter.
+        
+        Returns:
+            An InterviewStatusDictionary with counts for each task status
+            
+        Notes:
+            - Used for monitoring interview progress
+            - Helps identify bottlenecks (many tasks waiting for capacity)
+            - Tracks cache utilization via the 'number_from_cache' count
+            - Useful for status dashboards and progress reporting
+            
+        Example:
+            >>> t = TaskCreators()
+            >>> status = t.interview_status
+            >>> 'number_from_cache' in status
+            True
+            >>> status['number_from_cache']  # Check the cache count
+            0
+        """
+        from ..interviews import InterviewStatusDictionary
+        status_dict = InterviewStatusDictionary()
+        for task_creator in self.values():
+            status_dict[task_creator.task_status] += 1
+            status_dict["number_from_cache"] += task_creator.from_cache
+        return status_dict
+
+    def status_logs(self) -> List['TaskStatusLog']:
+        """
+        Get all task status logs for the interview.
+        
+        This method collects the status logs from all task creators, providing
+        a complete history of status changes for every task in the interview.
+        The resulting list can be used for detailed analysis of task execution
+        patterns, timing, and visualization.
+        
+        Returns:
+            A list of TaskStatusLog objects, one for each task in the interview
+            
+        Notes:
+            - Used by visualization tools like task timeline charts
+            - Helpful for performance analysis and bottleneck identification
+            - Provides data for execution reports and dashboards
+        """
+        return [task_creator.status_log for task_creator in self.values()]
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)