edsl 0.1.38.dev4__py3-none-any.whl → 0.1.39__py3-none-any.whl

edsl/agents/Agent.py

@@ -4,20 +4,46 @@ from __future__ import annotations
 import copy
 import inspect
 import types
-from typing import Callable, Optional, Union, Any, TYPE_CHECKING
+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 import Cache, Survey, Scenario
+    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
+
+
+@runtime_checkable
+class DirectAnswerMethod(Protocol):
+    """Protocol defining the required signature for direct answer methods."""
+
+    def __call__(self, self_: A, question: QuestionBase, scenario: Scenario) -> Any: ...
+
 
 from uuid import uuid4
 
 from edsl.Base import Base
-from edsl.prompts import Prompt
-from edsl.exceptions import QuestionScenarioRenderError
+from edsl.exceptions.questions import QuestionScenarioRenderError
 
 from edsl.exceptions.agents import (
     AgentErrors,
@@ -34,17 +60,25 @@ from edsl.agents.descriptors import (
 )
 from edsl.utilities.decorators import (
     sync_wrapper,
-    add_edsl_version,
-    remove_edsl_version,
 )
+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.scenarios.Scenario import Scenario
+
+
+class AgentTraits(Scenario):
+    """A class representing the traits of an agent."""
+
+    def __repr__(self):
+        return f"{self.data}"
+
 
 class Agent(Base):
     """An class representing an agent that can answer questions."""
 
-    __doc__ = "https://docs.expectedparrot.com/en/latest/agents.html"
+    __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."""
 
@@ -77,6 +111,8 @@ class Agent(Base):
         :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.
@@ -119,17 +155,59 @@ class Agent(Base):
         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
+
+    def _initialize_basic_attributes(self, traits, name, codebook) -> None:
+        """Initialize the basic attributes of the agent."""
         self.name = name
-        self._traits = traits or dict()
+        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 or self.default_instruction
-        self.dynamic_traits_function = dynamic_traits_function
+            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
+
         if self.dynamic_traits_function:
             self.dynamic_traits_function_name = self.dynamic_traits_function.__name__
             self.has_dynamic_traits_function = True
@@ -142,7 +220,11 @@ class Agent(Base):
                 dynamic_traits_function_name, dynamic_traits_function
             )
 
-        # Deal with direct answer 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
@@ -154,18 +236,56 @@ class Agent(Base):
             bound_method = types.MethodType(protected_method, self)
             setattr(self, "answer_question_directly", bound_method)
 
-        self._check_dynamic_traits_function()
-
-        self.current_question = None
-
+    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:
@@ -241,59 +361,111 @@ class Agent(Base):
             else:
                 return self.dynamic_traits_function()
         else:
-            return self._traits
-
-    def _repr_html_(self):
-        # d = self.to_dict(add_edsl_version=False)
-        d = self.traits
-        data = [[k, v] for k, v in d.items()]
-        from tabulate import tabulate
+            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.",
+            )
 
-        table = str(tabulate(data, headers=["keys", "values"], tablefmt="html"))
-        return f"<pre>{table}</pre>"
+    @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], new_name: Optional[str] = None
+        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})
-        >>> a.rename("age", "years") == Agent(traits = {'years': 10, 'hair': 'brown', 'height': 5.5})
+        >>> newa = a.rename("age", "years")
+        >>> newa == Agent(traits = {'years': 10, 'hair': 'brown', 'height': 5.5})
         True
 
-        >>> a.rename({'years': 'smage'})
-        Agent(traits = {'hair': 'brown', 'height': 5.5, 'smage': 10})
+        >>> newa.rename({'years': 'smage'}) == Agent(traits = {'smage': 10, 'hair': 'brown', 'height': 5.5})
+        True
 
         """
-        if isinstance(old_name_or_dict, dict) and new_name is None:
-            for old_name, new_name in old_name_or_dict.items():
-                self = self._rename(old_name, new_name)
-            return self
-
+        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):
-            self._rename(old_name_or_dict, new_name)
-            return self
+            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("age", "years") == Agent(traits = {'years': 10, 'hair': 'brown', 'height': 5.5})
-        True
+        >>> a._rename(old_name="age", new_name="years")
+        Agent(traits = {'years': 10, 'hair': 'brown', 'height': 5.5})
+
         """
-        self.traits[new_name] = self.traits.pop(old_name)
-        return self
+        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.
@@ -324,7 +496,7 @@ class Agent(Base):
 
     def add_direct_question_answering_method(
         self,
-        method: Callable,
+        method: DirectAnswerMethod,
         validate_response: bool = False,
         translate_response: bool = False,
     ) -> None:
@@ -353,6 +525,12 @@ class Agent(Base):
         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"]:
             if argument not in signature.parameters:
@@ -371,12 +549,11 @@ class Agent(Base):
         survey: Optional["Survey"] = None,
         scenario: Optional["Scenario"] = None,
         model: Optional["LanguageModel"] = None,
-        debug: bool = False,
         memory_plan: Optional["MemoryPlan"] = None,
         current_answers: Optional[dict] = None,
         iteration: int = 1,
-        sidecar_model=None,
         raise_validation_errors: bool = True,
+        key_lookup: Optional["KeyLookup"] = None,
     ) -> "InvigilatorBase":
         """Create an Invigilator.
 
@@ -391,7 +568,9 @@ class Agent(Base):
         An invigator is an object that is responsible for administering a question to an agent and
         recording the responses.
         """
-        from edsl import Model, Scenario
+        from edsl.language_models.model import Model
+
+        from edsl.scenarios.Scenario import Scenario
 
         cache = cache
         self.current_question = question
@@ -402,13 +581,12 @@ class Agent(Base):
             scenario=scenario,
             survey=survey,
             model=model,
-            debug=debug,
             memory_plan=memory_plan,
             current_answers=current_answers,
             iteration=iteration,
             cache=cache,
-            sidecar_model=sidecar_model,
             raise_validation_errors=raise_validation_errors,
+            key_lookup=key_lookup,
         )
         if hasattr(self, "validate_response"):
             invigilator.validate_response = self.validate_response
@@ -428,6 +606,7 @@ class Agent(Base):
         memory_plan: Optional[MemoryPlan] = None,
         current_answers: Optional[dict] = None,
         iteration: int = 0,
+        key_lookup: Optional["KeyLookup"] = None,
     ) -> AgentResponseDict:
         """
         Answer a posed question.
@@ -442,7 +621,7 @@ class Agent(Base):
 
         >>> a = Agent(traits = {})
         >>> a.add_direct_question_answering_method(lambda self, question, scenario: "I am a direct answer.")
-        >>> from edsl import QuestionFreeText
+        >>> from edsl.questions.QuestionFreeText import QuestionFreeText
         >>> q = QuestionFreeText.example()
         >>> a.answer_question(question = q, cache = False).answer
         'I am a direct answer.'
@@ -457,16 +636,35 @@ class Agent(Base):
             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: QuestionBase,
@@ -474,53 +672,25 @@ class Agent(Base):
         scenario: Optional[Scenario] = None,
         model: Optional[LanguageModel] = None,
         survey: Optional[Survey] = None,
-        debug: bool = False,
         memory_plan: Optional[MemoryPlan] = None,
         current_answers: Optional[dict] = None,
         iteration: int = 0,
-        sidecar_model=None,
         raise_validation_errors: bool = True,
+        key_lookup: Optional["KeyLookup"] = None,
     ) -> "InvigilatorBase":
         """Create an Invigilator."""
-        from edsl import Model
-        from edsl import Scenario
+        from edsl.language_models.model import Model
+        from edsl.scenarios.Scenario import Scenario
 
         model = model or Model()
         scenario = scenario or Scenario()
 
-        from edsl.agents.Invigilator import (
-            InvigilatorHuman,
-            InvigilatorFunctional,
-            InvigilatorAI,
-            InvigilatorBase,
-        )
-
         if cache is None:
             from edsl.data.Cache import Cache
 
             cache = Cache()
 
-        if debug:
-            raise NotImplementedError("Debug mode is not yet implemented.")
-            # use the question's _simulate_answer method
-            # invigilator_class = InvigilatorDebug
-        elif hasattr(question, "answer_question_directly"):
-            # It's 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 given to the
-            # 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 sidecar_model is not None:
-            # this is the case when a 'simple' model is being used
-            from edsl.agents.Invigilator import InvigilatorSidecar
-
-            invigilator_class = InvigilatorSidecar
+        invigilator_class = self._get_invigilator_class(question)
 
         invigilator = invigilator_class(
             self,
@@ -532,22 +702,24 @@ class Agent(Base):
             current_answers=current_answers,
             iteration=iteration,
             cache=cache,
-            sidecar_model=sidecar_model,
             raise_validation_errors=raise_validation_errors,
+            key_lookup=key_lookup,
         )
         return invigilator
 
     def select(self, *traits: str) -> Agent:
         """Selects agents with only the references traits
 
-        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
+        >>> 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})
+        Agent(traits = {'age': 10, 'height': 5.5}, codebook = {'age': 'Their age is'})
 
-        >>> a.select("age")
-        Agent(traits = {'age': 10})
+        >>> a.select("height")
+        Agent(traits = {'height': 5.5})
 
         """
 
@@ -556,7 +728,17 @@ class Agent(Base):
         else:
             traits_to_select = list(traits)
 
-        return Agent(traits={trait: self.traits[trait] for trait in traits_to_select})
+        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:
         """
@@ -575,6 +757,10 @@ class Agent(Base):
         ...
         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
@@ -583,9 +769,14 @@ 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:
         """Check if two agents are equal.
@@ -602,7 +793,11 @@ class Agent(Base):
         return self.data == other.data
 
     def __getattr__(self, name):
-        # This will be called only if 'name' is not found in the usual places
+        """
+        >>> 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
 
@@ -624,12 +819,6 @@ class Agent(Base):
         if "_traits" not in self.__dict__:
             self._traits = {}
 
-    def print(self) -> None:
-        from rich import print_json
-        import json
-
-        print_json(json.dumps(self.to_dict()))
-
     def __repr__(self) -> str:
         """Return representation of Agent."""
         class_name = self.__class__.__name__
@@ -640,14 +829,6 @@ class Agent(Base):
         ]
         return f"{class_name}({', '.join(items)})"
 
-    # def _repr_html_(self):
-    #     from edsl.utilities.utilities import data_to_html
-
-    #     return data_to_html(self.to_dict())
-
-    #######################
-    # SERIALIZATION METHODS
-    #######################
     @property
     def data(self) -> dict:
         """Format the data for serialization.
@@ -678,9 +859,9 @@ class Agent(Base):
             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
+                raw_data["dynamic_traits_function_name"] = (
+                    self.dynamic_traits_function_name
+                )
         if hasattr(self, "answer_question_directly"):
             raw_data.pop(
                 "answer_question_directly", None
@@ -694,18 +875,23 @@ class Agent(Base):
                 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["answer_question_directly_function_name"] = (
+                    self.answer_question_directly_function_name
+                )
+        raw_data["traits"] = dict(raw_data["traits"])
 
         return raw_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))
 
-    # @add_edsl_version
     def to_dict(self, add_edsl_version=True) -> dict[str, Union[dict, bool]]:
         """Serialize to a dictionary with EDSL info.
 
@@ -713,9 +899,22 @@ class Agent(Base):
 
         >>> a = Agent(name = "Steve", traits = {"age": 10, "hair": "brown", "height": 5.5})
         >>> a.to_dict()
-        {'name': 'Steve', 'traits': {'age': 10, 'hair': 'brown', 'height': 5.5}, 'edsl_version': '...', 'edsl_class_name': 'Agent'}
+        {'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 = copy.deepcopy(self.data)
+        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__
 
@@ -735,7 +934,18 @@ class Agent(Base):
         Agent(name = \"""Steve\""", traits = {'age': 10, 'hair': 'brown', 'height': 5.5})
 
         """
-        return cls(**agent_dict)
+        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."""
@@ -746,10 +956,15 @@ class Agent(Base):
         return table_data, column_names
 
     def add_trait(self, trait_name_or_dict: str, value: Optional[Any] = None) -> Agent:
-        """Adds a trait to an agent and returns that 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:
-            self.traits.update(trait_name_or_dict)
-            return self
+            newagent = self.duplicate()
+            newagent.traits = {**self.traits, **trait_name_or_dict}
+            return newagent
 
         if isinstance(trait_name_or_dict, dict) and value:
             raise AgentErrors(
@@ -757,9 +972,9 @@ class Agent(Base):
             )
 
         if isinstance(trait_name_or_dict, str):
-            trait = trait_name_or_dict
-            self.traits[trait] = value
-            return self
+            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")
 
@@ -772,8 +987,9 @@ class Agent(Base):
         >>> a.remove_trait("age")
         Agent(traits = {'hair': 'brown', 'height': 5.5})
         """
-        _ = self.traits.pop(trait)
-        return self
+        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.
@@ -784,32 +1000,15 @@ class Agent(Base):
 
         :param values_codebook: The new codebook.
         """
+        new_traits = {}
         for key, value in self.traits.items():
             if key in values_codebook:
-                self.traits[key] = values_codebook[key][value]
-        return self
-
-    def rich_print(self):
-        """Display an object as a rich table.
-
-        Example usage:
-
-        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
-        >>> a.rich_print()
-        <rich.table.Table object at ...>
-        """
-        from rich.table import 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")
-
-        for row in table_data:
-            row_data = [row[column] for column in column_names]
-            table.add_row(*row_data)
-
-        return table
+                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, randomize: bool = False) -> Agent:
@@ -817,6 +1016,9 @@ class Agent(Base):
         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})
         """
         addition = "" if not randomize else str(uuid4())
         return cls(traits={"age": 22, "hair": f"brown{addition}", "height": 5.5})
@@ -828,10 +1030,12 @@ class Agent(Base):
 
         >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
         >>> print(a.code())
-        from edsl import Agent
+        from edsl.agents.Agent import Agent
         agent = Agent(traits={'age': 10, 'hair': 'brown', 'height': 5.5})
         """
-        return f"from edsl import Agent\nagent = Agent(traits={self.traits})"
+        return (
+            f"from edsl.agents.Agent import Agent\nagent = Agent(traits={self.traits})"
+        )
 
 
 def main():