griptape

1
from tests.utils.structure_tester import StructureTester
2
import pytest
3

4

5
class TestTextQueryTask:
6
    @pytest.fixture(
7
        autouse=True,
8
        params=StructureTester.TEXT_SUMMARY_TASK_CAPABLE_PROMPT_DRIVERS,
9
        ids=StructureTester.prompt_driver_id_fn,
10
    )
11
    def structure_tester(self, request):
12
        from griptape.structures import Agent
13
        from griptape.tasks import TextQueryTask
14
        from griptape.drivers import LocalVectorStoreDriver, OpenAiEmbeddingDriver
15
        from griptape.engines import VectorQueryEngine
16
        from griptape.artifacts import TextArtifact
17

18
        vector_store_driver = LocalVectorStoreDriver(embedding_driver=OpenAiEmbeddingDriver())
19

20
        artifact = TextArtifact("John Doe works as as software engineer at Griptape.")
21

22
        vector_query_engine = VectorQueryEngine(prompt_driver=request.param, vector_store_driver=vector_store_driver)
23
        vector_query_engine.upsert_text_artifact(artifact=artifact)
24

25
        agent = Agent(prompt_driver=request.param)
26
        agent.add_task(
27
            TextQueryTask("Respond to the users following query: {{ args[0] }}", query_engine=vector_query_engine)
28
        )
29

30
        return StructureTester(agent)
31

32
    def test_text_query_task(self, structure_tester):
33
        structure_tester.run("What is the job role of John Doe?")
34