Streaming chunk accumulation

guardrails/classes/history/iteration.py

@@ -15,6 +15,7 @@
 from guardrails.utils.logs_utils import ValidatorLogs
 from guardrails.utils.pydantic_utils import ArbitraryModel
 from guardrails.utils.reask_utils import ReAsk
+from guardrails.validator_base import ErrorSpan
 
 
 class Iteration(ArbitraryModel):
@@ -155,6 +156,12 @@ def failed_validations(self) -> List[ValidatorLogs]:
         iteration."""
         return self.outputs.failed_validations
 
+    @property
+    def error_spans_in_output(self) -> List[ErrorSpan]:
+        """The error spans from the LLM response.
+        These indices are relative to the complete LLM output."""
+        return self.outputs.error_spans_in_output
+
     @property
     def status(self) -> str:
         """Representation of the end state of this iteration.

guardrails/classes/history/outputs.py

@@ -1,5 +1,4 @@
 from typing import Dict, List, Optional, Sequence, Union
-
 from pydantic import Field
 from typing_extensions import deprecated
 
@@ -8,7 +7,7 @@
 from guardrails.utils.logs_utils import ValidatorLogs
 from guardrails.utils.pydantic_utils import ArbitraryModel
 from guardrails.utils.reask_utils import ReAsk
-from guardrails.validator_base import FailResult
+from guardrails.validator_base import ErrorSpan, FailResult
 
 
 class Outputs(ArbitraryModel):
@@ -75,6 +74,28 @@ def failed_validations(self) -> List[ValidatorLogs]:
             ]
         )
 
+    @property
+    def error_spans_in_output(self) -> List[ErrorSpan]:
+        """The error spans from the LLM response.
+        These indices are relative to the complete LLM output."""
+        total_len = 0
+        spans_in_output = []
+        for log in self.validator_logs:
+            result = log.validation_result
+            if isinstance(result, FailResult):
+                if result.error_spans is not None:
+                    for error_span in result.error_spans:
+                        spans_in_output.append(
+                            ErrorSpan(
+                                start=error_span.start + total_len,
+                                end=error_span.end + total_len,
+                                reason=error_span.reason,
+                            )
+                        )
+            if result.validated_chunk is not None:
+                total_len += len(result.validated_chunk)
+        return spans_in_output
+
     @property
     def status(self) -> str:
         """Representation of the end state of the validation run.

guardrails/guard.py

@@ -1130,6 +1130,15 @@ async def _async_parse(
 
         return ValidationOutcome[OT].from_guard_history(call)
 
+    def error_spans_in_output(self):
+        try:
+            call = self.history[0]
+            iter = call.iterations[0]
+            llm_spans = iter.error_spans_in_output
+            return llm_spans
+        except (AttributeError, TypeError):
+            return []
+
     @deprecated(
         """The `with_prompt_validation` method is deprecated,
         and will be removed in 0.5.x. Instead, please use

guardrails/run/runner.py

@@ -547,17 +547,29 @@ def validate(
         index: int,
         parsed_output: Any,
         output_schema: Schema,
+        stream: Optional[bool] = False,
         **kwargs,
     ):
         """Validate the output."""
-        validated_output = output_schema.validate(
-            iteration,
-            parsed_output,
-            self.metadata,
-            attempt_number=index,
-            disable_tracer=self._disable_tracer,
-            **kwargs,
-        )
+        if isinstance(output_schema, StringSchema):
+            validated_output = output_schema.validate(
+                iteration,
+                parsed_output,
+                self.metadata,
+                index,
+                self._disable_tracer,
+                stream,
+                **kwargs,
+            )
+        else:
+            validated_output = output_schema.validate(
+                iteration,
+                parsed_output,
+                self.metadata,
+                attempt_number=index,
+                disable_tracer=self._disable_tracer,
+                **kwargs,
+            )
 
         return validated_output
 

guardrails/run/stream_runner.py

@@ -153,54 +153,149 @@ def step(
         verified = set()
         # Loop over the stream
         # and construct "fragments" of concatenated chunks
-        for chunk in stream:
-            # 1. Get the text from the chunk and append to fragment
-            chunk_text = self.get_chunk_text(chunk, api)
-            fragment += chunk_text
+        # for now, handle string and json schema differently
 
-            # 2. Parse the fragment
-            parsed_fragment, move_to_next = self.parse(
-                index, fragment, output_schema, verified
-            )
-            if move_to_next:
-                # Continue to next chunk
-                continue
+        if isinstance(output_schema, StringSchema):
+            stream_finished = False
+            last_chunk_text = ""
+            for chunk in stream:
+                # 1. Get the text from the chunk and append to fragment
+                chunk_text = self.get_chunk_text(chunk, api)
+                last_chunk_text = chunk_text
+                finished = self.is_last_chunk(chunk, api)
+                if finished:
+                    stream_finished = True
+                fragment += chunk_text
 
-            # 3. Run output validation
-            validated_fragment = self.validate(
-                iteration,
-                index,
-                parsed_fragment,
-                output_schema,
-                validate_subschema=True,
-            )
-            if isinstance(validated_fragment, SkeletonReAsk):
-                raise ValueError(
-                    "Received fragment schema is an invalid sub-schema "
-                    "of the expected output JSON schema."
+                # 2. Parse the chunk
+                parsed_chunk, move_to_next = self.parse(
+                    index, chunk_text, output_schema, verified
+                )
+                if move_to_next:
+                    # Continue to next chunk
+                    continue
+                validated_text = self.validate(
+                    iteration,
+                    index,
+                    parsed_chunk,
+                    output_schema,
+                    True,
+                    validate_subschema=True,
+                    # if it is the last chunk, validate everything that's left
+                    remainder=finished,
                 )
+                if isinstance(validated_text, SkeletonReAsk):
+                    raise ValueError(
+                        "Received fragment schema is an invalid sub-schema "
+                        "of the expected output JSON schema."
+                    )
 
-            # 4. Introspect: inspect the validated fragment for reasks
-            reasks, valid_op = self.introspect(index, validated_fragment, output_schema)
-            if reasks:
-                raise ValueError(
-                    "Reasks are not yet supported with streaming. Please "
-                    "remove reasks from schema or disable streaming."
+                # 4. Introspect: inspect the validated fragment for reasks
+                reasks, valid_op = self.introspect(index, validated_text, output_schema)
+                if reasks:
+                    raise ValueError(
+                        "Reasks are not yet supported with streaming. Please "
+                        "remove reasks from schema or disable streaming."
+                    )
+                # 5. Convert validated fragment to a pretty JSON string
+                yield ValidationOutcome(
+                    #  The chunk or the whole output?
+                    raw_llm_output=chunk_text,
+                    validated_output=validated_text,
+                    validation_passed=validated_text is not None,
+                )
+            # handle case where generator doesn't give finished status
+            if not stream_finished:
+                last_result = self.validate(
+                    iteration,
+                    index,
+                    "",
+                    output_schema,
+                    True,
+                    validate_subschema=True,
+                    remainder=True,
                 )
+                if len(last_result) > 0:
+                    yield ValidationOutcome(
+                        raw_llm_output=last_chunk_text,
+                        validated_output=last_result,
+                        validation_passed=last_result is not None,
+                    )
+        # handle non string schema
+        else:
+            for chunk in stream:
+                # 1. Get the text from the chunk and append to fragment
+                chunk_text = self.get_chunk_text(chunk, api)
+                fragment += chunk_text
 
-            # 5. Convert validated fragment to a pretty JSON string
-            yield ValidationOutcome(
-                raw_llm_output=fragment,
-                validated_output=validated_fragment,
-                validation_passed=validated_fragment is not None,
-            )
+                parsed_fragment, move_to_next = self.parse(
+                    index, fragment, output_schema, verified
+                )
+                if move_to_next:
+                    # Continue to next chunk
+                    continue
+                validated_fragment = self.validate(
+                    iteration,
+                    index,
+                    parsed_fragment,
+                    output_schema,
+                    validate_subschema=True,
+                )
+                if isinstance(validated_fragment, SkeletonReAsk):
+                    raise ValueError(
+                        "Received fragment schema is an invalid sub-schema "
+                        "of the expected output JSON schema."
+                    )
+
+                # 4. Introspect: inspect the validated fragment for reasks
+                reasks, valid_op = self.introspect(
+                    index, validated_fragment, output_schema
+                )
+                if reasks:
+                    raise ValueError(
+                        "Reasks are not yet supported with streaming. Please "
+                        "remove reasks from schema or disable streaming."
+                    )
+
+                # 5. Convert validated fragment to a pretty JSON string
+                yield ValidationOutcome(
+                    raw_llm_output=fragment,
+                    validated_output=validated_fragment,
+                    validation_passed=validated_fragment is not None,
+                )
 
         # Finally, add to logs
         iteration.outputs.raw_output = fragment
         iteration.outputs.parsed_output = parsed_fragment
         iteration.outputs.validation_response = validated_fragment
         iteration.outputs.guarded_output = valid_op
 
+    def is_last_chunk(self, chunk: Any, api: Union[PromptCallableBase, None]) -> bool:
+        """Detect if chunk is final chunk"""
+        if isinstance(api, OpenAICallable):
+            if OPENAI_VERSION.startswith("0"):
+                finished = chunk["choices"][0]["finish_reason"]
+                return finished is not None
+            else:
+                finished = chunk.choices[0].finish_reason
+                return finished is not None
+        elif isinstance(api, OpenAIChatCallable):
+            if OPENAI_VERSION.startswith("0"):
+                finished = chunk["choices"][0]["finish_reason"]
+                return finished is not None
+            else:
+                finished = chunk.choices[0].finish_reason
+                return finished is not None
+        elif isinstance(api, LiteLLMCallable):
+            finished = chunk.choices[0].finish_reason
+            return finished is not None
+        else:
+            try:
+                finished = chunk.choices[0].finish_reason
+                return finished is not None
+            except (AttributeError, TypeError):
+                return False
+
     def get_chunk_text(self, chunk: Any, api: Union[PromptCallableBase, None]) -> str:
         """Get the text from a chunk."""
         chunk_text = ""

guardrails/schema/string_schema.py

@@ -134,6 +134,7 @@ def validate(
         metadata: Dict,
         attempt_number: int = 0,
         disable_tracer: Optional[bool] = True,
+        stream: Optional[bool] = False,
         **kwargs,
     ) -> Any:
         """Validate a dictionary of data against the schema.
@@ -160,19 +161,20 @@ def validate(
                 dummy_key: data,
             },
         )
-
         validated_response, metadata = validator_service.validate(
             value=data,
             metadata=metadata,
             validator_setup=validation,
             iteration=iteration,
             disable_tracer=disable_tracer,
+            stream=stream,
+            **kwargs,
         )
 
         validated_response = {dummy_key: validated_response}
 
         if check_refrain_in_dict(validated_response):
-            # If the data contains a `Refain` value, we return an empty
+            # If the data contains a `Refrain` value, we return an empty
             # dictionary.
             logger.debug("Refrain detected.")
             validated_response = {}