From c79f15b131cc9cd215499ad104214117380a4425 Mon Sep 17 00:00:00 2001 From: Collin Dutter Date: Mon, 24 Feb 2025 12:20:03 -0800 Subject: [PATCH] fix(agent): update `output_schema` type hint to include pydantic model --- griptape/structures/agent.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/griptape/structures/agent.py b/griptape/structures/agent.py index 34cdd48da..a82db0e75 100644 --- a/griptape/structures/agent.py +++ b/griptape/structures/agent.py @@ -12,6 +12,7 @@ from griptape.tasks import PromptTask if TYPE_CHECKING: + from pydantic import BaseModel from schema import Schema from griptape.artifacts import BaseArtifact @@ -27,7 +28,7 @@ class Agent(Structure): ) stream: bool = field(default=None, kw_only=True) prompt_driver: BasePromptDriver = field(default=None, kw_only=True) - output_schema: Optional[Schema] = field(default=None, kw_only=True) + output_schema: Optional[Union[Schema, type[BaseModel]]] = field(default=None, kw_only=True) tools: list[BaseTool] = field(factory=list, kw_only=True) max_meta_memory_entries: Optional[int] = field(default=20, kw_only=True) fail_fast: bool = field(default=False, kw_only=True)