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

edsl/surveys/MemoryPlan.py

@@ -1,16 +1,26 @@
+"""A survey has a memory plan that specifies what the agent should remember when answering a question."""
+
 from collections import UserDict, defaultdict
-from edsl.surveys.Memory import Memory
+from typing import Optional
 
-from edsl.prompts.Prompt import Prompt
-from edsl.surveys.DAG import DAG
+# from edsl.surveys.Memory import Memory
+# from edsl.prompts.Prompt import Prompt
+# from edsl.surveys.DAG import DAG
 
 
 class MemoryPlan(UserDict):
     """A survey has a memory plan that specifies what the agent should remember when answering a question.
-    {focal_question: [prior_questions], focal_question: [prior_questions]}
+
+    The basic structure of a memory plan is a dictionary of focal questions to memories.
+
+    {focal_question1: [prior_question1, prior_question2, ...], focal_question: [prior_question3]}
     """
 
-    def __init__(self, survey: "Survey" = None, data=None):
+    def __init__(self, survey: Optional["Survey"] = None, data: Optional[dict] = None):
+        """Initialize a memory plan.
+
+        The actual 'data' attributes of the memory plan are a dictionary of focal questions to memories.
+        """
         if survey is not None:
             self.survey = survey
             self.survey_question_names = [q.question_name for q in survey.questions]
@@ -18,24 +28,42 @@ class MemoryPlan(UserDict):
         super().__init__(data or {})
 
     @property
-    def name_to_text(self):
-        "Returns a dictionary mapping question names to question texts"
+    def name_to_text(self) -> dict:
+        """Return a dictionary mapping question names to question texts."""
         return dict(zip(self.survey_question_names, self.question_texts))
 
-    def add_question(self, question):
+    def add_question(self, question: "QuestionBase") -> None:
+        """Add a question to the survey.
+
+        :param question: A question to add to the survey
+
+        """
         self.survey_question_names.append(question.question_name)
         self.question_texts.append(question.question_text)
 
-    def check_valid_question_name(self, question_name):
-        "Make sure a passed question name is valid"
+    def _check_valid_question_name(self, question_name: str) -> None:
+        """Ensure a passed question name is valid.
+
+        :param question_name: The name of the question to check.
+
+        """
         if question_name not in self.survey_question_names:
             raise ValueError(
                 f"{question_name} is not in the survey. Current names are {self.survey_question_names}"
             )
 
-    def get_memory_prompt_fragment(self, focal_question, answers) -> "Prompt":
-        "Generates the prompt fragment"
-        self.check_valid_question_name(focal_question)
+    def get_memory_prompt_fragment(
+        self, focal_question: str, answers: dict
+    ) -> "Prompt":
+        """Generate the prompt fragment descripting that past question and answer.
+
+        :param focal_question: The current question being answered.
+        :param answers: A dictionary of question names to answers.
+
+        """
+        from edsl.prompts.Prompt import Prompt
+
+        self._check_valid_question_name(focal_question)
 
         if focal_question not in self:
             return Prompt("")
@@ -50,7 +78,7 @@ class MemoryPlan(UserDict):
         """
 
         def gen_line(question_text, answer):
-            "Returns a line of memory"
+            """Return a line of memory."""
             return f"\tQuestion: {question_text}\n\tAnswer: {answer}\n"
 
         lines = [gen_line(*pair) for pair in q_and_a_pairs]
@@ -61,16 +89,41 @@ class MemoryPlan(UserDict):
         else:
             return Prompt("")
 
-    def check_order(self, focal_question, prior_question):
+    def _check_order(self, focal_question: str, prior_question: str) -> None:
+        """Ensure the prior question comes before the focal question."""
         focal_index = self.survey_question_names.index(focal_question)
         prior_index = self.survey_question_names.index(prior_question)
         if focal_index <= prior_index:
             raise ValueError(f"{prior_question} must come before {focal_question}.")
 
-    def add_single_memory(self, focal_question: str, prior_question: str):
-        self.check_valid_question_name(focal_question)
-        self.check_valid_question_name(prior_question)
-        self.check_order(focal_question, prior_question)
+    def add_single_memory(self, focal_question: str, prior_question: str) -> None:
+        """Add a single memory to the memory plan.
+
+        :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.
+
+        >>> 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']
+
+        >>> 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']
+        """
+        self._check_valid_question_name(focal_question)
+        self._check_valid_question_name(prior_question)
+        self._check_order(focal_question, prior_question)
+        from edsl.surveys.Memory import Memory
 
         if focal_question not in self:
             memory = Memory()
@@ -79,27 +132,53 @@ class MemoryPlan(UserDict):
         else:
             self[focal_question].add_prior_question(prior_question)
 
-    def add_memory_collection(self, focal_question, prior_questions: list[str]):
+    def add_memory_collection(
+        self, focal_question: str, prior_questions: list[str]
+    ) -> None:
+        """Add a collection of prior questions to the memory plan.
+
+        :param focal_question: The current question being answered.
+        :param prior_questions: A list of questions that were answered before the focal question that should be remembered.
+        """
         for question in prior_questions:
             self.add_single_memory(focal_question, question)
 
-    def to_dict(self):
+    def to_dict(self, add_edsl_version=True) -> dict:
+        """Serialize the memory plan to a dictionary.
+
+        >>> mp = MemoryPlan.example()
+        >>> mp.to_dict()
+        {'survey_question_names': ['q0', 'q1', 'q2'], 'survey_question_texts': ['Do you like school?', 'Why not?', 'Why?'], 'data': {'q1': {'prior_questions': ['q0']}}}
+        """
+        newdata = {}
+        for question_name, memory in self.items():
+            newdata[question_name] = memory.to_dict()
+
         return {
             "survey_question_names": self.survey_question_names,
             "survey_question_texts": self.question_texts,
-            "data": {k: v.to_dict() for k, v in self.items()},
+            "data": newdata,
         }
 
     @classmethod
-    def from_dict(cls, data):
-        # we avoid serializing the survey
-        memory_plan = cls(survey=None, data=data["data"])
+    def from_dict(cls, data) -> "MemoryPlan":
+        """Deserialize a memory plan from a dictionary."""
+        from edsl.surveys.Memory import Memory
+
+        newdata = {}
+        for question_name, memory in data["data"].items():
+            newdata[question_name] = Memory.from_dict(memory)
+
+        memory_plan = cls(survey=None, data=newdata)
         memory_plan.survey_question_names = data["survey_question_names"]
         memory_plan.question_texts = data["survey_question_texts"]
-        # memory_plan.data = data
         return memory_plan
 
-    def _indexify(self, d):
+    def _indexify(self, d: dict):
+        """Convert a dictionary of question names to a dictionary of question indices.
+
+        :param d: A dictionary of question names to indices.
+        """
         new_d = {}
         for k, v in d.items():
             key = self.survey_question_names.index(k)
@@ -108,16 +187,58 @@ class MemoryPlan(UserDict):
         return new_d
 
     @property
-    def dag(self):
+    def dag(self) -> "DAG":
+        """Return a directed acyclic graph of the memory plan.
+
+        >>> mp = MemoryPlan.example()
+        >>> mp.dag
+        {1: {0}}
+        """
+        from edsl.surveys.DAG import DAG
+
         d = defaultdict(set)
-        "Returns a directed acyclic graph of the memory plan"
         for focal_question, memory in self.items():
             for prior_question in memory:
                 d[focal_question].add(prior_question)
-                # if focal_question not in d:
-                #     d[focal_question] = set({prior_question})
-                # else:
-                #     d[focal_question].add(prior_question)
-        # focal_index = self.survey_question_names.index(focal_question)
-        # prior_index = self.survey_question_names.index(prior_question)
         return DAG(self._indexify(d))
+
+    @classmethod
+    def example(cls):
+        """Return an example memory plan."""
+        from edsl import Survey
+
+        mp = cls(survey=Survey.example())
+        mp.add_single_memory("q1", "q0")
+        return mp
+
+    def remove_question(self, question_name: str) -> None:
+        """Remove a question from the memory plan.
+
+        :param question_name: The name of the question to remove.
+        """
+        self._check_valid_question_name(question_name)
+
+        # Remove the question from survey_question_names and question_texts
+        index = self.survey_question_names.index(question_name)
+        self.survey_question_names.pop(index)
+        self.question_texts.pop(index)
+
+        # Remove the question from the memory plan if it's a focal question
+        self.pop(question_name, None)
+
+        # Remove the question from all memories where it appears as a prior question
+        for focal_question, memory in self.items():
+            memory.remove_prior_question(question_name)
+
+        # Update the DAG
+        self.dag.remove_node(index)
+
+    def remove_prior_question(self, question_name: str) -> None:
+        """Remove a prior question from the memory."""
+        self.prior_questions = [q for q in self.prior_questions if q != question_name]
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/surveys/Rule.py

@@ -1,9 +1,33 @@
+"""The Rule class defines a rule for determining the next question presented to an agent.
+
+The key component is an expression specifiying the logic of the rule, which can include any combination of logical operators ('and', 'or', 'not'), e.g.:
+
+.. code-block:: python
+
+    "q1 == 'yes' or q2 == 'no'"
+
+The expression must be about questions "before" the current question.
+
+Only one rule should apply at each priority level.
+If multiple rules apply, the one with the highest priority is used. 
+If there are conflicting rules, an exception is raised.
+
+If no rule is specified, the next question is given as the default.
+When a question is added with index, it is always given a rule the next question is index + 1, but
+with a low (-1) priority.
+"""
+
 import ast
-from collections import namedtuple
-from rich import print
-from simpleeval import EvalWithCompoundTypes
+import random
 from typing import Any, Union, List
-from edsl.exceptions import (
+
+
+# from rich import print
+from simpleeval import EvalWithCompoundTypes
+
+from edsl.exceptions.surveys import SurveyError
+
+from edsl.exceptions.surveys import (
     SurveyRuleCannotEvaluateError,
     SurveyRuleCollectionHasNoRulesAtNodeError,
     SurveyRuleRefersToFutureStateError,
@@ -13,48 +37,39 @@ from edsl.exceptions import (
 )
 from edsl.surveys.base import EndOfSurvey
 from edsl.utilities.ast_utilities import extract_variable_names
-from edsl.utilities.interface import print_table_with_rich
+from edsl.utilities.remove_edsl_version import remove_edsl_version
 
 
-class Rule:
-    """
-    The class defines a a "rule" for what question an agent should be presented next.
+class QuestionIndex:
+    def __set_name__(self, owner, name):
+        self.name = f"_{name}"
 
-    The key structure of a rule is:
-    - current_q: The question at which the rule is potentially applied (an index)
-    - expression: A string that if evaluates to true, then next_q (an index) is next
-    - next_q: The question is true
-    - priority: an integer that determines which rule is applied if multiple rules apply
+    def __get__(self, obj, objtype=None):
+        return getattr(obj, self.name)
 
-    The key part of this is expression specifiying the skip logic.
-    An expression can be any combination of and/or/not/parenthesis etc.
+    def __set__(self, obj, value):
+        if not isinstance(value, (int, EndOfSurvey.__class__)):
+            raise SurveyError(f"{self.name} must be an integer or EndOfSurvey")
+        if self.name == "_next_q" and isinstance(value, int):
+            current_q = getattr(obj, "_current_q")
+            if value <= current_q:
+                raise SurveyError("next_q must be greater than current_q")
+        setattr(obj, self.name, value)
 
-            "q1 == 'yes' or q2 == 'no'"
 
-    The skip-logic expression has to be about questions "before" the current node.
-    To determine the value of the expression, we currently, we use eval.
-    Eventually, we'll use the AST of the expression to make it safer.
-
-    If multiple rules apply, the one with the highest priority is used.
-    This is to deal with the fact that when we create a survey, we give the
-    next question as the default.
-    So when a question is added with index, it is always given a rule the next question is index + 1, but
-    given a low (-1) priority.
+class Rule:
+    """The Rule class defines a "rule" for determining the next question presented to an agent."""
 
-    Only one rule should apply at each priority level.
-    If there are conflicting rules, an exception is raised.
-    Ideally, we'd have a way to resolve this ex ante, perhaps to traversing the implied tree
-    each time a rule is added, but for now, we'll let the error emerge at run-time.
+    current_q = QuestionIndex()
+    next_q = QuestionIndex()
 
-    ## Not implemented but nice to have
-    We could potentially use the question pydantic models to check for rule conflicts, as
-    they define the potential trees through a survey.
+    # Not implemented but nice to have:
+    # We could potentially use the question pydantic models to check for rule conflicts, as
+    # they define the potential trees through a survey.
 
-    We could also use the AST to check for conflicts by inspecting the types of a rule.
-    For example, if we know the answer to a question is a string, we could check that
-    the expression only contains string comparisons.
-    This would be a lot of work.
-    """
+    # We could also use the AST to check for conflicts by inspecting the types of a rule.
+    # For example, if we know the answer to a question is a string, we could check that
+    # the expression only contains string comparisons.
 
     def __init__(
         self,
@@ -63,24 +78,41 @@ class Rule:
         next_q: Union[int, EndOfSurvey.__class__],
         question_name_to_index: dict[str, int],
         priority: int,
+        before_rule: bool = False,
     ):
-        """Questions are represented by int indices."""
+        """Represent a rule for determining the next question presented to an agent.
 
+        Questions are represented by int indices.
+
+        :param current_q: The question at which the rule is potentially applied.
+        :param expression: A string that evaluates to true or false. If true, then next_q is next.
+        :param next_q: The next question if the expression is true.
+        :param question_name_to_index: A dictionary mapping question names to indices.
+        :param priority: An integer that determines which rule is applied, if multiple rules apply.
+        """
         self.current_q = current_q
         self.expression = expression
         self.next_q = next_q
-        self.priority = priority
         self.question_name_to_index = question_name_to_index
+        self.priority = priority
+        self.before_rule = before_rule
 
-        if not next_q == EndOfSurvey and current_q > next_q:
-            raise SurveyRuleSendsYouBackwardsError
+        if not self.next_q == EndOfSurvey:
+            if self.next_q <= self.current_q:
+                raise SurveyRuleSendsYouBackwardsError
 
-        # get the AST for the expression - used to extract
-        # the variables referenced in the expression
+        if not self.next_q == EndOfSurvey and self.current_q > self.next_q:
+            raise SurveyRuleSendsYouBackwardsError(
+                f"current_q: {self.current_q}, next_q: {self.next_q}"
+            )
+
+        # get the AST for the expression - used to extract the variables referenced in the expression
         try:
             self.ast_tree = ast.parse(self.expression)
         except SyntaxError:
-            raise SurveyRuleSkipLogicSyntaxError
+            raise SurveyRuleSkipLogicSyntaxError(
+                f"The expression {self.expression} is not valid Python syntax."
+            )
 
         # get the names of the variables in the expression
         # e.g., q1 == 'yes' -> ['q1']
@@ -90,91 +122,206 @@ class Rule:
         try:
             assert all([q in question_name_to_index for q in extracted_question_names])
         except AssertionError:
-            raise SurveyRuleReferenceInRuleToUnknownQuestionError
+            pass
+            # import warnings
+            # warnings.warn(f"There is an extracted field in the expression that is not a known question. It could be a scenario variable. That's fine! But it also could be a typo or mistake.")
+            # print(f"Question name to index: {question_name_to_index}")
+            # print(f"Extracted question names: {extracted_question_names}")
+            # raise SurveyRuleReferenceInRuleToUnknownQuestionError
 
         # get the indices of the questions mentioned in the expression
         self.named_questions_by_index = [
-            question_name_to_index[q] for q in extracted_question_names
+            question_name_to_index[q]
+            for q in extracted_question_names
+            if q in question_name_to_index
         ]
 
         # A rule should only refer to questions that have already been asked.
         # so the named questions in the expression should not be higher than the current question
         if self.named_questions_by_index:
             if max(self.named_questions_by_index) > self.current_q:
+                print(
+                    "A rule refers to a future question, the answer to which would not be available here."
+                )
                 raise SurveyRuleRefersToFutureStateError
 
-    def to_dict(self):
+    def _checks(self):
+        pass
+
+    def to_dict(self, add_edsl_version=True):
+        """Convert the rule to a dictionary for serialization.
+
+        >>> r = Rule.example()
+        >>> r.to_dict()
+        {'current_q': 1, 'expression': "q1 == 'yes'", 'next_q': 2, 'priority': 0, 'question_name_to_index': {'q1': 1}, 'before_rule': False}
+        """
         return {
             "current_q": self.current_q,
             "expression": self.expression,
             "next_q": "EndOfSurvey" if self.next_q == EndOfSurvey else self.next_q,
             "priority": self.priority,
             "question_name_to_index": self.question_name_to_index,
+            "before_rule": self.before_rule,
         }
 
     @classmethod
+    @remove_edsl_version
     def from_dict(self, rule_dict):
+        """Create a rule from a dictionary."""
         if rule_dict["next_q"] == "EndOfSurvey":
             rule_dict["next_q"] = EndOfSurvey
 
+        if "before_rule" not in rule_dict:
+            rule_dict["before_rule"] = False
+
         return Rule(
             current_q=rule_dict["current_q"],
             expression=rule_dict["expression"],
             next_q=rule_dict["next_q"],
             priority=rule_dict["priority"],
             question_name_to_index=rule_dict["question_name_to_index"],
+            before_rule=rule_dict["before_rule"],
         )
 
     def __repr__(self):
-        return f'Rule(current_q={self.current_q}, expression="{self.expression}", next_q={self.next_q}, priority={self.priority}, question_name_to_index={self.question_name_to_index})'
+        """Pretty-print the rule."""
+        return f'Rule(current_q={self.current_q}, expression="{self.expression}", next_q={self.next_q}, priority={self.priority}, question_name_to_index={self.question_name_to_index}, before_rule={self.before_rule})'
 
     def __str__(self):
+        """Return a string representation of the rule."""
         return self.__repr__()
 
     @property
     def question_index_to_name(self):
-        """Reverses the dictionary do we can look up questions by name"""
+        """Reverse the dictionary do we can look up questions by name.
+
+        >>> r = Rule.example()
+        >>> r.question_index_to_name
+        {1: 'q1'}
+
+        """
         return {v: k for k, v in self.question_name_to_index.items()}
 
     def show_ast_tree(self):
-        """Pretty-prints the AST tree to the terminal"""
+        """Pretty-print the AST tree to the terminal.
+
+        >>> r = Rule.example()
+        >>> r.show_ast_tree()
+        Module(...)
+        """
         print(
             ast.dump(
                 self.ast_tree, annotate_fields=True, indent=4, include_attributes=True
             )
         )
 
-    def evaluate(self, answers: dict[int, Any]):
-        """
-        Computes the value of the expression, given a dictionary of known questions answers.
+    @staticmethod
+    def _prepare_replacement(current_info_env: dict[int, Any]):
+        d = {}
+        for var, value in current_info_env.items():
+            if isinstance(value, str):
+                replacement = f"'{value}'"
+            else:
+                replacement = str(value)
+            d[var] = replacement
+        return d
+
+    def evaluate(self, current_info_env: dict[int, Any]):
+        """Compute the value of the expression, given a dictionary of known questions answers.
+
+        :param current_info_env: A dictionary mapping question, scenario, and agent names to their values.
+
         If the expression cannot be evaluated, it raises a CannotEvaluate exception.
+
+        >>> r = Rule.example()
+        >>> r.evaluate({'q1' : 'yes'})
+        True
+        >>> r.evaluate({'q1' : 'no'})
+        False
+
+        >>> r = Rule.example(jinja2=True)
+        >>> r.evaluate({'q1' : 'yes'})
+        True
+
+        >>> r = Rule.example(jinja2=True)
+        >>> r.evaluate({'q1' : 'This is q1'})
+        False
+
+        >>> r = Rule.example(jinja2=False, bad = True)
+        >>> r.evaluate({'q1' : 'yes'})
+        Traceback (most recent call last):
+        ...
+        edsl.exceptions.surveys.SurveyRuleCannotEvaluateError...
         """
+        from jinja2 import Template
+
+        def substitute_in_answers(expression, current_info_env):
+            """Take the dictionary of answers and substitute them into the expression."""
 
-        def substitute_in_answers(expression, answers):
-            "Take the dictionary of answers and substitute them into the expression"
-            for var, value in answers.items():
-                # If it's a string, add quotes; otherwise, just convert to string
-                if isinstance(value, str):
-                    replacement = f"'{value}'"
-                else:
-                    replacement = str(value)
+            current_info = self._prepare_replacement(current_info_env)
 
-                expression = expression.replace(var, replacement)
-            return expression
+            if "{{" in expression and "}}" in expression:
+                template_expression = Template(self.expression)
+                to_evaluate = template_expression.render(current_info)
+            else:
+                # import warnings
+                # import textwrap
+                # warnings.warn(textwrap.dedent("""\
+                # The expression is not a Jinja2 template with {{ }}. This is not recommended.
+                # You can re-write your expression say "q1 == 'yes'" as "{{ q1 }} == 'yes'".
+                # """))
+                to_evaluate = expression
+                for var, value in current_info.items():
+                    to_evaluate = to_evaluate.replace(var, value)
+
+            return to_evaluate
 
         try:
-            to_evaluate = substitute_in_answers(self.expression, answers)
-            return EvalWithCompoundTypes().eval(to_evaluate)
+            to_evaluate = substitute_in_answers(self.expression, current_info_env)
         except Exception as e:
-            print(f"Exception in evaluation: {e}")
-            raise SurveyRuleCannotEvaluateError
+            msg = f"""Exception in evaluation: {e}. The expression is: {self.expression}. The current info env trying to substitute in is: {current_info_env}. After the substition, the expression was: {to_evaluate}."""
+            raise SurveyRuleCannotEvaluateError(msg)
+
+        random_functions = {
+            "randint": random.randint,
+            "choice": random.choice,
+            "random": random.random,
+            "uniform": random.uniform,
+            # Add any other random functions you want to allow
+        }
+
+        try:
+            return EvalWithCompoundTypes(functions=random_functions).eval(to_evaluate)
+        except Exception as e:
+            msg = f"""Exception in evaluation: {e}. The expression is: {self.expression}. The current info env trying to substitute in is: {current_info_env}. After the substition, the expression was: {to_evaluate}."""
+            raise SurveyRuleCannotEvaluateError(msg)
+
+    @classmethod
+    def example(cls, jinja2=False, bad=False):
+        if jinja2:
+            # a rule written in jinja2 style with {{ }}
+            expression = "{{ q1 }} == 'yes'"
+        else:
+            expression = "q1 == 'yes'"
+
+        if bad and jinja2:
+            # a rule written in jinja2 style with {{ }} but with a 'bad' expression
+            expression = "{{ q1 }} == 'This is q1'"
+
+        if bad and not jinja2:
+            expression = "q1 == 'This is q1'"
+
+        r = Rule(
+            current_q=1,
+            expression=expression,
+            next_q=2,
+            question_name_to_index={"q1": 1},
+            priority=0,
+        )
+        return r
 
 
 if __name__ == "__main__":
-    r = Rule(
-        current_q=1,
-        expression="q1 == 'yes'",
-        next_q=2,
-        question_name_to_index={"q1": 1},
-        priority=0,
-    )
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)