Fix tests

griptape/drivers/prompt/amazon_bedrock_prompt_driver.py

@@ -3,7 +3,7 @@
 import logging
 from typing import TYPE_CHECKING, Any, Literal
 
-from attrs import Factory, define, field
+from attrs import Attribute, Factory, define, field
 from schema import Schema
 
 from griptape.artifacts import (
@@ -55,12 +55,20 @@ class AmazonBedrockPromptDriver(BasePromptDriver):
         kw_only=True,
     )
     use_native_tools: bool = field(default=True, kw_only=True, metadata={"serializable": True})
+    use_native_structured_output: bool = field(default=True, kw_only=True, metadata={"serializable": True})
     native_structured_output_mode: Literal["native", "tool"] = field(
         default="tool", kw_only=True, metadata={"serializable": True}
     )
     tool_choice: dict = field(default=Factory(lambda: {"auto": {}}), kw_only=True, metadata={"serializable": True})
     _client: Any = field(default=None, kw_only=True, alias="client", metadata={"serializable": False})
 
+    @native_structured_output_mode.validator  # pyright: ignore[reportAttributeAccessIssue, reportOptionalMemberAccess]
+    def validate_native_structured_output_mode(self, attribute: Attribute, value: str) -> str:
+        if value == "native":
+            raise ValueError("AmazonBedrockPromptDriver does not support `native` structured output mode.")
+
+        return value
+
     @lazy_property()
     def client(self) -> Any:
         return self.session.client("bedrock-runtime")
@@ -105,8 +113,6 @@ def try_stream(self, prompt_stack: PromptStack) -> Iterator[DeltaMessage]:
             raise Exception("model response is empty")
 
     def _base_params(self, prompt_stack: PromptStack) -> dict:
-        from griptape.tools.structured_output.tool import StructuredOutputTool
-
         system_messages = [{"text": message.to_text()} for message in prompt_stack.system_messages]
         messages = self.__to_bedrock_messages([message for message in prompt_stack.messages if not message.is_system()])
 
@@ -119,23 +125,21 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
                 **({"maxTokens": self.max_tokens} if self.max_tokens is not None else {}),
             },
             "additionalModelRequestFields": self.additional_model_request_fields,
+            **(
+                {"toolConfig": {"tools": self.__to_bedrock_tools(prompt_stack.tools), "toolChoice": self.tool_choice}}
+                if prompt_stack.tools and self.use_native_tools
+                else {}
+            ),
             **self.extra_params,
         }
 
-        if prompt_stack.output_schema is not None and self.use_native_structured_output:
-            if self.native_structured_output_mode == "tool":
-                structured_output_tool = StructuredOutputTool(output_schema=prompt_stack.output_schema)
-                params["toolConfig"] = {
-                    "toolChoice": {"any": {}},
-                }
-                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}")
-
-        if prompt_stack.tools and self.use_native_tools:
+        if (
+            prompt_stack.output_schema is not None
+            and self.use_native_structured_output
+            and self.native_structured_output_mode == "tool"
+        ):
             params["toolConfig"] = {
-                "tools": self.__to_bedrock_tools(prompt_stack.tools),
+                "toolChoice": {"any": {}},
             }
 
         return params

griptape/drivers/prompt/anthropic_prompt_driver.py

@@ -68,6 +68,7 @@ class AnthropicPromptDriver(BasePromptDriver):
     top_k: int = field(default=250, kw_only=True, metadata={"serializable": True})
     tool_choice: dict = field(default=Factory(lambda: {"type": "auto"}), kw_only=True, metadata={"serializable": False})
     use_native_tools: bool = field(default=True, kw_only=True, metadata={"serializable": True})
+    use_native_structured_output: bool = field(default=True, kw_only=True, metadata={"serializable": True})
     native_structured_output_mode: Literal["native", "tool"] = field(
         default="tool", kw_only=True, metadata={"serializable": True}
     )
@@ -81,7 +82,7 @@ def client(self) -> Client:
     @native_structured_output_mode.validator  # pyright: ignore[reportAttributeAccessIssue, reportOptionalMemberAccess]
     def validate_native_structured_output_mode(self, attribute: Attribute, value: str) -> str:
         if value == "native":
-            raise ValueError("Anthropic does not support native structured output mode.")
+            raise ValueError("AnthropicPromptDriver does not support `native` structured output mode.")
 
         return value
 
@@ -120,7 +121,7 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
         system_messages = prompt_stack.system_messages
         system_message = system_messages[0].to_text() if system_messages else None
 
-        return {
+        params = {
             "model": self.model,
             "temperature": self.temperature,
             "stop_sequences": self.tokenizer.stop_sequences,
@@ -137,6 +138,15 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
             **self.extra_params,
         }
 
+        if (
+            prompt_stack.output_schema is not None
+            and self.use_native_structured_output
+            and self.native_structured_output_mode == "tool"
+        ):
+            params["tool_choice"] = {"type": "any"}
+
+        return params
+
     def __to_anthropic_messages(self, messages: list[Message]) -> list[dict]:
         return [
             {"role": self.__to_anthropic_role(message), "content": self.__to_anthropic_content(message)}

griptape/drivers/prompt/base_prompt_driver.py

@@ -77,18 +77,11 @@ def after_run(self, result: Message) -> None:
 
     @observable(tags=["PromptDriver.run()"])
     def run(self, prompt_input: PromptStack | BaseArtifact) -> Message:
-        from griptape.tools.structured_output.tool import StructuredOutputTool
-
         if isinstance(prompt_input, BaseArtifact):
             prompt_stack = PromptStack.from_artifact(prompt_input)
         else:
             prompt_stack = prompt_input
 
-        if not self.use_native_structured_output and prompt_stack.output_schema is not None:
-            structured_ouptut_tool = StructuredOutputTool(output_schema=prompt_stack.output_schema)
-            if structured_ouptut_tool not in prompt_stack.tools:
-                prompt_stack.tools.append(structured_ouptut_tool)
-
         for attempt in self.retrying():
             with attempt:
                 self.before_run(prompt_stack)

griptape/drivers/prompt/cohere_prompt_driver.py

@@ -53,6 +53,7 @@ class CoherePromptDriver(BasePromptDriver):
     model: str = field(metadata={"serializable": True})
     force_single_step: bool = field(default=False, kw_only=True, metadata={"serializable": True})
     use_native_tools: bool = field(default=True, kw_only=True, metadata={"serializable": True})
+    use_native_structured_output: bool = field(default=True, kw_only=True, metadata={"serializable": True})
     _client: ClientV2 = field(default=None, kw_only=True, alias="client", metadata={"serializable": False})
     tokenizer: BaseTokenizer = field(
         default=Factory(lambda self: CohereTokenizer(model=self.model, client=self.client), takes_self=True),
@@ -97,8 +98,6 @@ def try_stream(self, prompt_stack: PromptStack) -> Iterator[DeltaMessage]:
                 yield DeltaMessage(content=self.__to_prompt_stack_delta_message_content(event))
 
     def _base_params(self, prompt_stack: PromptStack) -> dict:
-        from griptape.tools.structured_output.tool import StructuredOutputTool
-
         tool_results = []
 
         messages = self.__to_cohere_messages(prompt_stack.messages)
@@ -110,23 +109,23 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
             "stop_sequences": self.tokenizer.stop_sequences,
             "max_tokens": self.max_tokens,
             **({"tool_results": tool_results} if tool_results else {}),
+            **(
+                {"tools": self.__to_cohere_tools(prompt_stack.tools)}
+                if prompt_stack.tools and self.use_native_tools
+                else {}
+            ),
             **self.extra_params,
         }
 
-        if prompt_stack.output_schema is not None:
-            if self.use_native_structured_output:
-                params["response_format"] = {
-                    "type": "json_object",
-                    "schema": prompt_stack.output_schema.json_schema("Output"),
-                }
-            else:
-                # This does not work great since Cohere does not support forced tool use.
-                structured_output_tool = StructuredOutputTool(output_schema=prompt_stack.output_schema)
-                if structured_output_tool not in prompt_stack.tools:
-                    prompt_stack.tools.append(structured_output_tool)
-
-        if prompt_stack.tools and self.use_native_tools:
-            params["tools"] = self.__to_cohere_tools(prompt_stack.tools)
+        if (
+            prompt_stack.output_schema is not None
+            and self.use_native_structured_output
+            and self.native_structured_output_mode == "native"
+        ):
+            params["response_format"] = {
+                "type": "json_object",
+                "schema": prompt_stack.output_schema.json_schema("Output"),
+            }
 
         return params
 

griptape/drivers/prompt/google_prompt_driver.py

@@ -63,6 +63,7 @@ class GooglePromptDriver(BasePromptDriver):
     top_p: Optional[float] = field(default=None, kw_only=True, metadata={"serializable": True})
     top_k: Optional[int] = field(default=None, kw_only=True, metadata={"serializable": True})
     use_native_tools: bool = field(default=True, kw_only=True, metadata={"serializable": True})
+    use_native_structured_output: bool = field(default=True, kw_only=True, metadata={"serializable": True})
     tool_choice: str = field(default="auto", kw_only=True, metadata={"serializable": True})
     _client: GenerativeModel = field(default=None, kw_only=True, alias="client", metadata={"serializable": False})
 
@@ -125,8 +126,6 @@ def try_stream(self, prompt_stack: PromptStack) -> Iterator[DeltaMessage]:
                 )
 
     def _base_params(self, prompt_stack: PromptStack) -> dict:
-        from griptape.tools.structured_output.tool import StructuredOutputTool
-
         types = import_optional_dependency("google.generativeai.types")
         protos = import_optional_dependency("google.generativeai.protos")
 
@@ -150,17 +149,23 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
                     **self.extra_params,
                 },
             ),
+            **(
+                {
+                    "tools": self.__to_google_tools(prompt_stack.tools),
+                    "tool_config": {"function_calling_config": {"mode": self.tool_choice}},
+                }
+                if prompt_stack.tools and self.use_native_tools
+                else {}
+            ),
         }
-        if prompt_stack.output_schema is not None:
-            structured_ouptut_tool = StructuredOutputTool(output_schema=prompt_stack.output_schema)
-            params["tool_config"] = {
-                "function_calling_config": {"mode": self.tool_choice},
-            }
-            if structured_ouptut_tool not in prompt_stack.tools:
-                prompt_stack.tools.append(structured_ouptut_tool)
-
-        if prompt_stack.tools and self.use_native_tools:
-            params["tools"] = self.__to_google_tools(prompt_stack.tools)
+        if prompt_stack.output_schema is not None and self.use_native_structured_output:
+            if self.native_structured_output_mode == "tool":
+                params["tool_config"] = {
+                    "function_calling_config": {"mode": "auto"},
+                }
+            elif self.native_structured_output_mode == "native":
+                # TODO: Add support for native structured output
+                ...
 
         return params
 

griptape/drivers/prompt/huggingface_hub_prompt_driver.py

@@ -111,9 +111,14 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
             **self.extra_params,
         }
 
-        if prompt_stack.output_schema and self.use_native_structured_output:
+        if (
+            prompt_stack.output_schema
+            and self.use_native_structured_output
+            and self.native_structured_output_mode == "native"
+        ):
             # https://huggingface.co/learn/cookbook/en/structured_generation#-constrained-decoding
             params["grammar"] = {"type": "json", "value": prompt_stack.output_schema.json_schema("Output Schema")}
+            # Grammar does not support $schema and $id
             del params["grammar"]["value"]["$schema"]
             del params["grammar"]["value"]["$id"]
 

griptape/drivers/prompt/ollama_prompt_driver.py

@@ -101,8 +101,6 @@ def try_stream(self, prompt_stack: PromptStack) -> Iterator[DeltaMessage]:
             raise Exception("invalid model response")
 
     def _base_params(self, prompt_stack: PromptStack) -> dict:
-        from griptape.tools.structured_output.tool import StructuredOutputTool
-
         messages = self._prompt_stack_to_messages(prompt_stack)
 
         params = {
@@ -124,13 +122,12 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
             **self.extra_params,
         }
 
-        if prompt_stack.output_schema is not None:
-            if self.use_native_structured_output:
-                params["format"] = prompt_stack.output_schema.json_schema("Output")
-            else:
-                structured_ouptut_tool = StructuredOutputTool(output_schema=prompt_stack.output_schema)
-                if structured_ouptut_tool not in prompt_stack.tools:
-                    prompt_stack.tools.append(structured_ouptut_tool)
+        if (
+            prompt_stack.output_schema is not None
+            and self.use_native_structured_output
+            and self.native_structured_output_mode == "tool"
+        ):
+            params["format"] = prompt_stack.output_schema.json_schema("Output")
 
         return params
 

griptape/drivers/prompt/openai_chat_prompt_driver.py

@@ -149,7 +149,15 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
             "temperature": self.temperature,
             "user": self.user,
             "seed": self.seed,
-            "stop": self.tokenizer.stop_sequences,
+            **(
+                {
+                    "tools": self.__to_openai_tools(prompt_stack.tools),
+                    "tool_choice": self.tool_choice,
+                    "parallel_tool_calls": self.parallel_tool_calls,
+                }
+                if prompt_stack.tools and self.use_native_tools
+                else {}
+            ),
             **({"stop": self.tokenizer.stop_sequences} if self.tokenizer.stop_sequences else {}),
             **({"max_tokens": self.max_tokens} if self.max_tokens is not None else {}),
             **({"stream_options": {"include_usage": True}} if self.stream else {}),
@@ -166,7 +174,7 @@ def _base_params(self, prompt_stack: PromptStack) -> dict:
                         "strict": True,
                     },
                 }
-            else:
+            elif self.native_structured_output_mode == "tool":
                 params["tool_choice"] = "required"
 
         if prompt_stack.tools and self.use_native_tools:

griptape/mixins/rule_mixin.py

@@ -28,6 +28,3 @@ def rulesets(self) -> list[Ruleset]:
             rulesets.append(Ruleset(id=self._default_ruleset_id, name=self._default_ruleset_name, rules=self.rules))
 
         return rulesets
-
-    def get_rules_for_type(self, rule_type: type[T]) -> list[T]:
-        return [rule for ruleset in self.rulesets for rule in ruleset.rules if isinstance(rule, rule_type)]

griptape/schemas/base_schema.py

@@ -151,6 +151,8 @@ def _resolve_types(cls, attrs_cls: type) -> None:
         from collections.abc import Sequence
         from typing import Any
 
+        from schema import Schema
+
         from griptape.artifacts import BaseArtifact
         from griptape.common import (
             BaseDeltaMessageContent,
@@ -215,6 +217,7 @@ def _resolve_types(cls, attrs_cls: type) -> None:
                 "BaseRule": BaseRule,
                 "Ruleset": Ruleset,
                 # Third party modules
+                "Schema": Schema,
                 "Client": import_optional_dependency("cohere").Client if is_dependency_installed("cohere") else Any,
                 "ClientV2": import_optional_dependency("cohere").ClientV2 if is_dependency_installed("cohere") else Any,
                 "GenerativeModel": import_optional_dependency("google.generativeai").GenerativeModel

griptape/structures/agent.py

@@ -3,7 +3,6 @@
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
 from attrs import Attribute, Factory, define, evolve, field
-from schema import Schema
 
 from griptape.artifacts.text_artifact import TextArtifact
 from griptape.common import observable
@@ -12,6 +11,8 @@
 from griptape.tasks import PromptTask
 
 if TYPE_CHECKING:
+    from schema import Schema
+
     from griptape.artifacts import BaseArtifact
     from griptape.drivers import BasePromptDriver
     from griptape.tasks import BaseTask
@@ -51,7 +52,6 @@ def __attrs_post_init__(self) -> None:
                 prompt_driver=self.prompt_driver,
                 tools=self.tools,
                 max_meta_memory_entries=self.max_meta_memory_entries,
-                output_schema=self._build_schema_from_type(self.output_type) if self.output_type is not None else None,
             )
             self.add_task(task)
 
@@ -78,9 +78,3 @@ def try_run(self, *args) -> Agent:
         self.task.run()
 
         return self
-
-    def _build_schema_from_type(self, output_type: type | Schema) -> Schema:
-        if isinstance(output_type, Schema):
-            return output_type
-        else:
-            return Schema(output_type)