added extraction task

2024-11-16 12:38:51 +13:00
parent ea6239c142
commit 393678482b
1 changed files with 103 additions and 0 deletions
@@ -0,0 +1,103 @@
+from griptape.engines import CsvExtractionEngine, JsonExtractionEngine
+from griptape.tasks import ExtractionTask
+
+from ..agent.gtComfyAgent import gtComfyAgent as Agent
+from .gtUIBaseTask import gtUIBaseTask
+
+extraction_engines = ["csv", "json"]
+
+
+class gtUIExtractionTask(gtUIBaseTask):
+    """
+    The Griptape Text Extraction Task
+    """
+
+    @classmethod
+    def INPUT_TYPES(s):
+        inputs = super().INPUT_TYPES()
+        inputs["optional"].update(
+            {
+                "extraction_type": (
+                    extraction_engines,
+                    {
+                        "default": "json",
+                        "tooltip": "The type of extraction to perform.",
+                    },
+                ),
+                "column_names": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "Comma separated list of column names to extract. Example:\n name, age",
+                    },
+                ),
+                "template_schema": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "multiline": True,
+                        "tooltip": 'Schema template for the json extraction. Example:\n {"name": str, "age": int}',
+                    },
+                ),
+            }
+        )
+        del inputs["required"]["STRING"]
+        inputs["optional"].update(
+            {
+                "STRING": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "multiline": True,
+                    },
+                ),
+            }
+        )
+        return inputs
+
+    DESCRIPTION = "Extract text using csv or json."
+    CATEGORY = "Griptape/Agent"
+
+    def run(self, **kwargs):
+        STRING = kwargs.get("STRING")
+        input_string = kwargs.get("input_string", None)
+        agent = kwargs.get("agent", None)
+        extraction_type = kwargs.get("extraction_type", "json")
+        column_names_string = kwargs.get("column_names", "")
+        column_names = [
+            column_name.strip() for column_name in column_names_string.split(",")
+        ]
+        template_schema = kwargs.get("template_schema", "")
+        params = {}
+
+        if not agent:
+            agent = Agent()
+        prompt_driver = agent.drivers_config.prompt_driver
+        if extraction_type == "csv":
+            print("CSV Extraction")
+            print(column_names)
+            engine = CsvExtractionEngine(
+                prompt_driver=prompt_driver, column_names=column_names
+            )
+        elif extraction_type == "json":
+            engine = JsonExtractionEngine(
+                prompt_driver=prompt_driver, template_schema=template_schema
+            )
+
+        prompt_text = self.get_prompt_text(STRING, input_string)
+
+        task = ExtractionTask(extraction_engine=engine)
+        try:
+            agent.add_task(task)
+        except Exception as e:
+            print(e)
+        result = agent.run(prompt_text)
+        artifacts = result.output_task.output.value
+        if extraction_type == "csv":
+            output = [artifact.value.split("\n") for artifact in artifacts]
+        elif extraction_type == "json":
+            output = [artifact.value for artifact in artifacts]
+        return (
+            output,
+            agent,
+        )