From 66c8f4a81d7b5db3ad6c45b10ca33a444f65f5a3 Mon Sep 17 00:00:00 2001 From: Valentin De Matos <43698357+Thytu@users.noreply.github.com> Date: Sat, 4 Jan 2025 12:26:27 +0100 Subject: [PATCH] v0.3.0 (#2) * build: add dill as req * git: ignore .DS_Store * git: ignore .dill & .pkl * refactor(demo): remove useless print * refactor(demo): name the agents * !refactor: separate actions from agent and switch pickl to dill * refactor(demo): use new custom action system * doc: update docstrings * chore: bump version to 0.3.0 * test: add unit tests * test: remove 3.11-related tests (not needed) * chore(checkpoint): remove unused code * fix(action): think action no longer trigger talk action * fix(interaction): check if agent exist before storing interaction * fix(checkpoint): check if user tryies to create a new checkpoint manager * fix(test): resolve failing UTs * doc: README --------- Signed-off-by: Valentin De Matos --- .github/workflows/tests.yml | 33 +++ .gitignore | 6 +- README.md | 105 +++++++--- agentarium/Agent.py | 233 ++++++++-------------- agentarium/AgentInteractionManager.py | 7 + agentarium/CheckpointManager.py | 37 ++-- agentarium/__init__.py | 10 +- agentarium/actions/Action.py | 138 +++++++++++++ agentarium/actions/default_actions.py | 110 ++++++++++ agentarium/constant.py | 4 + agentarium/utils.py | 4 +- examples/1_basic_chat/demo.py | 3 +- examples/2_checkpointing/demo.py | 4 +- examples/3_adding_a_custom_action/demo.py | 48 ++--- pyproject.toml | 14 +- requirements.txt | 1 + tests/conftest.py | 56 ++++++ tests/test_action.py | 119 +++++++++++ tests/test_agent.py | 112 +++++++++++ tests/test_checkpoint_manager.py | 79 ++++++++ tests/test_interaction_manager.py | 107 ++++++++++ 21 files changed, 992 insertions(+), 238 deletions(-) create mode 100644 .github/workflows/tests.yml create mode 100644 agentarium/actions/Action.py create mode 100644 agentarium/actions/default_actions.py create mode 100644 agentarium/constant.py create mode 100644 tests/conftest.py create mode 100644 tests/test_action.py create mode 100644 tests/test_agent.py create mode 100644 tests/test_checkpoint_manager.py create mode 100644 tests/test_interaction_manager.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..5ba7f86 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,33 @@ +name: Run Tests + +on: + pull_request: + branches: [ main ] + push: + branches: [ main ] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[test]" + + - name: Run tests with pytest + env: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + run: | + pytest \ No newline at end of file diff --git a/.gitignore b/.gitignore index 5d12fc8..b2a8661 100644 --- a/.gitignore +++ b/.gitignore @@ -174,5 +174,9 @@ cython_debug/ *.pickle *.pkl +*.dill +*.pkl + # VS Code -*.code-workspace \ No newline at end of file +*.code-workspace +.DS_Store diff --git a/README.md b/README.md index e7bfb0e..199f51a 100644 --- a/README.md +++ b/README.md @@ -32,20 +32,22 @@ from agentarium import Agent agent1 = Agent(name="agent1") agent2 = Agent(name="agent2") -agent1.talk_to(agent2, "Hello, how are you?") -agent2.talk_to(agent1, "I'm fine, thank you!") +# Direct communication between agents +alice.talk_to(bob, "Hello Bob! I heard you're working on some interesting ML projects.") -agent1.act() # Same as agent.talk_to but it's the agent who decides what to do +# Agent autonomously decides its next action based on context +bob.act() ``` ## ✨ Features - **🤖 Advanced Agent Management**: Create and orchestrate multiple AI agents with different roles and capabilities -- **🔄 Robust Interaction Management**: Coordinate complex interactions between agents -- **💾 Checkpoint System**: Save and restore agent states and interactions -- **📊 Data Generation**: Generate synthetic data through agent interactions +- **🔄 Autonomous Decision Making**: Agents can make decisions and take actions based on their context +- **💾 Checkpoint System**: Save and restore agent states and interactions for reproducibility +- **🎭 Customizable Actions**: Define custom actions beyond the default talk/think capabilities +- **🧠 Memory & Context**: Agents maintain memory of past interactions for contextual responses +- **⚡ AI Integration**: Seamless integration with various AI providers through aisuite - **⚡ Performance Optimized**: Built for efficiency and scalability -- **🌍 Flexible Environment Configuration**: Define custom environments with YAML configuration files - **🛠️ Extensible Architecture**: Easy to extend and customize for your specific needs ## 📚 Examples @@ -54,32 +56,62 @@ agent1.act() # Same as agent.talk_to but it's the agent who decides what to do Create a simple chat interaction between agents: ```python -# examples/basic_chat/demo.py from agentarium import Agent -alice = Agent.create_agent() -bob = Agent.create_agent() +# Create agents with specific characteristics +alice = Agent.create_agent(name="Alice", occupation="Software Engineer") +bob = Agent.create_agent(name="Bob", occupation="Data Scientist") + +# Direct communication +alice.talk_to(bob, "Hello Bob! I heard you're working on some interesting projects.") -alice.talk_to(bob, "Hello Bob! I heard you're working on some interesting data science projects.") +# Let Bob autonomously decide how to respond bob.act() ``` -### Synthetic Data Generation -Generate synthetic data through agent interactions: +### Adding Custom Actions +Add new capabilities to your agents: + +```python +from agentarium import Agent, Action + +# Define a simple greeting action +def greet(name: str, **kwargs) -> str: + return f"Hello, {name}!" + +# Create an agent and add the greeting action +agent = Agent.create_agent(name="Alice") +agent.add_action( + Action( + name="GREET", + description="Greet someone by name", + parameters=["name"], + function=greet + ) +) + +# Use the custom action +agent.execute_action("GREET", "Bob") +``` + +### Using Checkpoints +Save and restore agent states: ```python -# examples/synthetic_data/demo.py from agentarium import Agent from agentarium.CheckpointManager import CheckpointManager +# Initialize checkpoint manager checkpoint = CheckpointManager("demo") -alice = Agent.create_agent() -bob = Agent.create_agent() +# Create and interact with agents +alice = Agent.create_agent(name="Alice") +bob = Agent.create_agent(name="Bob") alice.talk_to(bob, "What a beautiful day!") checkpoint.update(step="interaction_1") +# Save the current state checkpoint.save() ``` @@ -87,23 +119,38 @@ More examples can be found in the [examples/](examples/) directory. ## 📖 Documentation -### Environment Configuration -Configure your environment using YAML files: +### Agent Creation +Create agents with custom characteristics: + +```python +agent = Agent.create_agent( + name="Alice", + age=28, + occupation="Software Engineer", + location="San Francisco", + bio="A passionate developer who loves AI" +) +``` + +### LLM Configuration +Configure your LLM provider and credentials using a YAML file: ```yaml llm: - provider: "openai" # any provider supported by aisuite - model: "gpt-4o-mini" # any model supported by the provider + provider: "openai" # The LLM provider to use (any provider supported by aisuite) + model: "gpt-4" # The specific model to use from the provider -aisuite: # optional, credentials for aisuite - openai: - api_key: "sk-..." +aisuite: # (optional) Credentials for aisuite + openai: # Provider-specific configuration + api_key: "sk-..." # Your API key ``` ### Key Components -- **Agent**: Base class for creating AI agents -- **CheckpointManager**: Handles saving and loading of agent states +- **Agent**: Core class for creating AI agents with personalities and autonomous behavior +- **CheckpointManager**: Handles saving and loading of agent states and interactions +- **Action**: Base class for defining custom agent actions +- **AgentInteractionManager**: Manages and tracks all agent interactions ## 🤝 Contributing @@ -122,11 +169,5 @@ This project is licensed under the Apache 2.0 License - see the [LICENSE](LICENS ## 🙏 Acknowledgments -- Thanks to all contributors who have helped shape Agentarium -- Special thanks to the open-source community - ---- +Thanks to all contributors who have helped shape Agentarium 🫶 -
-Made with ❤️ by thytu -
\ No newline at end of file diff --git a/agentarium/Agent.py b/agentarium/Agent.py index 956cfc0..2fec016 100644 --- a/agentarium/Agent.py +++ b/agentarium/Agent.py @@ -4,18 +4,16 @@ import logging import aisuite as ai -from enum import Enum from copy import deepcopy -from typing import List, Callable +from .actions.Action import Action +from typing import List, Dict, Any from faker import Faker from .Interaction import Interaction from .AgentInteractionManager import AgentInteractionManager from .Config import Config from .utils import cache_w_checkpoint_manager - - -class DefaultValue(Enum): - NOT_PROVIDED = "NOT_PROVIDED" +from .actions.default_actions import default_actions +from .constant import DefaultValue faker = Faker() @@ -58,19 +56,6 @@ class Agent: Bio: [Bio of the person] """ - _default_actions = { - "THINK": { - "format": "[THINK][CONTENT]", - "prompt": "Think about something.", - "example": "[THINK][It's such a beautiful day!]", - }, - "TALK": { - "format": "[TALK][AGENT_ID][CONTENT]", - "prompt": "Talk to the agent with the given ID. Note that you must the agent ID, not the agent name.", - "example": "[TALK][123][Hello!]", - }, - } - _default_self_introduction_prompt = """ Informations about yourself: {agent_informations} @@ -134,12 +119,14 @@ def __init__(self, **kwargs): self._self_introduction_prompt = deepcopy(Agent._default_self_introduction_prompt) self._act_prompt = deepcopy(Agent._default_act_prompt) - self._actions = deepcopy(Agent._default_actions) - self._actions["TALK"]["function"] = self._talk_action_function - self._actions["THINK"]["function"] = self._think_action_function + self._actions = deepcopy(default_actions) if kwargs["actions"] == DefaultValue.NOT_PROVIDED else kwargs["actions"] self.storage = {} # Useful for storing data between actions. Note: not used by the agentarium system. + def __setstate__(self, state): + self.__dict__.update(state) # default __setstate__ + self._interaction_manager.register_agent(self) # add the agent to the interaction manager + @staticmethod @cache_w_checkpoint_manager def create_agent( @@ -150,6 +137,7 @@ def create_agent( occupation: str = DefaultValue.NOT_PROVIDED, location: str = DefaultValue.NOT_PROVIDED, bio: str = DefaultValue.NOT_PROVIDED, + actions: dict = DefaultValue.NOT_PROVIDED, **kwargs, ) -> Agent: """ @@ -167,7 +155,7 @@ def create_agent( - Occupation (randomly selected job) - Location (randomly selected city) - Bio (generated based on other characteristics) - + - Actions, by default "talk" and "think" actions are available (default actions or custom actions) Args: **kwargs: Dictionary of agent characteristics to use instead of generating them. Any characteristic not provided will be @@ -175,7 +163,17 @@ def create_agent( """ try: Agent._allow_init = True - return Agent(agent_id=agent_id, gender=gender, name=name, age=age, occupation=occupation, location=location, bio=bio, **kwargs) + return Agent( + agent_id=agent_id, + gender=gender, + name=name, + age=age, + occupation=occupation, + location=location, + bio=bio, + actions=actions, + **kwargs + ) finally: Agent._allow_init = False @@ -301,24 +299,30 @@ def _generate_agent_bio(agent_informations: dict) -> str: @cache_w_checkpoint_manager def act(self) -> str: """ - Generate and execute the agent's next action. + Generate and execute the agent's next action based on their current state. + + This method: + 1. Generates a self-introduction based on agent's characteristics and history + 2. Creates a prompt combining the self-introduction and available actions + 3. Uses the language model to decide the next action + 4. Parses and executes the chosen action - Uses the language model to determine and perform the agent's next - action based on their characteristics and interaction history. - The agent can either think to themselves or talk to another agent. + The agent's decision is based on: + - Their characteristics (personality, role, etc.) + - Their interaction history + - Available actions in their action space Returns: - str: The complete response from the language model, including - the agent's thoughts and chosen action. + Dict[str, Any]: A dictionary containing the action results, including: + - 'action': The name of the executed action + - Additional keys depending on the specific action executed Raises: - RuntimeError: If the action format is invalid or if the target - agent for a TALK action is not found. + RuntimeError: If no actions are available or if the chosen action is invalid """ - # [THINK][CONTENT]: Think about something. (i.e [THINK][It's such a beautiful day!]) - # [TALK][AGENT_ID][CONTENT]: Talk to the agent with the given ID. Note that you must the agent ID, not the agent name. (i.e [TALK][123][Hello!]) - # [CHATGPT][CONTENT]: Use ChatGPT. (i.e [CHATGPT][Hello!]) + if len(self._actions) == 0: + raise RuntimeError("No actions available for the agent to perform") self_introduction = self._self_introduction_prompt.format( agent_informations=self.agent_informations, @@ -327,7 +331,7 @@ def act(self) -> str: prompt = self._act_prompt.format( self_introduction=self_introduction, - actions="\n".join([f"{action['format']}: {action['prompt']} (i.e {action['example']})" for action in self._actions.values()]), + actions="\n".join([f"{action.get_format()}: {action.description if action.description else ''}" for action in self._actions.values()]), list_of_actions=list(self._actions.keys()), ) @@ -339,13 +343,14 @@ def act(self) -> str: temperature=0.4, ) - actions = re.search(r"(.*?)", response.choices[0].message.content, re.DOTALL).group(1).strip().split("]") - actions = [action.replace("[", "").replace("]", "").strip() for action in actions] + regex_result = re.search(r"(.*?)", response.choices[0].message.content, re.DOTALL).group(1).strip() + action_name, *args = [value.replace("[", "").replace("]", "").strip() for value in regex_result.split("]")] - if actions[0] in self._actions: - return self._actions[actions[0]]["function"](*actions[1:]) + if action_name not in self._actions: + logging.error(f"Received an invalid action: '{action_name=}' in the output: {response.choices[0].message.content}") + raise RuntimeError(f"Invalid action received: '{action_name}'.") - raise RuntimeError(f"Invalid action: '{actions[0]}'. Received output: {response.choices[0].message.content}") + return self._actions[action_name].function(*args, agent=self) # we pass the agent to the action function to allow it to use the agent's methods def dump(self) -> dict: """ @@ -386,68 +391,6 @@ def get_interactions(self) -> List[Interaction]: """ return self._interaction_manager.get_agent_interactions(self) - def _talk_action_function(self, *params) -> None: - """ - Send a message to another agent. - """ - if len(params) < 2: - raise RuntimeError(f"Received a TALK action with less than 2 arguments: {params}") - - if (receiver := self._interaction_manager.get_agent(params[0])) is None: - raise RuntimeError(f"Received a TALK action with an invalid agent ID: {params[0]}. {params=}") - - return self.talk_to(receiver, params[1]) - - @cache_w_checkpoint_manager - def talk_to(self, receiver: Agent, message: str) -> None: - """ - Send a message to another agent. - - Records an interaction where this agent sends a message to another agent. - - Args: - receiver (Agent): The agent to send the message to. - message (str): The content of the message to send. - """ - - self._interaction_manager.record_interaction(self, receiver, message) - - return { - "action": "TALK", - "sender": self.agent_id, - "receiver": receiver.agent_id, - "message": message, - } - - def _think_action_function(self, *params) -> None: - """ - Think about a message. - """ - if len(params) < 1: - raise RuntimeError(f"Received a THINK action with less than 1 argument: {params}") - - return self.think(params[0]) - - @cache_w_checkpoint_manager - def think(self, message: str) -> None: - """ - Think about a message. - - Records an interaction where this agent thinks about a message. - - Args: - message (str): The content of the message to think about. - """ - - self._interaction_manager.record_interaction(self, self, message) - - return { - "action": "THINK", - "sender": self.agent_id, - "receiver": self.agent_id, - "message": message, - } - def ask(self, message: str) -> None: """ Ask the agent a question and receive a contextually aware response. @@ -481,7 +424,7 @@ def ask(self, message: str) -> None: return response.choices[0].message.content - def add_new_action(self, action_descriptor: dict[str, str], action_function: Callable[[Agent, str], str | dict]) -> None: + def add_action(self, action: Action) -> None: """ Add a new action to the agent's capabilities. @@ -509,7 +452,7 @@ def add_new_action(self, action_descriptor: dict[str, str], action_function: Cal Example: ```python # Adding a ChatGPT integration action - agent.add_new_action( + agent.add_action( action_descriptor={ "format": "[CHATGPT][Your message here]", "prompt": "Use ChatGPT to have a conversation or ask questions.", @@ -533,59 +476,57 @@ def use_chatgpt(agent: Agent, message: str) -> dict: - Any 'action' key in the function's output dictionary will be overwritten """ - if "format" not in action_descriptor: - raise RuntimeError(f"Invalid action: {action_descriptor}, missing 'format'") + if action.name in self._actions: + raise RuntimeError(f"Invalid action: {action.name}, action already exists in the agent's action space") - name = action_descriptor["format"].split("[")[1].split("]")[0] + self._actions[action.name] = action - if len(name) < 1: - raise RuntimeError(f"Invalid action format: {action_descriptor['format']}, first name must be a single word") - - if name in self._actions: - raise RuntimeError(f"Action {action_descriptor} already exists") - - if "prompt" not in action_descriptor: - raise RuntimeError(f"Invalid action: {action_descriptor}, missing 'prompt'") + def execute_action(self, action_name: str, *args, **kwargs) -> Dict[str, Any]: + """ + Manually execute an action by name. - if "example" not in action_descriptor: - raise RuntimeError(f"Invalid action: {action_descriptor}, missing 'example'") + Args: + name: Name of the action to execute + *args, **kwargs: Arguments to pass to the action function - def standardize_action_output(fun): - """ - Standardizes the output format of action functions. + Returns: + Dict containing the action results + """ - This decorator ensures that all action functions return a dictionary with a - consistent format, containing at minimum an 'action' key with the action name. - If the original function returns a dictionary, its contents are preserved - (except for the 'action' key which may be overwritten). If it returns a - non-dictionary value, it is stored under the 'output' key. + if action_name not in self._actions: + raise RuntimeError(f"Invalid action: {action_name}, action does not exist in the agent's action space") - Args: - fun (Callable): The action function to decorate. + return self._actions[action_name].function(*args, **kwargs, agent=self) - Returns: - Callable: A wrapped function that standardizes the output format. - """ + def remove_action(self, action_name: str) -> None: + """ + Remove an action from the agent's action space. + """ + if action_name not in self._actions: + raise RuntimeError(f"Invalid action: {action_name}, action does not exist in the agent's action space") - def wrapper(*args, **kwargs): - fn_output = fun(self, *args, **kwargs) + del self._actions[action_name] - output = fn_output if type(fn_output) == dict else {"output": fn_output} + def talk_to(self, agent: Agent, message: str) -> None: + """ + Send a message from one agent to another and record the interaction. + """ - if "action" in output: - logging.warning(( - f"The action '{name}' returned an output with an 'action' key. " - "This is not allowed, it will be overwritten." - )) + if "talk" not in self._actions: + # Did you really removed the default "talk" action and expect the talk_to method to work? + raise RuntimeError("Talk action not found in the agent's action space.") - output["action"] = name + self.execute_action("talk", agent.agent_id, message) - return output + def think(self, message: str) -> None: + """ + Make the agent think about a message. + """ - return wrapper + if "think" not in self._actions: + raise RuntimeError("Think action not found in the agent's action space.") - self._actions[name] = action_descriptor - self._actions[name]["function"] = standardize_action_output(action_function) + self.execute_action("think", message) def reset(self) -> None: """ @@ -598,8 +539,8 @@ def reset(self) -> None: if __name__ == "__main__": interaction = Interaction( - sender=Agent(name="Alice", bio="Alice is a software engineer."), - receiver=Agent(name="Bob", bio="Bob is a data scientist."), + sender=Agent.create_agent(name="Alice", bio="Alice is a software engineer."), + receiver=Agent.create_agent(name="Bob", bio="Bob is a data scientist."), message="Hello Bob! I heard you're working on some interesting data science projects." ) diff --git a/agentarium/AgentInteractionManager.py b/agentarium/AgentInteractionManager.py index 421fa66..908bda4 100644 --- a/agentarium/AgentInteractionManager.py +++ b/agentarium/AgentInteractionManager.py @@ -90,6 +90,13 @@ def record_interaction(self, sender: Agent, receiver: Agent, message: str) -> No receiver (Agent): The agent receiving the interaction. message (str): The content of the interaction. """ + + if sender.agent_id not in self._agents: + raise ValueError(f"Sender agent {sender.agent_id} is not registered in the interaction manager.") + + if receiver.agent_id not in self._agents: + raise ValueError(f"Receiver agent {receiver.agent_id} is not registered in the interaction manager.") + interaction = Interaction(sender=sender, receiver=receiver, message=message) # Record in global interactions diff --git a/agentarium/CheckpointManager.py b/agentarium/CheckpointManager.py index 21d238b..a445ed4 100644 --- a/agentarium/CheckpointManager.py +++ b/agentarium/CheckpointManager.py @@ -1,9 +1,14 @@ import os -import pickle +import dill from collections import OrderedDict from .AgentInteractionManager import AgentInteractionManager +# NOTE: if someone knows how to fix this, please do it :D +import warnings +warnings.filterwarnings("ignore", category=dill.PicklingWarning) + + class CheckpointManager: """ @@ -26,6 +31,9 @@ def __new__(cls, name: str = None): if cls._instance is None: cls._instance = super(CheckpointManager, cls).__new__(cls) cls._instance._initialized = False + elif name and name != cls._instance.name: + # Allow only one instance of CheckpointManager + raise RuntimeError(f"CheckpointManager instance already exists with a different name: {name}") return cls._instance def __init__(self, name: str = "default"): @@ -40,10 +48,9 @@ def __init__(self, name: str = "default"): return self.name = name - self.path = f"{self.name}.pickle" + self.path = f"{self.name}.dill" self._interaction_manager = AgentInteractionManager() - self._state = OrderedDict() self._action_idx = 0 self.recorded_actions = [] @@ -53,38 +60,18 @@ def __init__(self, name: str = "default"): self._initialized = True - def update(self, step: str = None, save: bool = False): - """ - Update the checkpoint with the current state of the simulation. - """ - - self._state[step] = { - "agents": [agent.dump() for agent in self._interaction_manager._agents.values()], - "interactions": [interaction.dump() for interaction in self._interaction_manager._interactions], - } - - if save: - self.save() - - def _set_internal_state(self, state: OrderedDict): - self._state = state - def save(self) -> None: """ Save the current state of a simulation. """ - pickle.dump({"state": self._state, "actions": self.recorded_actions}, open(self.path, "wb")) - # should be imported before - # creating chatbots the next week? + dill.dump({"actions": self.recorded_actions}, open(self.path, "wb"), byref=True) def load(self) -> None: """ Load a simulation from a checkpoint. """ - from agentarium import Agent - env_data = pickle.load(open(self.path, "rb")) + env_data = dill.load(open(self.path, "rb")) - self._state = OrderedDict(env_data["state"]) self.recorded_actions = env_data["actions"] diff --git a/agentarium/__init__.py b/agentarium/__init__.py index 35de920..9f4e0a8 100644 --- a/agentarium/__init__.py +++ b/agentarium/__init__.py @@ -1,13 +1,13 @@ from .Agent import Agent -# from .Environment import Environment from .AgentInteractionManager import AgentInteractionManager +from .actions.Action import Action from .Interaction import Interaction -__version__ = "0.2.8" +__version__ = "0.3.0" __all__ = [ "Agent", - "Environment", "AgentInteractionManager", - "Interaction" -] \ No newline at end of file + "Interaction", + "Action", +] diff --git a/agentarium/actions/Action.py b/agentarium/actions/Action.py new file mode 100644 index 0000000..e381b41 --- /dev/null +++ b/agentarium/actions/Action.py @@ -0,0 +1,138 @@ +import logging + +from functools import wraps +from dataclasses import dataclass +from typing import Callable, Any, Dict, Union, Optional, List + + +# TODO: rename the output key "action" to "_action_name" +def standardize_action_output(action): + """ + Decorator that standardizes the output format of action functions. + + This decorator ensures that all action functions return a dictionary with a + consistent format, containing at minimum an 'action' key with the action name. + If the original function returns a dictionary, its contents are preserved + (except for the 'action' key which may be overwritten). If it returns a + non-dictionary value, it is stored under the 'output' key. + + Args: + fun (Callable): The action function to decorate. + + Returns: + Callable: A wrapped function that standardizes the output format. + """ + + def decorator(function): + @wraps(function) + def wrapper(*args, **kwargs): + + fn_output = function(*args, **kwargs) + + output = fn_output if isinstance(fn_output, dict) else {"output": fn_output} + + if "action" in output: + logging.warning(( + f"The action '{action.name}' returned an output with an 'action' key. " + "This is not allowed, it will be overwritten." + )) + + output["action"] = action.name + + return output + + return wrapper + + return decorator + + +def verify_agent_in_kwargs(function): + """ + Decorator that verifies the presence of an 'agent' key in kwargs. + + This decorator ensures that any action function receives an agent instance + through its kwargs. This is crucial for actions that need to interact with + or modify the agent's state. + + Args: + function (Callable): The action function to decorate. + + Returns: + Callable: A wrapped function that verifies the presence of 'agent' in kwargs. + + Raises: + RuntimeError: If 'agent' key is not found in kwargs. + """ + @wraps(function) + def wrapper(*args, **kwargs): + if "agent" not in kwargs: + raise RuntimeError("Action functions must receive an agent instance through kwargs") + return function(*args, **kwargs) + return wrapper + + +@dataclass +class Action: + """ + Represents an action that an agent can perform. + + Attributes: + name (str): Unique identifier for the action. Also used to remove the action from the agent's action space. + description (str, optional): Short description of what the action does. + parameters (Union[List[str], str]): List of parameters names or a single parameter name. + function: Callable that implements the action. + + Example: + ```python + action = Action( + name="CHATGPT", + description="Use ChatGPT", + parameters=["prompt"], + function=use_chatgpt, + ) + ``` + """ + + name: str + description: Optional[str] + parameters: Union[List[str], str] + function: Callable[[Any, ...], Dict[str, Any]] + + def __post_init__(self): + """Validate action attributes after initialization.""" + + if len(self.name) < 1: + raise ValueError("Action name must be non-empty") + + if isinstance(self.parameters, str): + self.parameters = [self.parameters] + + if not isinstance(self.parameters, list): + raise ValueError("Parameters must be a list of strings or a single string") + + if any(not isinstance(p, str) for p in self.parameters): + raise ValueError("Parameter names must be strings") + + if any(len(p) < 1 for p in self.parameters): + raise ValueError("Parameter names must be non-empty") + + # Store the function's module and name for pickling + if hasattr(self.function, '__module__') and hasattr(self.function, '__name__'): + self._function_module = self.function.__module__ + self._function_name = self.function.__name__ + else: + raise ValueError("Function must be a named function (not a lambda) defined at module level") + + self.function = standardize_action_output(self)(self.function) + + def get_format(self): + """Returns a string representation of the action's format for use in prompts. + + The format consists of the action name followed by its parameters, all enclosed in square brackets. + For example, an action named "TALK" with parameters ["agent_id", "message"] would return: + "[TALK][agent_id][message]" + + Returns: + str: The formatted string showing how to use the action in prompts. + """ + return f"[{self.name}]" + "".join([f"[{p}]" for p in self.parameters]) diff --git a/agentarium/actions/default_actions.py b/agentarium/actions/default_actions.py new file mode 100644 index 0000000..5b7ff40 --- /dev/null +++ b/agentarium/actions/default_actions.py @@ -0,0 +1,110 @@ +import logging + +from agentarium.utils import cache_w_checkpoint_manager +from agentarium.actions.Action import Action, verify_agent_in_kwargs + +logger = logging.getLogger(__name__) + +@verify_agent_in_kwargs +@cache_w_checkpoint_manager +def talk_action_function(*args, **kwargs): + """ + Send a message from one agent to another and record the interaction. + + Args: + *args: Variable length argument list where: + - First argument (str): The ID of the agent to send the message to + - Second argument (str): The content of the message to send + **kwargs: Arbitrary keyword arguments. Must contain 'agent' key with the Agent instance. + + Returns: + dict: A dictionary containing: + - sender (str): The ID of the sending agent + - receiver (str): The ID of the receiving agent + - message (str): The content of the message + + Raises: + RuntimeError: If less than 1 argument is provided, if 'agent' is not in kwargs, + or if the receiver agent ID is invalid. + """ + + if len(args) < 1: + raise RuntimeError(f"Received a TALK action with less than 1 argument: {args}") + + if "agent" not in kwargs: + raise RuntimeError(f"Couldn't find agent in {kwargs=} for TALK action") + + receiver = kwargs["agent"]._interaction_manager.get_agent(args[0]) + + if receiver is None: + logger.error(f"Received a TALK action with an invalid agent ID {args[0]=} {args=} {kwargs['agent']._interaction_manager._agents=}") + raise RuntimeError(f"Received a TALK action with an invalid agent ID: {args[0]}. {args=}") + + message = args[1] + + kwargs["agent"]._interaction_manager.record_interaction(kwargs["agent"], receiver, message) + + return { + "sender": kwargs["agent"].agent_id, + "receiver": receiver.agent_id, + "message": message, + } + +@verify_agent_in_kwargs +@cache_w_checkpoint_manager +def think_action_function(*args, **kwargs): + """ + Record an agent's internal thought. + + Args: + *params: Variable length argument list where the first argument is the thought content. + **kwargs: Arbitrary keyword arguments. Must contain 'agent' key with the Agent instance. + + Returns: + dict: A dictionary containing: + - sender (str): The ID of the thinking agent + - receiver (str): Same as sender (since it's an internal thought) + - message (str): The content of the thought + + Raises: + RuntimeError: If no parameters are provided for the thought content. + """ + + if len(args) < 1: + raise RuntimeError(f"Received a TALK action with less than 1 argument: {args}") + + if "agent" not in kwargs: + raise RuntimeError(f"Couldn't find agent in {kwargs=} for TALK action") + + message = args[0] + + kwargs["agent"]._interaction_manager.record_interaction(kwargs["agent"], kwargs["agent"], message) + + return { + "sender": kwargs["agent"].agent_id, + "receiver": kwargs["agent"].agent_id, + "message": message, + } + + +# Create action instances at module level +talk_action = Action( + name="talk", + description="Talk to the agent with the given ID.", + parameters=["agent_id", "message"], + function=talk_action_function, +) + +think_action = Action( + name="think", + description="Think about something.", + parameters=["content"], + function=think_action_function, +) + +default_actions = { + talk_action.name: talk_action, + think_action.name: think_action, +} + +# __all__ = ["talk_action", "think_action", "default_actions"] diff --git a/agentarium/constant.py b/agentarium/constant.py new file mode 100644 index 0000000..543718d --- /dev/null +++ b/agentarium/constant.py @@ -0,0 +1,4 @@ +from enum import Enum + +class DefaultValue(Enum): + NOT_PROVIDED = "NOT_PROVIDED" diff --git a/agentarium/utils.py b/agentarium/utils.py index 4f36de7..819f7de 100644 --- a/agentarium/utils.py +++ b/agentarium/utils.py @@ -2,7 +2,7 @@ import hashlib from typing import Dict, Any -from .CheckpointManager import CheckpointManager +# from .CheckpointManager import CheckpointManager def dict_hash(dictionary: Dict[str, Any]) -> str: @@ -33,6 +33,8 @@ def cache_w_checkpoint_manager(function): def wrapper(*args, **kwargs): + from .CheckpointManager import CheckpointManager + checkpoint_manager = CheckpointManager() expected_hash = checkpoint_manager.recorded_actions[checkpoint_manager._action_idx]["hash"] if len(checkpoint_manager.recorded_actions) > checkpoint_manager._action_idx else None diff --git a/examples/1_basic_chat/demo.py b/examples/1_basic_chat/demo.py index 74ac1b2..f292c2a 100644 --- a/examples/1_basic_chat/demo.py +++ b/examples/1_basic_chat/demo.py @@ -7,10 +7,9 @@ alice_agent.talk_to(bob_agent, "Hello Bob! I heard you're working on some interesting data science projects.") bob_agent.talk_to(alice_agent, "Hi Alice! Yes, I'm currently working on a machine learning model for natural language processing.") -alice_agent.act() # Here it's the agent that decides to talk to Bob +alice_agent.act() # Let the agents decide what to do :D bob_agent.act() -print("\n === Interactions ===\n") print("Alice's interactions:") print(alice_agent.get_interactions()) diff --git a/examples/2_checkpointing/demo.py b/examples/2_checkpointing/demo.py index 8c55b56..369b172 100644 --- a/examples/2_checkpointing/demo.py +++ b/examples/2_checkpointing/demo.py @@ -15,8 +15,8 @@ # Create two agents - their states will be tracked by the checkpoint manager # Even agent creation is cached - if these agents were created before, # they'll be loaded from cache with their exact same properties - alice = Agent.create_agent() - bob = Agent.create_agent() + alice = Agent.create_agent(name="Alice") + bob = Agent.create_agent(name="Bob") # When this interaction happens: # - First run: Actually calls the LLM and stores the result diff --git a/examples/3_adding_a_custom_action/demo.py b/examples/3_adding_a_custom_action/demo.py index 09351ff..af37826 100644 --- a/examples/3_adding_a_custom_action/demo.py +++ b/examples/3_adding_a_custom_action/demo.py @@ -5,16 +5,17 @@ 2. Adding a new action (ChatGPT integration) to the agent 3. Using both direct action calls and autonomous agent behavior -The demo specifically shows how to use the `add_new_action` method to extend an agent's capabilities: +The demo specifically shows how to use the `add_action` method to extend an agent's capabilities. Action Definition Format: - { - "prompt": str, # Description of what the action does - "format": str, # Format string showing how to use the action: [ACTION_NAME][ARG1][ARG2]... - "example": str, # Concrete example of the action's usage - } + Action( + name: str, # Unique identifier for the action + description: str, # Description of what the action does + parameters: List[str], # List of parameter names the action accepts + function: Callable # Function that implements the action behavior + ) -The action_function parameter should be a callable that takes (agent, *args) and returns +The action function should take (*args, **kwargs) as parameters and return either a string or a dictionary. If a dictionary is returned, it can contain any keys except 'action' which is reserved. The function's output will be automatically standardized to include the action name. @@ -22,14 +23,14 @@ import aisuite as ai from typing import Dict, Any -from agentarium.agent import Agent +from agentarium import Agent, Action # Initialize the AI client for ChatGPT interactions llm_client = ai.Client() -def use_chatgpt(agent: Agent, prompt: str, *args, **kwargs) -> Dict[str, str]: +def use_chatgpt(prompt: str, *args, **kwargs) -> Dict[str, str]: """ A custom action that allows an agent to interact with ChatGPT. This function demonstrates how to: @@ -53,6 +54,9 @@ def use_chatgpt(agent: Agent, prompt: str, *args, **kwargs) -> Dict[str, str]: - chatgpt_response: ChatGPT's response Note: The 'action' key will be automatically added by the framework when called by agent.act() """ + + agent: Agent = kwargs["agent"] # Every action receives the agent in kwargs + # Initialize chat history if it doesn't exist if "chatgpt_history" not in agent.storage: agent.storage["chatgpt_history"] = [] @@ -114,31 +118,29 @@ def main(): ) # Add the ChatGPT action to the agent's capabilities - # This demonstrates how to use add_new_action: + # This demonstrates how to use add_action: # 1. Define the action descriptor with prompt, format, and example # 2. Provide an action function that implements the behavior - agent.add_new_action( + agent.add_action( # Action descriptor - tells the agent how to use this action - { - "prompt": "Use ChatGPT to have a conversation or ask questions.", - "format": "[CHATGPT][Your message here]", # Format must start with action name in caps - "example": "[CHATGPT][What's the weather like today?]", - }, - # Action function - implements the actual behavior - action_function=use_chatgpt + Action( + name="CHATGPT", + description="Use ChatGPT to have a conversation or ask questions.", + parameters=["prompt"], + function=use_chatgpt + ) ) # Initialize the agent's thoughts agent.think("I've just been created and I can use ChatGPT to interact!") # Example 1: Direct action call - output = use_chatgpt(agent, "Hello! Can you help me learn about artificial intelligence?") - print_agent_output(output | {"action": "CHATGPT"}) # Add the action name to the output + output = agent.execute_action("CHATGPT", "Hello! Can you help me learn about artificial intelligence?") + print_agent_output(output) # Add the action name to the output # Example 2: Autonomous behavior - agent can now choose to use ChatGPT on its own - for _ in range(3): - output = agent.act() - print_agent_output(output) + output = agent.act() + print_agent_output(output) if __name__ == '__main__': diff --git a/pyproject.toml b/pyproject.toml index bea31cb..e68032c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "agentarium" -version = "0.2.8" +version = "0.3.0" authors = [ { name = "thytu" }, ] @@ -22,7 +22,19 @@ dependencies = [ "PyYAML>=6.0.1", "boto3>=1.35.86", "aisuite>=0.1.7", + "dill>=0.3.8", ] +[project.optional-dependencies] +test = [ + "pytest>=7.0.0", + "pytest-cov>=4.0.0", +] + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +addopts = "-v --cov=agentarium --cov-report=term-missing --cov-fail-under=80" + [project.urls] Homepage = "https://github.com/thytu/Agentarium" \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 6367431..190da0b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ faker>=33.1.0 PyYAML>=6.0.1 boto3>=1.35.86 aisuite>=0.1.7 +dill>=0.3.8 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..03a7fa7 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,56 @@ +import os +import pytest + +from agentarium import Agent +from agentarium.CheckpointManager import CheckpointManager +from agentarium.AgentInteractionManager import AgentInteractionManager + + +@pytest.fixture(autouse=True) +def cleanup_checkpoint_files(): + """Automatically clean up any checkpoint files after each test.""" + yield + # Clean up any .dill files in the current directory + for file in os.listdir(): + if file.endswith('.dill'): + os.remove(file) + + +@pytest.fixture +def base_agent(): + """Create a basic agent for testing.""" + agent = Agent.create_agent( + name="TestAgent", + age=25, + occupation="Software Engineer", + location="Test City", + bio="A test agent" + ) + yield agent + # Reset the agent after each test that uses this fixture + agent.reset() + + +@pytest.fixture +def interaction_manager(): + """Create a fresh interaction manager for testing.""" + return AgentInteractionManager() + + +@pytest.fixture +def checkpoint_manager(): + """Create a test checkpoint manager.""" + manager = CheckpointManager("test_checkpoint") + yield manager + # Cleanup is handled by cleanup_checkpoint_files fixture + + +@pytest.fixture +def agent_pair(): + """Create a pair of agents for interaction testing.""" + alice = Agent.create_agent(name="Alice", bio="Alice is a test agent") + bob = Agent.create_agent(name="Bob", bio="Bob is a test agent") + yield alice, bob + # Reset both agents after each test that uses this fixture + alice.reset() + bob.reset() diff --git a/tests/test_action.py b/tests/test_action.py new file mode 100644 index 0000000..5e1e31d --- /dev/null +++ b/tests/test_action.py @@ -0,0 +1,119 @@ +import pytest +from agentarium import Action, Agent + +def test_action_creation(): + """Test creating an action with valid parameters.""" + def test_function(*args, **kwargs): + return {"result": "success"} + + action = Action( + name="test", + description="A test action", + parameters=["param1", "param2"], + function=test_function + ) + + assert action.name == "test" + assert action.description == "A test action" + assert action.parameters == ["param1", "param2"] + assert action.function is not None + +def test_action_single_parameter(): + """Test creating an action with a single parameter string.""" + def test_function(*args, **kwargs): + return {"result": "success"} + + # Test with a single parameter string + action = Action( + name="test", + description="A test action", + parameters="param1", + function=test_function + ) + + assert action.parameters == ["param1"] + + # Test with a list of parameters + action = Action( + name="test", + description="A test action", + parameters=["param1"], + function=test_function + ) + + assert action.parameters == ["param1"] + + +def test_invalid_action_name(): + """Test that action creation fails with invalid name.""" + def test_function(*args, **kwargs): + return {"result": "success"} + + with pytest.raises(ValueError): + Action( + name="", + description="A test action", + parameters=["param1"], + function=test_function + ) + + +def test_invalid_parameters(): + """Test that action creation fails with invalid parameters.""" + def test_function(*args, **kwargs): + return {"result": "success"} + + # Test empty parameter name + with pytest.raises(ValueError): + Action( + name="test", + description="A test action", + parameters=[""], + function=test_function + ) + + # Test non-string parameter + with pytest.raises(ValueError): + Action( + name="test", + description="A test action", + parameters=[123], + function=test_function + ) + + +def test_action_format(): + """Test the action format string generation.""" + def test_function(*args, **kwargs): + return {"result": "success"} + + action = Action( + name="test", + description="A test action", + parameters=["param1", "param2"], + function=test_function + ) + + expected_format = "[test][param1][param2]" + assert action.get_format() == expected_format + + +def test_action_execution(base_agent): + """Test executing an action with an agent.""" + + def test_function(*args, **kwargs): + agent = kwargs["agent"] + return {"message": f"Action executed by {agent.name}"} + + action = Action( + name="test", + description="A test action", + parameters=["param"], + function=test_function + ) + + result = action.function("test_param", agent=base_agent) + + assert result["action"] == "test" + assert "message" in result + assert "TestAgent" in result["message"] diff --git a/tests/test_agent.py b/tests/test_agent.py new file mode 100644 index 0000000..0e066f4 --- /dev/null +++ b/tests/test_agent.py @@ -0,0 +1,112 @@ +import pytest + +from agentarium import Agent, Action +from agentarium.constant import DefaultValue + + +def test_agent_creation(base_agent): + """Test basic agent creation with default and custom values.""" + # Test with default values + assert base_agent.agent_id is not None + assert base_agent.name is not None + assert base_agent.age is not None + assert base_agent.occupation is not None + assert base_agent.location is not None + assert base_agent.bio is not None + + # Test without default values (nor bio) + custom_agent = Agent.create_agent( + name="Alice", + age=25, + occupation="Software Engineer", + location="San Francisco", + ) + assert isinstance(custom_agent.name, str) + assert isinstance(custom_agent.age, int) + assert isinstance(custom_agent.occupation, str) + assert isinstance(custom_agent.location, str) + assert isinstance(custom_agent.bio, str) + + +def test_agent_default_actions(agent_pair): + """Test that agents are created with default actions.""" + alice, _ = agent_pair + assert "talk" in alice._actions + assert "think" in alice._actions + + +def test_agent_custom_actions(base_agent): + """Test adding and using custom actions.""" + def custom_action(*args, **kwargs): + agent = kwargs["agent"] + return {"message": f"Custom action by {agent.name}"} + + custom_action_obj = Action( + name="custom", + description="A custom action", + parameters=["message"], + function=custom_action + ) + + base_agent.add_action(custom_action_obj) + + assert "custom" in base_agent._actions + result = base_agent.execute_action("custom", "test message") + assert result["action"] == "custom" + assert "message" in result + + +def test_agent_interaction(agent_pair): + """Test basic interaction between agents.""" + + alice, bob = agent_pair + + message = "Hello Bob!" + alice.talk_to(bob, message) + + # Check Alice's interactions + alice_interactions = alice.get_interactions() + assert len(alice_interactions) == 1 + assert alice_interactions[0].sender == alice + assert alice_interactions[0].receiver == bob + assert alice_interactions[0].message == message + + # Check Bob's interactions + bob_interactions = bob.get_interactions() + assert len(bob_interactions) == 1 + assert bob_interactions[0].sender == alice + assert bob_interactions[0].receiver == bob + assert bob_interactions[0].message == message + + +def test_agent_think(agent_pair): + """Test agent's ability to think.""" + + agent, _ = agent_pair + thought = "I should learn more about AI" + + agent.think(thought) + + interactions = agent.get_interactions() + assert len(interactions) == 1 + assert interactions[0].sender == agent + assert interactions[0].receiver == agent + assert interactions[0].message == thought + + +def test_invalid_action(base_agent): + """Test handling of invalid actions.""" + + with pytest.raises(RuntimeError): + base_agent.execute_action("nonexistent_action", "test") + + +def test_agent_reset(base_agent): + """Test resetting agent state.""" + base_agent.think("Initial thought") + + assert len(base_agent.get_interactions()) == 1 + + base_agent.reset() + assert len(base_agent.get_interactions()) == 0 + assert len(base_agent.storage) == 0 diff --git a/tests/test_checkpoint_manager.py b/tests/test_checkpoint_manager.py new file mode 100644 index 0000000..bbb1398 --- /dev/null +++ b/tests/test_checkpoint_manager.py @@ -0,0 +1,79 @@ +import os +import pytest +from agentarium import Agent +from agentarium.CheckpointManager import CheckpointManager + +@pytest.fixture +def checkpoint_manager(): + """Create a test checkpoint manager.""" + # Reset the singleton instance + CheckpointManager._instance = None + manager = CheckpointManager("test_checkpoint") + yield manager + # Cleanup after tests + if os.path.exists(manager.path): + os.remove(manager.path) + + +def test_checkpoint_manager_initialization(checkpoint_manager): + """Test checkpoint manager initialization.""" + assert checkpoint_manager.name == "test_checkpoint" + assert checkpoint_manager.path == "test_checkpoint.dill" + assert checkpoint_manager._action_idx == 0 + assert len(checkpoint_manager.recorded_actions) == 0 + + +def test_checkpoint_save_load(checkpoint_manager, agent_pair): + """Test saving and loading checkpoint data.""" + # Create agents and perform actions + alice, bob = agent_pair + + alice.talk_to(bob, "Hello Bob!") + bob.talk_to(alice, "Hi Alice!") + + # Save checkpoint + checkpoint_manager.save() + + # Create a new checkpoint manager and load data + new_manager = CheckpointManager("test_checkpoint") + new_manager.load() + + # Verify recorded actions were loaded + assert len(new_manager.recorded_actions) > 0 + assert new_manager.recorded_actions == checkpoint_manager.recorded_actions + + # Reset agents after test + alice.reset() + bob.reset() + + +def test_checkpoint_recording(checkpoint_manager, agent_pair): + """Test that actions are properly recorded.""" + alice, bob = agent_pair + + # Check that actions were recorded + assert len(checkpoint_manager.recorded_actions) == 2 # two agents created + + # Perform some actions + alice.talk_to(bob, "Hello!") + bob.think("I should respond...") + bob.talk_to(alice, "Hi there!") + + # Check that actions were recorded + assert len(checkpoint_manager.recorded_actions) == 3 + 2 # three actions + two agents created + + # Reset agents after test + alice.reset() + bob.reset() + +def test_singleton_behavior(): + """Test that CheckpointManager behaves like a singleton.""" + manager1 = CheckpointManager() + manager2 = CheckpointManager() + + # Both instances should reference the same object + assert manager1 is manager2 + + # Clean up + if os.path.exists(manager1.path): + os.remove(manager1.path) diff --git a/tests/test_interaction_manager.py b/tests/test_interaction_manager.py new file mode 100644 index 0000000..8a948f0 --- /dev/null +++ b/tests/test_interaction_manager.py @@ -0,0 +1,107 @@ +import pytest +from agentarium import Agent, Interaction +from agentarium.AgentInteractionManager import AgentInteractionManager + + +@pytest.fixture +def interaction_manager(): + """Create a test interaction manager.""" + return AgentInteractionManager() + + +def test_agent_registration(interaction_manager, agent_pair): + """Test registering agents with the interaction manager.""" + alice, bob = agent_pair + + # Verify agents are registered + assert alice.agent_id in interaction_manager._agents + assert bob.agent_id in interaction_manager._agents + + # Verify correct agent objects are stored + assert interaction_manager._agents[alice.agent_id] is alice + assert interaction_manager._agents[bob.agent_id] is bob + + +def test_interaction_recording(interaction_manager, agent_pair): + """Test recording interactions between agents.""" + alice, bob = agent_pair + message = "Hello Bob!" + + # Record an interaction + interaction_manager.record_interaction(alice, bob, message) + + # Verify interaction was recorded + assert len(alice.get_interactions()) == 1 + + interaction = alice.get_interactions()[0] + assert interaction.sender is alice + assert interaction.receiver is bob + assert interaction.message == message + +def test_get_agent_interactions(interaction_manager, agent_pair): + """Test retrieving agent interactions.""" + alice, bob = agent_pair + + # Record multiple interactions + interaction_manager.record_interaction(alice, bob, "Hello!") + interaction_manager.record_interaction(bob, alice, "Hi there!") + interaction_manager.record_interaction(alice, bob, "How are you?") + + # Get Alice's interactions + alice_interactions = interaction_manager.get_agent_interactions(alice) + assert len(alice_interactions) == 3 + + # Get Bob's interactions + bob_interactions = interaction_manager.get_agent_interactions(bob) + assert len(bob_interactions) == 3 + +def test_get_agent(interaction_manager, agent_pair): + """Test retrieving agents by ID.""" + alice, bob = agent_pair + + # Test getting existing agents + assert interaction_manager.get_agent(alice.agent_id) is alice + assert interaction_manager.get_agent(bob.agent_id) is bob + + # Test getting non-existent agent + assert interaction_manager.get_agent("nonexistent_id") is None + +def test_interaction_order(interaction_manager, agent_pair): + """Test that interactions are recorded in order.""" + alice, bob = agent_pair + + messages = ["First", "Second", "Third"] + + for msg in messages: + interaction_manager.record_interaction(alice, bob, msg) + + # Verify interactions are in order + interactions = interaction_manager.get_agent_interactions(alice) + for i, interaction in enumerate(interactions): + assert interaction.message == messages[i] + +def test_self_interaction(interaction_manager, agent_pair): + """Test recording self-interactions (thoughts).""" + alice, _ = agent_pair + thought = "I should learn more" + + interaction_manager.record_interaction(alice, alice, thought) + + interactions = interaction_manager.get_agent_interactions(alice) + assert len(interactions) == 1 + assert interactions[0].sender is alice + assert interactions[0].receiver is alice + assert interactions[0].message == thought + +def test_interaction_validation(interaction_manager, agent_pair): + """Test validation of interaction recording.""" + + alice, _ = agent_pair + + unregistered_agent = type("UnregisteredAgent", (object,), {"agent_id": "some_unregistered_id"}) + + with pytest.raises(ValueError): + interaction_manager.record_interaction(unregistered_agent, alice, "Hello") + + with pytest.raises(ValueError): + interaction_manager.record_interaction(alice, unregistered_agent, "Hello")