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

edsl/agents/Agent.py

@@ -1,67 +1,84 @@
+"""An Agent is an AI agent that can reference a set of traits in answering questions."""
+
 from __future__ import annotations
 import copy
 import inspect
 import types
-import io
-from typing import Any, Callable, Optional, Union, Dict
+from typing import (
+    Callable,
+    Optional,
+    Union,
+    Any,
+    TYPE_CHECKING,
+    Protocol,
+    runtime_checkable,
+    TypeVar,
+)
+from contextlib import contextmanager
+from dataclasses import dataclass
+
+# Type variable for the Agent class
+A = TypeVar("A", bound="Agent")
+
+if TYPE_CHECKING:
+    from edsl.data.Cache import Cache
+    from edsl.surveys.Survey import Survey
+    from edsl.scenarios.Scenario import Scenario
+    from edsl.language_models import LanguageModel
+    from edsl.surveys.MemoryPlan import MemoryPlan
+    from edsl.questions import QuestionBase
+    from edsl.agents.Invigilator import InvigilatorBase
+    from edsl.prompts import Prompt
+    from edsl.questions.QuestionBase import QuestionBase
+    from edsl.scenarios.Scenario import Scenario
+
 
-from jinja2 import Template
+@runtime_checkable
+class DirectAnswerMethod(Protocol):
+    """Protocol defining the required signature for direct answer methods."""
 
-from rich.console import Console
-from rich.table import Table
+    def __call__(self, self_: A, question: QuestionBase, scenario: Scenario) -> Any: ...
+
+
+from uuid import uuid4
 
 from edsl.Base import Base
+from edsl.exceptions.questions import QuestionScenarioRenderError
 
 from edsl.exceptions.agents import (
+    AgentErrors,
     AgentCombinationError,
     AgentDirectAnswerFunctionError,
     AgentDynamicTraitsFunctionError,
 )
 
-from edsl.agents.Invigilator import (
-    InvigilatorDebug,
-    InvigilatorHuman,
-    InvigilatorFunctional,
-    InvigilatorAI,
-)
-
-from edsl.language_models.registry import Model
-from edsl.scenarios import Scenario
-from edsl.enums import LanguageModelType
-
-# from edsl.utilities import (
-#     dict_to_html,
-#     print_dict_as_html_table,
-#     print_dict_with_rich,
-# )
-
 from edsl.agents.descriptors import (
     TraitsDescriptor,
     CodebookDescriptor,
     InstructionDescriptor,
     NameDescriptor,
 )
-
-from edsl.utilities.decorators import sync_wrapper
-
+from edsl.utilities.decorators import (
+    sync_wrapper,
+)
+from edsl.utilities.remove_edsl_version import remove_edsl_version
 from edsl.data_transfer_models import AgentResponseDict
+from edsl.utilities.restricted_python import create_restricted_function
 
-from edsl.prompts.library.agent_persona import AgentPersona
+from edsl.scenarios.Scenario import Scenario
 
 
-class Agent(Base):
-    """An agent that can answer questions.
+class AgentTraits(Scenario):
+    """A class representing the traits of an agent."""
 
-    Parameters
-    ----------
-    traits : dict, optional - A dictionary of traits that the agent has. The keys need to be
-    valid python variable names. The values can be any python object that has a valid __str__ method.
-    codebook : dict, optional - A codebook mapping trait keys to trait descriptions.
-    instruction : str, optional - Instructions for the agent.
+    def __repr__(self):
+        return f"{self.data}"
 
-    dynamic_traits_function : Callable, optional - A function that returns a dictionary of traits.
 
-    """
+class Agent(Base):
+    """An class representing an agent that can answer questions."""
+
+    __documentation__ = "https://docs.expectedparrot.com/en/latest/agents.html"
 
     default_instruction = """You are answering questions as if you were a human. Do not break character."""
 
@@ -69,36 +86,250 @@ class Agent(Base):
     codebook = CodebookDescriptor()
     instruction = InstructionDescriptor()
     name = NameDescriptor()
+    dynamic_traits_function_name = ""
+    answer_question_directly_function_name = ""
+    has_dynamic_traits_function = False
 
     def __init__(
         self,
-        *,
-        traits: dict = None,
-        name: str = None,
-        codebook: dict = None,
-        instruction: str = None,
-        trait_presentation_template: str = None,
-        dynamic_traits_function: Callable = None,
+        traits: Optional[dict] = None,
+        name: Optional[str] = None,
+        codebook: Optional[dict] = None,
+        instruction: Optional[str] = None,
+        traits_presentation_template: Optional[str] = None,
+        dynamic_traits_function: Optional[Callable] = None,
+        dynamic_traits_function_source_code: Optional[str] = None,
+        dynamic_traits_function_name: Optional[str] = None,
+        answer_question_directly_source_code: Optional[str] = None,
+        answer_question_directly_function_name: Optional[str] = None,
     ):
-        self.name = name
-        self._traits = traits or dict()
-        self.codebook = codebook or dict()
-        self.instruction = instruction or self.default_instruction
-        self.dynamic_traits_function = dynamic_traits_function
+        """Initialize a new instance of Agent.
+
+        :param traits: A dictionary of traits that the agent has. The keys need to be valid identifiers.
+        :param name: A name for the agent
+        :param codebook: A codebook mapping trait keys to trait descriptions.
+        :param instruction: Instructions for the agent in how to answer questions.
+        :param trait_presentation_template: A template for how to present the agent's traits.
+        :param dynamic_traits_function: A function that returns a dictionary of traits.
+        :param dynamic_traits_function_source_code: The source code for the dynamic traits function.
+        :param dynamic_traits_function_name: The name of the dynamic traits function.
+
+        The `traits` parameter is a dictionary of traits that the agent has.
+        These traits are used to construct a prompt that is presented to the LLM.
+        In the absence of a `traits_presentation_template`, the default is used.
+        This is a template that is used to present the agent's traits to the LLM.
+        See :py:class:`edsl.prompts.library.agent_persona.AgentPersona` for more information.
+
+        Example usage:
+
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> a.traits
+        {'age': 10, 'hair': 'brown', 'height': 5.5}
+
+        These traits are used to construct a prompt that is presented to the LLM.
+
+        In the absence of a `traits_presentation_template`, the default is used.
+
+        >>> a = Agent(traits = {"age": 10}, traits_presentation_template = "I am a {{age}} year old.")
+        >>> repr(a.agent_persona)
+        'Prompt(text=\"""I am a {{age}} year old.\""")'
+
+        When this is rendered for presentation to the LLM, it will replace the `{{age}}` with the actual age.
+        it is also possible to use the `codebook` to provide a more human-readable description of the trait.
+        Here is an example where we give a prefix to the age trait (namely the age):
+
+        >>> traits = {"age": 10, "hair": "brown", "height": 5.5}
+        >>> codebook = {'age': 'Their age is'}
+        >>> a = Agent(traits = traits, codebook = codebook, traits_presentation_template = "This agent is Dave. {{codebook['age']}} {{age}}")
+        >>> d = a.traits | {'codebook': a.codebook}
+        >>> a.agent_persona.render(d)
+        Prompt(text=\"""This agent is Dave. Their age is 10\""")
+
+        Instructions
+        ------------
+        The agent can also have instructions. These are instructions that are given to the agent when answering questions.
+
+        >>> Agent.default_instruction
+        'You are answering questions as if you were a human. Do not break character.'
+
+        See see how these are used to actually construct the prompt that is presented to the LLM, see :py:class:`edsl.agents.Invigilator.InvigilatorBase`.
+
+        """
+        self._initialize_basic_attributes(traits, name, codebook)
+        self._initialize_instruction(instruction)
+        self._initialize_dynamic_traits_function(
+            dynamic_traits_function,
+            dynamic_traits_function_source_code,
+            dynamic_traits_function_name,
+        )
+        self._initialize_answer_question_directly(
+            answer_question_directly_source_code, answer_question_directly_function_name
+        )
         self._check_dynamic_traits_function()
+        self._initialize_traits_presentation_template(traits_presentation_template)
         self.current_question = None
 
-        if trait_presentation_template is not None:
-            self.trait_presentation_template = trait_presentation_template
-            self.agent_persona = AgentPersona(text=self.trait_presentation_template)
+    def _initialize_basic_attributes(self, traits, name, codebook) -> None:
+        """Initialize the basic attributes of the agent."""
+        self.name = name
+        self._traits = AgentTraits(traits or dict())
+        self.codebook = codebook or dict()
+
+    def _initialize_instruction(self, instruction) -> None:
+        """Initialize the instruction for the agent i.e., how the agent should answer questions."""
+        if instruction is None:
+            self.instruction = self.default_instruction
+            self._instruction = self.default_instruction
+            self.set_instructions = False
+        else:
+            self.instruction = instruction
+            self._instruction = instruction
+            self.set_instructions = True
+
+    def _initialize_traits_presentation_template(
+        self, traits_presentation_template
+    ) -> None:
+        """Initialize the traits presentation template. How the agent's traits are presented to the LLM."""
+        if traits_presentation_template is not None:
+            self._traits_presentation_template = traits_presentation_template
+            self.traits_presentation_template = traits_presentation_template
+            self.set_traits_presentation_template = True
+        else:
+            self.traits_presentation_template = "Your traits: {{traits}}"
+            self.set_traits_presentation_template = False
+
+    def _initialize_dynamic_traits_function(
+        self,
+        dynamic_traits_function,
+        dynamic_traits_function_source_code,
+        dynamic_traits_function_name,
+    ) -> None:
+        """Initialize the dynamic traits function i.e., a function that returns a dictionary of traits based on the question."""
+        # Deal with dynamic traits function
+        self.dynamic_traits_function = dynamic_traits_function
 
-    def _check_dynamic_traits_function(self):
         if self.dynamic_traits_function:
+            self.dynamic_traits_function_name = self.dynamic_traits_function.__name__
+            self.has_dynamic_traits_function = True
+        else:
+            self.has_dynamic_traits_function = False
+
+        if dynamic_traits_function_source_code:
+            self.dynamic_traits_function_name = dynamic_traits_function_name
+            self.dynamic_traits_function = create_restricted_function(
+                dynamic_traits_function_name, dynamic_traits_function
+            )
+
+    def _initialize_answer_question_directly(
+        self,
+        answer_question_directly_source_code,
+        answer_question_directly_function_name,
+    ) -> None:
+        if answer_question_directly_source_code:
+            self.answer_question_directly_function_name = (
+                answer_question_directly_function_name
+            )
+            protected_method = create_restricted_function(
+                answer_question_directly_function_name,
+                answer_question_directly_source_code,
+            )
+            bound_method = types.MethodType(protected_method, self)
+            setattr(self, "answer_question_directly", bound_method)
+
+    def _initialize_traits_presentation_template(
+        self, traits_presentation_template
+    ) -> None:
+        if traits_presentation_template is not None:
+            self._traits_presentation_template = traits_presentation_template
+            self.traits_presentation_template = traits_presentation_template
+            self.set_traits_presentation_template = True
+        else:
+            self.traits_presentation_template = "Your traits: {{traits}}"
+            self.set_traits_presentation_template = False
+
+    def duplicate(self) -> Agent:
+        """Return a duplicate of the agent.
+
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5}, codebook = {'age': 'Their age is'})
+        >>> a2 = a.duplicate()
+        >>> a2 == a
+        True
+        >>> id(a) == id(a2)
+        False
+        >>> def f(self, question, scenario): return "I am a direct answer."
+        >>> a.add_direct_question_answering_method(f)
+        >>> hasattr(a, "answer_question_directly")
+        True
+        >>> a2 = a.duplicate()
+        >>> a2.answer_question_directly(None, None)
+        'I am a direct answer.'
+
+        >>> a = Agent(traits = {'age': 10}, instruction = "Have fun!")
+        >>> a2 = a.duplicate()
+        >>> a2.instruction
+        'Have fun!'
+        """
+        new_agent = Agent.from_dict(self.to_dict())
+        if hasattr(self, "answer_question_directly"):
+            answer_question_directly = self.answer_question_directly
+            newf = lambda self, question, scenario: answer_question_directly(
+                question, scenario
+            )
+            new_agent.add_direct_question_answering_method(newf)
+        if hasattr(self, "dynamic_traits_function"):
+            dynamic_traits_function = self.dynamic_traits_function
+            new_agent.dynamic_traits_function = dynamic_traits_function
+        return new_agent
+
+    @property
+    def agent_persona(self) -> Prompt:
+        """Return the agent persona template."""
+        from edsl.prompts.Prompt import Prompt
+
+        return Prompt(text=self.traits_presentation_template)
+
+    def prompt(self) -> str:
+        """Return the prompt for the agent.
+
+        Example usage:
+
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> a.prompt()
+        Prompt(text=\"""Your traits: {'age': 10, 'hair': 'brown', 'height': 5.5}\""")
+        """
+        replacement_dict = (
+            self.traits | {"traits": self.traits} | {"codebook": self.codebook}
+        )
+        if undefined := self.agent_persona.undefined_template_variables(
+            replacement_dict
+        ):
+            raise QuestionScenarioRenderError(
+                f"Agent persona still has variables that were not rendered: {undefined}"
+            )
+        else:
+            return self.agent_persona.render(replacement_dict)
+
+    def _check_dynamic_traits_function(self) -> None:
+        """Check whether dynamic trait function is valid.
+
+        This checks whether the dynamic traits function is valid.
+
+        >>> def f(question): return {"age": 10, "hair": "brown", "height": 5.5}
+        >>> a = Agent(dynamic_traits_function = f)
+        >>> a._check_dynamic_traits_function()
+
+        >>> def g(question, poo): return {"age": 10, "hair": "brown", "height": 5.5}
+        >>> a = Agent(dynamic_traits_function = g)
+        Traceback (most recent call last):
+        ...
+        edsl.exceptions.agents.AgentDynamicTraitsFunctionError: ...
+        """
+        if self.has_dynamic_traits_function:
             sig = inspect.signature(self.dynamic_traits_function)
             if "question" in sig.parameters:
                 if len(sig.parameters) > 1:
                     raise AgentDynamicTraitsFunctionError(
-                        f"The dynamic traits function {self.dynamic_traits_function} has too many parameters. It should only have one parameter: 'question'."
+                        message=f"The dynamic traits function {self.dynamic_traits_function} has too many parameters. It should only have one parameter: 'question'."
                     )
             else:
                 if len(sig.parameters) > 0:
@@ -109,35 +340,196 @@ class Agent(Base):
 
     @property
     def traits(self) -> dict[str, str]:
-        """A agent's traits, which is a dictionary.
+        """An agent's traits, which is a dictionary.
 
-        >> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
-        >> a.traits
-        {'age': 10, 'hair': 'brown', 'height': 5.5}
+        The agent could have a a dynamic traits function (`dynamic_traits_function`) that returns a dictionary of traits
+        when called. This function can also take a `question` as an argument.
+        If so, the dynamic traits function is called and the result is returned.
+        Otherwise, the traits are returned.
 
-        >> a = Agent()
-        >> a.add_direct_question_answering_method(lambda question, scenario, self: {"age": 10, "hair": "brown", "height": 5.5})
-        >> a.traits
+        Example:
+
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> a.traits
         {'age': 10, 'hair': 'brown', 'height': 5.5}
 
-        The agent could have a a dynamic traits function (dynamic_traits_function) that returns a dictionary of traits
-        when called. This function can also take a question as an argument.
-        If so, the dynamic traits function is called and the result is returned.
-        Otherwise, the traits are returned.
         """
-        if self.dynamic_traits_function:
+        if self.has_dynamic_traits_function:
             sig = inspect.signature(self.dynamic_traits_function)
             if "question" in sig.parameters:
                 return self.dynamic_traits_function(question=self.current_question)
             else:
                 return self.dynamic_traits_function()
         else:
-            return self._traits
+            return dict(self._traits)
+
+    @contextmanager
+    def modify_traits_context(self):
+        self._check_before_modifying_traits()
+        try:
+            yield
+        finally:
+            self._traits = AgentTraits(self._traits)
+
+    def _check_before_modifying_traits(self):
+        """Check before modifying traits."""
+        if self.has_dynamic_traits_function:
+            raise AgentErrors(
+                "You cannot modify the traits of an agent that has a dynamic traits function.",
+                "If you want to modify the traits, you should remove the dynamic traits function.",
+            )
+
+    @traits.setter
+    def traits(self, traits: dict[str, str]):
+        with self.modify_traits_context():
+            self._traits = traits
+        # self._check_before_modifying_traits()
+        # self._traits = AgentTraits(traits)
+
+    def rename(
+        self,
+        old_name_or_dict: Union[str, dict[str, str]],
+        new_name: Optional[str] = None,
+    ) -> Agent:
+        """Rename a trait.
+
+        :param old_name_or_dict: The old name of the trait or a dictionary of old names and new names.
+        :param new_name: The new name of the trait.
+
+        Example usage:
+
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> newa = a.rename("age", "years")
+        >>> newa == Agent(traits = {'years': 10, 'hair': 'brown', 'height': 5.5})
+        True
+
+        >>> newa.rename({'years': 'smage'}) == Agent(traits = {'smage': 10, 'hair': 'brown', 'height': 5.5})
+        True
+
+        """
+        self._check_before_modifying_traits()
+        if isinstance(old_name_or_dict, dict) and new_name:
+            raise AgentErrors(
+                f"You passed a dict: {old_name_or_dict} and a new name: {new_name}. You should pass only a dict."
+            )
+
+        if isinstance(old_name_or_dict, dict) and new_name is None:
+            return self._rename_dict(old_name_or_dict)
+
+        if isinstance(old_name_or_dict, str):
+            return self._rename(old_name_or_dict, new_name)
+
+        raise AgentErrors("Something is not right with Agent renaming")
+
+    def _rename_dict(self, renaming_dict: dict[str, str]) -> Agent:
+        """
+        Internal method to rename traits using a dictionary.
+        The keys should all be old names and the values should all be new names.
+
+        Example usage:
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> a._rename_dict({"age": "years", "height": "feet"})
+        Agent(traits = {'years': 10, 'hair': 'brown', 'feet': 5.5})
+
+        """
+        try:
+            assert all(k in self.traits for k in renaming_dict.keys())
+        except AssertionError:
+            raise AgentErrors(
+                f"The trait(s) {set(renaming_dict.keys()) - set(self.traits.keys())} do not exist in the agent's traits, which are {self.traits}."
+            )
+        new_agent = self.duplicate()
+        new_agent.traits = {renaming_dict.get(k, k): v for k, v in self.traits.items()}
+        return new_agent
+
+    def _rename(self, old_name: str, new_name: str) -> Agent:
+        """Rename a trait.
+
+        Example usage:
+
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> a._rename(old_name="age", new_name="years")
+        Agent(traits = {'years': 10, 'hair': 'brown', 'height': 5.5})
+
+        """
+        try:
+            assert old_name in self.traits
+        except AssertionError:
+            raise AgentErrors(
+                f"The trait '{old_name}' does not exist in the agent's traits, which are {self.traits}."
+            )
+        newagent = self.duplicate()
+        newagent.traits = {
+            new_name if k == old_name else k: v for k, v in self.traits.items()
+        }
+        newagent.codebook = {
+            new_name if k == old_name else k: v for k, v in self.codebook.items()
+        }
+        return newagent
+
+    def __getitem__(self, key):
+        """Allow for accessing traits using the bracket notation.
+
+        Example:
+
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> a['traits']['age']
+        10
+
+        """
+        return getattr(self, key)
+
+    def remove_direct_question_answering_method(self) -> None:
+        """Remove the direct question answering method.
 
-    def add_direct_question_answering_method(self, method: Callable):
-        """Adds a method to the agent that can answer a particular question type."""
+        Example usage:
+
+        >>> a = Agent()
+        >>> def f(self, question, scenario): return "I am a direct answer."
+        >>> a.add_direct_question_answering_method(f)
+        >>> a.remove_direct_question_answering_method()
+        >>> hasattr(a, "answer_question_directly")
+        False
+        """
+        if hasattr(self, "answer_question_directly"):
+            delattr(self, "answer_question_directly")
+
+    def add_direct_question_answering_method(
+        self,
+        method: DirectAnswerMethod,
+        validate_response: bool = False,
+        translate_response: bool = False,
+    ) -> None:
+        """Add a method to the agent that can answer a particular question type.
+        https://docs.expectedparrot.com/en/latest/agents.html#agent-direct-answering-methods
+
+        :param method: A method that can answer a question directly.
+        :param validate_response: Whether to validate the response.
+        :param translate_response: Whether to translate the response.
+
+        Example usage:
+
+        >>> a = Agent()
+        >>> def f(self, question, scenario): return "I am a direct answer."
+        >>> a.add_direct_question_answering_method(f)
+        >>> a.answer_question_directly(question = None, scenario = None)
+        'I am a direct answer.'
+        """
         if hasattr(self, "answer_question_directly"):
-            print("Warning: overwriting existing answer_question_directly method")
+            import warnings
+
+            warnings.warn(
+                "Warning: overwriting existing answer_question_directly method"
+            )
+
+        self.validate_response = validate_response
+        self.translate_response = translate_response
+
+        # if not isinstance(method, DirectAnswerMethod):
+        #     raise AgentDirectAnswerFunctionError(
+        #         f"Method {method} does not match required signature. "
+        #         "Must take (self, question, scenario) parameters."
+        #     )
 
         signature = inspect.signature(method)
         for argument in ["question", "scenario", "self"]:
@@ -147,93 +539,215 @@ class Agent(Base):
                 )
         bound_method = types.MethodType(method, self)
         setattr(self, "answer_question_directly", bound_method)
+        self.answer_question_directly_function_name = bound_method.__name__
 
     def create_invigilator(
         self,
-        question: Question,
-        scenario: Optional[Scenario] = None,
-        model: Optional[LanguageModel] = None,
-        debug: bool = False,
-        memory_plan: Optional[MemoryPlan] = None,
+        *,
+        question: "QuestionBase",
+        cache: "Cache",
+        survey: Optional["Survey"] = None,
+        scenario: Optional["Scenario"] = None,
+        model: Optional["LanguageModel"] = None,
+        memory_plan: Optional["MemoryPlan"] = None,
         current_answers: Optional[dict] = None,
-    ) -> "Invigilator":
-        """
-        An invigator is an object that is responsible administering a question to an agent and
+        iteration: int = 1,
+        raise_validation_errors: bool = True,
+        key_lookup: Optional["KeyLookup"] = None,
+    ) -> "InvigilatorBase":
+        """Create an Invigilator.
+
+        An invigilator is an object that is responsible for administering a question to an agent.
+        There are several different types of invigilators, depending on the type of question and the agent.
+        For example, there are invigilators for functional questions (i.e., question is of type :class:`edsl.questions.QuestionFunctional`:), for direct questions, and for LLM questions.
+
+        >>> a = Agent(traits = {})
+        >>> a.create_invigilator(question = None, cache = False)
+        InvigilatorAI(...)
+
+        An invigator is an object that is responsible for administering a question to an agent and
         recording the responses.
         """
+        from edsl.language_models.model import Model
+
+        from edsl.scenarios.Scenario import Scenario
+
+        cache = cache
         self.current_question = question
-        model = model or Model(LanguageModelType.GPT_4.value, use_cache=True)
+        model = model or Model()
         scenario = scenario or Scenario()
         invigilator = self._create_invigilator(
-            question, scenario, model, debug, memory_plan, current_answers
+            question=question,
+            scenario=scenario,
+            survey=survey,
+            model=model,
+            memory_plan=memory_plan,
+            current_answers=current_answers,
+            iteration=iteration,
+            cache=cache,
+            raise_validation_errors=raise_validation_errors,
+            key_lookup=key_lookup,
         )
+        if hasattr(self, "validate_response"):
+            invigilator.validate_response = self.validate_response
+        if hasattr(self, "translate_response"):
+            invigilator.translate_response = self.translate_response
         return invigilator
 
     async def async_answer_question(
         self,
-        question: Question,
+        *,
+        question: QuestionBase,
+        cache: Cache,
         scenario: Optional[Scenario] = None,
+        survey: Optional[Survey] = None,
         model: Optional[LanguageModel] = None,
         debug: bool = False,
         memory_plan: Optional[MemoryPlan] = None,
         current_answers: Optional[dict] = None,
+        iteration: int = 0,
+        key_lookup: Optional["KeyLookup"] = None,
     ) -> AgentResponseDict:
         """
+        Answer a posed question.
+
+        :param question: The question to answer.
+        :param scenario: The scenario in which the question is asked.
+        :param model: The language model to use.
+        :param debug: Whether to run in debug mode.
+        :param memory_plan: The memory plan to use.
+        :param current_answers: The current answers.
+        :param iteration: The iteration number.
+
+        >>> a = Agent(traits = {})
+        >>> a.add_direct_question_answering_method(lambda self, question, scenario: "I am a direct answer.")
+        >>> from edsl.questions.QuestionFreeText import QuestionFreeText
+        >>> q = QuestionFreeText.example()
+        >>> a.answer_question(question = q, cache = False).answer
+        'I am a direct answer.'
+
         This is a function where an agent returns an answer to a particular question.
         However, there are several different ways an agent can answer a question, so the
-        actual functionality is delegated to an Invigilator object.
+        actual functionality is delegated to an :class:`edsl.agents.InvigilatorBase`: object.
         """
         invigilator = self.create_invigilator(
             question=question,
+            cache=cache,
             scenario=scenario,
+            survey=survey,
             model=model,
-            debug=debug,
             memory_plan=memory_plan,
             current_answers=current_answers,
+            iteration=iteration,
+            key_lookup=key_lookup,
         )
         response: AgentResponseDict = await invigilator.async_answer_question()
         return response
 
     answer_question = sync_wrapper(async_answer_question)
 
+    def _get_invigilator_class(self, question: QuestionBase) -> Type[InvigilatorBase]:
+        """Get the invigilator class for a question.
+
+        This method returns the invigilator class that should be used to answer a question.
+        The invigilator class is determined by the type of question and the type of agent.
+        """
+        from edsl.agents.Invigilator import (
+            InvigilatorHuman,
+            InvigilatorFunctional,
+            InvigilatorAI,
+        )
+
+        if hasattr(question, "answer_question_directly"):
+            return InvigilatorFunctional
+        elif hasattr(self, "answer_question_directly"):
+            return InvigilatorHuman
+        else:
+            return InvigilatorAI
+
     def _create_invigilator(
         self,
-        question: Question,
+        question: QuestionBase,
+        cache: Optional[Cache] = None,
         scenario: Optional[Scenario] = None,
         model: Optional[LanguageModel] = None,
-        debug: bool = False,
+        survey: Optional[Survey] = None,
         memory_plan: Optional[MemoryPlan] = None,
         current_answers: Optional[dict] = None,
-    ):
-        model = model or Model(LanguageModelType.GPT_4.value, use_cache=True)
+        iteration: int = 0,
+        raise_validation_errors: bool = True,
+        key_lookup: Optional["KeyLookup"] = None,
+    ) -> "InvigilatorBase":
+        """Create an Invigilator."""
+        from edsl.language_models.model import Model
+        from edsl.scenarios.Scenario import Scenario
+
+        model = model or Model()
         scenario = scenario or Scenario()
 
-        if debug:
-            # use the question's simulate_answer method
-            invigilator_class = InvigilatorDebug
-        elif hasattr(question, "answer_question_directly"):
-            # it is a functional question and the answer only depends on the agent's traits & the scenario
-            invigilator_class = InvigilatorFunctional
-        elif hasattr(self, "answer_question_directly"):
-            # this of the case where the agent has a method that can answer the question directly
-            # this occurrs when 'answer_question_directly' has been monkey-patched onto the agent
-            # which happens when the agent is created from an existing survey
-            invigilator_class = InvigilatorHuman
-        else:
-            # this means an LLM agent will be used. This is the standard case.
-            invigilator_class = InvigilatorAI
+        if cache is None:
+            from edsl.data.Cache import Cache
+
+            cache = Cache()
+
+        invigilator_class = self._get_invigilator_class(question)
 
         invigilator = invigilator_class(
-            self, question, scenario, model, memory_plan, current_answers
+            self,
+            question=question,
+            scenario=scenario,
+            survey=survey,
+            model=model,
+            memory_plan=memory_plan,
+            current_answers=current_answers,
+            iteration=iteration,
+            cache=cache,
+            raise_validation_errors=raise_validation_errors,
+            key_lookup=key_lookup,
         )
         return invigilator
 
-    ################
-    # Dunder Methods
-    ################
-    def __add__(self, other_agent: Agent = None) -> Agent:
+    def select(self, *traits: str) -> Agent:
+        """Selects agents with only the references traits
+
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5}, codebook = {'age': 'Their age is'})
+        >>> a
+        Agent(traits = {'age': 10, 'hair': 'brown', 'height': 5.5}, codebook = {'age': 'Their age is'})
+
+
+        >>> a.select("age", "height")
+        Agent(traits = {'age': 10, 'height': 5.5}, codebook = {'age': 'Their age is'})
+
+        >>> a.select("height")
+        Agent(traits = {'height': 5.5})
+
+        """
+
+        if len(traits) == 1:
+            traits_to_select = [list(traits)[0]]
+        else:
+            traits_to_select = list(traits)
+
+        def _remove_none(d):
+            return {k: v for k, v in d.items() if v is not None}
+
+        newagent = self.duplicate()
+        newagent.traits = {
+            trait: self.traits.get(trait, None) for trait in traits_to_select
+        }
+        newagent.codebook = _remove_none(
+            {trait: self.codebook.get(trait, None) for trait in traits_to_select}
+        )
+        return newagent
+
+    def __add__(self, other_agent: Optional[Agent] = None) -> Agent:
         """
-        Combines two agents by joining their traits.The agents must not have overlapping traits.
+        Combine two agents by joining their traits.
+
+        The agents must not have overlapping traits.
+
+        Example usage:
+
         >>> a1 = Agent(traits = {"age": 10})
         >>> a2 = Agent(traits = {"height": 5.5})
         >>> a1 + a2
@@ -242,6 +756,11 @@ class Agent(Base):
         Traceback (most recent call last):
         ...
         edsl.exceptions.agents.AgentCombinationError: The agents have overlapping traits: {'age'}.
+        ...
+        >>> a1 = Agent(traits = {"age": 10}, codebook = {"age": "Their age is"})
+        >>> a2 = Agent(traits = {"height": 5.5}, codebook = {"height": "Their height is"})
+        >>> a1 + a2
+        Agent(traits = {'age': 10, 'height': 5.5}, codebook = {'age': 'Their age is', 'height': 'Their height is'})
         """
         if other_agent is None:
             return self
@@ -250,12 +769,19 @@ class Agent(Base):
                 f"The agents have overlapping traits: {common_traits}."
             )
         else:
-            new_agent = Agent(traits=copy.deepcopy(self.traits))
-            new_agent.traits.update(other_agent.traits)
-            return new_agent
+            new_codebook = copy.deepcopy(self.codebook) | copy.deepcopy(
+                other_agent.codebook
+            )
+            d = self.traits | other_agent.traits
+            newagent = self.duplicate()
+            newagent.traits = d
+            newagent.codebook = new_codebook
+            return newagent
 
     def __eq__(self, other: Agent) -> bool:
-        """Checks if two agents are equal. Only checks the traits.
+        """Check if two agents are equal.
+
+        This only checks the traits.
         >>> a1 = Agent(traits = {"age": 10})
         >>> a2 = Agent(traits = {"age": 10})
         >>> a1 == a2
@@ -266,25 +792,57 @@ class Agent(Base):
         """
         return self.data == other.data
 
-    def __repr__(self):
+    def __getattr__(self, name):
+        """
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> a.age
+        10
+        """
+        if name == "has_dynamic_traits_function":
+            return self.has_dynamic_traits_function
+
+        if name in self._traits:
+            return self._traits[name]
+
+        raise AttributeError(
+            f"'{type(self).__name__}' object has no attribute '{name}'"
+        )
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        # Include any additional state that needs to be serialized
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        # Ensure _traits is initialized if it's missing
+        if "_traits" not in self.__dict__:
+            self._traits = {}
+
+    def __repr__(self) -> str:
+        """Return representation of Agent."""
         class_name = self.__class__.__name__
         items = [
-            f"{k} = '{v}'" if isinstance(v, str) else f"{k} = {v}"
+            f'{k} = """{v}"""' if isinstance(v, str) else f"{k} = {v}"
             for k, v in self.data.items()
             if k != "question_type"
         ]
         return f"{class_name}({', '.join(items)})"
 
-    ################
-    # SERIALIZATION METHODS
-    ################
     @property
-    def data(self):
+    def data(self) -> dict:
+        """Format the data for serialization.
+
+        TODO: Warn if has dynamic traits function or direct answer function that cannot be serialized.
+        TODO: Add ability to have coop-hosted functions that are serializable.
+        """
+
         raw_data = {
             k.replace("_", "", 1): v
             for k, v in self.__dict__.items()
             if k.startswith("_")
         }
+
         if hasattr(self, "set_instructions"):
             if not self.set_instructions:
                 raw_data.pop("instruction")
@@ -292,20 +850,102 @@ class Agent(Base):
             raw_data.pop("codebook")
         if self.name == None:
             raw_data.pop("name")
+
+        if hasattr(self, "dynamic_traits_function"):
+            raw_data.pop(
+                "dynamic_traits_function", None
+            )  # in case dynamic_traits_function will appear with _ in self.__dict__
+            dynamic_traits_func = self.dynamic_traits_function
+            if dynamic_traits_func:
+                func = inspect.getsource(dynamic_traits_func)
+                raw_data["dynamic_traits_function_source_code"] = func
+                raw_data["dynamic_traits_function_name"] = (
+                    self.dynamic_traits_function_name
+                )
+        if hasattr(self, "answer_question_directly"):
+            raw_data.pop(
+                "answer_question_directly", None
+            )  # in case answer_question_directly will appear with _ in self.__dict__
+            answer_question_directly_func = self.answer_question_directly
+
+            if (
+                answer_question_directly_func
+                and raw_data.get("answer_question_directly_source_code", None) != None
+            ):
+                raw_data["answer_question_directly_source_code"] = inspect.getsource(
+                    answer_question_directly_func
+                )
+                raw_data["answer_question_directly_function_name"] = (
+                    self.answer_question_directly_function_name
+                )
+        raw_data["traits"] = dict(raw_data["traits"])
+
         return raw_data
 
-    def to_dict(self) -> dict[str, Union[dict, bool]]:
-        """Serializes to a dictionary."""
-        return self.data
+    def __hash__(self) -> int:
+        """Return a hash of the agent.
+
+        >>> hash(Agent.example())
+        2067581884874391607
+        """
+        from edsl.utilities.utilities import dict_hash
+
+        return dict_hash(self.to_dict(add_edsl_version=False))
+
+    def to_dict(self, add_edsl_version=True) -> dict[str, Union[dict, bool]]:
+        """Serialize to a dictionary with EDSL info.
+
+        Example usage:
+
+        >>> a = Agent(name = "Steve", traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> a.to_dict()
+        {'traits': {'age': 10, 'hair': 'brown', 'height': 5.5}, 'name': 'Steve', 'edsl_version': '...', 'edsl_class_name': 'Agent'}
+
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5}, instruction = "Have fun.")
+        >>> a.to_dict()
+        {'traits': {'age': 10, 'hair': 'brown', 'height': 5.5}, 'instruction': 'Have fun.', 'edsl_version': '...', 'edsl_class_name': 'Agent'}
+        """
+        d = {}
+        d["traits"] = copy.deepcopy(dict(self._traits))
+        if self.name:
+            d["name"] = self.name
+        if self.set_instructions:
+            d["instruction"] = self.instruction
+        if self.set_traits_presentation_template:
+            d["traits_presentation_template"] = self.traits_presentation_template
+        if self.codebook:
+            d["codebook"] = self.codebook
+        if add_edsl_version:
+            from edsl import __version__
+
+            d["edsl_version"] = __version__
+            d["edsl_class_name"] = self.__class__.__name__
+
+        return d
 
     @classmethod
+    @remove_edsl_version
     def from_dict(cls, agent_dict: dict[str, Union[dict, bool]]) -> Agent:
-        """Deserializes from a dictionary."""
-        return cls(**agent_dict)
+        """Deserialize from a dictionary.
+
+        Example usage:
 
-    ################
-    # DISPLAY Methods
-    ################
+        >>> Agent.from_dict({'name': "Steve", 'traits': {'age': 10, 'hair': 'brown', 'height': 5.5}})
+        Agent(name = \"""Steve\""", traits = {'age': 10, 'hair': 'brown', 'height': 5.5})
+
+        """
+        if "traits" in agent_dict:
+            return cls(
+                traits=agent_dict["traits"],
+                name=agent_dict.get("name", None),
+                instruction=agent_dict.get("instruction", None),
+                traits_presentation_template=agent_dict.get(
+                    "traits_presentation_template", None
+                ),
+                codebook=agent_dict.get("codebook", None),
+            )
+        else:  # old-style agent - we used to only store the traits
+            return cls(**agent_dict)
 
     def _table(self) -> tuple[dict, list]:
         """Prepare generic table data."""
@@ -315,35 +955,95 @@ class Agent(Base):
         column_names = ["Attribute", "Value"]
         return table_data, column_names
 
-    def rich_print(self):
-        """Displays an object as a rich table."""
-        table_data, column_names = self._table()
-        table = Table(title=f"{self.__class__.__name__} Attributes")
-        for column in column_names:
-            table.add_column(column, style="bold")
+    def add_trait(self, trait_name_or_dict: str, value: Optional[Any] = None) -> Agent:
+        """Adds a trait to an agent and returns that agent
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> a.add_trait("weight", 150)
+        Agent(traits = {'age': 10, 'hair': 'brown', 'height': 5.5, 'weight': 150})
+        """
+        if isinstance(trait_name_or_dict, dict) and value is None:
+            newagent = self.duplicate()
+            newagent.traits = {**self.traits, **trait_name_or_dict}
+            return newagent
+
+        if isinstance(trait_name_or_dict, dict) and value:
+            raise AgentErrors(
+                f"You passed a dict: {trait_name_or_dict} and a value: {value}. You should pass only a dict."
+            )
+
+        if isinstance(trait_name_or_dict, str):
+            newagent = self.duplicate()
+            newagent.traits = {**self.traits, **{trait_name_or_dict: value}}
+            return newagent
+
+        raise AgentErrors("Something is not right with adding a trait to an Agent")
 
-        for row in table_data:
-            row_data = [row[column] for column in column_names]
-            table.add_row(*row_data)
+    def remove_trait(self, trait: str) -> Agent:
+        """Remove a trait from the agent.
 
-        return table
+        Example usage:
+
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> a.remove_trait("age")
+        Agent(traits = {'hair': 'brown', 'height': 5.5})
+        """
+        newagent = self.duplicate()
+        newagent.traits = {k: v for k, v in self.traits.items() if k != trait}
+        return newagent
+
+    def translate_traits(self, values_codebook: dict) -> Agent:
+        """Translate traits to a new codebook.
+
+        >>> a = Agent(traits = {"age": 10, "hair": 1, "height": 5.5})
+        >>> a.translate_traits({"hair": {1:"brown"}})
+        Agent(traits = {'age': 10, 'hair': 'brown', 'height': 5.5})
+
+        :param values_codebook: The new codebook.
+        """
+        new_traits = {}
+        for key, value in self.traits.items():
+            if key in values_codebook:
+                new_traits[key] = values_codebook[key].get(value, value)
+            else:
+                new_traits[key] = value
+        newagent = self.duplicate()
+        newagent.traits = new_traits
+        return newagent
 
     @classmethod
-    def example(cls) -> Agent:
-        """Returns an example agent.
+    def example(cls, randomize: bool = False) -> Agent:
+        """
+        Returns an example Agent instance.
+
+        :param randomize: If True, adds a random string to the value of an example key.
 
         >>> Agent.example()
         Agent(traits = {'age': 22, 'hair': 'brown', 'height': 5.5})
         """
-        return cls(traits={"age": 22, "hair": "brown", "height": 5.5})
+        addition = "" if not randomize else str(uuid4())
+        return cls(traits={"age": 22, "hair": f"brown{addition}", "height": 5.5})
 
     def code(self) -> str:
-        """Returns the code for the agent."""
-        return f"Agent(traits={self.traits})"
+        """Return the code for the agent.
+
+        Example usage:
+
+        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> print(a.code())
+        from edsl.agents.Agent import Agent
+        agent = Agent(traits={'age': 10, 'hair': 'brown', 'height': 5.5})
+        """
+        return (
+            f"from edsl.agents.Agent import Agent\nagent = Agent(traits={self.traits})"
+        )
 
 
 def main():
-    """Consumes API credits"""
+    """
+    Give an example of usage.
+
+    WARNING: Consume API credits
+    """
     from edsl.agents import Agent
     from edsl.questions import QuestionMultipleChoice
 
@@ -361,13 +1061,11 @@ def main():
         question_options=["Yes", "No"],
         question_name="food_preference",
     )
-    job = agent.to(question)
-    # run the job
+    job = question.by(agent)
     results = job.run()
-    # results
 
 
 if __name__ == "__main__":
     import doctest
 
-    doctest.testmod()
+    doctest.testmod(optionflags=doctest.ELLIPSIS)