Add rules to all Tasks (#415)

Author: collindutter

griptape/engines/extraction/base_extraction_engine.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
+from typing import Optional
 from abc import ABC, abstractmethod
 from attr import define, field, Factory
 from griptape.artifacts import ListArtifact
 from griptape.chunkers import BaseChunker, TextChunker
 from griptape.drivers import BasePromptDriver, OpenAiChatPromptDriver
+from griptape.rules import Ruleset
 from griptape.tokenizers import OpenAiTokenizer
 
 
@@ -54,5 +56,10 @@ def min_response_tokens(self) -> int:
         )
 
     @abstractmethod
-    def extract(self, text: str | ListArtifact, **kwargs) -> ListArtifact:
+    def extract(
+        self,
+        text: str | ListArtifact,
+        rulesets: Optional[list[Ruleset]] = None,
+        **kwargs,
+    ) -> ListArtifact:
         ...

griptape/engines/extraction/csv_extraction_engine.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from typing import Optional
 import csv
 import io
 from attr import field, Factory, define
@@ -11,6 +12,7 @@
 from griptape.utils import PromptStack
 from griptape.engines import BaseExtractionEngine
 from griptape.utils import J2
+from griptape.rules import Ruleset, rule
 
 
 @define
@@ -21,7 +23,10 @@ class CsvExtractionEngine(BaseExtractionEngine):
     )
 
     def extract(
-        self, text: str | ListArtifact, column_names: list[str], **kwargs
+        self,
+        text: str | ListArtifact,
+        column_names: list[str],
+        rulesets: Optional[Ruleset] = None,
     ) -> ListArtifact | ErrorArtifact:
         try:
             return ListArtifact(
@@ -31,6 +36,7 @@ def extract(
                     else [TextArtifact(text)],
                     column_names,
                     [],
+                    rulesets=rulesets,
                 ),
                 item_separator="\n",
             )
@@ -57,10 +63,13 @@ def _extract_rec(
         artifacts: list[TextArtifact],
         column_names: list[str],
         rows: list[CsvRowArtifact],
+        rulesets: Optional[Ruleset] = None,
     ) -> list[CsvRowArtifact]:
         artifacts_text = self.chunk_joiner.join([a.value for a in artifacts])
         full_text = self.template_generator.render(
-            column_names=column_names, text=artifacts_text
+            column_names=column_names,
+            text=artifacts_text,
+            rulesets=J2("rulesets/rulesets.j2").render(rulesets=rulesets),
         )
 
         if (
@@ -86,7 +95,9 @@ def _extract_rec(
         else:
             chunks = self.chunker.chunk(artifacts_text)
             partial_text = self.template_generator.render(
-                column_names=column_names, text=chunks[0].value
+                column_names=column_names,
+                text=chunks[0].value,
+                rulesets=J2("rulesets/rulesets.j2").render(rulesets=rulesets),
             )
 
             rows.extend(
@@ -104,4 +115,6 @@ def _extract_rec(
                 )
             )
 
-            return self._extract_rec(chunks[1:], column_names, rows)
+            return self._extract_rec(
+                chunks[1:], column_names, rows, rulesets=rulesets
+            )

griptape/engines/extraction/json_extraction_engine.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
+from typing import Optional
 import json
 from attr import field, Factory, define
 from griptape.artifacts import TextArtifact, ListArtifact, ErrorArtifact
 from griptape.engines import BaseExtractionEngine
 from griptape.utils import J2
 from griptape.utils import PromptStack
+from griptape.rules import Ruleset, rule
 
 
 @define
@@ -15,7 +17,10 @@ class JsonExtractionEngine(BaseExtractionEngine):
     )
 
     def extract(
-        self, text: str | ListArtifact, template_schema: dict, **kwargs
+        self,
+        text: str | ListArtifact,
+        template_schema: dict,
+        rulesets: Optional[Ruleset] = None,
     ) -> ListArtifact | ErrorArtifact:
         try:
             json_schema = json.dumps(template_schema)
@@ -27,6 +32,7 @@ def extract(
                     else [TextArtifact(text)],
                     json_schema,
                     [],
+                    rulesets=rulesets,
                 ),
                 item_separator="\n",
             )
@@ -41,10 +47,13 @@ def _extract_rec(
         artifacts: list[TextArtifact],
         json_template_schema: str,
         extractions: list[TextArtifact],
+        rulesets: Optional[Ruleset] = None,
     ) -> list[TextArtifact]:
         artifacts_text = self.chunk_joiner.join([a.value for a in artifacts])
         full_text = self.template_generator.render(
-            json_template_schema=json_template_schema, text=artifacts_text
+            json_template_schema=json_template_schema,
+            text=artifacts_text,
+            rulesets=J2("rulesets/rulesets.j2").render(rulesets=rulesets),
         )
 
         if (
@@ -69,7 +78,9 @@ def _extract_rec(
         else:
             chunks = self.chunker.chunk(artifacts_text)
             partial_text = self.template_generator.render(
-                template_schema=json_template_schema, text=chunks[0].value
+                template_schema=json_template_schema,
+                text=chunks[0].value,
+                rulesets=J2("rulesets/rulesets.j2").render(rulesets=rulesets),
             )
 
             extractions.extend(
@@ -87,5 +98,5 @@ def _extract_rec(
             )
 
             return self._extract_rec(
-                chunks[1:], json_template_schema, extractions
+                chunks[1:], json_template_schema, extractions, rulesets=rulesets
             )

griptape/engines/query/base_query_engine.py

@@ -2,13 +2,18 @@
 from typing import Optional
 from attr import define
 from griptape.artifacts import TextArtifact, ListArtifact
+from griptape.rules import Ruleset
 
 
 @define
 class BaseQueryEngine(ABC):
     @abstractmethod
     def query(
-        self, query: str, namespace: Optional[str] = None, **kwargs
+        self,
+        query: str,
+        namespace: Optional[str] = None,
+        rulesets: Optional[list[Ruleset]] = None,
+        **kwargs
     ) -> TextArtifact:
         ...
 

griptape/engines/query/vector_query_engine.py

@@ -35,6 +35,7 @@ def query(
         metadata: Optional[str] = None,
         top_n: Optional[int] = None,
         namespace: Optional[str] = None,
+        rulesets: Optional[str] = None,
     ) -> TextArtifact:
         tokenizer = self.prompt_driver.tokenizer
         result = self.vector_store_driver.query(query, top_n, namespace)
@@ -52,7 +53,10 @@ def query(
             text_segments.append(artifact.value)
 
             message = self.template_generator.render(
-                metadata=metadata, query=query, text_segments=text_segments
+                metadata=metadata,
+                query=query,
+                text_segments=text_segments,
+                rulesets=J2("rulesets/rulesets.j2").render(rulesets=rulesets),
             )
             message_token_count = self.prompt_driver.token_count(
                 PromptStack(
@@ -69,7 +73,12 @@ def query(
                 text_segments.pop()
 
                 message = self.template_generator.render(
-                    metadata=metadata, query=query, text_segments=text_segments
+                    metadata=metadata,
+                    query=query,
+                    text_segments=text_segments,
+                    rulesets=J2("rulesets/rulesets.j2").render(
+                        rulesets=rulesets
+                    ),
                 )
 
                 break

griptape/engines/summary/base_summary_engine.py

@@ -1,15 +1,21 @@
 from abc import ABC, abstractmethod
+from typing import Optional
 from attr import define
 from griptape.artifacts import TextArtifact, ListArtifact
+from griptape.rules import Ruleset
 
 
 @define
 class BaseSummaryEngine(ABC):
-    def summarize_text(self, text: str) -> str:
+    def summarize_text(
+        self, text: str, rulesets: Optional[list[Ruleset]] = None
+    ) -> str:
         return self.summarize_artifacts(
-            ListArtifact([TextArtifact(text)])
+            ListArtifact([TextArtifact(text)]), rulesets=rulesets
         ).value
 
     @abstractmethod
-    def summarize_artifacts(self, artifacts: ListArtifact) -> TextArtifact:
+    def summarize_artifacts(
+        self, artifacts: ListArtifact, rulesets: Optional[list[Ruleset]] = None
+    ) -> TextArtifact:
         ...

griptape/engines/summary/prompt_summary_engine.py

@@ -7,6 +7,7 @@
 from griptape.engines import BaseSummaryEngine
 from griptape.utils import J2
 from griptape.tokenizers import OpenAiTokenizer
+from griptape.rules import Ruleset
 
 
 @define
@@ -57,18 +58,27 @@ def min_response_tokens(self) -> int:
             * self.max_token_multiplier
         )
 
-    def summarize_artifacts(self, artifacts: ListArtifact) -> TextArtifact:
-        return self.summarize_artifacts_rec(artifacts.value, None)
+    def summarize_artifacts(
+        self, artifacts: ListArtifact, rulesets: Optional[Ruleset] = None
+    ) -> TextArtifact:
+        return self.summarize_artifacts_rec(
+            artifacts.value, None, rulesets=rulesets
+        )
 
     def summarize_artifacts_rec(
-        self, artifacts: list[BaseArtifact], summary: Optional[str]
+        self,
+        artifacts: list[BaseArtifact],
+        summary: Optional[str],
+        rulesets: Optional[Ruleset] = None,
     ) -> TextArtifact:
         artifacts_text = self.chunk_joiner.join(
             [a.to_text() for a in artifacts]
         )
 
         full_text = self.template_generator.render(
-            summary=summary, text=artifacts_text
+            summary=summary,
+            text=artifacts_text,
+            rulesets=J2("rulesets/rulesets.j2").render(rulesets=rulesets),
         )
 
         if (
@@ -86,7 +96,9 @@ def summarize_artifacts_rec(
             chunks = self.chunker.chunk(artifacts_text)
 
             partial_text = self.template_generator.render(
-                summary=summary, text=chunks[0].value
+                summary=summary,
+                text=chunks[0].value,
+                rulesets=J2("rulesets/rulesets.j2").render(rulesets=rulesets),
             )
 
             return self.summarize_artifacts_rec(
@@ -100,4 +112,5 @@ def summarize_artifacts_rec(
                         ]
                     )
                 ).value,
+                rulesets=rulesets,
             )

griptape/tasks/base_text_input_task.py

@@ -3,16 +3,21 @@
 
 from attr import define, field
 from griptape.artifacts import TextArtifact
+from griptape.rules import Ruleset, Rule
 from griptape.tasks import BaseTask
 from griptape.utils import J2
 
 
 @define
 class BaseTextInputTask(BaseTask, ABC):
     DEFAULT_INPUT_TEMPLATE = "{{ args[0] }}"
+    DEFAULT_RULESET_NAME = "Default Ruleset"
+    ADDITIONAL_RULESET_NAME = "Additional Ruleset"
 
     input_template: str = field(default=DEFAULT_INPUT_TEMPLATE)
     context: dict[str, Any] = field(factory=dict, kw_only=True)
+    rulesets: list[Ruleset] = field(factory=list, kw_only=True)
+    rules: list[Rule] = field(factory=list, kw_only=True)
 
     @property
     def input(self) -> TextArtifact:
@@ -31,6 +36,50 @@ def full_context(self) -> dict[str, Any]:
         else:
             return {}
 
+    @rulesets.validator
+    def validate_rulesets(self, _, rulesets: list[Ruleset]) -> None:
+        if not rulesets:
+            return
+
+        if self.rules:
+            raise ValueError("Can't have both rulesets and rules specified.")
+
+    @rules.validator
+    def validate_rules(self, _, rules: list[Rule]) -> None:
+        if not rules:
+            return
+
+        if self.rulesets:
+            raise ValueError("Can't have both rules and rulesets specified.")
+
+    @property
+    def all_rulesets(self) -> list[Ruleset]:
+        structure_rulesets = []
+
+        if self.structure:
+            if self.structure.rulesets:
+                structure_rulesets = self.structure.rulesets
+            elif self.structure.rules:
+                structure_rulesets = [
+                    Ruleset(
+                        name=self.DEFAULT_RULESET_NAME,
+                        rules=self.structure.rules,
+                    )
+                ]
+
+        task_rulesets = []
+        if self.rulesets:
+            task_rulesets = self.rulesets
+        elif self.rules:
+            if structure_rulesets:
+                task_ruleset_name = self.ADDITIONAL_RULESET_NAME
+            else:
+                task_ruleset_name = self.DEFAULT_RULESET_NAME
+
+            task_rulesets = [Ruleset(name=task_ruleset_name, rules=self.rules)]
+
+        return structure_rulesets + task_rulesets
+
     def before_run(self) -> None:
         super().before_run()
 

griptape/tasks/extraction_task.py

@@ -10,4 +10,6 @@ class ExtractionTask(BaseTextInputTask):
     args: dict = field(kw_only=True)
 
     def run(self) -> ListArtifact:
-        return self.extraction_engine.extract(self.input.to_text(), **self.args)
+        return self.extraction_engine.extract(
+            self.input.to_text(), rulesets=self.all_rulesets, **self.args
+        )