Skip to content

Commit

Permalink
Fix saving results of csv extraction (#1105)
Browse files Browse the repository at this point in the history
  • Loading branch information
collindutter authored Aug 26, 2024
1 parent ad500c1 commit 1d1f0c0
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 1 deletion.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

### Fixed
- Ensure thread safety when publishing events by adding a thread lock to batch operations in `BaseEventListenerDriver`.
- `FileManagerTool` failing to save Artifacts created by `ExtractionTool` with a `CsvExtractionEngine`.

## [0.30.1] - 2024-08-21

Expand Down
3 changes: 2 additions & 1 deletion griptape/tools/file_manager/tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,8 @@ def save_memory_artifacts_to_disk(self, params: dict) -> ErrorArtifact | InfoArt

for artifact in list_artifact.value:
formatted_file_name = f"{artifact.name}-{file_name}" if len(list_artifact) > 1 else file_name
result = self.file_manager_driver.save_file(os.path.join(dir_name, formatted_file_name), artifact.value)
value = artifact.value if isinstance(artifact.value, (str, bytes)) else artifact.to_text()
result = self.file_manager_driver.save_file(os.path.join(dir_name, formatted_file_name), value)
if isinstance(result, ErrorArtifact):
return result

Expand Down
24 changes: 24 additions & 0 deletions tests/unit/tools/test_file_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import pytest

from griptape.artifacts import ListArtifact, TextArtifact
from griptape.artifacts.csv_row_artifact import CsvRowArtifact
from griptape.artifacts.error_artifact import ErrorArtifact
from griptape.drivers.file_manager.local_file_manager_driver import LocalFileManagerDriver
from griptape.loaders.text_loader import TextLoader
Expand Down Expand Up @@ -108,6 +109,29 @@ def test_save_memory_artifacts_to_disk_for_multiple_artifacts(self, temp_dir):
assert Path(os.path.join(temp_dir, "test", f"{artifacts[1].name}-{file_name}")).read_text() == "baz"
assert result.value == "Successfully saved memory artifacts to disk"

def test_save_memory_artifacts_to_disk_for_non_string_artifact(self, temp_dir):
memory = defaults.text_task_memory("Memory1")
artifact = CsvRowArtifact({"foo": "bar"})

memory.store_artifact("foobar", artifact)

file_manager = FileManagerTool(
input_memory=[memory], file_manager_driver=LocalFileManagerDriver(workdir=temp_dir)
)
result = file_manager.save_memory_artifacts_to_disk(
{
"values": {
"dir_name": "test",
"file_name": "foobar.txt",
"memory_name": memory.name,
"artifact_namespace": "foobar",
}
}
)

assert Path(os.path.join(temp_dir, "test", "foobar.txt")).read_text() == "bar"
assert result.value == "Successfully saved memory artifacts to disk"

def test_save_content_to_file(self, temp_dir):
file_manager = FileManagerTool(file_manager_driver=LocalFileManagerDriver(workdir=temp_dir))
result = file_manager.save_content_to_file(
Expand Down

0 comments on commit 1d1f0c0

Please sign in to comment.