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

edsl/agents/QuestionTemplateReplacementsBuilder.py

@@ -0,0 +1,137 @@
+from jinja2 import Environment, meta
+from typing import Any, Set, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from edsl.agents.PromptConstructor import PromptConstructor
+    from edsl.scenarios.Scenario import Scenario
+
+
+class QuestionTemplateReplacementsBuilder:
+    def __init__(self, prompt_constructor: "PromptConstructor"):
+        self.prompt_constructor = prompt_constructor
+
+    def question_file_keys(self):
+        question_text = self.prompt_constructor.question.question_text
+        file_keys = self._find_file_keys(self.prompt_constructor.scenario)
+        return self._extract_file_keys_from_question_text(question_text, file_keys)
+
+    def scenario_file_keys(self):
+        return self._find_file_keys(self.prompt_constructor.scenario)
+
+    def get_jinja2_variables(template_str: str) -> Set[str]:
+        """
+        Extracts all variable names from a Jinja2 template using Jinja2's built-in parsing.
+
+        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)
+
+    @staticmethod
+    def _find_file_keys(scenario: "Scenario") -> 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 import Scenario
+        >>> from edsl.scenarios.FileStore import FileStore
+        >>> import tempfile
+        >>> with tempfile.NamedTemporaryFile() as f:
+        ...     _ = f.write(b"Hello, world!")
+        ...     _ = f.seek(0)
+        ...     fs = FileStore(f.name)
+        ...     scenario = Scenario({"fs_file": fs, 'a': 1})
+        ...     QuestionTemplateReplacementsBuilder._find_file_keys(scenario)
+        ['fs_file']
+        """
+        from edsl.scenarios.FileStore import FileStore
+
+        file_entries = []
+        for key, value in scenario.items():
+            if isinstance(value, FileStore):
+                file_entries.append(key)
+        return file_entries
+
+    @staticmethod
+    def _extract_file_keys_from_question_text(
+        question_text: str, scenario_file_keys: list
+    ) -> list:
+        """
+        Extracts the file keys from a question text.
+
+        >>> from edsl import Scenario
+        >>> from edsl.scenarios.FileStore import FileStore
+        >>> import tempfile
+        >>> with tempfile.NamedTemporaryFile() as f:
+        ...     _ = f.write(b"Hello, world!")
+        ...     _ = f.seek(0)
+        ...     fs = FileStore(f.name)
+        ...     scenario = Scenario({"fs_file": fs, 'a': 1})
+        ...     QuestionTemplateReplacementsBuilder._extract_file_keys_from_question_text("{{ fs_file }}", ['fs_file'])
+        ['fs_file']
+        """
+        variables = QuestionTemplateReplacementsBuilder.get_jinja2_variables(
+            question_text
+        )
+        question_file_keys = []
+        for var in variables:
+            if var in scenario_file_keys:
+                question_file_keys.append(var)
+        return question_file_keys
+
+    def _scenario_replacements(self) -> dict[str, Any]:
+        # File references dictionary
+        file_refs = {key: f"<see file {key}>" for key in self.scenario_file_keys()}
+
+        # Scenario items excluding file keys
+        scenario_items = {
+            k: v
+            for k, v in self.prompt_constructor.scenario.items()
+            if k not in self.scenario_file_keys()
+        }
+        return {**file_refs, **scenario_items}
+
+    @staticmethod
+    def _question_data_replacements(
+        question: dict, question_data: dict
+    ) -> dict[str, Any]:
+        """Builds a dictionary of replacement values for rendering a prompt by combining multiple data sources.
+
+        >>> from edsl import QuestionMultipleChoice
+        >>> q = QuestionMultipleChoice(question_text="Do you like school?", question_name = "q0", question_options = ["yes", "no"])
+        >>> QuestionTemplateReplacementsBuilder._question_data_replacements(q, q.data)
+        {'use_code': False, 'include_comment': True, 'question_name': 'q0', 'question_text': 'Do you like school?', 'question_options': ['yes', 'no']}
+
+        """
+        question_settings = {
+            "use_code": getattr(question, "_use_code", True),
+            "include_comment": getattr(question, "_include_comment", False),
+        }
+        return {**question_settings, **question_data}
+
+    def build_replacement_dict(self, question_data: dict) -> dict[str, Any]:
+        """Builds a dictionary of replacement values for rendering a prompt by combining multiple data sources."""
+        rpl = {}
+        rpl["scenario"] = self._scenario_replacements()
+        rpl["question"] = self._question_data_replacements(
+            self.prompt_constructor.question, question_data
+        )
+        rpl["prior_answers"] = self.prompt_constructor.prior_answers_dict()
+        rpl["agent"] = {"agent": self.prompt_constructor.agent}
+
+        # Combine all dictionaries using dict.update() for clarity
+        replacement_dict = {}
+        for r in rpl.values():
+            replacement_dict.update(r)
+
+        return replacement_dict
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

edsl/agents/descriptors.py

@@ -1,42 +1,68 @@
+"""This module contains the descriptors used to set the attributes of the Agent class."""
+
 from typing import Dict
-from edsl.utilities.utilities import is_valid_variable_name
-from edsl.exceptions.agents import AgentNameError
+from edsl.exceptions.agents import AgentNameError, AgentTraitKeyError
+
+
+def convert_agent_name(x):
+    # potentially a numpy int64
+    import numpy as np
+
+    if isinstance(x, np.int64):
+        return int(x)
+    elif x is None:
+        return None
+    elif isinstance(x, int):
+        return x
+    else:
+        return str(x)
 
 
 class NameDescriptor:
     """ABC for something."""
 
     def __get__(self, instance, owner):
-        """"""
+        """Return the value of the attribute."""
         return instance.__dict__[self.name]
 
     def __set__(self, instance, name: str) -> None:
-        instance.__dict__[self.name] = name
+        """Set the value of the attribute."""
+        instance.__dict__[self.name] = convert_agent_name(name)
 
     def __set_name__(self, owner, name: str) -> None:
+        """Set the name of the attribute."""
         self.name = "_" + name
 
 
 class TraitsDescriptor:
-    """ABC for something."""
+    """Traits descriptor."""
 
     def __get__(self, instance, owner):
-        """"""
+        """Return the value of the attribute."""
         return instance.__dict__[self.name]
 
     def __set__(self, instance, traits_dict: Dict[str, str]) -> None:
+        """Set the value of the attribute."""
+        from edsl.utilities.utilities import is_valid_variable_name
+
         for key, value in traits_dict.items():
-            if not is_valid_variable_name(key):
-                raise AgentNameError("Trait keys must be a valid variable name!")
             if key == "name":
                 raise AgentNameError(
-                    """Trait keys cannot be 'name'!. Instead, use the 'name' attribute directly e.g., 
-                                >>> Agent(name="my_agent", traits={"trait1": "value1", "trait2": "value2"})
-                                """
+                    "Trait keys cannot be 'name'. Instead, use the 'name' attribute directly e.g.,\n"
+                    'Agent(name="my_agent", traits={"trait1": "value1", "trait2": "value2"})'
                 )
+
+            if not is_valid_variable_name(key):
+                raise AgentTraitKeyError(
+                    f"""Trait keys must be valid Python identifiers (must be alphanumeric, cannot start with a number and must use underscores instead of spaces).
+                    You passed: {key}
+                    """
+                )
+
         instance.__dict__[self.name] = traits_dict
 
     def __set_name__(self, owner, name: str) -> None:
+        """Set the name of the attribute."""
         self.name = name
 
 
@@ -44,13 +70,15 @@ class CodebookDescriptor:
     """ABC for something."""
 
     def __get__(self, instance, owner):
-        """"""
+        """Return the value of the attribute."""
         return instance.__dict__[self.name]
 
     def __set__(self, instance, codebook_dict: Dict[str, str]) -> None:
+        """Set the value of the attribute."""
         instance.__dict__[self.name] = codebook_dict
 
     def __set_name__(self, owner, name: str) -> None:
+        """Set the name of the attribute."""
         self.name = "_" + name
 
 
@@ -58,12 +86,14 @@ class InstructionDescriptor:
     """ABC for something."""
 
     def __get__(self, instance, owner):
-        """"""
+        """Return the value of the attribute."""
         return instance.__dict__[self.name]
 
     def __set__(self, instance, instruction) -> None:
+        """Set the value of the attribute."""
         instance.__dict__[self.name] = instruction
         instance.set_instructions = instruction != instance.default_instruction
 
     def __set_name__(self, owner, name: str) -> None:
+        """Set the name of the attribute."""
         self.name = "_" + name

edsl/agents/prompt_helpers.py

@@ -0,0 +1,129 @@
+import enum
+from typing import Dict, Optional
+from collections import UserList
+from edsl.prompts.Prompt 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(),
+        }

edsl/agents/question_option_processor.py

@@ -0,0 +1,172 @@
+from jinja2 import Environment, meta
+from typing import List, Optional, Union
+
+
+class QuestionOptionProcessor:
+    """
+    Class that manages the processing of question options.
+    These can be provided directly, as a template string, or fetched from prior answers or the scenario.
+    """
+
+    def __init__(self, prompt_constructor):
+        self.prompt_constructor = prompt_constructor
+
+    @staticmethod
+    def _get_default_options() -> list:
+        """Return default placeholder options."""
+        return [f"<< Option {i} - Placeholder >>" for i in range(1, 4)]
+
+    @staticmethod
+    def _parse_template_variable(template_str: str) -> str:
+        """
+        Extract the variable name from a template string.
+
+        Args:
+            template_str (str): Jinja template string
+
+        Returns:
+            str: Name of the first undefined variable in the template
+
+        >>> QuestionOptionProcessor._parse_template_variable("Here are some {{ options }}")
+        'options'
+        >>> QuestionOptionProcessor._parse_template_variable("Here are some {{ options }} and {{ other }}")
+        Traceback (most recent call last):
+        ...
+        ValueError: Multiple variables found in template string
+        >>> QuestionOptionProcessor._parse_template_variable("Here are some")
+        Traceback (most recent call last):
+        ...
+        ValueError: No variables found in template string
+        """
+        env = Environment()
+        parsed_content = env.parse(template_str)
+        undeclared_variables = list(meta.find_undeclared_variables(parsed_content))
+        if not undeclared_variables:
+            raise ValueError("No variables found in template string")
+        if len(undeclared_variables) > 1:
+            raise ValueError("Multiple variables found in template string")
+        return undeclared_variables[0]
+
+    @staticmethod
+    def _get_options_from_scenario(
+        scenario: dict, option_key: str
+    ) -> Union[list, None]:
+        """
+        Try to get options from scenario data.
+
+        >>> from edsl import Scenario
+        >>> scenario = Scenario({"options": ["Option 1", "Option 2"]})
+        >>> QuestionOptionProcessor._get_options_from_scenario(scenario, "options")
+        ['Option 1', 'Option 2']
+
+
+        Returns:
+            list | None: List of options if found in scenario, None otherwise
+        """
+        scenario_options = scenario.get(option_key)
+        return scenario_options if isinstance(scenario_options, list) else None
+
+    @staticmethod
+    def _get_options_from_prior_answers(
+        prior_answers: dict, option_key: str
+    ) -> Union[list, None]:
+        """
+        Try to get options from prior answers.
+
+        prior_answers (dict): Dictionary of prior answers
+        option_key (str): Key to look up in prior answers
+
+        >>> from edsl import QuestionList as Q
+        >>> q = Q.example()
+        >>> q.answer = ["Option 1", "Option 2"]
+        >>> prior_answers = {"options": q}
+        >>> QuestionOptionProcessor._get_options_from_prior_answers(prior_answers, "options")
+        ['Option 1', 'Option 2']
+        >>> QuestionOptionProcessor._get_options_from_prior_answers(prior_answers, "wrong_key") is None
+        True
+
+        Returns:
+            list | None: List of options if found in prior answers, None otherwise
+        """
+        prior_answer = prior_answers.get(option_key)
+        if prior_answer and hasattr(prior_answer, "answer"):
+            if isinstance(prior_answer.answer, list):
+                return prior_answer.answer
+        return None
+
+    def get_question_options(self, question_data: dict) -> list:
+        """
+        Extract and process question options from question data.
+
+        Args:
+            question_data (dict): Dictionary containing question configuration
+
+        Returns:
+            list: List of question options. Returns default placeholders if no valid options found.
+
+        >>> class MockPromptConstructor:
+        ...     pass
+        >>> mpc = MockPromptConstructor()
+        >>> from edsl import Scenario
+        >>> mpc.scenario = Scenario({"options": ["Option 1", "Option 2"]})
+        >>> processor = QuestionOptionProcessor(mpc)
+
+        The basic case where options are directly provided:
+
+        >>> question_data = {"question_options": ["Option 1", "Option 2"]}
+        >>> processor.get_question_options(question_data)
+        ['Option 1', 'Option 2']
+
+        The case where options are provided as a template string:
+
+        >>> question_data = {"question_options": "{{ options }}"}
+        >>> processor.get_question_options(question_data)
+        ['Option 1', 'Option 2']
+
+        The case where there is a templace string but it's in the prior answers:
+
+        >>> class MockQuestion:
+        ...     pass
+        >>> q0 = MockQuestion()
+        >>> q0.answer = ["Option 1", "Option 2"]
+        >>> mpc.prior_answers_dict = lambda: {'q0': q0}
+        >>> processor = QuestionOptionProcessor(mpc)
+        >>> question_data = {"question_options": "{{ q0 }}"}
+        >>> processor.get_question_options(question_data)
+        ['Option 1', 'Option 2']
+
+        The case we're no options are found:
+        >>> processor.get_question_options({"question_options": "{{ poop }}"})
+        ['<< Option 1 - Placeholder >>', '<< Option 2 - Placeholder >>', '<< Option 3 - Placeholder >>']
+
+        """
+        options_entry = question_data.get("question_options")
+
+        # If not a template string, return as is or default
+        if not isinstance(options_entry, str):
+            return options_entry if options_entry else self._get_default_options()
+
+        # Parse template to get variable name
+        option_key = self._parse_template_variable(options_entry)
+
+        # Try getting options from scenario
+        scenario_options = self._get_options_from_scenario(
+            self.prompt_constructor.scenario, option_key
+        )
+        if scenario_options:
+            return scenario_options
+
+        # Try getting options from prior answers
+        prior_answer_options = self._get_options_from_prior_answers(
+            self.prompt_constructor.prior_answers_dict(), option_key
+        )
+        if prior_answer_options:
+            return prior_answer_options
+
+        return self._get_default_options()
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

edsl/auto/AutoStudy.py

@@ -0,0 +1,130 @@
+from typing import Optional, TYPE_CHECKING
+
+from edsl import Model
+from edsl.auto.StageQuestions import StageQuestions
+from edsl.auto.StagePersona import StagePersona
+from edsl.auto.StagePersonaDimensions import StagePersonaDimensions
+from edsl.auto.StagePersonaDimensionValues import StagePersonaDimensionValues
+from edsl.auto.StagePersonaDimensionValueRanges import (
+    StagePersonaDimensionValueRanges,
+)
+from edsl.auto.StageLabelQuestions import StageLabelQuestions
+from edsl.auto.StageGenerateSurvey import StageGenerateSurvey
+
+from edsl.auto.utilities import agent_generator, create_agents, gen_pipeline
+
+if TYPE_CHECKING:
+    from edsl.surveys.Survey import Survey
+    from edsl.agents.AgentList import AgentList
+
+
+class AutoStudy:
+    def __init__(
+        self,
+        overall_question: str,
+        population: str,
+        model: Optional["Model"] = None,
+        survey: Optional["Survey"] = None,
+        agent_list: Optional["AgentList"] = None,
+        default_num_agents: int = 11,
+    ):
+        """AutoStudy class for generating surveys and agents."""
+
+        self.overall_question = overall_question
+        self.population = population
+        self._survey = survey
+        self._agent_list = agent_list
+        self._agent_list_generator = None
+        self._persona_mapping = None
+        self._results = None
+        self.default_num_agents = default_num_agents
+        self.model = model or Model()
+
+    def to_dict(self):
+        return {
+            "overall_question": self.overall_question,
+            "population": self.population,
+            "survey": self.survey.to_dict(),
+            "persona_mapping": self.persona_mapping.to_dict(),
+            "results": self.results.to_dict(),
+        }
+
+    @property
+    def survey(self):
+        if self._survey is None:
+            self._survey = self._create_survey()
+        return self._survey
+
+    @property
+    def persona_mapping(self):
+        if self._persona_mapping is None:
+            self._persona_mapping = self._create_persona_mapping()
+        return self._persona_mapping
+
+    @property
+    def agent_list_generator(self):
+        if self._agent_list_generator is None:
+            self._agent_list_generator = self._create_agent_list_generator()
+        return self._agent_list_generator
+
+    @property
+    def results(self):
+        if self._results is None:
+            self._results = self._create_results()
+        return self._results
+
+    def _create_survey(self):
+        survey_pipline_stages = [
+            StageQuestions,
+            StageLabelQuestions,
+            StageGenerateSurvey,
+        ]
+        survey_pipeline = gen_pipeline(survey_pipline_stages)
+        return survey_pipeline.process(
+            data=survey_pipeline.input(
+                overall_question=self.overall_question, population=self.population
+            )
+        ).survey
+
+    def _create_persona_mapping(self):
+        persona_pipeline_stages = [
+            StageQuestions,
+            StagePersona,
+            StagePersonaDimensions,
+            StagePersonaDimensionValues,
+            StagePersonaDimensionValueRanges,
+        ]
+
+        persona_pipeline = gen_pipeline(persona_pipeline_stages)
+        sample_agent_results = persona_pipeline.process(
+            persona_pipeline.input(
+                overall_question=overall_question, population=self.population
+            )
+        )
+        return sample_agent_results
+
+    def _create_agent_list_generator(self):
+        return agent_generator(
+            persona=self.persona_mapping.persona,
+            dimension_dict=self.persona_mapping.mapping,
+        )
+
+    def agent_list(self, num_agents):
+        return create_agents(
+            agent_generator=self.agent_list_generator,
+            survey=self.survey,
+            num_agents=num_agents,
+        )
+
+    def _create_results(self, num_agents=None):
+        if num_agents is None:
+            num_agents = self.default_num_agents
+        agent_list = self.agent_list(num_agents)
+        return self.survey.by(agent_list).by(self.model).run()
+
+
+if __name__ == "__main__":
+    overall_question = "I have an open source Python library for working with LLMs. What are some ways we can market this to others?"
+    auto_study = AutoStudy(overall_question, population="US Adults")
+
+    results = auto_study.results