griptape

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

4

5
class TestCsvExtractionTask:
6
    @pytest.fixture(
7
        autouse=True,
8
        params=StructureTester.CSV_EXTRACTION_TASK_CAPABLE_PROMPT_DRIVERS,
9
        ids=StructureTester.prompt_driver_id_fn,
10
    )
11
    def structure_tester(self, request):
12
        from griptape.tasks import ExtractionTask
13
        from griptape.structures import Agent
14
        from griptape.engines import CsvExtractionEngine
15

16
        columns = ["Name", "Age", "Address"]
17

18
        # Create an agent and add the ExtractionTask to it
19
        agent = Agent(prompt_driver=request.param)
20
        agent.add_task(
21
            ExtractionTask(
22
                extraction_engine=CsvExtractionEngine(prompt_driver=request.param), args={"column_names": columns}
23
            )
24
        )
25

26
        return StructureTester(agent)
27

28
    def test_csv_extraction_task(self, structure_tester):
29
        structure_tester.run(
30
            """
31
            John (Age 25) lives at 123 Main St
32
            Jane (Age 30) lives at 456 Elm St
33
        """
34
        )
35