@@ -146,8 +146,8 @@ def test_guardrail_sequence_with_template_and_toxicity(self):
146
146
def test_fn (chain : GuardrailSequence ):
147
147
output = chain .run ("cats" , num_generations = 5 )
148
148
self .assertIsInstance (output , GuardrailIO )
149
- self .assertIsInstance (output .data , list )
150
- self .assertEqual (len ( output .data ), 1 )
149
+ self .assertIsInstance (output .data , str )
150
+ self .assertEqual (output .data , "Tell me a joke about cats" )
151
151
self .assertIsInstance (output .info , list )
152
152
self .assertEqual (len (output .info ), len (chain .steps ))
153
153
@@ -166,9 +166,8 @@ def test_guardrail_sequence_with_filtering(self):
166
166
def test_fn (chain : GuardrailSequence ):
167
167
output = chain .run (self .TOXIC_CONTENT )
168
168
self .assertIsInstance (output , GuardrailIO )
169
- self .assertIsInstance (output .data , list )
170
- self .assertEqual (len (output .data ), 1 )
171
- self .assertEqual (output .data [0 ], message )
169
+ self .assertIsInstance (output .data , str )
170
+ self .assertEqual (output .data , message )
172
171
self .assertIsInstance (output .info , list )
173
172
self .assertEqual (len (output .info ), len (chain .steps ))
174
173
0 commit comments