Update typing and overload

guardrails-ai · Dec 12, 2023 · ec9d5b9 · ec9d5b9
1 parent a30b72c
commit ec9d5b9
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/guardrails/guard.py b/guardrails/guard.py
@@ -238,9 +238,10 @@ def __call__(
         msg_history: Optional[List[Dict]] = None,
         metadata: Optional[Dict] = None,
         full_schema_reask: Optional[bool] = None,
+        stream: Optional[bool] = False,
         *args,
         **kwargs,
-    ) -> ValidationOutcome[OT]:
+    ) -> Union[ValidationOutcome[OT], Iterable[str]]:
         ...
 
     @overload
@@ -272,7 +273,7 @@ def __call__(
         *args,
         **kwargs,
     ) -> Union[
-        Union[ValidationOutcome[OT], Iterable], Awaitable[ValidationOutcome[OT]]
+        Union[ValidationOutcome[OT], Iterable[str]], Awaitable[ValidationOutcome[OT]]
     ]:
         """Call the LLM and validate the output. Pass an async LLM API to
         return a coroutine.