Use generics for loaders, remove instances of path

MIGRATION.md

@@ -77,14 +77,14 @@ TextArtifact("name: John\nAge: 30")
 #### Before
 
 ```python
-results = CsvLoader().load(Path("people.csv").read_text())
+results = CsvLoader().load("people.csv")
 
 print(results[0].value) # {"name": "John", "age": 30}
 ```
 
 #### After
 ```python
-results = CsvLoader().load(Path("people.csv").read_text())
+results = CsvLoader().load("people.csv")
 
 print(results[0].value) # name: John\nAge: 30
 print(results[0].meta["row"]) # 0

docs/griptape-framework/drivers/src/image_generation_drivers_8.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.artifacts import TextArtifact
 from griptape.drivers import (
     HuggingFacePipelineImageGenerationDriver,
@@ -11,7 +9,7 @@
 from griptape.tasks import VariationImageGenerationTask
 
 prompt_artifact = TextArtifact("landscape photograph, verdant, countryside, 8k")
-input_image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+input_image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
 image_variation_task = VariationImageGenerationTask(
     input=(prompt_artifact, input_image_artifact),

docs/griptape-framework/drivers/src/image_generation_drivers_9.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.artifacts import TextArtifact
 from griptape.drivers import (
     HuggingFacePipelineImageGenerationDriver,
@@ -11,7 +9,7 @@
 from griptape.tasks import VariationImageGenerationTask
 
 prompt_artifact = TextArtifact("landscape photograph, verdant, countryside, 8k")
-control_image_artifact = ImageLoader().load(Path("canny_control_image.png").read_bytes())
+control_image_artifact = ImageLoader().load("canny_control_image.png")
 
 controlnet_task = VariationImageGenerationTask(
     input=(prompt_artifact, control_image_artifact),

docs/griptape-framework/drivers/src/image_query_drivers_1.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.drivers import AnthropicImageQueryDriver
 from griptape.engines import ImageQueryEngine
 from griptape.loaders import ImageLoader
@@ -13,6 +11,6 @@
     image_query_driver=driver,
 )
 
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
 engine.run("Describe the weather in the image", [image_artifact])

docs/griptape-framework/drivers/src/image_query_drivers_2.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.drivers import AnthropicImageQueryDriver
 from griptape.engines import ImageQueryEngine
 from griptape.loaders import ImageLoader
@@ -13,9 +11,9 @@
     image_query_driver=driver,
 )
 
-image_artifact1 = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact1 = ImageLoader().load("tests/resources/mountain.png")
 
-image_artifact2 = ImageLoader().load(Path("tests/resources/cow.png").read_bytes())
+image_artifact2 = ImageLoader().load("tests/resources/cow.png")
 
 result = engine.run("Describe the weather in the image", [image_artifact1, image_artifact2])
 

docs/griptape-framework/drivers/src/image_query_drivers_3.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.drivers import OpenAiImageQueryDriver
 from griptape.engines import ImageQueryEngine
 from griptape.loaders import ImageLoader
@@ -13,6 +11,6 @@
     image_query_driver=driver,
 )
 
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
 engine.run("Describe the weather in the image", [image_artifact])

docs/griptape-framework/drivers/src/image_query_drivers_4.py

@@ -1,5 +1,4 @@
 import os
-from pathlib import Path
 
 from griptape.drivers import AzureOpenAiImageQueryDriver
 from griptape.engines import ImageQueryEngine
@@ -17,6 +16,6 @@
     image_query_driver=driver,
 )
 
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
 engine.run("Describe the weather in the image", [image_artifact])

docs/griptape-framework/drivers/src/image_query_drivers_5.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 import boto3
 
 from griptape.drivers import AmazonBedrockImageQueryDriver, BedrockClaudeImageQueryModelDriver
@@ -16,7 +14,7 @@
 
 engine = ImageQueryEngine(image_query_driver=driver)
 
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
 
 result = engine.run("Describe the weather in the image", [image_artifact])

docs/griptape-framework/engines/src/image_generation_engines_3.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.drivers import AmazonBedrockImageGenerationDriver, BedrockStableDiffusionImageGenerationModelDriver
 from griptape.engines import VariationImageGenerationEngine
 from griptape.loaders import ImageLoader
@@ -15,7 +13,7 @@
     image_generation_driver=driver,
 )
 
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
 engine.run(
     prompts=["A photo of a mountain landscape in winter"],

docs/griptape-framework/engines/src/image_generation_engines_4.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.drivers import AmazonBedrockImageGenerationDriver, BedrockStableDiffusionImageGenerationModelDriver
 from griptape.engines import InpaintingImageGenerationEngine
 from griptape.loaders import ImageLoader
@@ -15,9 +13,9 @@
     image_generation_driver=driver,
 )
 
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
-mask_artifact = ImageLoader().load(Path("tests/resources/mountain-mask.png").read_bytes())
+mask_artifact = ImageLoader().load("tests/resources/mountain-mask.png")
 
 engine.run(
     prompts=["A photo of a castle built into the side of a mountain"],

docs/griptape-framework/engines/src/image_generation_engines_5.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.drivers import AmazonBedrockImageGenerationDriver, BedrockStableDiffusionImageGenerationModelDriver
 from griptape.engines import OutpaintingImageGenerationEngine
 from griptape.loaders import ImageLoader
@@ -15,9 +13,9 @@
     image_generation_driver=driver,
 )
 
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
-mask_artifact = ImageLoader().load(Path("tests/resources/mountain-mask.png").read_bytes())
+mask_artifact = ImageLoader().load("tests/resources/mountain-mask.png")
 
 engine.run(
     prompts=["A photo of a mountain shrouded in clouds"],

docs/griptape-framework/engines/src/image_query_engines_1.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.drivers import OpenAiImageQueryDriver
 from griptape.engines import ImageQueryEngine
 from griptape.loaders import ImageLoader
@@ -8,6 +6,6 @@
 
 engine = ImageQueryEngine(image_query_driver=driver)
 
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
 engine.run("Describe the weather in the image", [image_artifact])

docs/griptape-framework/structures/src/tasks_12.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.drivers import AmazonBedrockImageGenerationDriver, BedrockStableDiffusionImageGenerationModelDriver
 from griptape.engines import VariationImageGenerationEngine
 from griptape.loaders import ImageLoader
@@ -18,7 +16,7 @@
 )
 
 # Load input image artifact.
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
 # Instantiate a pipeline.
 pipeline = Pipeline()

docs/griptape-framework/structures/src/tasks_13.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.drivers import AmazonBedrockImageGenerationDriver, BedrockStableDiffusionImageGenerationModelDriver
 from griptape.engines import InpaintingImageGenerationEngine
 from griptape.loaders import ImageLoader
@@ -18,9 +16,9 @@
 )
 
 # Load input image artifacts.
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
-mask_artifact = ImageLoader().load(Path("tests/resources/mountain-mask.png").read_bytes())
+mask_artifact = ImageLoader().load("tests/resources/mountain-mask.png")
 
 # Instantiate a pipeline.
 pipeline = Pipeline()

docs/griptape-framework/structures/src/tasks_14.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.drivers import AmazonBedrockImageGenerationDriver, BedrockStableDiffusionImageGenerationModelDriver
 from griptape.engines import OutpaintingImageGenerationEngine
 from griptape.loaders import ImageLoader
@@ -18,9 +16,9 @@
 )
 
 # Load input image artifacts.
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
-mask_artifact = ImageLoader().load(Path("tests/resources/mountain-mask.png").read_bytes())
+mask_artifact = ImageLoader().load("tests/resources/mountain-mask.png")
 
 # Instantiate a pipeline.
 pipeline = Pipeline()

docs/griptape-framework/structures/src/tasks_15.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 from griptape.drivers import OpenAiImageQueryDriver
 from griptape.engines import ImageQueryEngine
 from griptape.loaders import ImageLoader
@@ -18,7 +16,7 @@
 )
 
 # Load the input image artifact.
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.png").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.png")
 
 # Instantiate a pipeline.
 pipeline = Pipeline()

docs/griptape-framework/structures/src/tasks_3.py

@@ -1,10 +1,8 @@
-from pathlib import Path
-
 from griptape.loaders import ImageLoader
 from griptape.structures import Agent
 
 agent = Agent()
 
-image_artifact = ImageLoader().load(Path("tests/resources/mountain.jpg").read_bytes())
+image_artifact = ImageLoader().load("tests/resources/mountain.jpg")
 
 agent.run([image_artifact, "What's in this image?"])

griptape/loaders/audio_loader.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from typing import Any, cast
-
 import filetype
 from attrs import define
 
@@ -10,11 +8,8 @@
 
 
 @define
-class AudioLoader(BaseFileLoader):
+class AudioLoader(BaseFileLoader[AudioArtifact]):
     """Loads audio content into audio artifacts."""
 
-    def load(self, source: Any, *args, **kwargs) -> AudioArtifact:
-        return cast(AudioArtifact, super().load(source, *args, **kwargs))
-
-    def parse(self, source: bytes, *args, **kwargs) -> AudioArtifact:
+    def parse(self, source: bytes) -> AudioArtifact:
         return AudioArtifact(source, format=filetype.guess(source).extension)

griptape/loaders/base_file_loader.py

@@ -1,24 +1,29 @@
 from __future__ import annotations
 
 from abc import ABC
-from typing import TYPE_CHECKING
+from os import PathLike
+from typing import TypeVar, Union
 
 from attrs import Factory, define, field
 
+from griptape.artifacts import BaseArtifact
 from griptape.drivers import BaseFileManagerDriver, LocalFileManagerDriver
 from griptape.loaders import BaseLoader
 
-if TYPE_CHECKING:
-    from os import PathLike
+A = TypeVar("A", bound=BaseArtifact)
 
 
 @define
-class BaseFileLoader(BaseLoader, ABC):
+class BaseFileLoader(BaseLoader[Union[str, PathLike], bytes, A], ABC):
     file_manager_driver: BaseFileManagerDriver = field(
         default=Factory(lambda: LocalFileManagerDriver(workdir=None)),
         kw_only=True,
     )
     encoding: str = field(default="utf-8", kw_only=True)
 
-    def fetch(self, source: str | PathLike, *args, **kwargs) -> str | bytes:
-        return self.file_manager_driver.load_file(str(source), *args, **kwargs).value
+    def fetch(self, source: str | PathLike) -> bytes:
+        data = self.file_manager_driver.load_file(str(source)).value
+        if isinstance(data, str):
+            return data.encode(self.encoding)
+        else:
+            return data

griptape/loaders/base_loader.py

@@ -1,42 +1,50 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
 
 from attrs import define, field
 
+from griptape.artifacts import BaseArtifact
 from griptape.mixins import FuturesExecutorMixin
 from griptape.utils.futures import execute_futures_dict
 from griptape.utils.hash import bytes_to_hash, str_to_hash
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
 
-    from griptape.artifacts import BaseArtifact
     from griptape.common import Reference
 
+S = TypeVar("S")
+F = TypeVar("F")
+A = TypeVar("A", bound=BaseArtifact)
+
 
 @define
-class BaseLoader(FuturesExecutorMixin, ABC):
+class BaseLoader(FuturesExecutorMixin, ABC, Generic[S, F, A]):
     reference: Optional[Reference] = field(default=None, kw_only=True)
 
-    def load(self, source: Any, *args, **kwargs) -> BaseArtifact:
+    def load(self, source: S, *args, **kwargs) -> A:
         data = self.fetch(source)
 
-        return self.parse(data)
+        artifact = self.parse(data)
+
+        artifact.reference = self.reference
+
+        return artifact
 
     @abstractmethod
-    def fetch(self, source: Any, *args, **kwargs) -> Any: ...
+    def fetch(self, source: S) -> F: ...
 
     @abstractmethod
-    def parse(self, source: Any, *args, **kwargs) -> BaseArtifact: ...
+    def parse(self, source: F) -> A: ...
 
     def load_collection(
         self,
         sources: list[Any],
         *args,
         **kwargs,
-    ) -> Mapping[str, BaseArtifact]:
+    ) -> Mapping[str, A]:
         # Create a dictionary before actually submitting the jobs to the executor
         # to avoid duplicate work.
         sources_by_key = {self.to_key(source): source for source in sources}