diff --git a/tests/workflows/pydantic_output_reference/__init__.py b/tests/workflows/pydantic_output_reference/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/workflows/pydantic_output_reference/tests/__init__.py b/tests/workflows/pydantic_output_reference/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/workflows/pydantic_output_reference/tests/test_workflow.py b/tests/workflows/pydantic_output_reference/tests/test_workflow.py new file mode 100644 index 000000000..93d96393a --- /dev/null +++ b/tests/workflows/pydantic_output_reference/tests/test_workflow.py @@ -0,0 +1,20 @@ +from tests.workflows.pydantic_output_reference.workflow import PydanticOutputReferenceWorkflow + + +def test_run_workflow__pydantic_field_reference(): + """ + Tests that a node can reference a specific field from a Pydantic model output of a previous node. + """ + + # GIVEN a workflow where FirstNode outputs a Pydantic model + # AND SecondNode references a field from that model + workflow = PydanticOutputReferenceWorkflow() + + # WHEN we run the workflow + terminal_event = workflow.run() + + # THEN the workflow should have completed successfully + assert terminal_event.name == "workflow.execution.fulfilled", terminal_event + + # AND the output should contain the greeting using the referenced field + assert terminal_event.outputs == {"greeting": "Hello, Alice!"} diff --git a/tests/workflows/pydantic_output_reference/workflow.py b/tests/workflows/pydantic_output_reference/workflow.py new file mode 100644 index 000000000..8a3c8e77c --- /dev/null +++ b/tests/workflows/pydantic_output_reference/workflow.py @@ -0,0 +1,50 @@ +from pydantic import BaseModel + +from vellum.workflows import BaseWorkflow +from vellum.workflows.nodes.bases import BaseNode +from vellum.workflows.outputs import BaseOutputs + + +class UserProfile(BaseModel): + """A Pydantic model representing a user profile.""" + + name: str + age: int + email: str + + +class FirstNode(BaseNode): + """A node that outputs a Pydantic model.""" + + class Outputs(BaseOutputs): + profile: UserProfile + + def run(self) -> BaseOutputs: + return self.Outputs( + profile=UserProfile( + name="Alice", + age=30, + email="alice@example.com", + ) + ) + + +class SecondNode(BaseNode): + """A node that references a Pydantic model output and operates on its field.""" + + profile = FirstNode.Outputs.profile + + class Outputs(BaseOutputs): + greeting: str + + def run(self) -> BaseOutputs: + return self.Outputs(greeting=f"Hello, {self.profile.name}!") + + +class PydanticOutputReferenceWorkflow(BaseWorkflow): + """A workflow demonstrating Pydantic model output field references between nodes.""" + + graph = FirstNode >> SecondNode + + class Outputs(BaseOutputs): + greeting = SecondNode.Outputs.greeting