Shared state Initialization

Hello,

I am trying to understand how tools function with an easy example and while exploring the shared_state fundamentals, got into the following error using the below code.

from agency_swarm.tools import BaseTool
from pydantic import Field, model_validator

class MyCustomTool(BaseTool):
    """
    A brief description of what the custom tool does. 
    The docstring should clearly explain the tool's purpose and functionality.
    It will be used by the agent to determine when to use this tool.
    """

    # Define the fields with descriptions using Pydantic Field
    example_field: str = Field(
        ..., description="Description of the example field, explaining its purpose and usage for the Agent."
    )

    # Additional Pydantic fields as required
    # ...

    def run(self):
        """
        The implementation of the run method, where the tool's main functionality is executed.
        This method should utilize the fields defined above to perform the task.
        Doc string is not required for this method and will not be used by your agent.
        """
        self.shared_state.set("context", "This is my context")

        # Your custom tool logic goes here
        # do_something(self.example_field)

        # Return the result of the tool's operation as a string
        return "Result of MyCustomTool operation"

    # Has access to the complete model attributes after pydantic pre-validation.
    @model_validator(mode="after")
    def validate_field(self):
        if len(self.example_field) < 10:
            raise ValueError("The example field must be at least 10 characters long.")

if __name__ == "__main__":
    print(MyCustomTool(example_field="example field greater than 10 chars").run())

Error:

AttributeError: 'NoneType' object has no attribute 'set'

On further inspection, I could not find the initialization of the shared_state variable in the BaseTool class in order for the extended tools to use this variable. As we can see below, it is set to None by default.

class BaseTool(OpenAISchema, ABC):
    shared_state: ClassVar[SharedState] = None
    caller_agent: Any = None
    event_handler: Any = None
    one_call_at_a_time: bool = False

    @classmethod
    @property
    def openai_schema(cls):
        schema = super(BaseTool, cls).openai_schema

        properties = schema.get("parameters", {}).get("properties", {})

        properties.pop("caller_agent", None)
        properties.pop("shared_state", None)
        properties.pop("event_handler", None)
        properties.pop("one_call_at_a_time", None)

        required = schema.get("parameters", {}).get("required", [])
        if "caller_agent" in required:
            required.remove("caller_agent")
        if "shared_state" in required:
            required.remove("shared_state")
        if "event_handler" in required:
            required.remove("event_handler")
        if "one_call_at_a_time" in required:
            required.remove("one_call_at_a_time")

        return schema

    def model_dump(self, exclude=None, **kwargs):
        if exclude is None:
            exclude = {"caller_agent", "shared_state", "event_handler", "one_call_at_a_time"}
        else:
            exclude.update({"caller_agent", "shared_state", "event_handler", "one_call_at_a_time"})
        return super().model_dump(exclude=exclude, **kwargs)

    @abstractmethod
    def run(self, **kwargs):
        pass

Can someone guide on how to resolve this error or the right approach to use the shared_state variable ? btw, I have already referred the Advanced Usage section of the documentation

VRSEN / agency-swarm

Shared state Initialization #165