526 lines
18 KiB
Python
526 lines
18 KiB
Python
|
import copy
|
||
|
|
from functools import partial
|
||
|
|
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
|
||
|
|
|
||
|
|
|
||
|
|
from guardrails import validator_service
|
||
|
|
from guardrails.actions.reask import get_reask_setup
|
||
|
|
from guardrails.classes.execution.guard_execution_options import GuardExecutionOptions
|
||
|
|
from guardrails.classes.history import Call, Inputs, Iteration, Outputs
|
||
|
|
from guardrails.classes.output_type import OutputTypes
|
||
|
|
from guardrails.constants import fail_status
|
||
|
|
from guardrails.errors import ValidationError
|
||
|
|
from guardrails.llm_providers import (
|
||
|
|
AsyncPromptCallableBase,
|
||
|
|
PromptCallableBase,
|
||
|
|
)
|
||
|
|
from guardrails.logger import set_scope
|
||
|
|
from guardrails.prompt import Prompt
|
||
|
|
from guardrails.prompt.messages import Messages
|
||
|
|
from guardrails.run.utils import messages_source
|
||
|
|
from guardrails.schema.rail_schema import json_schema_to_rail_output
|
||
|
|
from guardrails.schema.validator import schema_validation
|
||
|
|
from guardrails.hub_telemetry.hub_tracing import trace
|
||
|
|
from guardrails.types import ModelOrListOfModels, ValidatorMap, MessageHistory
|
||
|
|
from guardrails.utils.exception_utils import UserFacingException
|
||
|
|
from guardrails.utils.hub_telemetry_utils import HubTelemetry
|
||
|
|
from guardrails.classes.llm.llm_response import LLMResponse
|
||
|
|
from guardrails.utils.parsing_utils import (
|
||
|
|
coerce_types,
|
||
|
|
parse_llm_output,
|
||
|
|
prune_extra_keys,
|
||
|
|
)
|
||
|
|
from guardrails.utils.prompt_utils import (
|
||
|
|
prompt_content_for_schema,
|
||
|
|
)
|
||
|
|
from guardrails.actions.reask import NonParseableReAsk, ReAsk, introspect
|
||
|
|
from guardrails.telemetry import trace_call, trace_step
|
||
|
|
|
||
|
|
|
||
|
|
class Runner:
|
||
|
|
"""Runner class that calls an LLM API with a prompt, and performs input and
|
||
|
|
output validation.
|
||
|
|
|
||
|
|
This class will repeatedly call the API until the
|
||
|
|
reask budget is exhausted, or the output is valid.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
prompt: The prompt to use.
|
||
|
|
api: The LLM API to call, which should return a string.
|
||
|
|
output_schema: The output schema to use for validation.
|
||
|
|
num_reasks: The maximum number of times to reask the LLM in case of
|
||
|
|
validation failure, defaults to 0.
|
||
|
|
output: The output to use instead of calling the API, used in cases
|
||
|
|
where the output is already known.
|
||
|
|
"""
|
||
|
|
|
||
|
|
# Validation Inputs
|
||
|
|
output_schema: Dict[str, Any]
|
||
|
|
output_type: OutputTypes
|
||
|
|
validation_map: ValidatorMap = {}
|
||
|
|
metadata: Dict[str, Any]
|
||
|
|
|
||
|
|
# LLM Inputs
|
||
|
|
messages: Optional[List[Dict[str, Union[Prompt, str]]]] = None
|
||
|
|
base_model: Optional[ModelOrListOfModels]
|
||
|
|
exec_options: Optional[GuardExecutionOptions]
|
||
|
|
|
||
|
|
# LLM Calling Details
|
||
|
|
api: Optional[PromptCallableBase] = None
|
||
|
|
output: Optional[str] = None
|
||
|
|
num_reasks: int
|
||
|
|
full_schema_reask: bool = False
|
||
|
|
|
||
|
|
# Internal Metrics Collection
|
||
|
|
disable_tracer: Optional[bool] = True
|
||
|
|
|
||
|
|
# QUESTION: Are any of these init args actually necessary for initialization?
|
||
|
|
# ANSWER: _Maybe_ messages for Prompt initialization
|
||
|
|
# but even that can happen at execution time.
|
||
|
|
# TODO: In versions >=0.6.x, remove this class and just execute a Guard functionally
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
output_type: OutputTypes,
|
||
|
|
output_schema: Dict[str, Any],
|
||
|
|
num_reasks: int,
|
||
|
|
validation_map: ValidatorMap,
|
||
|
|
*,
|
||
|
|
messages: Optional[List[Dict]] = None,
|
||
|
|
api: Optional[PromptCallableBase] = None,
|
||
|
|
metadata: Optional[Dict[str, Any]] = None,
|
||
|
|
output: Optional[str] = None,
|
||
|
|
base_model: Optional[ModelOrListOfModels] = None,
|
||
|
|
full_schema_reask: bool = False,
|
||
|
|
disable_tracer: Optional[bool] = True,
|
||
|
|
exec_options: Optional[GuardExecutionOptions] = None,
|
||
|
|
):
|
||
|
|
# Validation Inputs
|
||
|
|
self.output_type = output_type
|
||
|
|
self.output_schema = output_schema
|
||
|
|
self.validation_map = validation_map
|
||
|
|
self.metadata = metadata or {}
|
||
|
|
self.exec_options = copy.deepcopy(exec_options) or GuardExecutionOptions()
|
||
|
|
|
||
|
|
# LLM Inputs
|
||
|
|
|
||
|
|
stringified_output_schema = prompt_content_for_schema(
|
||
|
|
output_type, output_schema, validation_map
|
||
|
|
)
|
||
|
|
xml_output_schema = json_schema_to_rail_output(
|
||
|
|
json_schema=output_schema, validator_map=validation_map
|
||
|
|
)
|
||
|
|
|
||
|
|
if messages:
|
||
|
|
self.exec_options.messages = messages
|
||
|
|
messages_copy = []
|
||
|
|
for msg in messages:
|
||
|
|
msg_copy = copy.deepcopy(msg)
|
||
|
|
msg_copy["content"] = Prompt(
|
||
|
|
msg_copy["content"],
|
||
|
|
output_schema=stringified_output_schema,
|
||
|
|
xml_output_schema=xml_output_schema,
|
||
|
|
)
|
||
|
|
messages_copy.append(msg_copy)
|
||
|
|
self.messages = messages_copy
|
||
|
|
|
||
|
|
self.base_model = base_model
|
||
|
|
|
||
|
|
# LLM Calling Details
|
||
|
|
self.api = api
|
||
|
|
self.output = output
|
||
|
|
self.num_reasks = num_reasks
|
||
|
|
self.full_schema_reask = full_schema_reask
|
||
|
|
|
||
|
|
# Internal Metrics Collection
|
||
|
|
# Get metrics opt-out from credentials
|
||
|
|
self._disable_tracer = disable_tracer
|
||
|
|
|
||
|
|
# Get the HubTelemetry singleton
|
||
|
|
self._hub_telemetry = HubTelemetry()
|
||
|
|
self._hub_telemetry._enabled = not self._disable_tracer
|
||
|
|
|
||
|
|
@trace(name="/reasks", origin="Runner.__call__")
|
||
|
|
def __call__(self, call_log: Call, prompt_params: Optional[Dict] = None) -> Call:
|
||
|
|
"""Execute the runner by repeatedly calling step until the reask budget
|
||
|
|
is exhausted.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
prompt_params: Parameters to pass to the prompt in order to
|
||
|
|
generate the prompt string.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The Call log for this run.
|
||
|
|
"""
|
||
|
|
prompt_params = prompt_params or {}
|
||
|
|
try:
|
||
|
|
# NOTE: At first glance this seems gratuitous,
|
||
|
|
# but these local variables are reassigned after
|
||
|
|
# calling self.prepare_to_loop
|
||
|
|
(
|
||
|
|
messages,
|
||
|
|
output_schema,
|
||
|
|
) = (
|
||
|
|
self.messages,
|
||
|
|
self.output_schema,
|
||
|
|
)
|
||
|
|
|
||
|
|
index = 0
|
||
|
|
for index in range(self.num_reasks + 1):
|
||
|
|
# Run a single step.
|
||
|
|
iteration = self.step(
|
||
|
|
index=index,
|
||
|
|
api=self.api,
|
||
|
|
messages=messages,
|
||
|
|
prompt_params=prompt_params,
|
||
|
|
output_schema=output_schema,
|
||
|
|
output=self.output if index == 0 else None,
|
||
|
|
call_log=call_log,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Loop again?
|
||
|
|
if not self.do_loop(index, iteration.reasks):
|
||
|
|
break
|
||
|
|
|
||
|
|
# Get new prompt and output schema.
|
||
|
|
(output_schema, messages) = self.prepare_to_loop(
|
||
|
|
iteration.reasks,
|
||
|
|
output_schema,
|
||
|
|
parsed_output=iteration.outputs.parsed_output,
|
||
|
|
validated_output=call_log.validation_response,
|
||
|
|
prompt_params=prompt_params,
|
||
|
|
)
|
||
|
|
|
||
|
|
except UserFacingException as e:
|
||
|
|
# Because Pydantic v1 doesn't respect property setters
|
||
|
|
call_log.exception = e.original_exception
|
||
|
|
raise e.original_exception
|
||
|
|
except Exception as e:
|
||
|
|
# Because Pydantic v1 doesn't respect property setters
|
||
|
|
call_log.exception = e
|
||
|
|
raise e
|
||
|
|
return call_log
|
||
|
|
|
||
|
|
@trace(name="/step", origin="Runner.step")
|
||
|
|
@trace_step
|
||
|
|
def step(
|
||
|
|
self,
|
||
|
|
index: int,
|
||
|
|
output_schema: Dict[str, Any],
|
||
|
|
call_log: Call,
|
||
|
|
*,
|
||
|
|
api: Optional[PromptCallableBase],
|
||
|
|
messages: Optional[List[Dict]] = None,
|
||
|
|
prompt_params: Optional[Dict] = None,
|
||
|
|
output: Optional[str] = None,
|
||
|
|
) -> Iteration:
|
||
|
|
"""Run a full step."""
|
||
|
|
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,
|
||
|
|
)
|
||
|
|
outputs = Outputs()
|
||
|
|
iteration = Iteration(
|
||
|
|
callId=call_log.id, index=index, inputs=inputs, outputs=outputs
|
||
|
|
)
|
||
|
|
set_scope(str(id(iteration)))
|
||
|
|
call_log.iterations.push(iteration)
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Prepare: run pre-processing, and input validation.
|
||
|
|
if output is not None:
|
||
|
|
messages = None
|
||
|
|
else:
|
||
|
|
messages = self.prepare(
|
||
|
|
call_log,
|
||
|
|
messages=messages,
|
||
|
|
prompt_params=prompt_params,
|
||
|
|
api=api,
|
||
|
|
attempt_number=index,
|
||
|
|
)
|
||
|
|
|
||
|
|
iteration.inputs.messages = messages
|
||
|
|
|
||
|
|
# Call: run the API.
|
||
|
|
llm_response = self.call(messages, api, output)
|
||
|
|
|
||
|
|
iteration.outputs.llm_response_info = llm_response
|
||
|
|
raw_output = llm_response.output
|
||
|
|
|
||
|
|
# Parse: parse the output.
|
||
|
|
parsed_output, parsing_error = self.parse(raw_output, output_schema)
|
||
|
|
if parsing_error or isinstance(parsed_output, ReAsk):
|
||
|
|
iteration.outputs.exception = parsing_error # type: ignore
|
||
|
|
iteration.outputs.error = str(parsing_error)
|
||
|
|
iteration.outputs.reasks.append(parsed_output) # type: ignore
|
||
|
|
else:
|
||
|
|
iteration.outputs.parsed_output = parsed_output
|
||
|
|
|
||
|
|
# Validate: run output validation.
|
||
|
|
if parsing_error and isinstance(parsed_output, NonParseableReAsk):
|
||
|
|
reasks, _ = self.introspect(parsed_output)
|
||
|
|
else:
|
||
|
|
# Validate: run output validation.
|
||
|
|
validated_output = self.validate(
|
||
|
|
iteration, index, parsed_output, output_schema
|
||
|
|
)
|
||
|
|
iteration.outputs.validation_response = validated_output
|
||
|
|
|
||
|
|
# Introspect: inspect validated output for reasks.
|
||
|
|
reasks, valid_output = self.introspect(validated_output)
|
||
|
|
iteration.outputs.guarded_output = valid_output
|
||
|
|
|
||
|
|
iteration.outputs.reasks = list(reasks)
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
error_message = str(e)
|
||
|
|
iteration.outputs.error = error_message
|
||
|
|
iteration.outputs.exception = e
|
||
|
|
raise e
|
||
|
|
return iteration
|
||
|
|
|
||
|
|
@trace(name="/input_validation", origin="Runner.validate_messages")
|
||
|
|
def validate_messages(
|
||
|
|
self, call_log: Call, messages: MessageHistory, attempt_number: int
|
||
|
|
) -> None:
|
||
|
|
for msg in messages:
|
||
|
|
content = (
|
||
|
|
msg["content"].source
|
||
|
|
if isinstance(msg["content"], Prompt)
|
||
|
|
else msg["content"]
|
||
|
|
)
|
||
|
|
inputs = Inputs(
|
||
|
|
llm_output=content,
|
||
|
|
)
|
||
|
|
iteration = Iteration(
|
||
|
|
callId=call_log.id, index=attempt_number, inputs=inputs
|
||
|
|
)
|
||
|
|
call_log.iterations.insert(0, iteration)
|
||
|
|
value, _metadata = validator_service.validate(
|
||
|
|
value=content,
|
||
|
|
metadata=self.metadata,
|
||
|
|
validator_map=self.validation_map,
|
||
|
|
iteration=iteration,
|
||
|
|
disable_tracer=self._disable_tracer,
|
||
|
|
path="messages",
|
||
|
|
)
|
||
|
|
|
||
|
|
validated_msg = validator_service.post_process_validation(
|
||
|
|
value, attempt_number, iteration, OutputTypes.STRING
|
||
|
|
)
|
||
|
|
|
||
|
|
iteration.outputs.validation_response = validated_msg
|
||
|
|
|
||
|
|
if isinstance(validated_msg, ReAsk):
|
||
|
|
raise ValidationError(f"Messages validation failed: {validated_msg}")
|
||
|
|
elif not validated_msg or iteration.status == fail_status:
|
||
|
|
raise ValidationError("Messages validation failed")
|
||
|
|
|
||
|
|
msg["content"] = cast(str, validated_msg)
|
||
|
|
|
||
|
|
return messages # type: ignore
|
||
|
|
|
||
|
|
def prepare_messages(
|
||
|
|
self,
|
||
|
|
call_log: Call,
|
||
|
|
messages: MessageHistory,
|
||
|
|
prompt_params: Dict,
|
||
|
|
attempt_number: int,
|
||
|
|
) -> MessageHistory:
|
||
|
|
formatted_messages: MessageHistory = []
|
||
|
|
# Format any variables in the message history with the prompt params.
|
||
|
|
for msg in messages:
|
||
|
|
msg_copy = copy.deepcopy(msg)
|
||
|
|
if attempt_number == 0:
|
||
|
|
msg_copy["content"] = msg_copy["content"].format(**prompt_params)
|
||
|
|
formatted_messages.append(msg_copy)
|
||
|
|
|
||
|
|
# validate messages
|
||
|
|
if "messages" in self.validation_map:
|
||
|
|
self.validate_messages(call_log, formatted_messages, attempt_number)
|
||
|
|
|
||
|
|
return formatted_messages
|
||
|
|
|
||
|
|
@trace(name="/input_validation", origin="Runner.validate_prompt")
|
||
|
|
def validate_prompt(self, call_log: Call, prompt: Prompt, attempt_number: int):
|
||
|
|
inputs = Inputs(
|
||
|
|
llm_output=prompt.source,
|
||
|
|
)
|
||
|
|
iteration = Iteration(callId=call_log.id, index=attempt_number, inputs=inputs)
|
||
|
|
call_log.iterations.insert(0, iteration)
|
||
|
|
value, _metadata = validator_service.validate(
|
||
|
|
value=prompt.source,
|
||
|
|
metadata=self.metadata,
|
||
|
|
validator_map=self.validation_map,
|
||
|
|
iteration=iteration,
|
||
|
|
disable_tracer=self._disable_tracer,
|
||
|
|
path="prompt",
|
||
|
|
)
|
||
|
|
|
||
|
|
validated_prompt = validator_service.post_process_validation(
|
||
|
|
value, attempt_number, iteration, OutputTypes.STRING
|
||
|
|
)
|
||
|
|
|
||
|
|
iteration.outputs.validation_response = validated_prompt
|
||
|
|
|
||
|
|
if isinstance(validated_prompt, ReAsk):
|
||
|
|
raise ValidationError(f"Prompt validation failed: {validated_prompt}")
|
||
|
|
elif not validated_prompt or iteration.status == fail_status:
|
||
|
|
raise ValidationError("Prompt validation failed")
|
||
|
|
return Prompt(cast(str, validated_prompt))
|
||
|
|
|
||
|
|
@trace(name="/input_prep", origin="Runner.prepare")
|
||
|
|
def prepare(
|
||
|
|
self,
|
||
|
|
call_log: Call,
|
||
|
|
attempt_number: int,
|
||
|
|
*,
|
||
|
|
messages: Optional[MessageHistory],
|
||
|
|
prompt_params: Optional[Dict] = None,
|
||
|
|
api: Optional[Union[PromptCallableBase, AsyncPromptCallableBase]],
|
||
|
|
) -> Optional[MessageHistory]:
|
||
|
|
"""Prepare by running pre-processing and input validation.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The message history.
|
||
|
|
"""
|
||
|
|
prompt_params = prompt_params or {}
|
||
|
|
if api is None:
|
||
|
|
raise UserFacingException(ValueError("API must be provided."))
|
||
|
|
|
||
|
|
if messages:
|
||
|
|
messages = self.prepare_messages(
|
||
|
|
call_log, messages, prompt_params, attempt_number
|
||
|
|
)
|
||
|
|
|
||
|
|
return messages
|
||
|
|
|
||
|
|
@trace(name="/llm_call", origin="Runner.call")
|
||
|
|
@trace_call
|
||
|
|
def call(
|
||
|
|
self,
|
||
|
|
messages: Optional[MessageHistory],
|
||
|
|
api: Optional[PromptCallableBase],
|
||
|
|
output: Optional[str] = None,
|
||
|
|
) -> LLMResponse:
|
||
|
|
"""Run a step.
|
||
|
|
|
||
|
|
1. Query the LLM API,
|
||
|
|
2. Convert the response string to a dict,
|
||
|
|
3. Log the output
|
||
|
|
"""
|
||
|
|
|
||
|
|
# If the API supports a base model, pass it in.
|
||
|
|
api_fn = api
|
||
|
|
if api is not None:
|
||
|
|
supports_base_model = getattr(api, "supports_base_model", False)
|
||
|
|
if supports_base_model:
|
||
|
|
api_fn = partial(api, base_model=self.base_model)
|
||
|
|
|
||
|
|
if output is not None:
|
||
|
|
llm_response = LLMResponse(output=output)
|
||
|
|
elif api_fn is None:
|
||
|
|
raise ValueError("API or output must be provided.")
|
||
|
|
elif messages:
|
||
|
|
llm_response = api_fn(messages=messages_source(messages))
|
||
|
|
else:
|
||
|
|
llm_response = api_fn()
|
||
|
|
|
||
|
|
return llm_response
|
||
|
|
|
||
|
|
def parse(self, output: str, output_schema: Dict[str, Any], **kwargs):
|
||
|
|
parsed_output, error = parse_llm_output(output, self.output_type, **kwargs)
|
||
|
|
if parsed_output and not error and not isinstance(parsed_output, ReAsk):
|
||
|
|
parsed_output = prune_extra_keys(parsed_output, output_schema)
|
||
|
|
parsed_output = coerce_types(parsed_output, output_schema)
|
||
|
|
return parsed_output, error
|
||
|
|
|
||
|
|
@trace(name="/validation", origin="Runner.validate")
|
||
|
|
def validate(
|
||
|
|
self,
|
||
|
|
iteration: Iteration,
|
||
|
|
attempt_number: int,
|
||
|
|
parsed_output: Any,
|
||
|
|
output_schema: Dict[str, Any],
|
||
|
|
stream: Optional[bool] = False,
|
||
|
|
**kwargs,
|
||
|
|
):
|
||
|
|
"""Validate the output."""
|
||
|
|
# Break early if empty
|
||
|
|
if parsed_output is None:
|
||
|
|
return None
|
||
|
|
|
||
|
|
skeleton_reask = schema_validation(parsed_output, output_schema, **kwargs)
|
||
|
|
if skeleton_reask:
|
||
|
|
return skeleton_reask
|
||
|
|
|
||
|
|
if self.output_type != OutputTypes.STRING:
|
||
|
|
stream = None
|
||
|
|
|
||
|
|
validated_output, metadata = validator_service.validate(
|
||
|
|
value=parsed_output,
|
||
|
|
metadata=self.metadata,
|
||
|
|
validator_map=self.validation_map,
|
||
|
|
iteration=iteration,
|
||
|
|
disable_tracer=self._disable_tracer,
|
||
|
|
path="$",
|
||
|
|
stream=stream,
|
||
|
|
**kwargs,
|
||
|
|
)
|
||
|
|
self.metadata.update(metadata)
|
||
|
|
validated_output = validator_service.post_process_validation(
|
||
|
|
validated_output, attempt_number, iteration, self.output_type
|
||
|
|
)
|
||
|
|
|
||
|
|
return validated_output
|
||
|
|
|
||
|
|
def introspect(
|
||
|
|
self,
|
||
|
|
validated_output: Any,
|
||
|
|
) -> Tuple[Sequence[ReAsk], Optional[Union[str, Dict, List]]]:
|
||
|
|
"""Introspect the validated output."""
|
||
|
|
if validated_output is None:
|
||
|
|
return [], None
|
||
|
|
reasks, valid_output = introspect(validated_output)
|
||
|
|
|
||
|
|
return reasks, valid_output
|
||
|
|
|
||
|
|
def do_loop(self, attempt_number: int, reasks: Sequence[ReAsk]) -> bool:
|
||
|
|
"""Determine if we should loop again."""
|
||
|
|
if reasks and attempt_number < self.num_reasks:
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
def prepare_to_loop(
|
||
|
|
self,
|
||
|
|
reasks: Sequence[ReAsk],
|
||
|
|
output_schema: Dict[str, Any],
|
||
|
|
*,
|
||
|
|
parsed_output: Optional[Union[str, List, Dict, ReAsk]] = None,
|
||
|
|
validated_output: Optional[Union[str, List, Dict, ReAsk]] = None,
|
||
|
|
prompt_params: Optional[Dict] = None,
|
||
|
|
) -> Tuple[
|
||
|
|
Dict[str, Any],
|
||
|
|
Optional[Union[List[Dict], Messages]],
|
||
|
|
]:
|
||
|
|
"""Prepare to loop again."""
|
||
|
|
prompt_params = prompt_params or {}
|
||
|
|
output_schema, messages = get_reask_setup(
|
||
|
|
output_type=self.output_type,
|
||
|
|
output_schema=output_schema,
|
||
|
|
validation_map=self.validation_map,
|
||
|
|
reasks=reasks,
|
||
|
|
parsing_response=parsed_output,
|
||
|
|
validation_response=validated_output,
|
||
|
|
use_full_schema=self.full_schema_reask,
|
||
|
|
prompt_params=prompt_params,
|
||
|
|
exec_options=self.exec_options,
|
||
|
|
)
|
||
|
|
|
||
|
|
return output_schema, messages
|