griptape

1
from griptape.artifacts import TextArtifact, BlobArtifact
2
from griptape.drivers import LocalVectorStoreDriver
3
from griptape.engines import VectorQueryEngine, PromptSummaryEngine, CsvExtractionEngine, JsonExtractionEngine
4
from griptape.memory import TaskMemory
5
from griptape.memory.task.storage import TextArtifactStorage, BlobArtifactStorage
6
from tests.mocks.mock_embedding_driver import MockEmbeddingDriver
7
from tests.mocks.mock_prompt_driver import MockPromptDriver
8

9

10
def text_tool_artifact_storage():
11
    return TextArtifactStorage(
12
        query_engine=VectorQueryEngine(
13
            vector_store_driver=LocalVectorStoreDriver(embedding_driver=MockEmbeddingDriver()),
14
            prompt_driver=MockPromptDriver(),
15
        ),
16
        summary_engine=PromptSummaryEngine(prompt_driver=MockPromptDriver()),
17
        csv_extraction_engine=CsvExtractionEngine(prompt_driver=MockPromptDriver()),
18
        json_extraction_engine=JsonExtractionEngine(prompt_driver=MockPromptDriver()),
19
    )
20

21

22
def text_task_memory(name):
23
    return TaskMemory(
24
        name=name, artifact_storages={TextArtifact: text_tool_artifact_storage(), BlobArtifact: BlobArtifactStorage()}
25
    )
26