edsl 0.1.33.dev2__py3-none-any.whl → 0.1.34__py3-none-any.whl

edsl/agents/{PromptConstructionMixin.py → PromptConstructor.py}

@@ -1,145 +1,35 @@
 from __future__ import annotations
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, Set
 from collections import UserList
+import pdb
 
-# from functools import reduce
-from edsl.prompts.Prompt import Prompt
+from jinja2 import Environment, meta
 
-# from edsl.utilities.decorators import sync_wrapper, jupyter_nb_handler
+from edsl.prompts.Prompt import Prompt
+from edsl.data_transfer_models import ImageInfo
 from edsl.prompts.registry import get_classes as prompt_lookup
 from edsl.exceptions import QuestionScenarioRenderError
 
-import enum
-
-
-class PromptComponent(enum.Enum):
-    AGENT_INSTRUCTIONS = "agent_instructions"
-    AGENT_PERSONA = "agent_persona"
-    QUESTION_INSTRUCTIONS = "question_instructions"
-    PRIOR_QUESTION_MEMORY = "prior_question_memory"
-
-
-class PromptList(UserList):
-    separator = Prompt(" ")
-
-    def reduce(self):
-        """Reduce the list of prompts to a single prompt.
-
-        >>> p = PromptList([Prompt("You are a happy-go lucky agent."), Prompt("You are an agent with the following persona: {'age': 22, 'hair': 'brown', 'height': 5.5}")])
-        >>> p.reduce()
-        Prompt(text=\"""You are a happy-go lucky agent. You are an agent with the following persona: {'age': 22, 'hair': 'brown', 'height': 5.5}\""")
-
-        """
-        p = self[0]
-        for prompt in self[1:]:
-            if len(prompt) > 0:
-                p = p + self.separator + prompt
-        return p
-
-
-class PromptPlan:
-    """A plan for constructing prompts for the LLM call.
-    Every prompt plan has a user prompt order and a system prompt order.
-    It must contain each of the values in the PromptComponent enum.
-
-
-    >>> p = PromptPlan(user_prompt_order=(PromptComponent.AGENT_INSTRUCTIONS, PromptComponent.AGENT_PERSONA),system_prompt_order=(PromptComponent.QUESTION_INSTRUCTIONS, PromptComponent.PRIOR_QUESTION_MEMORY))
-    >>> p._is_valid_plan()
-    True
+from edsl.agents.prompt_helpers import PromptComponent, PromptList, PromptPlan
 
-    >>> p.arrange_components(agent_instructions=1, agent_persona=2, question_instructions=3, prior_question_memory=4)
-    {'user_prompt': ..., 'system_prompt': ...}
-
-    >>> p = PromptPlan(user_prompt_order=("agent_instructions", ), system_prompt_order=("question_instructions", "prior_question_memory"))
-    Traceback (most recent call last):
-    ...
-    ValueError: Invalid plan: must contain each value of PromptComponent exactly once.
 
+def get_jinja2_variables(template_str: str) -> Set[str]:
     """
+    Extracts all variable names from a Jinja2 template using Jinja2's built-in parsing.
 
-    def __init__(
-        self,
-        user_prompt_order: Optional[tuple] = None,
-        system_prompt_order: Optional[tuple] = None,
-    ):
-        """Initialize the PromptPlan."""
-
-        if user_prompt_order is None:
-            user_prompt_order = (
-                PromptComponent.QUESTION_INSTRUCTIONS,
-                PromptComponent.PRIOR_QUESTION_MEMORY,
-            )
-        if system_prompt_order is None:
-            system_prompt_order = (
-                PromptComponent.AGENT_INSTRUCTIONS,
-                PromptComponent.AGENT_PERSONA,
-            )
-
-        # very commmon way to screw this up given how python treats single strings as iterables
-        if isinstance(user_prompt_order, str):
-            user_prompt_order = (user_prompt_order,)
-
-        if isinstance(system_prompt_order, str):
-            system_prompt_order = (system_prompt_order,)
-
-        if not isinstance(user_prompt_order, tuple):
-            raise TypeError(
-                f"Expected a tuple, but got {type(user_prompt_order).__name__}"
-            )
-
-        if not isinstance(system_prompt_order, tuple):
-            raise TypeError(
-                f"Expected a tuple, but got {type(system_prompt_order).__name__}"
-            )
-
-        self.user_prompt_order = self._convert_to_enum(user_prompt_order)
-        self.system_prompt_order = self._convert_to_enum(system_prompt_order)
-        if not self._is_valid_plan():
-            raise ValueError(
-                "Invalid plan: must contain each value of PromptComponent exactly once."
-            )
-
-    def _convert_to_enum(self, prompt_order: tuple):
-        """Convert string names to PromptComponent enum values."""
-        return tuple(
-            PromptComponent(component) if isinstance(component, str) else component
-            for component in prompt_order
-        )
-
-    def _is_valid_plan(self):
-        """Check if the plan is valid."""
-        combined = self.user_prompt_order + self.system_prompt_order
-        return set(combined) == set(PromptComponent)
-
-    def arrange_components(self, **kwargs) -> Dict[PromptComponent, Prompt]:
-        """Arrange the components in the order specified by the plan."""
-        # check is valid components passed
-        component_strings = set([pc.value for pc in PromptComponent])
-        if not set(kwargs.keys()) == component_strings:
-            raise ValueError(
-                f"Invalid components passed: {set(kwargs.keys())} but expected {PromptComponent}"
-            )
-
-        user_prompt = PromptList(
-            [kwargs[component.value] for component in self.user_prompt_order]
-        )
-        system_prompt = PromptList(
-            [kwargs[component.value] for component in self.system_prompt_order]
-        )
-        return {"user_prompt": user_prompt, "system_prompt": system_prompt}
-
-    def get_prompts(self, **kwargs) -> Dict[str, Prompt]:
-        """Get both prompts for the LLM call."""
-        prompts = self.arrange_components(**kwargs)
-        return {
-            "user_prompt": prompts["user_prompt"].reduce(),
-            "system_prompt": prompts["system_prompt"].reduce(),
-        }
+    Args:
+    template_str (str): The Jinja2 template string
 
+    Returns:
+    Set[str]: A set of variable names found in the template
+    """
+    env = Environment()
+    ast = env.parse(template_str)
+    return meta.find_undeclared_variables(ast)
 
-class PromptConstructorMixin:
-    """Mixin for constructing prompts for the LLM call.
 
+class PromptConstructor:
+    """
     The pieces of a prompt are:
     - The agent instructions - "You are answering questions as if you were a human. Do not break character."
     - The persona prompt - "You are an agent with the following persona: {'age': 22, 'hair': 'brown', 'height': 5.5}"
@@ -149,16 +39,42 @@ class PromptConstructorMixin:
     This is mixed into the Invigilator class.
     """
 
-    prompt_plan = PromptPlan()
+    def __init__(self, invigilator):
+        self.invigilator = invigilator
+        self.agent = invigilator.agent
+        self.question = invigilator.question
+        self.scenario = invigilator.scenario
+        self.survey = invigilator.survey
+        self.model = invigilator.model
+        self.current_answers = invigilator.current_answers
+        self.memory_plan = invigilator.memory_plan
+        self.prompt_plan = PromptPlan()
+
+    @property
+    def scenario_file_keys(self) -> list:
+        """We need to find all the keys in the scenario that refer to FileStore objects.
+        These will be used to append to the prompt a list of files that are part of the scenario.
+        """
+        from edsl.scenarios.FileStore import FileStore
+
+        file_entries = []
+        for key, value in self.scenario.items():
+            if isinstance(value, FileStore):
+                file_entries.append(key)
+        return file_entries
 
     @property
     def agent_instructions_prompt(self) -> Prompt:
         """
         >>> from edsl.agents.InvigilatorBase import InvigilatorBase
         >>> i = InvigilatorBase.example()
-        >>> i.agent_instructions_prompt
+        >>> i.prompt_constructor.agent_instructions_prompt
         Prompt(text=\"""You are answering questions as if you were a human. Do not break character.\""")
         """
+        from edsl import Agent
+
+        if self.agent == Agent():  # if agent is empty, then return an empty prompt
+            return Prompt(text="")
         if not hasattr(self, "_agent_instructions_prompt"):
             applicable_prompts = prompt_lookup(
                 component_type="agent_instructions",
@@ -176,47 +92,56 @@ class PromptConstructorMixin:
         """
         >>> from edsl.agents.InvigilatorBase import InvigilatorBase
         >>> i = InvigilatorBase.example()
-        >>> i.agent_persona_prompt
+        >>> i.prompt_constructor.agent_persona_prompt
         Prompt(text=\"""You are an agent with the following persona:
         {'age': 22, 'hair': 'brown', 'height': 5.5}\""")
 
         """
-        if not hasattr(self, "_agent_persona_prompt"):
-            if not hasattr(self.agent, "agent_persona"):
-                applicable_prompts = prompt_lookup(
-                    component_type="agent_persona",
-                    model=self.model.model,
-                )
-                persona_prompt_template = applicable_prompts[0]()
-            else:
-                persona_prompt_template = self.agent.agent_persona
-
-            # TODO: This multiple passing of agent traits - not sure if it is necessary. Not harmful.
-            if undefined := persona_prompt_template.undefined_template_variables(
-                self.agent.traits
-                | {"traits": self.agent.traits}
-                | {"codebook": self.agent.codebook}
-                | {"traits": self.agent.traits}
-            ):
-                raise QuestionScenarioRenderError(
-                    f"Agent persona still has variables that were not rendered: {undefined}"
-                )
+        from edsl import Agent
+
+        if hasattr(self, "_agent_persona_prompt"):
+            return self._agent_persona_prompt
 
-            persona_prompt = persona_prompt_template.render(
-                self.agent.traits | {"traits": self.agent.traits},
-                codebook=self.agent.codebook,
-                traits=self.agent.traits,
+        if self.agent == Agent():  # if agent is empty, then return an empty prompt
+            return Prompt(text="")
+
+        if not hasattr(self.agent, "agent_persona"):
+            applicable_prompts = prompt_lookup(
+                component_type="agent_persona",
+                model=self.model.model,
             )
-            if persona_prompt.has_variables:
-                raise QuestionScenarioRenderError(
-                    "Agent persona still has variables that were not rendered."
-                )
-            self._agent_persona_prompt = persona_prompt
+            persona_prompt_template = applicable_prompts[0]()
+        else:
+            persona_prompt_template = self.agent.agent_persona
+
+        # TODO: This multiple passing of agent traits - not sure if it is necessary. Not harmful.
+        template_parameter_dictionary = (
+            self.agent.traits
+            | {"traits": self.agent.traits}
+            | {"codebook": self.agent.codebook}
+            | {"traits": self.agent.traits}
+        )
+
+        if undefined := persona_prompt_template.undefined_template_variables(
+            template_parameter_dictionary
+        ):
+            raise QuestionScenarioRenderError(
+                f"Agent persona still has variables that were not rendered: {undefined}"
+            )
+
+        persona_prompt = persona_prompt_template.render(template_parameter_dictionary)
+        if persona_prompt.has_variables:
+            raise QuestionScenarioRenderError(
+                "Agent persona still has variables that were not rendered."
+            )
+
+        self._agent_persona_prompt = persona_prompt
 
         return self._agent_persona_prompt
 
     def prior_answers_dict(self) -> dict:
         d = self.survey.question_names_to_questions()
+        # This attaches the answer to the question
         for question, answer in self.current_answers.items():
             if question in d:
                 d[question].answer = answer
@@ -226,51 +151,70 @@ class PromptConstructorMixin:
                     d[new_question].comment = answer
         return d
 
+    @property
+    def question_file_keys(self):
+        raw_question_text = self.question.question_text
+        variables = get_jinja2_variables(raw_question_text)
+        question_file_keys = []
+        for var in variables:
+            if var in self.scenario_file_keys:
+                question_file_keys.append(var)
+        return question_file_keys
+
     @property
     def question_instructions_prompt(self) -> Prompt:
         """
         >>> from edsl.agents.InvigilatorBase import InvigilatorBase
         >>> i = InvigilatorBase.example()
-        >>> i.question_instructions_prompt
+        >>> i.prompt_constructor.question_instructions_prompt
         Prompt(text=\"""...
         ...
         """
+        # The user might have passed a custom prompt, which would be stored in _question_instructions_prompt
         if not hasattr(self, "_question_instructions_prompt"):
+            # Gets the instructions for the question - this is how the question should be answered
             question_prompt = self.question.get_instructions(model=self.model.model)
 
-            # TODO: Try to populate the answers in the question object if they are available
-            # d = self.survey.question_names_to_questions()
-            # for question, answer in self.current_answers.items():
-            #     if question in d:
-            #         d[question].answer = answer
-            #     else:
-            #         # adds a comment to the question
-            #         if (new_question := question.split("_comment")[0]) in d:
-            #             d[new_question].comment = answer
-
+            # Get the data for the question - this is a dictionary of the question data
+            # e.g., {'question_text': 'Do you like school?', 'question_name': 'q0', 'question_options': ['yes', 'no']}
             question_data = self.question.data.copy()
 
             # check to see if the question_options is actually a string
-            # This is used when the user is using the question_options as a variable from a sceario
+            # This is used when the user is using the question_options as a variable from a scenario
             # if "question_options" in question_data:
             if isinstance(self.question.data.get("question_options", None), str):
-                from jinja2 import Environment, meta
-
                 env = Environment()
                 parsed_content = env.parse(self.question.data["question_options"])
                 question_option_key = list(
                     meta.find_undeclared_variables(parsed_content)
                 )[0]
 
+                # look to see if the question_option_key is in the scenario
                 if isinstance(
                     question_options := self.scenario.get(question_option_key), list
                 ):
                     question_data["question_options"] = question_options
                     self.question.question_options = question_options
 
+                # might be getting it from the prior answers
+                if self.prior_answers_dict().get(question_option_key) is not None:
+                    if isinstance(
+                        question_options := self.prior_answers_dict()
+                        .get(question_option_key)
+                        .answer,
+                        list,
+                    ):
+                        question_data["question_options"] = question_options
+                        self.question.question_options = question_options
+
             replacement_dict = (
-                question_data
-                | self.scenario
+                {key: f"<see file {key}>" for key in self.scenario_file_keys}
+                | question_data
+                | {
+                    k: v
+                    for k, v in self.scenario.items()
+                    if k not in self.scenario_file_keys
+                }  # don't include images in the replacement dict
                 | self.prior_answers_dict()
                 | {"agent": self.agent}
                 | {
@@ -280,9 +224,10 @@ class PromptConstructorMixin:
                     ),
                 }
             )
-            # breakpoint()
+
             rendered_instructions = question_prompt.render(replacement_dict)
-            # breakpoint()
+
+            # is there anything left to render?
             undefined_template_variables = (
                 rendered_instructions.undefined_template_variables({})
             )
@@ -296,11 +241,14 @@ class PromptConstructorMixin:
                     )
 
             if undefined_template_variables:
+                # breakpoint()
                 raise QuestionScenarioRenderError(
                     f"Question instructions still has variables: {undefined_template_variables}."
                 )
 
-            # Check if question has an instructions
+            ####################################
+            # Check if question has instructions - these are instructions in a Survey that can apply to multiple follow-on questions
+            ####################################
             relevant_instructions = self.survey.relevant_instructions(
                 self.question.question_name
             )
@@ -329,6 +277,23 @@ class PromptConstructorMixin:
             self._prior_question_memory_prompt = memory_prompt
         return self._prior_question_memory_prompt
 
+    def create_memory_prompt(self, question_name: str) -> Prompt:
+        """Create a memory for the agent.
+
+        The returns a memory prompt for the agent.
+
+        >>> from edsl.agents.InvigilatorBase import InvigilatorBase
+        >>> i = InvigilatorBase.example()
+        >>> i.current_answers = {"q0": "Prior answer"}
+        >>> i.memory_plan.add_single_memory("q1", "q0")
+        >>> p = i.prompt_constructor.create_memory_prompt("q1")
+        >>> p.text.strip().replace("\\n", " ").replace("\\t", " ")
+        'Before the question you are now answering, you already answered the following question(s):          Question: Do you like school?  Answer: Prior answer'
+        """
+        return self.memory_plan.get_memory_prompt_fragment(
+            question_name, self.current_answers
+        )
+
     def construct_system_prompt(self) -> Prompt:
         """Construct the system prompt for the LLM call."""
         import warnings
@@ -357,15 +322,18 @@ class PromptConstructorMixin:
         >>> i.get_prompts()
         {'user_prompt': ..., 'system_prompt': ...}
         """
+        # breakpoint()
         prompts = self.prompt_plan.get_prompts(
             agent_instructions=self.agent_instructions_prompt,
             agent_persona=self.agent_persona_prompt,
             question_instructions=self.question_instructions_prompt,
             prior_question_memory=self.prior_question_memory_prompt,
         )
-
-        if hasattr(self.scenario, "has_image") and self.scenario.has_image:
-            prompts["encoded_image"] = self.scenario["encoded_image"]
+        if self.question_file_keys:
+            files_list = []
+            for key in self.question_file_keys:
+                files_list.append(self.scenario[key])
+            prompts["files_list"] = files_list
         return prompts
 
     def _get_scenario_with_image(self) -> Scenario:

edsl/agents/prompt_helpers.py

@@ -0,0 +1,129 @@
+import enum
+from typing import Dict, Optional
+from collections import UserList
+from edsl.prompts import Prompt
+
+
+class PromptComponent(enum.Enum):
+    AGENT_INSTRUCTIONS = "agent_instructions"
+    AGENT_PERSONA = "agent_persona"
+    QUESTION_INSTRUCTIONS = "question_instructions"
+    PRIOR_QUESTION_MEMORY = "prior_question_memory"
+
+
+class PromptList(UserList):
+    separator = Prompt(" ")
+
+    def reduce(self):
+        """Reduce the list of prompts to a single prompt.
+
+        >>> p = PromptList([Prompt("You are a happy-go lucky agent."), Prompt("You are an agent with the following persona: {'age': 22, 'hair': 'brown', 'height': 5.5}")])
+        >>> p.reduce()
+        Prompt(text=\"""You are a happy-go lucky agent. You are an agent with the following persona: {'age': 22, 'hair': 'brown', 'height': 5.5}\""")
+
+        """
+        p = self[0]
+        for prompt in self[1:]:
+            if len(prompt) > 0:
+                p = p + self.separator + prompt
+        return p
+
+
+class PromptPlan:
+    """A plan for constructing prompts for the LLM call.
+    Every prompt plan has a user prompt order and a system prompt order.
+    It must contain each of the values in the PromptComponent enum.
+
+
+    >>> p = PromptPlan(user_prompt_order=(PromptComponent.AGENT_INSTRUCTIONS, PromptComponent.AGENT_PERSONA),system_prompt_order=(PromptComponent.QUESTION_INSTRUCTIONS, PromptComponent.PRIOR_QUESTION_MEMORY))
+    >>> p._is_valid_plan()
+    True
+
+    >>> p.arrange_components(agent_instructions=1, agent_persona=2, question_instructions=3, prior_question_memory=4)
+    {'user_prompt': ..., 'system_prompt': ...}
+
+    >>> p = PromptPlan(user_prompt_order=("agent_instructions", ), system_prompt_order=("question_instructions", "prior_question_memory"))
+    Traceback (most recent call last):
+    ...
+    ValueError: Invalid plan: must contain each value of PromptComponent exactly once.
+
+    """
+
+    def __init__(
+        self,
+        user_prompt_order: Optional[tuple] = None,
+        system_prompt_order: Optional[tuple] = None,
+    ):
+        """Initialize the PromptPlan."""
+
+        if user_prompt_order is None:
+            user_prompt_order = (
+                PromptComponent.QUESTION_INSTRUCTIONS,
+                PromptComponent.PRIOR_QUESTION_MEMORY,
+            )
+        if system_prompt_order is None:
+            system_prompt_order = (
+                PromptComponent.AGENT_INSTRUCTIONS,
+                PromptComponent.AGENT_PERSONA,
+            )
+
+        # very commmon way to screw this up given how python treats single strings as iterables
+        if isinstance(user_prompt_order, str):
+            user_prompt_order = (user_prompt_order,)
+
+        if isinstance(system_prompt_order, str):
+            system_prompt_order = (system_prompt_order,)
+
+        if not isinstance(user_prompt_order, tuple):
+            raise TypeError(
+                f"Expected a tuple, but got {type(user_prompt_order).__name__}"
+            )
+
+        if not isinstance(system_prompt_order, tuple):
+            raise TypeError(
+                f"Expected a tuple, but got {type(system_prompt_order).__name__}"
+            )
+
+        self.user_prompt_order = self._convert_to_enum(user_prompt_order)
+        self.system_prompt_order = self._convert_to_enum(system_prompt_order)
+        if not self._is_valid_plan():
+            raise ValueError(
+                "Invalid plan: must contain each value of PromptComponent exactly once."
+            )
+
+    def _convert_to_enum(self, prompt_order: tuple):
+        """Convert string names to PromptComponent enum values."""
+        return tuple(
+            PromptComponent(component) if isinstance(component, str) else component
+            for component in prompt_order
+        )
+
+    def _is_valid_plan(self):
+        """Check if the plan is valid."""
+        combined = self.user_prompt_order + self.system_prompt_order
+        return set(combined) == set(PromptComponent)
+
+    def arrange_components(self, **kwargs) -> Dict[PromptComponent, Prompt]:
+        """Arrange the components in the order specified by the plan."""
+        # check is valid components passed
+        component_strings = set([pc.value for pc in PromptComponent])
+        if not set(kwargs.keys()) == component_strings:
+            raise ValueError(
+                f"Invalid components passed: {set(kwargs.keys())} but expected {PromptComponent}"
+            )
+
+        user_prompt = PromptList(
+            [kwargs[component.value] for component in self.user_prompt_order]
+        )
+        system_prompt = PromptList(
+            [kwargs[component.value] for component in self.system_prompt_order]
+        )
+        return {"user_prompt": user_prompt, "system_prompt": system_prompt}
+
+    def get_prompts(self, **kwargs) -> Dict[str, Prompt]:
+        """Get both prompts for the LLM call."""
+        prompts = self.arrange_components(**kwargs)
+        return {
+            "user_prompt": prompts["user_prompt"].reduce(),
+            "system_prompt": prompts["system_prompt"].reduce(),
+        }