Repository: griptape-ai/griptape

Issue:
Tool Drivers don't seem to be serializing properly
Some Tool Drivers aren't serializing properly.

For example, the `AudioTranscriptionTool` takes an `audio_transcription_driver` - but when you use `to_dict()` on it, the `audio_transcription_driver` is not included.

```python
from griptape.drivers.audio_transcription.openai import OpenAiAudioTranscriptionDriver
from griptape.tools import AudioTranscriptionTool
from rich import print

driver = OpenAiAudioTranscriptionDriver(model="whisper-1")
tool = AudioTranscriptionTool(audio_transcription_driver=driver)

print("Original tool")
print(tool)
tool_dict = tool.to_dict()
print("\nTool as dict")
print(tool_dict)
```

Result:
```
Original tool
AudioTranscriptionTool(
    on_before_run=None,
    on_after_run=None,
    type='AudioTranscriptionTool',
    module_name='griptape.tools.audio_transcription.tool',
    allowlist=None,
    denylist=None,
    extra_schema_properties=None,
    name='AudioTranscriptionTool',
    input_memory=None,
    output_memory=None,
    install_dependencies_on_init=True,
    dependencies_install_directory=None,
    verbose=False,
    off_prompt=False,
    audio_transcription_driver=OpenAiAudioTranscriptionDriver(
        min_retry_delay=2,
        max_retry_delay=10,
        max_attempts=2,
        after_hook=<function ExponentialBackoffMixin.<lambda> at 0x1076149a0>,
        ignored_exception_types=(),
        type='OpenAiAudioTranscriptionDriver',
        module_name='griptape.drivers.audio_transcription.openai_audio_transcription_driver',
        model='whisper-1',
        api_type=None,
        api_version=None,
        base_url=None,
        api_key=None,
        organization=None,
        _client=None
    ),
    audio_loader=AudioLoader(
        create_futures_executor=<function FuturesExecutorMixin.<lambda>.<locals>.<lambda> at 0x1030351c0>,
        _futures_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x1030e4f50>,
        reference=None,
        file_manager_driver=LocalFileManagerDriver(encoding=None, _workdir='/Users/jason/Documents/GitHub/griptape-intro-demos'),
        encoding='utf-8'
    )
)

Tool as dict
{
    'type': 'AudioTranscriptionTool',
    'name': 'AudioTranscriptionTool',
    'input_memory': None,
    'output_memory': None,
    'install_dependencies_on_init': True,
    'dependencies_install_directory': None,
    'verbose': False,
    'off_prompt': False
}```

As you can see, there's no `audio_transcription_driver` in the `tool_dict`.

Discussion:
@shhlife: We're also seeing this type of error a lot where we get a notification about some missing `Base` driver.. for example:

```python
from griptape.drivers.web_search.duck_duck_go import DuckDuckGoWebSearchDriver
from griptape.tools import WebSearchTool

""" Web Search Tool """
web_driver = DuckDuckGoWebSearchDriver()
web_tool = WebSearchTool(web_search_driver=web_driver)

web_tool_dict = web_tool.to_dict()
```

This gives the following error:

```
NameError: name 'BaseWebSearchDriver' is not defined
```@shhlife: Other areas we are seeing it:

KnowledgeBaseTool
FileManagerTool
ExtractionTool@collindutter: #1587 strikes again

Solution patch:
diff --git a/docs/griptape-framework/misc/serialization.md b/docs/griptape-framework/misc/serialization.md
new file mode 100644
index 00000000..9f88aa56
--- /dev/null
+++ b/docs/griptape-framework/misc/serialization.md
@@ -0,0 +1,41 @@
+---
+search:
+  boost: 2
+---
+
+## Overview
+
+Many components in Griptape can be serialized and deserialized using the [to_dict](../../reference/griptape/mixins/serializable_mixin.md#griptape.mixins.serializable_mixin.SerializableMixin.to_dict) and [from_dict](../../reference/griptape/mixins/serializable_mixin.md#griptape.mixins.serializable_mixin.SerializableMixin.from_dict) methods.
+There are also [to_json](../../reference/griptape/mixins/serializable_mixin.md#griptape.mixins.serializable_mixin.SerializableMixin.to_json) and [from_json](../../reference/griptape/mixins/serializable_mixin.md#griptape.mixins.serializable_mixin.SerializableMixin.from_json) as a convenience.
+
+Here is how we can serialize an `Agent` and then deserialize it back:
+
+```python
+--8<-- "docs/griptape-framework/misc/src/serialization_1.py"
+```
+
+## Serialization Overrides
+
+All classes that implement the [SerializableMixin](../../reference/griptape/mixins/serializable_mixin.md#griptape.mixins.serializable_mixin.SerializableMixin) can be serialized using the above methods.
+However, only fields marked with `metadata={"serializable": True}` will be included in the serialization process.
+If you need to add or remove fields in the serialization process, you can pass [serialization_overrides](../../reference/griptape/mixins/serializable_mixin.md#griptape.mixins.serializable_mixin.SerializableMixin.serialization_overrides) to any of the serialization methods.
+
+```python
+--8<-- "docs/griptape-framework/misc/src/serialization_2.py"
+```
+
+## Types Overrides
+
+Due to some unfortunate internals of the Griptape's serialization process, you may occasionally run into a `NameError` when serializing. It will look something like this:
+
+```text
+NameError: name 'BaseWebSearchDriver' is not defined
+```
+
+This is something we're [actively working](https://github.com/griptape-ai/griptape/issues/1587) on fixing, but in the meantime, you can use the [types_overrides](../../reference/griptape/mixins/serializable_mixin.md#griptape.mixins.serializable_mixin.SerializableMixin.types_overrides) parameter to pass in a dictionary of types that need to be overridden during serialization.
+
+Here is an example of how you can use `types_overrides`:
+
+```python
+--8<-- "docs/griptape-framework/misc/src/serialization_3.py"
+```
diff --git a/docs/griptape-framework/misc/src/serialization_1.py b/docs/griptape-framework/misc/src/serialization_1.py
new file mode 100644
index 00000000..71d2d57c
--- /dev/null
+++ b/docs/griptape-framework/misc/src/serialization_1.py
@@ -0,0 +1,10 @@
+from griptape.structures import Agent
+
+agent = Agent()
+agent.run("My name is Collin")
+
+agent_dict = agent.to_dict()
+
+new_agent = Agent.from_dict(agent_dict)
+
+new_agent.run("What's my name?")
diff --git a/docs/griptape-framework/misc/src/serialization_2.py b/docs/griptape-framework/misc/src/serialization_2.py
new file mode 100644
index 00000000..2a3b06c4
--- /dev/null
+++ b/docs/griptape-framework/misc/src/serialization_2.py
@@ -0,0 +1,10 @@
+from rich.pretty import pprint
+
+from griptape.structures import Agent
+
+agent = Agent()
+agent.run("My name is Collin")
+
+agent_dict = agent.to_dict(serializable_overrides={"max_meta_memory_entries": True, "id": False})
+# `max_meta_memory_entries` will be included, `id` will not
+pprint(agent_dict)
diff --git a/docs/griptape-framework/misc/src/serialization_3.py b/docs/griptape-framework/misc/src/serialization_3.py
new file mode 100644
index 00000000..adde2cc4
--- /dev/null
+++ b/docs/griptape-framework/misc/src/serialization_3.py
@@ -0,0 +1,16 @@
+from duckduckgo_search import DDGS
+from rich.pretty import pprint
+
+from griptape.drivers.web_search import BaseWebSearchDriver
+from griptape.drivers.web_search.duck_duck_go import DuckDuckGoWebSearchDriver
+from griptape.tools import WebSearchTool
+
+web_driver = DuckDuckGoWebSearchDriver()
+web_tool = WebSearchTool(web_search_driver=web_driver)
+
+web_tool_dict = web_tool.to_dict(
+    serializable_overrides={"web_search_driver": True},
+    types_overrides={"BaseWebSearchDriver": BaseWebSearchDriver, "DDGS": DDGS},
+)
+
+pprint(web_tool_dict)
diff --git a/griptape/mixins/serializable_mixin.py b/griptape/mixins/serializable_mixin.py
index 35269b36..4be7cd36 100644
--- a/griptape/mixins/serializable_mixin.py
+++ b/griptape/mixins/serializable_mixin.py
@@ -29,12 +29,21 @@ class SerializableMixin(Generic[T]):
     )
 
     @classmethod
-    def get_schema(cls: type[T], subclass_name: Optional[str] = None, *, module_name: Optional[str] = None) -> Schema:
+    def get_schema(
+        cls: type[T],
+        subclass_name: Optional[str] = None,
+        *,
+        module_name: Optional[str] = None,
+        types_overrides: Optional[dict[str, type]] = None,
+        serializable_overrides: Optional[dict[str, bool]] = None,
+    ) -> Schema:
         """Generates a Marshmallow schema for the class.
 
         Args:
             subclass_name: An optional subclass name. Required if the class is abstract.
             module_name: An optional module name. Defaults to the class's module.
+            types_overrides: An optional dictionary of field names to override type.
+            serializable_overrides: An optional dictionary of field names to override serializable status.
         """
         if ABC in cls.__bases__:
             if subclass_name is None:
@@ -43,28 +52,66 @@ class SerializableMixin(Generic[T]):
             module_name = module_name or cls.__module__
             subclass_cls = cls._import_cls_rec(module_name, subclass_name)
 
-            schema_class = BaseSchema.from_attrs_cls(subclass_cls)
+            schema_class = BaseSchema.from_attrs_cls(
+                subclass_cls, types_overrides=types_overrides, serializable_overrides=serializable_overrides
+            )
         else:
-            schema_class = BaseSchema.from_attrs_cls(cls)
+            schema_class = BaseSchema.from_attrs_cls(
+                cls, types_overrides=types_overrides, serializable_overrides=serializable_overrides
+            )
 
         return schema_class()
 
     @classmethod
-    def from_dict(cls: type[T], data: dict) -> T:
-        return cast(T, cls.get_schema(subclass_name=data.get("type"), module_name=data.get("module_name")).load(data))
+    def from_dict(
+        cls: type[T],
+        data: dict,
+        *,
+        types_overrides: Optional[dict[str, type]] = None,
+        serializable_overrides: Optional[dict[str, bool]] = None,
+    ) -> T:
+        return cast(
+            T,
+            cls.get_schema(
+                subclass_name=data.get("type"),
+                module_name=data.get("module_name"),
+                types_overrides=types_overrides,
+                serializable_overrides=serializable_overrides,
+            ).load(data),
+        )
 
     @classmethod
-    def from_json(cls: type[T], data: str) -> T:
-        return cls.from_dict(json.loads(data))
+    def from_json(
+        cls: type[T],
+        data: str,
+        *,
+        types_overrides: Optional[dict[str, type]] = None,
+        serializable_overrides: Optional[dict[str, bool]] = None,
+    ) -> T:
+        return cls.from_dict(
+            json.loads(data), types_overrides=types_overrides, serializable_overrides=serializable_overrides
+        )
 
     def __str__(self) -> str:
         return json.dumps(self.to_dict())
 
-    def to_json(self) -> str:
-        return json.dumps(self.to_dict())
-
-    def to_dict(self) -> dict:
-        schema = BaseSchema.from_attrs_cls(self.__class__)
+    def to_json(
+        self,
+        *,
+        types_overrides: Optional[dict[str, type]] = None,
+        serializable_overrides: Optional[dict[str, bool]] = None,
+    ) -> str:
+        return json.dumps(self.to_dict(types_overrides=types_overrides, serializable_overrides=serializable_overrides))
+
+    def to_dict(
+        self,
+        *,
+        types_overrides: Optional[dict[str, type]] = None,
+        serializable_overrides: Optional[dict[str, bool]] = None,
+    ) -> dict:
+        schema = BaseSchema.from_attrs_cls(
+            self.__class__, types_overrides=types_overrides, serializable_overrides=serializable_overrides
+        )
 
         return dict(schema().dump(self))
 
diff --git a/griptape/schemas/base_schema.py b/griptape/schemas/base_schema.py
index 4d9ca557..c1161af8 100644
--- a/griptape/schemas/base_schema.py
+++ b/griptape/schemas/base_schema.py
@@ -27,15 +27,24 @@ class BaseSchema(Schema):
     }
 
     @classmethod
-    def from_attrs_cls(cls, attrs_cls: type) -> type:
+    def from_attrs_cls(
+        cls,
+        attrs_cls: type,
+        *,
+        types_overrides: Optional[dict[str, type]] = None,
+        serializable_overrides: Optional[dict[str, bool]] = None,
+    ) -> type:
         """Generate a Schema from an attrs class.
 
         Args:
             attrs_cls: An attrs class.
+            types_overrides: A dictionary of types to override when resolving types.
+            serializable_overrides: A dictionary of field names to whether they are serializable.
         """
         from marshmallow import post_load
 
-        from griptape.mixins.serializable_mixin import SerializableMixin
+        if serializable_overrides is None:
+            serializable_overrides = {}
 
         class SubSchema(cls):
             @post_load
@@ -50,30 +59,31 @@ class BaseSchema(Schema):
 
                 return attrs_cls(**data)
 
-        if issubclass(attrs_cls, SerializableMixin):
-            cls._resolve_types(attrs_cls)
-            return SubSchema.from_dict(
-                {
-                    a.alias or a.name: cls._get_field_for_type(
-                        a.type, serialization_key=a.metadata.get("serialization_key")
-                    )
-                    for a in attrs.fields(attrs_cls)
-                    if a.metadata.get("serializable")
-                },
-                name=f"{attrs_cls.__name__}Schema",
-            )
-        else:
-            raise ValueError(f"Class must implement SerializableMixin: {attrs_cls}")
+        cls._resolve_types(attrs_cls, types_override=types_overrides)
+        fields = {}
+        for field in attrs.fields(attrs_cls):
+            field_key = field.alias or field.name
+            if serializable_overrides.get(field_key, field.metadata.get("serializable", False)):
+                fields[field_key] = cls._get_field_for_type(
+                    field.type,
+                    serialization_key=field.metadata.get("serialization_key"),
+                    types_overrides=types_overrides,
+                )
+        return SubSchema.from_dict(fields, name=f"{attrs_cls.__name__}Schema")
 
     @classmethod
     def _get_field_for_type(
-        cls, field_type: type, serialization_key: Optional[str] = None
+        cls,
+        field_type: type,
+        serialization_key: Optional[str] = None,
+        types_overrides: Optional[dict[str, type]] = None,
     ) -> fields.Field | fields.Nested:
         """Generate a marshmallow Field instance from a Python type.
 
         Args:
             field_type: A field type.
             serialization_key: The key to pull the data from before serializing.
+            types_overrides: A dictionary of types to override when resolving types.
         """
         from griptape.schemas.polymorphic_schema import PolymorphicSchema
 
@@ -95,8 +105,11 @@ class BaseSchema(Schema):
                 serialization_key=serialization_key,
             )
         elif attrs.has(field_class):
-            schema = PolymorphicSchema if ABC in field_class.__bases__ else cls.from_attrs_cls
-            return fields.Nested(schema(field_class), allow_none=optional, attribute=serialization_key)
+            if ABC in field_class.__bases__:
+                schema = PolymorphicSchema(field_class, types_overrides=types_overrides)
+            else:
+                schema = cls.from_attrs_cls(field_class, types_overrides=types_overrides)
+            return fields.Nested(schema, allow_none=optional, attribute=serialization_key)
         elif cls._is_enum(field_type):
             return fields.String(allow_none=optional, attribute=serialization_key)
         elif cls._is_list_sequence(field_class):
@@ -191,11 +204,12 @@ class BaseSchema(Schema):
         return origin, args, optional
 
     @classmethod
-    def _resolve_types(cls, attrs_cls: type) -> None:
+    def _resolve_types(cls, attrs_cls: type, types_override: Optional[dict[str, type]] = None) -> None:
         """Resolve types in an attrs class.
 
         Args:
             attrs_cls: An attrs class.
+            types_override: A dictionary of types to override.
         """
         from collections.abc import Sequence
         from typing import Any
@@ -225,19 +239,26 @@ class BaseSchema(Schema):
             Reference,
             ToolAction,
         )
+        from griptape.drivers.assistant import BaseAssistantDriver
         from griptape.drivers.audio_transcription import BaseAudioTranscriptionDriver
         from griptape.drivers.embedding import BaseEmbeddingDriver
+        from griptape.drivers.file_manager import BaseFileManagerDriver
         from griptape.drivers.image_generation import BaseImageGenerationDriver, BaseMultiModelImageGenerationDriver
         from griptape.drivers.image_generation_model import BaseImageGenerationModelDriver
         from griptape.drivers.memory.conversation import BaseConversationMemoryDriver
+        from griptape.drivers.observability import BaseObservabilityDriver
         from griptape.drivers.prompt import BasePromptDriver
         from griptape.drivers.prompt.base_prompt_driver import StructuredOutputStrategy
         from griptape.drivers.ruleset import BaseRulesetDriver
+        from griptape.drivers.sql import BaseSqlDriver
         from griptape.drivers.text_to_speech import BaseTextToSpeechDriver
         from griptape.drivers.vector import BaseVectorStoreDriver
+        from griptape.drivers.web_scraper import BaseWebScraperDriver
+        from griptape.drivers.web_search import BaseWebSearchDriver
         from griptape.engines.rag import RagContext
         from griptape.events import EventListener
         from griptape.memory import TaskMemory
+        from griptape.memory.meta import BaseMetaEntry
         from griptape.memory.structure import BaseConversationMemory, Run
         from griptape.memory.task.storage import BaseArtifactStorage
         from griptape.rules.base_rule import BaseRule
@@ -248,6 +269,9 @@ class BaseSchema(Schema):
         from griptape.tools import BaseTool
         from griptape.utils import import_optional_dependency, is_dependency_installed
 
+        if types_override is None:
+            types_override = {}
+
         attrs.resolve_types(
             attrs_cls,
             localns={
@@ -262,7 +286,14 @@ class BaseSchema(Schema):
                 "BaseImageGenerationDriver": BaseImageGenerationDriver,
                 "BaseMultiModelImageGenerationDriver": BaseMultiModelImageGenerationDriver,
                 "BaseImageGenerationModelDriver": BaseImageGenerationModelDriver,
+                "BaseWebSearchDriver": BaseWebSearchDriver,
+                "BaseWebScraperDriver": BaseWebScraperDriver,
+                "BaseFileManagerDriver": BaseFileManagerDriver,
+                "BaseSqlDriver": BaseSqlDriver,
+                "BaseObservabilityDriver": BaseObservabilityDriver,
+                "BaseAssistantDriver": BaseAssistantDriver,
                 "BaseArtifact": BaseArtifact,
+                "BaseMetaEntry": BaseMetaEntry,
                 "PromptStack": PromptStack,
                 "EventListener": EventListener,
                 "BaseMessageContent": BaseMessageContent,
@@ -311,6 +342,7 @@ class BaseSchema(Schema):
                 "voyageai": import_optional_dependency("voyageai") if is_dependency_installed("voyageai") else Any,
                 "Schema": Schema,
                 "BaseModel": BaseModel,
+                **types_override,
             },
         )
 
diff --git a/griptape/schemas/polymorphic_schema.py b/griptape/schemas/polymorphic_schema.py
index 39749a43..bf2a1164 100644
--- a/griptape/schemas/polymorphic_schema.py
+++ b/griptape/schemas/polymorphic_schema.py
@@ -1,4 +1,6 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import Any, Optional
 
 from marshmallow import Schema, ValidationError
 
@@ -8,10 +10,11 @@ from griptape.schemas import BaseSchema
 class PolymorphicSchema(BaseSchema):
     """PolymorphicSchema is based on https://github.com/marshmallow-code/marshmallow-oneofschema."""
 
-    def __init__(self, inner_class: Any, **kwargs) -> None:
+    def __init__(self, inner_class: Any, types_overrides: Optional[dict[str, type]] = None, **kwargs) -> None:
         super().__init__(**kwargs)
 
         self.inner_class = inner_class
+        self.types_overrides = types_overrides
 
     type_field = "type"
     type_field_remove = True
@@ -58,7 +61,7 @@ class PolymorphicSchema(BaseSchema):
         if not obj_type:
             return (None, {"_schema": f"Unknown object class: {obj.__class__.__name__}"})
 
-        type_schema = BaseSchema.from_attrs_cls(obj.__class__)
+        type_schema = BaseSchema.from_attrs_cls(obj.__class__, types_overrides=self.types_overrides)
 
         if not type_schema:
             return None, {"_schema": f"Unsupported object type: {obj_type}"}
diff --git a/mkdocs.yml b/mkdocs.yml
index 2f3aad75..7ee51611 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -161,6 +161,7 @@ nav:
       - Misc:
           - Events: "griptape-framework/misc/events.md"
           - Tokenizers: "griptape-framework/misc/tokenizers.md"
+          - Serialization: "griptape-framework/misc/serialization.md"
   - Recipes:
       - Overview: "examples/index.md"
       - Agents:


Test patch:
diff --git a/tests/mocks/mock_meta_entry.py b/tests/mocks/mock_meta_entry.py
index 6f944883..82fa3995 100644
--- a/tests/mocks/mock_meta_entry.py
+++ b/tests/mocks/mock_meta_entry.py
@@ -1,6 +1,14 @@
+from __future__ import annotations
+
+import json
+from typing import Optional
+
 from griptape.memory.meta import BaseMetaEntry
 
 
 class MockMetaEntry(BaseMetaEntry):
-    def to_dict(self) -> dict:
+    def to_json(self, *, serializable_overrides: Optional[dict[str, bool]] = None) -> str:
+        return json.dumps(self.to_dict())
+
+    def to_dict(self, *, serializable_overrides: Optional[dict[str, bool]] = None) -> dict:
         return {"foo": "bar"}
diff --git a/tests/unit/schemas/test_base_schema.py b/tests/unit/schemas/test_base_schema.py
index 3f17aa75..b8487943 100644
--- a/tests/unit/schemas/test_base_schema.py
+++ b/tests/unit/schemas/test_base_schema.py
@@ -9,7 +9,6 @@ from marshmallow import fields
 from pydantic import BaseModel
 
 from griptape.artifacts import BaseArtifact, TextArtifact
-from griptape.loaders import TextLoader
 from griptape.schemas import PolymorphicSchema
 from griptape.schemas.base_schema import BaseSchema
 from griptape.schemas.bytes_field import Bytes
@@ -48,9 +47,6 @@ class TestBaseSchema:
         assert isinstance(schema.fields["baz"]._candidate_fields[0].inner, fields.Integer)
         assert schema.fields["baz"].allow_none is True
 
-        with pytest.raises(ValueError):
-            BaseSchema.from_attrs_cls(TextLoader)
-
     def test_get_field_for_type(self):
         assert isinstance(BaseSchema._get_field_for_type(BaseArtifact), fields.Nested)
 
@@ -170,3 +166,11 @@ class TestBaseSchema:
         field = BaseSchema._handle_union(Union[str, None], optional=True)
         assert isinstance(field, UnionField)
         assert field.allow_none is True
+
+    def test_serialization_override(self):
+        assert "bar" in MockSerializable().to_dict(serializable_overrides={"bar": True})
+        assert "bar" not in MockSerializable().to_dict(serializable_overrides={"bar": False})
+        assert MockSerializable().to_dict(serializable_overrides={"not_a_key": True})
+
+    def test_types_override(self):
+        assert MockSerializable().to_dict(types_overrides={"foo": int})