참고소스 수정본
This commit is contained in:
LASTA_DEV01\lasta
340
참고/guardrails-main/guardrails/run/async_stream_runner.py
Normal file
340
참고/guardrails-main/guardrails/run/async_stream_runner.py
Normal file
@@ -0,0 +1,340 @@
|
||||
from contextvars import ContextVar, copy_context
|
||||
import sys
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncIterator,
|
||||
Dict,
|
||||
List,
|
||||
Optional,
|
||||
cast,
|
||||
)
|
||||
|
||||
from guardrails.validator_service import AsyncValidatorService
|
||||
from guardrails.actions.reask import SkeletonReAsk
|
||||
from guardrails.classes import ValidationOutcome
|
||||
from guardrails.classes.history import Call, Inputs, Iteration, Outputs
|
||||
from guardrails.classes.output_type import OutputTypes
|
||||
from guardrails.llm_providers import (
|
||||
AsyncPromptCallableBase,
|
||||
)
|
||||
from guardrails.logger import set_scope
|
||||
from guardrails.run import StreamRunner
|
||||
from guardrails.run.async_runner import AsyncRunner
|
||||
from guardrails.telemetry import trace_async_stream_step
|
||||
from guardrails.hub_telemetry.hub_tracing import async_trace_stream
|
||||
from guardrails.types import OnFailAction
|
||||
from guardrails_ai.types import (
|
||||
PassResult,
|
||||
FailResult,
|
||||
)
|
||||
|
||||
|
||||
if sys.version_info.minor < 10:
|
||||
from guardrails.utils.polyfills import anext
|
||||
|
||||
|
||||
class AsyncStreamRunner(AsyncRunner, StreamRunner):
|
||||
# @async_trace_stream(name="/reasks", origin="AsyncStreamRunner.async_run")
|
||||
async def async_run(
|
||||
self, call_log: Call, prompt_params: Optional[Dict] = None
|
||||
) -> AsyncIterator[ValidationOutcome]:
|
||||
prompt_params = prompt_params or {}
|
||||
|
||||
(
|
||||
messages,
|
||||
output_schema,
|
||||
) = (
|
||||
self.messages,
|
||||
self.output_schema,
|
||||
)
|
||||
|
||||
result = await self.async_step(
|
||||
0,
|
||||
output_schema,
|
||||
call_log,
|
||||
api=self.api,
|
||||
messages=messages,
|
||||
prompt_params=prompt_params,
|
||||
output=self.output,
|
||||
)
|
||||
# FIXME: Where can this be moved to be less verbose? This is an await call on
|
||||
# the async generator.
|
||||
async for call in result:
|
||||
yield call
|
||||
|
||||
@async_trace_stream(name="/step", origin="AsyncStreamRunner.async_step")
|
||||
@trace_async_stream_step
|
||||
async def async_step(
|
||||
self,
|
||||
index: int,
|
||||
output_schema: Dict[str, Any],
|
||||
call_log: Call,
|
||||
*,
|
||||
api: Optional[AsyncPromptCallableBase],
|
||||
messages: Optional[List[Dict]] = None,
|
||||
prompt_params: Optional[Dict] = None,
|
||||
output: Optional[str] = None,
|
||||
) -> AsyncIterator[ValidationOutcome]:
|
||||
prompt_params = prompt_params or {}
|
||||
inputs = Inputs(
|
||||
llm_api=api,
|
||||
llm_output=output,
|
||||
messages=messages,
|
||||
prompt_params=prompt_params,
|
||||
num_reasks=self.num_reasks,
|
||||
metadata=self.metadata,
|
||||
full_schema_reask=self.full_schema_reask,
|
||||
stream=True,
|
||||
)
|
||||
outputs = Outputs()
|
||||
iteration = Iteration(
|
||||
callId=call_log.id, index=index, inputs=inputs, outputs=outputs
|
||||
)
|
||||
set_scope(str(id(iteration)))
|
||||
call_log.iterations.push(iteration)
|
||||
if output is not None:
|
||||
messages = None
|
||||
else:
|
||||
messages = await self.async_prepare(
|
||||
call_log,
|
||||
messages=messages,
|
||||
prompt_params=prompt_params,
|
||||
api=api,
|
||||
attempt_number=index,
|
||||
)
|
||||
|
||||
iteration.inputs.messages = messages
|
||||
|
||||
llm_response = await self.async_call(messages, api, output)
|
||||
iteration.outputs.llm_response_info = llm_response
|
||||
stream_output = llm_response.async_stream_output
|
||||
if stream_output is None:
|
||||
raise ValueError(
|
||||
"No async stream was returned from the API. Please check that "
|
||||
"the API is returning an async generator."
|
||||
)
|
||||
|
||||
fragment = ""
|
||||
parsed_fragment, validated_fragment, valid_op = None, None, None
|
||||
verified = set()
|
||||
validation_response = ""
|
||||
validation_progress = {}
|
||||
refrain_triggered = False
|
||||
validation_passed = True
|
||||
|
||||
context = copy_context()
|
||||
stream_context_vars: ContextVar[Dict[str, ContextVar[List[str]]]] = ContextVar(
|
||||
"stream_context"
|
||||
)
|
||||
context_vars: Dict[str, ContextVar[List[str]]] = {}
|
||||
for k, v in self.validation_map.items():
|
||||
if isinstance(v, list):
|
||||
for validator in v:
|
||||
property_validation_chunks = ContextVar(
|
||||
f"{k}_{validator.rail_alias}_chunks"
|
||||
)
|
||||
context.run(property_validation_chunks.set, [])
|
||||
context_vars[f"{k}_{validator.rail_alias}"] = (
|
||||
property_validation_chunks # noqa: E501
|
||||
)
|
||||
context.run(stream_context_vars.set, context_vars)
|
||||
|
||||
if self.output_type == OutputTypes.STRING:
|
||||
validator_service = AsyncValidatorService(self.disable_tracer)
|
||||
|
||||
next_exists = True
|
||||
while next_exists:
|
||||
try:
|
||||
chunk = await anext(stream_output)
|
||||
chunk_text = self.get_chunk_text(chunk, api)
|
||||
_ = self.is_last_chunk(chunk, api)
|
||||
|
||||
fragment += chunk_text
|
||||
|
||||
results = await validator_service.async_partial_validate(
|
||||
chunk_text,
|
||||
self.metadata,
|
||||
self.validation_map,
|
||||
iteration,
|
||||
"$",
|
||||
"$",
|
||||
True,
|
||||
context=context,
|
||||
context_vars=stream_context_vars,
|
||||
)
|
||||
validators = self.validation_map.get("$", [])
|
||||
|
||||
# collect the result validated_chunk into validation progress
|
||||
# per validator
|
||||
for result in results:
|
||||
validator_log = result.validator_logs # type: ignore
|
||||
validator = next(
|
||||
filter(
|
||||
lambda x: x.rail_alias == validator_log.registered_name,
|
||||
validators,
|
||||
),
|
||||
None,
|
||||
)
|
||||
if (
|
||||
validator_log.validation_result
|
||||
and validator_log.validation_result.validated_chunk
|
||||
):
|
||||
is_filter = (
|
||||
validator.on_fail_descriptor is OnFailAction.FILTER # type: ignore
|
||||
)
|
||||
is_refrain = (
|
||||
validator.on_fail_descriptor is OnFailAction.REFRAIN # type: ignore
|
||||
)
|
||||
if validator_log.validation_result.outcome == "fail":
|
||||
validation_passed = False
|
||||
reasks, valid_op = self.introspect(
|
||||
validator_log.validation_result
|
||||
)
|
||||
if reasks:
|
||||
raise ValueError(
|
||||
"Reasks are not yet supported with streaming. "
|
||||
"Please remove reasks from schema or disable"
|
||||
" streaming."
|
||||
)
|
||||
|
||||
if isinstance(validator_log.validation_result, PassResult):
|
||||
chunk = validator_log.validation_result.validated_chunk
|
||||
elif isinstance(
|
||||
validator_log.validation_result, FailResult
|
||||
):
|
||||
if is_filter or is_refrain:
|
||||
refrain_triggered = True
|
||||
chunk = ""
|
||||
else:
|
||||
chunk = validator_service.perform_correction(
|
||||
validator_log.validation_result,
|
||||
validator_log.validation_result.validated_chunk,
|
||||
validator, # type: ignore
|
||||
rechecked_value=None,
|
||||
) # type: ignore
|
||||
|
||||
if validator_log.validator_name not in validation_progress:
|
||||
validation_progress[validator_log.validator_name] = ""
|
||||
|
||||
validation_progress[validator_log.validator_name] += chunk
|
||||
# if there is an entry for every validator
|
||||
# run a merge and emit a validation outcome
|
||||
if (
|
||||
len(validation_progress) == len(validators)
|
||||
or len(validators) == 0
|
||||
):
|
||||
if refrain_triggered:
|
||||
current = ""
|
||||
else:
|
||||
merge_chunks = []
|
||||
for piece in validation_progress:
|
||||
merge_chunks.append(validation_progress[piece])
|
||||
|
||||
current = validator_service.multi_merge(
|
||||
fragment, merge_chunks
|
||||
)
|
||||
|
||||
vo = ValidationOutcome(
|
||||
callId=call_log.id,
|
||||
rawLlmOutput=fragment,
|
||||
validatedOutput=current,
|
||||
validationPassed=True,
|
||||
)
|
||||
fragment = ""
|
||||
validation_progress = {}
|
||||
refrain_triggered = False
|
||||
|
||||
yield vo
|
||||
|
||||
except StopIteration:
|
||||
next_exists = False
|
||||
except StopAsyncIteration:
|
||||
next_exists = False
|
||||
except Exception as e:
|
||||
raise e
|
||||
finally:
|
||||
# reset all context vars
|
||||
for context_var in context_vars.values():
|
||||
token = context.run(context_var.set, [])
|
||||
context.run(context_var.reset, token)
|
||||
token = context.run(stream_context_vars.set, {})
|
||||
context.run(stream_context_vars.reset, token)
|
||||
|
||||
# if theres anything left merge and emit a chunk
|
||||
if len(validation_progress) > 0:
|
||||
merge_chunks = []
|
||||
for piece in validation_progress:
|
||||
merge_chunks.append(validation_progress[piece])
|
||||
|
||||
current = validator_service.multi_merge(fragment, merge_chunks)
|
||||
yield ValidationOutcome(
|
||||
callId=call_log.id,
|
||||
rawLlmOutput=fragment,
|
||||
validatedOutput=current,
|
||||
validationPassed=validation_passed,
|
||||
)
|
||||
else:
|
||||
next_exists = True
|
||||
while next_exists:
|
||||
try:
|
||||
chunk = await anext(stream_output)
|
||||
chunk_text = self.get_chunk_text(chunk, api)
|
||||
fragment += chunk_text
|
||||
|
||||
parsed_fragment, move_to_next = self.parse(
|
||||
fragment, output_schema, verified=verified
|
||||
)
|
||||
if move_to_next:
|
||||
continue
|
||||
validated_fragment = await self.async_validate(
|
||||
iteration,
|
||||
index,
|
||||
parsed_fragment,
|
||||
output_schema,
|
||||
validate_subschema=True,
|
||||
context=context,
|
||||
context_vars=stream_context_vars,
|
||||
)
|
||||
if isinstance(validated_fragment, SkeletonReAsk):
|
||||
raise ValueError(
|
||||
"Received fragment schema is an invalid sub-schema "
|
||||
"of the expected output JSON schema."
|
||||
)
|
||||
|
||||
reasks, valid_op = self.introspect(validated_fragment)
|
||||
if reasks:
|
||||
raise ValueError(
|
||||
"Reasks are not yet supported with streaming. Please "
|
||||
"remove reasks from schema or disable streaming."
|
||||
)
|
||||
|
||||
if self.output_type == OutputTypes.LIST:
|
||||
validation_response = cast(list, validated_fragment)
|
||||
else:
|
||||
validation_response = cast(dict, validated_fragment)
|
||||
yield ValidationOutcome(
|
||||
callId=call_log.id,
|
||||
rawLlmOutput=fragment,
|
||||
validatedOutput=validated_fragment,
|
||||
validationPassed=validated_fragment is not None,
|
||||
)
|
||||
fragment = ""
|
||||
except StopIteration:
|
||||
next_exists = False
|
||||
except StopAsyncIteration:
|
||||
next_exists = False
|
||||
except Exception as e:
|
||||
raise e
|
||||
finally:
|
||||
# reset all context vars
|
||||
for context_var in context_vars.values():
|
||||
token = context.run(context_var.set, [])
|
||||
context.run(context_var.reset, token)
|
||||
token = context.run(stream_context_vars.set, {})
|
||||
context.run(stream_context_vars.reset, token)
|
||||
|
||||
iteration.outputs.raw_output = fragment
|
||||
# FIXME: Handle case where parsing continuously fails/is a reask
|
||||
iteration.outputs.parsed_output = parsed_fragment or fragment # type: ignore
|
||||
iteration.outputs.validation_response = validation_response
|
||||
iteration.outputs.guarded_output = valid_op
|
||||
Reference in New Issue
Block a user