edsl 0.1.49__py3-none-any.whl → 0.1.51__py3-none-any.whl

edsl/surveys/memory/memory_management.py

@@ -3,6 +3,7 @@ from typing import Callable, Union, List, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from edsl.questions.QuestionBase import QuestionBase
+    from ...surveys.survey import Survey
 
 
 class MemoryManagement:

edsl/surveys/memory/memory_plan.py

@@ -8,6 +8,7 @@ if TYPE_CHECKING:
     from ...surveys import Survey
     from ...prompts import Prompt
     from ..dag import DAG
+    from ...questions.question_base import QuestionBase
 
 class MemoryPlan(UserDict):
     """A survey has a memory plan that specifies what the agent should remember when answering a question.
@@ -103,23 +104,13 @@ class MemoryPlan(UserDict):
         :param focal_question: The current question being answered.
         :param prior_question: The question that was answered before the focal question that should be remembered.
 
-        >>> mp = MemoryPlan.example()
-        >>> mp.add_single_memory("q0", "q1")
-        Traceback (most recent call last):
-        ...
-        ValueError: q1 must come before q0.
+        # Prior question must come before focal question
+        # Example: adding "q0", "q1" would raise ValueError as q1 comes after q0
 
-        >>> mp = MemoryPlan.example()
-        >>> mp.add_single_memory("q0", "crap")
-        Traceback (most recent call last):
-        ...
-        ValueError: crap is not in the survey. Current names are ['q0', 'q1', 'q2']
+        # Question names must exist in the survey
+        # Example: using a non-existent question name "crap" would raise ValueError
 
-        >>> mp = MemoryPlan.example()
-        >>> mp.add_single_memory("crap", "q0")
-        Traceback (most recent call last):
-        ...
-        ValueError: crap is not in the survey. Current names are ['q0', 'q1', 'q2']
+        # Similarly, focal question must also exist in the survey
         """
         self._check_valid_question_name(focal_question)
         self._check_valid_question_name(prior_question)

edsl/surveys/rules/__init__.py

@@ -1,3 +1,5 @@
-from .rule import Rule
-from .rule_manager import RuleManager
-from .rule_collection import RuleCollection
+from .rule import Rule  # noqa: F401
+from .rule_manager import RuleManager  # noqa: F401
+from .rule_collection import RuleCollection  # noqa: F401
+
+__all__ = []

edsl/surveys/rules/rule.py

@@ -19,7 +19,7 @@ with a low (-1) priority.
 
 import ast
 import random
-from typing import Any, Union, List
+from typing import Any, Union
 from collections import defaultdict
 
 
@@ -27,7 +27,6 @@ from collections import defaultdict
 from simpleeval import EvalWithCompoundTypes
 
 from ..exceptions import SurveyError
-
 from ..exceptions import (
     SurveyRuleCannotEvaluateError,
     SurveyRuleRefersToFutureStateError,

edsl/surveys/rules/rule_collection.py

@@ -1,6 +1,6 @@
 """A collection of rules for a survey."""
 
-from typing import List, Union, Any, Optional
+from typing import List, Any, Optional
 from collections import defaultdict, UserList, namedtuple
 
 from ..exceptions import (

edsl/surveys/survey.py

@@ -25,10 +25,8 @@ from typing import (
     TYPE_CHECKING,
     Dict,
     Tuple,
-    Set,
-    Type,
 )
-from typing_extensions import Literal, TypeAlias
+from typing_extensions import Literal
 from ..base import Base
 from ..agents import Agent
 from ..scenarios import Scenario
@@ -45,7 +43,6 @@ if TYPE_CHECKING:
     from ..scenarios import ScenarioList
     from ..buckets.bucket_collection import BucketCollection
     from ..key_management.key_lookup import KeyLookup
-    from .memory import Memory
     
     # Define types for documentation purpose only  
     VisibilityType = Literal["unlisted", "public", "private"]
@@ -243,6 +240,14 @@ class Survey(Base):
 
         self._exporter = SurveyExport(self)
 
+
+    # In survey.py
+    @property
+    def ep(self):
+        """Return plugin host for this survey."""
+        from ..plugins.plugin_host import PluginHost
+        return PluginHost(self)
+
     def question_names_valid(self) -> bool:
         """Check if the question names are valid."""
         return all(q.is_valid_question_name() for q in self.questions)
@@ -350,11 +355,7 @@ class Survey(Base):
 
         This doesnt' work with questions that don't exist:
 
-        >>> s._get_question_index("poop")
-        Traceback (most recent call last):
-        ...
-        edsl.surveys.exceptions.SurveyError: Question name poop not found in survey. The current question names are {'q0': 0, 'q1': 1, 'q2': 2}.
-        ...
+        # Example with a non-existent question name would raise SurveyError
         """
         if q is EndOfSurvey:
             return EndOfSurvey
@@ -609,11 +610,6 @@ class Survey(Base):
         >>> s1 = Survey.example()
         >>> from edsl import QuestionFreeText
         >>> s2 = Survey([QuestionFreeText(question_text="What is your name?", question_name="yo")])
-        >>> s3 = s1 + s2
-        Traceback (most recent call last):
-        ...
-        edsl.surveys.exceptions.SurveyCreationError: ...
-        ...
         >>> s3 = s1.clear_non_default_rules() + s2
         >>> len(s3.questions)
         4
@@ -676,11 +672,7 @@ class Survey(Base):
         >>> q = QuestionMultipleChoice(question_text = "Do you like school?", question_options=["yes", "no"], question_name="q0")
         >>> s = Survey().add_question(q)
 
-        >>> s = Survey().add_question(q).add_question(q)
-        Traceback (most recent call last):
-        ...
-        edsl.surveys.exceptions.SurveyCreationError: Question name 'q0' already exists in survey. Existing names are ['q0'].
-        ...
+        # Adding a question with a duplicate name would raise SurveyCreationError
         """
         return EditSurvey(self).add_question(question, index)
 
@@ -972,11 +964,7 @@ class Survey(Base):
         >>> s.next_question("q0", {"q0.answer": "no"}).question_name
         'q1'
 
-        >>> s.add_stop_rule("q0", "{{ q1.answer }} <> 'yes'")
-        Traceback (most recent call last):
-        ...
-        edsl.surveys.exceptions.SurveyCreationError: The expression contains '<>', which is not allowed. You probably mean '!='.
-        ...
+        # Using invalid operators like '<>' would raise SurveyCreationError
         """
         return RuleManager(self).add_stop_rule(question, expression)
 

edsl/surveys/survey_css.py

@@ -1,5 +1,5 @@
 from typing import Optional
-from edsl.utilities.remove_edsl_version import remove_edsl_version
+from ..utilities.remove_edsl_version import remove_edsl_version
 
 
 class CSSRuleMeta(type):
@@ -73,7 +73,7 @@ class CSSRule(metaclass=CSSRuleMeta):
         d = {"selector": self.selector, "properties": self.properties}
 
         if add_esl_version:
-            from edsl import __version__
+            from .. import __version__
 
             d["edsl_version"] = __version__
             d["edsl_class_name"] = self.__class__.__name__
@@ -233,7 +233,7 @@ class SurveyCSS:
         """
         d = {"rules": [rule.to_dict() for rule in self.rules.values()]}
         if add_edsl_version:
-            from edsl import __version__
+            from .. import __version__
 
             d["edsl_version"] = __version__
             d["edsl_class_name"] = self.__class__.__name__

edsl/surveys/survey_export.py

@@ -1,12 +1,16 @@
 """A class for exporting surveys to different formats."""
 
-from typing import Union, Optional
+from typing import Union, Optional, TYPE_CHECKING
 
 import subprocess
 import platform
 import os
 import tempfile
 
+if TYPE_CHECKING:
+    from docx import Document
+    from ..scenarios import ScenarioList
+
 
 def open_docx(file_path):
     """
@@ -61,7 +65,7 @@ class SurveyExport:
 
         doc = Document()
         doc.add_heading("EDSL Survey")
-        doc.add_paragraph(f"\n")
+        doc.add_paragraph("\n")
         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
@@ -176,7 +180,6 @@ class SurveyExport:
         include_question_name=False,
     ):
         from IPython.display import display, HTML
-        import tempfile
         import os
         from edsl.utilities.utilities import is_notebook
 

edsl/surveys/survey_flow_visualization.py

@@ -1,10 +1,19 @@
 """A mixin for visualizing the flow of a survey with parameter nodes."""
 
-from typing import Optional, TYPE_CHECKING
+from typing import Optional
 from edsl.surveys.base import RulePriority, EndOfSurvey
 import tempfile
 
 
+# Import types for annotations
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..surveys.survey import Survey
+    from ..scenarios import Scenario
+    from ..agents import Agent
+
+
 class SurveyFlowVisualization:
     """A mixin for visualizing the flow of a survey with parameter visualization."""
 

edsl/surveys/survey_simulator.py

@@ -58,7 +58,8 @@ class Simulator:
 
             if num_passes > 100:
                 print("Too many passes.")
-                raise Exception("Too many passes.")
+                from .exceptions import SurveyError
+                raise SurveyError("Too many passes.")
         return self.survey.answers
 
     def create_agent(self) -> "Agent":

edsl/tasks/__init__.py

@@ -26,7 +26,29 @@ For most users, this module works behind the scenes, but understanding it can
 be helpful when debugging or optimizing complex EDSL workflows.
 """
 
+__all__ = [
+    'TaskHistory', 
+    'QuestionTaskCreator', 
+    'TaskCreators', 
+    'TaskStatus', 
+    'TaskStatusDescriptor',
+    'TaskError',
+    'TaskStatusError',
+    'TaskExecutionError',
+    'TaskDependencyError',
+    'TaskResourceError',
+    'TaskHistoryError'
+]
+
 from .task_history import TaskHistory
 from .question_task_creator import QuestionTaskCreator
 from .task_creators import TaskCreators
-from .task_status_enum import TaskStatus, TaskStatusDescriptor
+from .task_status_enum import TaskStatus, TaskStatusDescriptor
+from .exceptions import (
+    TaskError,
+    TaskStatusError,
+    TaskExecutionError,
+    TaskDependencyError,
+    TaskResourceError,
+    TaskHistoryError
+)

edsl/tasks/exceptions.py

@@ -0,0 +1,72 @@
+"""
+This module defines the exception hierarchy for the tasks module.
+
+All exceptions related to task creation, execution, and management are defined here.
+These exceptions provide detailed error information for debugging and error reporting.
+"""
+
+from ..base import BaseException
+
+
+class TaskError(BaseException):
+    """
+    Base exception for all tasks-related errors.
+    
+    This is the parent class for all exceptions raised within the tasks module.
+    It inherits from BaseException to ensure proper error tracking and reporting.
+    """
+    pass
+
+
+class TaskStatusError(TaskError):
+    """
+    Raised when a task encounters an invalid status transition.
+    
+    This exception is raised when a task attempts to transition to an invalid state
+    based on its current state, such as trying to set a completed task to running.
+    
+    Attributes:
+        current_status: The current status of the task
+        attempted_status: The status that could not be set
+    """
+    pass
+
+
+class TaskExecutionError(TaskError):
+    """
+    Raised when a task encounters an error during execution.
+    
+    This is a general exception for errors that occur while a task is running,
+    not specific to dependency resolution or resource allocation.
+    """
+    pass
+
+
+class TaskDependencyError(TaskError):
+    """
+    Raised when there is an issue with task dependencies.
+    
+    This exception is raised for dependency-related issues, such as circular
+    dependencies or errors in dependent tasks.
+    """
+    pass
+
+
+class TaskResourceError(TaskError):
+    """
+    Raised when a task cannot acquire necessary resources.
+    
+    This exception is used when a task cannot obtain required resources
+    such as tokens or request capacity, beyond normal waiting situations.
+    """
+    pass
+
+
+class TaskHistoryError(TaskError):
+    """
+    Raised for errors related to task history operations.
+    
+    This exception covers issues with recording, accessing, or analyzing
+    task execution history and logs.
+    """
+    pass

edsl/tasks/question_task_creator.py

@@ -8,7 +8,7 @@ execution unit in EDSL's task system.
 """
 
 import asyncio
-from typing import Callable, Union, List, Dict, Any, Optional, TYPE_CHECKING
+from typing import Callable, Optional, TYPE_CHECKING
 from collections import UserList, UserDict
 
 from ..jobs.exceptions import InterviewErrorPriorTaskCanceled
@@ -170,12 +170,12 @@ class QuestionTaskCreator(UserList):
         """
 
         requested_tokens = self.estimated_tokens()
-        if (estimated_wait_time := self.tokens_bucket.wait_time(requested_tokens)) > 0:
+        if (self.tokens_bucket.wait_time(requested_tokens)) > 0:
             self.task_status = TaskStatus.WAITING_FOR_TOKEN_CAPACITY
 
         await self.tokens_bucket.get_tokens(requested_tokens)
 
-        if (estimated_wait_time := self.model_buckets.requests_bucket.wait_time(1)) > 0:
+        if self.model_buckets.requests_bucket.wait_time(1) > 0:
             self.waiting = True  #  do we need this?
             self.task_status = TaskStatus.WAITING_FOR_REQUEST_CAPACITY
 

edsl/tasks/task_creators.py

@@ -6,15 +6,13 @@ for executing a single question within an interview. It aggregates status and to
 information across all tasks, providing a complete view of interview execution.
 """
 
-from typing import Callable, Union, List, Dict, TYPE_CHECKING
+from typing import List, 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):
     """

edsl/tasks/task_history.py

@@ -7,7 +7,10 @@ visualization and reporting to help users understand task execution patterns and
 issues.
 """
 
-from typing import List, Optional, Dict, Any, Union
+from typing import List, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..interviews import Interview
 from io import BytesIO
 import base64
 import os
@@ -89,12 +92,10 @@ class TaskHistory(RepresentationMixin):
     @classmethod
     def example(cls):
         """ """
-        from ..interviews import Interview
         from ..jobs import Jobs
 
         j = Jobs.example(throw_exception_probability=1, test_model=True)
 
-        from edsl.config import CONFIG
 
         results = j.run(
             print_exceptions=False,
@@ -145,7 +146,7 @@ class TaskHistory(RepresentationMixin):
             "include_traceback": self.include_traceback,
         }
         if add_edsl_version:
-            from edsl import __version__
+            from .. import __version__
 
             d["edsl_version"] = __version__
             d["edsl_class_name"] = "TaskHistory"
@@ -404,8 +405,8 @@ class TaskHistory(RepresentationMixin):
 
         models_used = set([i.model.model for index, i in self._interviews.items()])
 
-        from jinja2 import Environment, FileSystemLoader
-        from edsl.utilities import TemplateLoader
+        from jinja2 import Environment
+        from ..utilities import TemplateLoader
 
         env = Environment(loader=TemplateLoader("edsl", "templates/error_reporting"))
 
@@ -463,9 +464,8 @@ class TaskHistory(RepresentationMixin):
             - Exception details, status transitions, and timing are all included
         """
         from IPython.display import display, HTML
-        import tempfile
         import os
-        from edsl.utilities.utilities import is_notebook
+        from ..utilities.utilities import is_notebook
 
         output = self.generate_html_report(css)
 
@@ -509,8 +509,6 @@ class TaskHistory(RepresentationMixin):
         """Create a notebook with the HTML content embedded in the first cell, then delete the cell content while keeping the output."""
         from nbformat import v4 as nbf
         from nbconvert.preprocessors import ExecutePreprocessor
-        import nbformat
-        import os
 
         # Use the existing html method to generate the HTML content
         output_html = self.generate_html_report(css=None)

edsl/tasks/task_status_log.py

@@ -7,9 +7,8 @@ execution, including timing, state transitions, and status at any point in time.
 """
 
 from collections import UserList
-from typing import List, Optional, Union
 
-from .task_status_enum import TaskStatus, TaskStatusLogEntry
+from .task_status_enum import TaskStatus
 
 
 class TaskStatusLog(UserList):

edsl/tokens/__init__.py

@@ -1,2 +1,30 @@
+"""
+The tokens module provides functionality for tracking and analyzing token usage in EDSL.
+
+This module implements classes for tracking and reporting token usage across various
+components of EDSL, particularly for language model calls. It supports aggregation,
+cost calculation, and reporting of token usage metrics.
+
+Key components:
+1. TokenUsage - Tracks prompt and completion tokens for a single operation
+2. InterviewTokenUsage - Aggregates token usage across an entire interview
+3. Exception classes for handling token-related errors
+
+The token tracking system helps with:
+- Cost estimation and billing
+- Resource utilization analysis
+- Cache effectiveness measurement
+- API quota management
+"""
+
 from .token_usage import TokenUsage
-from .interview_token_usage import InterviewTokenUsage
+from .interview_token_usage import InterviewTokenUsage
+from .exceptions import TokenError, TokenUsageError, TokenCostError
+
+__all__ = [
+    "TokenUsage", 
+    "InterviewTokenUsage",
+    "TokenError",
+    "TokenUsageError", 
+    "TokenCostError"
+]

edsl/tokens/exceptions.py

@@ -0,0 +1,37 @@
+"""
+This module defines the exception hierarchy for the tokens module.
+
+All exceptions related to token usage tracking and cost calculations are defined here.
+"""
+
+from ..base import BaseException
+
+
+class TokenError(BaseException):
+    """
+    Base exception for all token-related errors.
+    
+    This is the parent class for all exceptions raised within the tokens module.
+    It inherits from BaseException to ensure proper error tracking and reporting.
+    """
+    pass
+
+
+class TokenUsageError(TokenError):
+    """
+    Raised when there is an error in token usage operations.
+    
+    This exception is raised for issues related to token usage tracking,
+    such as invalid token counts or incompatible token usage types.
+    """
+    pass
+
+
+class TokenCostError(TokenError):
+    """
+    Raised when there is an error in token cost calculations.
+    
+    This exception is used for issues with cost calculations, such as
+    missing or invalid pricing information.
+    """
+    pass

edsl/tokens/interview_token_usage.py

@@ -1,7 +1,8 @@
 from typing import Optional
 
 from .token_usage import TokenUsage
-from edsl.enums import TokenPricing
+from ..enums import TokenPricing
+from .exceptions import TokenUsageError
 
 class InterviewTokenUsage:
     """A class to represent the token usage of an interview."""
@@ -24,7 +25,7 @@ class InterviewTokenUsage:
         >>> usage3 = usage1 + usage2
         """
         if not isinstance(other, InterviewTokenUsage):
-            raise ValueError(f"Can't add {type(other)} to InterviewTokenSummary")
+            raise TokenUsageError(f"Can't add {type(other)} to InterviewTokenSummary")
         return InterviewTokenUsage(
             new_token_usage=self.new_token_usage + other.new_token_usage,
             cached_token_usage=self.cached_token_usage + other.cached_token_usage,

edsl/tokens/token_usage.py

@@ -1,4 +1,5 @@
-from edsl.enums import TokenPricing
+from ..enums import TokenPricing
+from .exceptions import TokenUsageError, TokenCostError
 
 
 class TokenUsage:
@@ -15,9 +16,9 @@ class TokenUsage:
 
     def __add__(self, other):
         if not isinstance(other, TokenUsage):
-            raise ValueError(f"Can't add {type(other)} to InterviewTokenUsage")
+            raise TokenUsageError(f"Can't add {type(other)} to InterviewTokenUsage")
         if self.from_cache != other.from_cache:
-            raise ValueError(f"Can't add token usages from different sources")
+            raise TokenUsageError("Can't add token usages from different sources")
         return TokenUsage(
             from_cache=self.from_cache,
             prompt_tokens=self.prompt_tokens + other.prompt_tokens,

edsl/utilities/__init__.py

@@ -26,5 +26,25 @@ from .is_valid_variable_name import is_valid_variable_name
 from .naming_utilities import sanitize_string
 
 # Interface module - note: print_results_long is imported directly in results.py
-from .interface import print_results_long
+
+__all__ = [
+    "TemplateLoader",
+    "PrettyList",
+    "create_restricted_function",
+    "remove_edsl_version",
+    "extract_variable_names",
+    "clean_json",
+    "dict_hash",
+    "hash_value",
+    "repair_json",
+    "create_valid_var_name",
+    "random_string",
+    "shorten_string", 
+    "is_gzipped",
+    "sync_wrapper",
+    "jupyter_nb_handler",
+    "is_notebook",
+    "is_valid_variable_name",
+    "sanitize_string"
+]
 

edsl/utilities/decorators.py

@@ -1,11 +1,10 @@
 import functools
 import asyncio
 import nest_asyncio
+from edsl import __version__ as edsl_version
 
 nest_asyncio.apply()
 
-from edsl import __version__ as edsl_version
-
 
 def add_edsl_version(func):
     """

edsl/utilities/markdown_to_docx.py

@@ -69,7 +69,7 @@ class MarkdownToDocx:
 
         try:
             # Run pandoc command
-            result = subprocess.run(
+            subprocess.run(
                 cmd,
                 input=self.markdown_content,
                 text=True,
@@ -115,7 +115,7 @@ class MarkdownToDocx:
             cmd = ["pandoc", "--print-default-data-file", "reference.docx"]
 
             with open(output_path, "wb") as f:
-                result = subprocess.run(cmd, stdout=f, check=True)
+                subprocess.run(cmd, stdout=f, check=True)
             return True
         except subprocess.CalledProcessError as e:
             print(f"Error creating template: {e.stderr}")

edsl/utilities/markdown_to_pdf.py

@@ -86,7 +86,7 @@ class MarkdownToPDF:
 
         try:
             # Run pandoc command
-            result = subprocess.run(
+            subprocess.run(
                 cmd,
                 input=self.markdown_content,
                 text=True,

edsl/utilities/repair_functions.py

@@ -1,5 +1,4 @@
 import json
-from .utilities import valid_json
 
 
 def extract_json_from_string(s):