edsl 0.1.46__py3-none-any.whl → 0.1.48__py3-none-any.whl

edsl/agents/{Agent.py → agent.py}

@@ -1,9 +1,53 @@
-"""An Agent is an AI agent that can reference a set of traits in answering questions."""
+"""Agent module for creating and managing AI agents with traits and question-answering capabilities.
+
+This module provides the Agent class, which represents an AI agent with customizable traits,
+instructions, and optional direct question-answering methods. Agents are the primary entities
+that answer questions in the EDSL framework, and they can be configured with:
+
+- Traits: Key-value pairs representing agent attributes (age, occupation, preferences, etc.)
+- Instructions: Directives for how the agent should answer questions
+- Codebooks: Human-readable descriptions for traits used in prompts
+- Dynamic traits: Functions that can modify traits based on questions
+- Direct answering methods: Functions that can answer specific questions without using an LLM
+
+Agents can be combined, modified, and used to answer various question types through different
+invigilator mechanisms.
+
+Codebook and Trait Rendering
+---------------------------
+One of the key features of the Agent class is its ability to use codebooks to improve
+how traits are presented to language models. A codebook is a dictionary that maps trait keys 
+to human-readable descriptions:
+
+```python
+traits = {"age": 30, "occupation": "doctor"}
+codebook = {"age": "Age in years", "occupation": "Current profession"}
+agent = Agent(traits=traits, codebook=codebook)
+```
+
+When an agent with a codebook generates prompts, it will use these descriptions
+instead of the raw trait keys, creating more natural and descriptive prompts:
+
+Without codebook: "Your traits: {'age': 30, 'occupation': 'doctor'}"
+With codebook: 
+```
+Your traits:
+Age in years: 30
+Current profession: doctor
+```
+
+This approach helps language models better understand the traits and can lead to
+more natural responses that properly interpret the agent's characteristics.
+"""
 
 from __future__ import annotations
 import copy
 import inspect
 import types
+import warnings
+from uuid import uuid4
+from contextlib import contextmanager
+
 from typing import (
     Callable,
     Optional,
@@ -14,23 +58,21 @@ from typing import (
     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 ..caching import Cache
+    from ..surveys import Survey
+    from ..scenarios import Scenario
+    from ..language_models import LanguageModel
+    from ..surveys.memory import MemoryPlan
+    from ..questions import QuestionBase
+    from ..invigilators import InvigilatorBase
+    from ..prompts import Prompt
+    from ..questions import QuestionBase
+    from ..key_management import KeyLookup
 
 
 @runtime_checkable
@@ -40,52 +82,92 @@ class DirectAnswerMethod(Protocol):
     def __call__(self, self_: A, question: QuestionBase, scenario: Scenario) -> Any: ...
 
 
-from uuid import uuid4
+from ..base import Base
+from ..scenarios import Scenario
+from ..questions import QuestionScenarioRenderError
+from ..data_transfer_models import AgentResponseDict
+from ..utilities import (
+    sync_wrapper,
+    create_restricted_function,
+    dict_hash,
+    remove_edsl_version,
+)
 
-from edsl.Base import Base
-from edsl.exceptions.questions import QuestionScenarioRenderError
 
-from edsl.exceptions.agents import (
+from .exceptions import (
     AgentErrors,
     AgentCombinationError,
     AgentDirectAnswerFunctionError,
     AgentDynamicTraitsFunctionError,
 )
 
-from edsl.agents.descriptors import (
+from .descriptors import (
     TraitsDescriptor,
     CodebookDescriptor,
     InstructionDescriptor,
     NameDescriptor,
 )
-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.scenarios.Scenario import Scenario
 
 
 class AgentTraits(Scenario):
-    """A class representing the traits of an agent."""
+    """A class representing the traits of an agent.
+    
+    AgentTraits inherits from Scenario to provide a structured way to store and
+    access agent characteristics. This allows agent traits to be handled consistently
+    throughout the EDSL framework, including for presentation in prompts.
+    
+    Attributes:
+        data: Dictionary containing the agent's traits as key-value pairs
+    """
 
     def __repr__(self):
+        """Generate a string representation of the agent traits.
+        
+        Returns:
+            str: String representation of the agent traits dictionary
+        """
         return f"{self.data}"
 
 
 class Agent(Base):
-    """An class representing an agent that can answer questions."""
+    """A class representing an AI agent with customizable traits that can answer questions.
+    
+    An Agent in EDSL represents an entity with specific characteristics (traits) that can
+    answer questions through various mechanisms. Agents can use language models to generate
+    responses based on their traits, directly answer questions through custom functions, or
+    dynamically adjust their traits based on the questions being asked.
+    
+    Key capabilities:
+    - Store and manage agent characteristics (traits)
+    - Provide instructions on how the agent should answer
+    - Support for custom codebooks to provide human-readable trait descriptions
+    - Integration with multiple question types and language models
+    - Combine agents to create more complex personas
+    - Customize agent behavior through direct answering methods
+    
+    Agents are used in conjunction with Questions, Scenarios, and Surveys to create
+    structured interactions with language models.
+    """
 
     __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."""
 
+    # Trait storage using descriptors for validation and management
     _traits = TraitsDescriptor()
+    
+    # Codebook maps trait keys to human-readable descriptions
+    # This improves prompt readability and LLM understanding
+    # Example: {'age': 'Age in years'} → "Age in years: 30" instead of "age: 30"
     codebook = CodebookDescriptor()
+    
+    # Instructions for how the agent should answer questions
     instruction = InstructionDescriptor()
+    
+    # Optional name identifier for the agent
     name = NameDescriptor()
+    
+    # Default values for function-related attributes
     dynamic_traits_function_name = ""
     answer_question_directly_function_name = ""
     has_dynamic_traits_function = False
@@ -103,57 +185,67 @@ class Agent(Base):
         answer_question_directly_source_code: Optional[str] = None,
         answer_question_directly_function_name: Optional[str] = None,
     ):
-        """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})
+        """Initialize a new Agent instance with specified traits and capabilities.
+
+        Args:
+            traits: Dictionary of agent characteristics (e.g., {"age": 30, "occupation": "doctor"})
+            name: Optional name identifier for the agent
+            codebook: Dictionary mapping trait keys to human-readable descriptions for prompts.
+                This provides more descriptive labels for traits when rendering prompts.
+                For example, {'age': 'Age in years'} would display "Age in years: 30" instead of "age: 30".
+            instruction: Directive for how the agent should answer questions
+            traits_presentation_template: Jinja2 template for formatting traits in prompts
+            dynamic_traits_function: Function that can modify traits based on questions
+            dynamic_traits_function_source_code: Source code string for the dynamic traits function
+            dynamic_traits_function_name: Name of the dynamic traits function
+            answer_question_directly_source_code: Source code for direct question answering method
+            answer_question_directly_function_name: Name of the direct answering function
+            
+        The Agent class brings together several key concepts:
+        
+        Traits
+        ------
+        Traits are key-value pairs that define an agent's characteristics. These are used
+        to construct a prompt that guides the language model on how to respond.
+        
+        Example:
+        >>> 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.")
+        
+        Traits Presentation
+        ------------------
+        The traits_presentation_template controls how traits are formatted in prompts.
+        It uses Jinja2 templating to insert trait values.
+        
+        Example:
+        >>> 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):
-
+        
+        Codebooks
+        ---------
+        Codebooks provide human-readable descriptions for traits in prompts.
+        
+        Example:
         >>> 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}}")
+        >>> 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.
-
+        Instructions guide how the agent should answer questions. If not provided,
+        a default instruction is used.
+        
         >>> 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`.
-
+        
+        For details on how these components are used to construct prompts, see
+        :py:class:`edsl.agents.Invigilator.InvigilatorBase`.
         """
         self._initialize_basic_attributes(traits, name, codebook)
         self._initialize_instruction(instruction)
@@ -170,13 +262,25 @@ class Agent(Base):
         self.current_question = None
 
     def _initialize_basic_attributes(self, traits, name, codebook) -> None:
-        """Initialize the basic attributes of the agent."""
+        """Initialize the basic attributes of the agent.
+        
+        Args:
+            traits: Dictionary of agent characteristics
+            name: Name identifier for the agent
+            codebook: Dictionary mapping trait keys to descriptions
+        """
         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."""
+        """Initialize the instruction for how the agent should answer questions.
+        
+        If no instruction is provided, uses the default instruction.
+        
+        Args:
+            instruction: Directive for how the agent should answer questions
+        """
         if instruction is None:
             self.instruction = self.default_instruction
             self._instruction = self.default_instruction
@@ -189,13 +293,71 @@ class Agent(Base):
     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."""
+        """Initialize the template for presenting agent traits in prompts.
+        
+        This method sets up how the agent's traits will be formatted in language model prompts.
+        The template is a Jinja2 template string that can reference trait values and other
+        agent properties.
+        
+        If no template is provided:
+        - If a codebook exists, the method creates a template that displays each trait with its 
+          codebook description instead of the raw key names (e.g., "Age in years: 30" instead of "age: 30")
+        - Without a codebook, it uses a default template that displays all traits as a dictionary
+        
+        Custom templates always take precedence over automatically generated ones, giving users
+        complete control over how traits are presented.
+        
+        Args:
+            traits_presentation_template: Optional Jinja2 template string for formatting traits.
+                If not provided, a default template will be generated.
+                
+        Examples:
+            With no template or codebook, traits are shown as a dictionary:
+            
+            >>> agent = Agent(traits={"age": 25, "occupation": "engineer"})
+            >>> str(agent.prompt().text)
+            'Your traits: {\'age\': 25, \'occupation\': \'engineer\'}'
+            
+            With a codebook but no custom template, traits are shown with descriptions:
+            
+            >>> codebook = {"age": "Age in years", "occupation": "Current profession"}
+            >>> agent = Agent(traits={"age": 25, "occupation": "engineer"}, codebook=codebook)
+            >>> print(agent.prompt().text)  # doctest: +NORMALIZE_WHITESPACE
+            Your traits:
+            Age in years: 25
+            Current profession: engineer
+            
+            With a custom template, that format is used regardless of codebook:
+            
+            >>> template = "Person: {{age}} years old, works as {{occupation}}"
+            >>> agent = Agent(traits={"age": 25, "occupation": "engineer"}, 
+            ...               codebook=codebook, traits_presentation_template=template)
+            >>> str(agent.prompt().text)
+            'Person: 25 years old, works as engineer'
+        """
         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}}"
+            # Set the default template based on whether a codebook exists
+            if self.codebook:
+                # Create a template that uses the codebook descriptions
+                traits_lines = []
+                for trait_key in self.traits.keys():
+                    if trait_key in self.codebook:
+                        # Use codebook description if available
+                        traits_lines.append(f"{self.codebook[trait_key]}: {{{{ {trait_key} }}}}")
+                    else:
+                        # Fall back to raw key for traits without codebook entries
+                        traits_lines.append(f"{trait_key}: {{{{ {trait_key} }}}}")
+                
+                # Join all trait lines with newlines
+                self.traits_presentation_template = "Your traits:\n" + "\n".join(traits_lines)
+            else:
+                # Use the standard dictionary format if no codebook
+                self.traits_presentation_template = "Your traits: {{traits}}"
+                
             self.set_traits_presentation_template = False
 
     def _initialize_dynamic_traits_function(
@@ -204,7 +366,17 @@ class Agent(Base):
         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."""
+        """Initialize a function that can dynamically modify agent traits based on questions.
+        
+        This allows traits to change based on the question being asked, enabling
+        more sophisticated agent behaviors. The function can be provided directly
+        or as source code that will be compiled.
+        
+        Args:
+            dynamic_traits_function: Function object that returns a dictionary of traits
+            dynamic_traits_function_source_code: Source code string for the function
+            dynamic_traits_function_name: Name to assign to the function
+        """
         # Deal with dynamic traits function
         self.dynamic_traits_function = dynamic_traits_function
 
@@ -225,6 +397,16 @@ class Agent(Base):
         answer_question_directly_source_code,
         answer_question_directly_function_name,
     ) -> None:
+        """Initialize a method for the agent to directly answer questions without using an LLM.
+        
+        This allows creating agents that answer programmatically rather than through
+        language model generation. The direct answering method can be provided as
+        source code that will be compiled and bound to this agent instance.
+        
+        Args:
+            answer_question_directly_source_code: Source code for the direct answering method
+            answer_question_directly_function_name: Name to assign to the method
+        """
         if answer_question_directly_source_code:
             self.answer_question_directly_function_name = (
                 answer_question_directly_function_name
@@ -248,58 +430,161 @@ class Agent(Base):
             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!'
+        """Create a deep copy of this agent with all its traits and capabilities.
+        
+        This method creates a completely independent copy of the agent, including
+        all its traits, codebook, instructions, and special functions like dynamic
+        traits and direct answering methods.
+        
+        Returns:
+            Agent: A new agent instance that is functionally identical to this one
+            
+        Examples:
+            Create a duplicate agent and verify it's equal but not the same object:
+            
+            >>> a = Agent(traits={"age": 10, "hair": "brown", "height": 5.5}, 
+            ...           codebook={'age': 'Their age is'})
+            >>> a2 = a.duplicate()
+            >>> a2 == a  # Functionally equivalent
+            True
+            >>> id(a) == id(a2)  # But different objects
+            False
+            
+            Duplicating preserves direct answering methods:
+            
+            >>> 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.'
+            
+            Duplicating preserves custom instructions:
+            
+            >>> a = Agent(traits={'age': 10}, instruction="Have fun!")
+            >>> a2 = a.duplicate()
+            >>> a2.instruction
+            'Have fun!'
         """
         new_agent = Agent.from_dict(self.to_dict())
+        
+        # Transfer direct answering method if present
         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)
+            
+        # Transfer dynamic traits function if present
         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
+        """Get the agent's persona template as a Prompt object.
+        
+        This property provides access to the template that formats the agent's traits
+        for presentation in prompts. The template is wrapped in a Prompt object
+        that supports rendering with variable substitution.
+        
+        Returns:
+            Prompt: A prompt object containing the traits presentation template
+        """
+        from ..prompts 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}\""")
+        """Generate a formatted prompt containing the agent's traits.
+        
+        This method renders the agent's traits presentation template with the
+        agent's traits and codebook, creating a formatted prompt that can be
+        used in language model requests.
+        
+        The method is dynamic and responsive to changes in the agent's state:
+        
+        1. If a custom template was explicitly set during initialization, it will be used
+        2. If using the default template and the codebook has been updated since
+           initialization, this method will recreate the template to reflect the current
+           codebook values
+        3. The template is rendered with access to all trait values, the complete traits
+           dictionary, and the codebook
+           
+        The template rendering makes the following variables available:
+        - All individual trait keys (e.g., {{age}}, {{occupation}})
+        - The full traits dictionary as {{traits}}
+        - The codebook as {{codebook}}
+        
+        Returns:
+            Prompt: A Prompt object containing the rendered template
+            
+        Raises:
+            QuestionScenarioRenderError: If any template variables remain undefined
+            
+        Examples:
+            Basic trait rendering without a codebook:
+            
+            >>> agent = Agent(traits={"age": 10, "hair": "brown", "height": 5.5})
+            >>> agent.prompt()
+            Prompt(text=\"""Your traits: {'age': 10, 'hair': 'brown', 'height': 5.5}\""")
+            
+            Trait rendering with a codebook (more readable format):
+            
+            >>> codebook = {"age": "Age in years", "hair": "Hair color"}
+            >>> agent = Agent(traits={"age": 10, "hair": "brown"}, codebook=codebook)
+            >>> print(agent.prompt().text)  # doctest: +NORMALIZE_WHITESPACE
+            Your traits:
+            Age in years: 10
+            Hair color: brown
+            
+            Adding a codebook after initialization updates the rendering:
+            
+            >>> agent = Agent(traits={"age": 30, "occupation": "doctor"})
+            >>> initial_prompt = agent.prompt()
+            >>> "Your traits: {" in initial_prompt.text
+            True
+            >>> agent.codebook = {"age": "Age", "occupation": "Profession"}
+            >>> updated_prompt = agent.prompt()
+            >>> "Age: 30" in updated_prompt.text
+            True
+            >>> "Profession: doctor" in updated_prompt.text
+            True
+            
+            Custom templates can reference any trait directly:
+            
+            >>> template = "Profile: {{age}} year old {{occupation}}"
+            >>> agent = Agent(traits={"age": 45, "occupation": "teacher"}, 
+            ...               traits_presentation_template=template)
+            >>> agent.prompt().text
+            'Profile: 45 year old teacher'
         """
+        # If using the default template and the codebook has been updated since initialization,
+        # recreate the template to use the current codebook
+        if not self.set_traits_presentation_template and self.codebook:
+            # Create a template that uses the codebook descriptions
+            traits_lines = []
+            for trait_key in self.traits.keys():
+                if trait_key in self.codebook:
+                    # Use codebook description if available
+                    traits_lines.append(f"{self.codebook[trait_key]}: {{{{ {trait_key} }}}}")
+                else:
+                    # Fall back to raw key for traits without codebook entries
+                    traits_lines.append(f"{trait_key}: {{{{ {trait_key} }}}}")
+            
+            # Join all trait lines with newlines
+            self.traits_presentation_template = "Your traits:\n" + "\n".join(traits_lines)
+            
+        # Create a dictionary with traits, a reference to all traits, and the codebook
         replacement_dict = (
             self.traits | {"traits": self.traits} | {"codebook": self.codebook}
         )
+        
+        # Check for any undefined variables in the template
         if undefined := self.agent_persona.undefined_template_variables(
             replacement_dict
         ):
@@ -310,28 +595,41 @@ class Agent(Base):
             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: ...
+        """Validate that the dynamic traits function has the correct signature.
+        
+        This method checks if the dynamic traits function (if present) has the correct
+        parameter list. The function should either take no parameters or a single
+        parameter named 'question'.
+        
+        Raises:
+            AgentDynamicTraitsFunctionError: If the function signature is invalid
+            
+        Examples:
+            Valid function with 'question' parameter:
+            
+            >>> def f(question): return {"age": 10, "hair": "brown", "height": 5.5}
+            >>> a = Agent(dynamic_traits_function=f)
+            >>> a._check_dynamic_traits_function()
+            
+            Invalid function with extra parameters:
+            
+            >>> def g(question, poo): return {"age": 10, "hair": "brown", "height": 5.5}
+            >>> a = Agent(dynamic_traits_function=g)
+            Traceback (most recent call last):
+            ...
+            edsl.agents.exceptions.AgentDynamicTraitsFunctionError: ...
         """
         if self.has_dynamic_traits_function:
             sig = inspect.signature(self.dynamic_traits_function)
+            
             if "question" in sig.parameters:
+                # If it has 'question' parameter, it should be the only one
                 if len(sig.parameters) > 1:
                     raise AgentDynamicTraitsFunctionError(
                         message=f"The dynamic traits function {self.dynamic_traits_function} has too many parameters. It should only have one parameter: 'question'."
                     )
             else:
+                # If it doesn't have 'question', it shouldn't have any parameters
                 if len(sig.parameters) > 0:
                     raise AgentDynamicTraitsFunctionError(
                         f"""The dynamic traits function {self.dynamic_traits_function} has too many parameters. It should have no parameters or 
@@ -340,27 +638,33 @@ class Agent(Base):
 
     @property
     def traits(self) -> dict[str, str]:
-        """An agent's traits, which is a dictionary.
-
-        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.
-
-        Example:
-
-        >>> a = Agent(traits = {"age": 10, "hair": "brown", "height": 5.5})
-        >>> a.traits
-        {'age': 10, 'hair': 'brown', 'height': 5.5}
-
+        """Get the agent's traits, potentially using dynamic generation.
+        
+        This property provides access to the agent's traits, either from the stored
+        traits dictionary or by calling a dynamic traits function if one is defined.
+        If a dynamic traits function is used, it may take the current question as a
+        parameter to generate context-aware traits.
+        
+        Returns:
+            dict: Dictionary of agent traits (key-value pairs)
+            
+        Examples:
+            >>> a = Agent(traits={"age": 10, "hair": "brown", "height": 5.5})
+            >>> a.traits
+            {'age': 10, 'hair': 'brown', 'height': 5.5}
         """
         if self.has_dynamic_traits_function:
+            # Check if the function expects a question parameter
             sig = inspect.signature(self.dynamic_traits_function)
+            
             if "question" in sig.parameters:
+                # Call with the current question
                 return self.dynamic_traits_function(question=self.current_question)
             else:
+                # Call without parameters
                 return self.dynamic_traits_function()
         else:
+            # Return the stored traits
             return dict(self._traits)
 
     @contextmanager
@@ -383,8 +687,6 @@ class Agent(Base):
     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,
@@ -516,7 +818,6 @@ class Agent(Base):
         'I am a direct answer.'
         """
         if hasattr(self, "answer_question_directly"):
-            import warnings
 
             warnings.warn(
                 "Warning: overwriting existing answer_question_directly method"
@@ -525,12 +826,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:
@@ -568,9 +863,8 @@ 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 ..language_models import Model
+        from ..scenarios import Scenario
 
         cache = cache
         self.current_question = question
@@ -621,7 +915,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.questions import QuestionFreeText
         >>> q = QuestionFreeText.example()
         >>> a.answer_question(question = q, cache = False).answer
         'I am a direct answer.'
@@ -652,7 +946,7 @@ class Agent(Base):
         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 (
+        from ..invigilators import (
             InvigilatorHuman,
             InvigilatorFunctional,
             InvigilatorAI,
@@ -679,14 +973,14 @@ class Agent(Base):
         key_lookup: Optional["KeyLookup"] = None,
     ) -> "InvigilatorBase":
         """Create an Invigilator."""
-        from edsl.language_models.model import Model
-        from edsl.scenarios.Scenario import Scenario
+        from ..language_models import Model
+        from ..scenarios import Scenario
 
         model = model or Model()
         scenario = scenario or Scenario()
 
         if cache is None:
-            from edsl.data.Cache import Cache
+            from ..caching import Cache
 
             cache = Cache()
 
@@ -755,7 +1049,7 @@ class Agent(Base):
         >>> a1 + a1
         Traceback (most recent call last):
         ...
-        edsl.exceptions.agents.AgentCombinationError: The agents have overlapping traits: {'age'}.
+        edsl.agents.exceptions.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"})
@@ -888,8 +1182,6 @@ class Agent(Base):
         >>> 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]]:
@@ -1030,12 +1322,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.agents 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.agents import Agent\nagent = Agent(traits={self.traits})"
 
 
 def main():
@@ -1044,8 +1334,8 @@ def main():
 
     WARNING: Consume API credits
     """
-    from edsl.agents import Agent
-    from edsl.questions import QuestionMultipleChoice
+    from ..agents import Agent
+    from ..questions import QuestionMultipleChoice
 
     # a simple agent
     agent = Agent(traits={"age": 10, "hair": "brown", "height": 5.5})