WIP

griptape-ai · Dec 23, 2024 · 56b8fff · 56b8fff
1 parent 9f00ec3
commit 56b8fff
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 9 deletions.
diff --git a/griptape/common/prompt_stack/prompt_stack.py b/griptape/common/prompt_stack/prompt_stack.py
@@ -3,7 +3,6 @@
 from typing import TYPE_CHECKING, Optional
 
 from attrs import define, field
-from schema import Schema
 
 from griptape.artifacts import (
     ActionArtifact,
@@ -25,6 +24,8 @@
 from griptape.mixins.serializable_mixin import SerializableMixin
 
 if TYPE_CHECKING:
+    from schema import Schema
+
     from griptape.tools import BaseTool
 
 

diff --git a/griptape/drivers/prompt/amazon_bedrock_prompt_driver.py b/griptape/drivers/prompt/amazon_bedrock_prompt_driver.py
@@ -124,12 +124,12 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
 
         if prompt_stack.output_schema is not None and self.use_native_structured_output:
             if self.native_structured_output_mode == "tool":
-                structured_ouptut_tool = StructuredOutputTool(output_schema=prompt_stack.output_schema)
+                structured_output_tool = StructuredOutputTool(output_schema=prompt_stack.output_schema)
                 params["toolConfig"] = {
                     "toolChoice": {"any": {}},
                 }
-                if structured_ouptut_tool not in prompt_stack.tools:
-                    prompt_stack.tools.append(structured_ouptut_tool)
+                if structured_output_tool not in prompt_stack.tools:
+                    prompt_stack.tools.append(structured_output_tool)
             else:
                 raise ValueError(f"Unsupported native structured output mode: {self.native_structured_output_mode}")
 

diff --git a/griptape/tasks/prompt_task.py b/griptape/tasks/prompt_task.py
@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
 from attrs import NOTHING, Attribute, Factory, NothingType, define, field
+from schema import Schema
 
 from griptape import utils
 from griptape.artifacts import ActionArtifact, BaseArtifact, ErrorArtifact, ListArtifact, TextArtifact
@@ -13,13 +14,12 @@
 from griptape.memory.structure import Run
 from griptape.mixins.actions_subtask_origin_mixin import ActionsSubtaskOriginMixin
 from griptape.mixins.rule_mixin import RuleMixin
-from griptape.rules import Ruleset
+from griptape.rules import JsonSchemaRule, Ruleset
 from griptape.tasks import ActionsSubtask, BaseTask
+from griptape.tools import StructuredOutputTool
 from griptape.utils import J2
 
 if TYPE_CHECKING:
-    from schema import Schema
-
     from griptape.drivers import BasePromptDriver
     from griptape.memory import TaskMemory
     from griptape.memory.structure.base_conversation_memory import BaseConversationMemory
@@ -158,7 +158,7 @@ def prompt_stack(self) -> PromptStack:
 
         if memory is not None:
             # inserting at index 1 to place memory right after system prompt
-            memory.add_to_prompt_stack(self.prompt_driver, stack, 1 if system_template else 0)
+            memory.add_to_prompt_stack(self.prompt_driver, stack, 1 if system_contents else 0)
 
         return stack
 

diff --git a/tests/unit/tasks/test_actions_subtask.py b/tests/unit/tasks/test_actions_subtask.py
@@ -7,7 +7,7 @@
 from griptape.artifacts.json_artifact import JsonArtifact
 from griptape.common import ToolAction
 from griptape.structures import Agent
-from griptape.tasks import ActionsSubtask, PromptTask
+from griptape.tasks import ActionsSubtask, PromptTask, ToolkitTask
 from tests.mocks.mock_tool.tool import MockTool