edsl 0.1.39__py3-none-any.whl → 0.1.39.dev1__py3-none-any.whl

edsl/agents/Agent.py

@@ -4,46 +4,20 @@ from __future__ import annotations
 import copy
 import inspect
 import types
-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")
+from typing import Callable, Optional, Union, Any, TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from edsl.data.Cache import Cache
-    from edsl.surveys.Survey import Survey
-    from edsl.scenarios.Scenario import Scenario
+    from edsl import Cache, Survey, 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.exceptions.questions import QuestionScenarioRenderError
+from edsl.prompts import Prompt
+from edsl.exceptions import QuestionScenarioRenderError
 
 from edsl.exceptions.agents import (
     AgentErrors,
@@ -60,25 +34,17 @@ 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."""
 
-    __documentation__ = "https://docs.expectedparrot.com/en/latest/agents.html"
+    __doc__ = "https://docs.expectedparrot.com/en/latest/agents.html"
 
     default_instruction = """You are answering questions as if you were a human. Do not break character."""
 
@@ -111,8 +77,6 @@ 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.
@@ -155,59 +119,17 @@ 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 = AgentTraits(traits or dict())
+        self._traits = 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.instruction = instruction or self.default_instruction
         self.dynamic_traits_function = dynamic_traits_function
 
+        # Deal with dynamic traits function
         if self.dynamic_traits_function:
             self.dynamic_traits_function_name = self.dynamic_traits_function.__name__
             self.has_dynamic_traits_function = True
@@ -220,11 +142,7 @@ class Agent(Base):
                 dynamic_traits_function_name, dynamic_traits_function
             )
 
-    def _initialize_answer_question_directly(
-        self,
-        answer_question_directly_source_code,
-        answer_question_directly_function_name,
-    ) -> None:
+        # Deal with direct answer function
         if answer_question_directly_source_code:
             self.answer_question_directly_function_name = (
                 answer_question_directly_function_name
@@ -236,56 +154,18 @@ class Agent(Base):
             bound_method = types.MethodType(protected_method, self)
             setattr(self, "answer_question_directly", bound_method)
 
-    def _initialize_traits_presentation_template(
-        self, traits_presentation_template
-    ) -> None:
+        self._check_dynamic_traits_function()
+
+        self.current_question = 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:
@@ -361,111 +241,59 @@ class Agent(Base):
             else:
                 return self.dynamic_traits_function()
         else:
-            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.",
-            )
+            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
 
-    @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)
+        table = str(tabulate(data, headers=["keys", "values"], tablefmt="html"))
+        return f"<pre>{table}</pre>"
 
     def rename(
-        self,
-        old_name_or_dict: Union[str, dict[str, str]],
-        new_name: Optional[str] = None,
+        self, old_name_or_dict: Union[str, dict], 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})
+        >>> a.rename("age", "years") == Agent(traits = {'years': 10, 'hair': 'brown', 'height': 5.5})
         True
 
-        >>> newa.rename({'years': 'smage'}) == Agent(traits = {'smage': 10, 'hair': 'brown', 'height': 5.5})
-        True
+        >>> a.rename({'years': 'smage'})
+        Agent(traits = {'hair': 'brown', 'height': 5.5, 'smage': 10})
 
         """
-        self._check_before_modifying_traits()
+        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
+
         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)
+            self._rename(old_name_or_dict, new_name)
+            return self
 
         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})
-
+        >>> a.rename("age", "years") == Agent(traits = {'years': 10, 'hair': 'brown', 'height': 5.5})
+        True
         """
-        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
+        self.traits[new_name] = self.traits.pop(old_name)
+        return self
 
     def __getitem__(self, key):
         """Allow for accessing traits using the bracket notation.
@@ -496,7 +324,7 @@ class Agent(Base):
 
     def add_direct_question_answering_method(
         self,
-        method: DirectAnswerMethod,
+        method: Callable,
         validate_response: bool = False,
         translate_response: bool = False,
     ) -> None:
@@ -525,12 +353,6 @@ 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:
@@ -549,11 +371,12 @@ 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.
 
@@ -568,9 +391,7 @@ class Agent(Base):
         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
+        from edsl import Model, Scenario
 
         cache = cache
         self.current_question = question
@@ -581,12 +402,13 @@ 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
@@ -606,7 +428,6 @@ 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.
@@ -621,7 +442,7 @@ class Agent(Base):
 
         >>> a = Agent(traits = {})
         >>> a.add_direct_question_answering_method(lambda self, question, scenario: "I am a direct answer.")
-        >>> from edsl.questions.QuestionFreeText import QuestionFreeText
+        >>> from edsl import QuestionFreeText
         >>> q = QuestionFreeText.example()
         >>> a.answer_question(question = q, cache = False).answer
         'I am a direct answer.'
@@ -636,35 +457,16 @@ 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,
@@ -672,25 +474,53 @@ 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.language_models.model import Model
-        from edsl.scenarios.Scenario import Scenario
+        from edsl import Model
+        from edsl 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()
 
-        invigilator_class = self._get_invigilator_class(question)
+        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 = invigilator_class(
             self,
@@ -702,24 +532,22 @@ 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}, codebook = {'age': 'Their age is'})
-        >>> 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})
 
 
         >>> a.select("age", "height")
-        Agent(traits = {'age': 10, 'height': 5.5}, codebook = {'age': 'Their age is'})
+        Agent(traits = {'age': 10, 'height': 5.5})
 
-        >>> a.select("height")
-        Agent(traits = {'height': 5.5})
+        >>> a.select("age")
+        Agent(traits = {'age': 10})
 
         """
 
@@ -728,17 +556,7 @@ class Agent(Base):
         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
+        return Agent(traits={trait: self.traits[trait] for trait in traits_to_select})
 
     def __add__(self, other_agent: Optional[Agent] = None) -> Agent:
         """
@@ -757,10 +575,6 @@ 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
@@ -769,14 +583,9 @@ class Agent(Base):
                 f"The agents have overlapping traits: {common_traits}."
             )
         else:
-            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
+            new_agent = Agent(traits=copy.deepcopy(self.traits))
+            new_agent.traits.update(other_agent.traits)
+            return new_agent
 
     def __eq__(self, other: Agent) -> bool:
         """Check if two agents are equal.
@@ -793,11 +602,7 @@ class Agent(Base):
         return self.data == other.data
 
     def __getattr__(self, name):
-        """
-        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
-        >>> a.age
-        10
-        """
+        # This will be called only if 'name' is not found in the usual places
         if name == "has_dynamic_traits_function":
             return self.has_dynamic_traits_function
 
@@ -819,6 +624,12 @@ 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__
@@ -829,6 +640,14 @@ 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.
@@ -859,9 +678,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
@@ -875,23 +694,18 @@ 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["traits"] = dict(raw_data["traits"])
+                raw_data[
+                    "answer_question_directly_function_name"
+                ] = self.answer_question_directly_function_name
 
         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.
 
@@ -899,22 +713,9 @@ class Agent(Base):
 
         >>> 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'}
+        {'name': 'Steve', 'traits': {'age': 10, 'hair': 'brown', 'height': 5.5}, '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
+        d = copy.deepcopy(self.data)
         if add_edsl_version:
             from edsl import __version__
 
@@ -934,18 +735,7 @@ class Agent(Base):
         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)
+        return cls(**agent_dict)
 
     def _table(self) -> tuple[dict, list]:
         """Prepare generic table data."""
@@ -956,15 +746,10 @@ 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
-        >>> 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})
-        """
+        """Adds a trait to an agent and returns that agent"""
         if isinstance(trait_name_or_dict, dict) and value is None:
-            newagent = self.duplicate()
-            newagent.traits = {**self.traits, **trait_name_or_dict}
-            return newagent
+            self.traits.update(trait_name_or_dict)
+            return self
 
         if isinstance(trait_name_or_dict, dict) and value:
             raise AgentErrors(
@@ -972,9 +757,9 @@ class Agent(Base):
             )
 
         if isinstance(trait_name_or_dict, str):
-            newagent = self.duplicate()
-            newagent.traits = {**self.traits, **{trait_name_or_dict: value}}
-            return newagent
+            trait = trait_name_or_dict
+            self.traits[trait] = value
+            return self
 
         raise AgentErrors("Something is not right with adding a trait to an Agent")
 
@@ -987,9 +772,8 @@ class Agent(Base):
         >>> 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
+        _ = self.traits.pop(trait)
+        return self
 
     def translate_traits(self, values_codebook: dict) -> Agent:
         """Translate traits to a new codebook.
@@ -1000,15 +784,32 @@ class Agent(Base):
 
         :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
+                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
 
     @classmethod
     def example(cls, randomize: bool = False) -> Agent:
@@ -1016,9 +817,6 @@ 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})
@@ -1030,12 +828,10 @@ class Agent(Base):
 
         >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
         >>> print(a.code())
-        from edsl.agents.Agent import Agent
+        from edsl 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})"
-        )
+        return f"from edsl import Agent\nagent = Agent(traits={self.traits})"
 
 
 def main():