219 lines
8.5 KiB
Python
219 lines
8.5 KiB
Python
from functools import wraps
|
|
from typing import (
|
|
Any,
|
|
Awaitable,
|
|
Callable,
|
|
Dict,
|
|
Optional,
|
|
)
|
|
|
|
from opentelemetry import context, trace
|
|
from opentelemetry.trace import StatusCode, Tracer, Span
|
|
|
|
try:
|
|
from openinference.semconv.trace import SpanAttributes # type: ignore
|
|
except ImportError:
|
|
SpanAttributes = None
|
|
|
|
from guardrails_ai.types import ValidationResult
|
|
|
|
from guardrails.settings import settings
|
|
from guardrails.telemetry.common import add_user_attributes, serialize
|
|
from guardrails.telemetry.open_inference import trace_operation
|
|
from guardrails.utils.casting_utils import to_string
|
|
from guardrails.utils.safe_get import safe_get
|
|
from guardrails.version import GUARDRAILS_VERSION
|
|
|
|
|
|
def add_validator_attributes(
|
|
*args,
|
|
validator_span: Span,
|
|
validator_name: str,
|
|
obj_id: int,
|
|
on_fail_descriptor: Optional[str] = None,
|
|
result: Optional[ValidationResult] = None,
|
|
init_kwargs: Dict[str, Any] = {},
|
|
validation_session_id: str,
|
|
**kwargs,
|
|
):
|
|
value_arg = serialize(safe_get(args, 0)) or ""
|
|
metadata_arg = serialize(safe_get(args, 1, {})) or "{}"
|
|
|
|
# Legacy Span Attributes
|
|
validator_span.set_attribute("on_fail_descriptor", on_fail_descriptor or "noop")
|
|
validator_span.set_attribute(
|
|
"args",
|
|
to_string({k: to_string(v) for k, v in init_kwargs.items()}) or "{}",
|
|
)
|
|
validator_span.set_attribute("instance_id", serialize(obj_id) or "")
|
|
validator_span.set_attribute("input", value_arg)
|
|
|
|
# New Span Attributes
|
|
validator_span.set_attribute("type", "guardrails/guard/step/validator")
|
|
validator_span.set_attribute("validation_session_id", validation_session_id)
|
|
|
|
### Validator.__init__ ###
|
|
validator_span.set_attribute("validator.name", validator_name or "unknown")
|
|
validator_span.set_attribute("validator.on_fail", on_fail_descriptor or "noop")
|
|
validator_span.set_attribute("validator.instance_id", serialize(obj_id) or "")
|
|
for k, v in init_kwargs.items():
|
|
if v is not None:
|
|
validator_span.set_attribute(f"validator.init.{k}", serialize(v) or "")
|
|
|
|
### Validator.validate ###
|
|
validator_span.set_attribute("validator.validate.input.value", value_arg)
|
|
validator_span.set_attribute("validator.validate.input.metadata", metadata_arg)
|
|
for k, v in kwargs.items():
|
|
if v is not None:
|
|
validator_span.set_attribute(
|
|
f"validator.validate.input.{k}", serialize(v) or ""
|
|
)
|
|
trace_operation(
|
|
input_value={"value": value_arg, "metadata": metadata_arg},
|
|
input_mime_type="application/json",
|
|
)
|
|
|
|
if result is not None:
|
|
output = result.model_dump()
|
|
trace_operation(
|
|
output_value=output,
|
|
output_mime_type="application/json",
|
|
)
|
|
for k, v in output.items():
|
|
if v is not None:
|
|
validator_span.set_attribute(
|
|
f"validator.validate.output.{k}", serialize(v) or ""
|
|
)
|
|
|
|
|
|
def trace_validator(
|
|
validator_name: str,
|
|
obj_id: int,
|
|
on_fail_descriptor: Optional[str] = None,
|
|
tracer: Optional[Tracer] = None,
|
|
*,
|
|
validation_session_id: str,
|
|
**init_kwargs,
|
|
):
|
|
def trace_validator_decorator(fn: Callable[..., Optional[ValidationResult]]):
|
|
@wraps(fn)
|
|
def trace_validator_wrapper(*args, **kwargs):
|
|
if not settings.disable_tracing:
|
|
current_otel_context = context.get_current()
|
|
_tracer = trace.get_tracer("guardrails-ai", GUARDRAILS_VERSION)
|
|
validator_span_name = f"{validator_name}.validate"
|
|
with _tracer.start_as_current_span(
|
|
name=validator_span_name, # type: ignore
|
|
context=current_otel_context, # type: ignore
|
|
) as validator_span:
|
|
if SpanAttributes is not None:
|
|
validator_span.set_attribute(
|
|
SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
|
|
)
|
|
|
|
try:
|
|
resp = fn(*args, **kwargs)
|
|
add_user_attributes(validator_span)
|
|
add_validator_attributes(
|
|
*args,
|
|
validator_span=validator_span,
|
|
validator_name=validator_name,
|
|
obj_id=obj_id,
|
|
on_fail_descriptor=on_fail_descriptor,
|
|
result=resp,
|
|
init_kwargs=init_kwargs,
|
|
validation_session_id=validation_session_id,
|
|
**kwargs,
|
|
)
|
|
return resp
|
|
except Exception as e:
|
|
validator_span.set_status(
|
|
status=StatusCode.ERROR, description=str(e)
|
|
)
|
|
add_user_attributes(validator_span)
|
|
add_validator_attributes(
|
|
*args,
|
|
validator_span=validator_span,
|
|
validator_name=validator_name,
|
|
obj_id=obj_id,
|
|
on_fail_descriptor=on_fail_descriptor,
|
|
result=None,
|
|
init_kwargs=init_kwargs,
|
|
validation_session_id=validation_session_id,
|
|
**kwargs,
|
|
)
|
|
raise e
|
|
else:
|
|
return fn(*args, **kwargs)
|
|
|
|
return trace_validator_wrapper
|
|
|
|
return trace_validator_decorator
|
|
|
|
|
|
def trace_async_validator(
|
|
validator_name: str,
|
|
obj_id: int,
|
|
on_fail_descriptor: Optional[str] = None,
|
|
tracer: Optional[Tracer] = None,
|
|
*,
|
|
validation_session_id: str,
|
|
**init_kwargs,
|
|
):
|
|
def trace_validator_decorator(
|
|
fn: Callable[..., Awaitable[Optional[ValidationResult]]],
|
|
):
|
|
@wraps(fn)
|
|
async def trace_validator_wrapper(*args, **kwargs):
|
|
if not settings.disable_tracing:
|
|
current_otel_context = context.get_current()
|
|
_tracer = trace.get_tracer("guardrails-ai", GUARDRAILS_VERSION)
|
|
validator_span_name = f"{validator_name}.validate"
|
|
with _tracer.start_as_current_span(
|
|
name=validator_span_name, # type: ignore
|
|
context=current_otel_context, # type: ignore
|
|
) as validator_span:
|
|
if SpanAttributes is not None:
|
|
validator_span.set_attribute(
|
|
SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
|
|
) # see here for a list of span kinds: https://github.com/Arize-ai/openinference/blob/main/python/openinference-semantic-conventions/src/openinference/semconv/trace/__init__.py#L271
|
|
|
|
try:
|
|
resp = await fn(*args, **kwargs)
|
|
add_user_attributes(validator_span)
|
|
add_validator_attributes(
|
|
*args,
|
|
validator_span=validator_span,
|
|
validator_name=validator_name,
|
|
obj_id=obj_id,
|
|
on_fail_descriptor=on_fail_descriptor,
|
|
result=resp,
|
|
init_kwargs=init_kwargs,
|
|
validation_session_id=validation_session_id,
|
|
**kwargs,
|
|
)
|
|
return resp
|
|
except Exception as e:
|
|
validator_span.set_status(
|
|
status=StatusCode.ERROR, description=str(e)
|
|
)
|
|
add_user_attributes(validator_span)
|
|
add_validator_attributes(
|
|
*args,
|
|
validator_span=validator_span,
|
|
validator_name=validator_name,
|
|
obj_id=obj_id,
|
|
on_fail_descriptor=on_fail_descriptor,
|
|
result=None,
|
|
init_kwargs=init_kwargs,
|
|
validation_session_id=validation_session_id,
|
|
**kwargs,
|
|
)
|
|
raise e
|
|
else:
|
|
return await fn(*args, **kwargs)
|
|
|
|
return trace_validator_wrapper
|
|
|
|
return trace_validator_decorator
|