257 lines
8.1 KiB
Python
257 lines
8.1 KiB
Python
from functools import wraps
|
|
from typing import (
|
|
Any,
|
|
Dict,
|
|
Optional,
|
|
AsyncGenerator,
|
|
)
|
|
|
|
from opentelemetry.trace import Span
|
|
from opentelemetry.trace.propagation import set_span_in_context
|
|
|
|
from guardrails_ai.types import ValidationResult
|
|
from guardrails.hub_token.token import VALIDATOR_HUB_SERVICE
|
|
from guardrails.types.primitives import PrimitiveTypes
|
|
from guardrails.utils.safe_get import safe_get
|
|
from guardrails.utils.hub_telemetry_utils import HubTelemetry
|
|
|
|
|
|
def get_guard_call_attributes(
|
|
attrs: Dict[str, Any], origin: str, *args, **kwargs
|
|
) -> Dict[str, Any]:
|
|
attrs["stream"] = kwargs.get("stream", False)
|
|
|
|
guard_self = safe_get(args, 0)
|
|
if guard_self is not None:
|
|
attrs["guard_id"] = guard_self.id
|
|
attrs["user_id"] = guard_self._user_id
|
|
attrs["custom_reask_messages"] = (
|
|
guard_self._exec_opts.reask_messages is not None
|
|
)
|
|
attrs["output_type"] = (
|
|
"unstructured"
|
|
if PrimitiveTypes.is_primitive(guard_self.output_schema.type)
|
|
else "structured"
|
|
)
|
|
return attrs
|
|
|
|
llm_api_str = "" # noqa
|
|
llm_api = kwargs.get("llm_api")
|
|
if origin in ["Guard.__call__", "AsyncGuard.__call__"]:
|
|
llm_api = safe_get(args, 1, llm_api)
|
|
|
|
if llm_api:
|
|
llm_api_module_name = (
|
|
llm_api.__module__ if hasattr(llm_api, "__module__") else ""
|
|
)
|
|
llm_api_name = (
|
|
llm_api.__name__ if hasattr(llm_api, "__name__") else type(llm_api).__name__
|
|
)
|
|
llm_api_str = f"{llm_api_module_name}.{llm_api_name}"
|
|
attrs["llm_api"] = llm_api_str if llm_api_str else "None"
|
|
|
|
return attrs
|
|
|
|
|
|
def get_validator_inference_attributes(
|
|
attrs: Dict[str, Any], *args, **kwargs
|
|
) -> Dict[str, Any]:
|
|
validator_self = safe_get(args, 0)
|
|
if validator_self is not None:
|
|
used_guardrails_endpoint = (
|
|
VALIDATOR_HUB_SERVICE in validator_self.validation_endpoint
|
|
and not validator_self.use_local
|
|
)
|
|
used_custom_endpoint = (
|
|
not validator_self.use_local and not used_guardrails_endpoint
|
|
)
|
|
attrs["validator_name"] = validator_self.rail_alias
|
|
attrs["used_remote_inference"] = not validator_self.use_local
|
|
attrs["used_local_inference"] = validator_self.use_local
|
|
attrs["used_guardrails_endpoint"] = used_guardrails_endpoint
|
|
attrs["used_custom_endpoint"] = used_custom_endpoint
|
|
return attrs
|
|
|
|
|
|
def get_validator_usage_attributes(
|
|
attrs: Dict[str, Any], response, *args, **kwargs
|
|
) -> Dict[str, Any]:
|
|
# We're wrapping a wrapped function,
|
|
# so the first arg is the validator service
|
|
validator_self = safe_get(args, 1)
|
|
if validator_self is not None:
|
|
attrs["validator_name"] = validator_self.rail_alias
|
|
attrs["validator_on_fail"] = validator_self.on_fail_descriptor
|
|
|
|
if response is not None:
|
|
attrs["validator_result"] = (
|
|
response.outcome if isinstance(response, ValidationResult) else None
|
|
)
|
|
|
|
return attrs
|
|
|
|
|
|
def add_attributes(
|
|
span: Span,
|
|
attrs: Dict[str, Any],
|
|
name: str,
|
|
origin: str,
|
|
*args,
|
|
response=None,
|
|
**kwargs,
|
|
):
|
|
attrs["origin"] = origin
|
|
if name == "/guard_call":
|
|
attrs = get_guard_call_attributes(attrs, origin, *args, **kwargs)
|
|
elif name == "/reasks":
|
|
if response is not None and hasattr(response, "iterations"):
|
|
attrs["reask_count"] = len(response.iterations) - 1
|
|
else:
|
|
attrs["reask_count"] = 0
|
|
elif name == "/validator_inference":
|
|
attrs = get_validator_inference_attributes(attrs, *args, **kwargs)
|
|
elif name == "/validator_usage":
|
|
attrs = get_validator_usage_attributes(attrs, response, *args, **kwargs)
|
|
|
|
for key, value in attrs.items():
|
|
if value is not None:
|
|
span.set_attribute(key, value)
|
|
|
|
|
|
def trace(
|
|
*,
|
|
name: str,
|
|
origin: Optional[str] = None,
|
|
**attrs,
|
|
):
|
|
def decorator(fn):
|
|
@wraps(fn)
|
|
def wrapper(*args, **kwargs):
|
|
hub_telemetry = HubTelemetry()
|
|
if hub_telemetry._enabled and hub_telemetry._tracer is not None:
|
|
with hub_telemetry._tracer.start_span(
|
|
name,
|
|
context=hub_telemetry.extract_current_context(),
|
|
set_status_on_exception=True,
|
|
) as span: # noqa
|
|
context = set_span_in_context(span)
|
|
hub_telemetry.inject_current_context(context=context)
|
|
nonlocal origin
|
|
origin = origin if origin is not None else name
|
|
|
|
resp = fn(*args, **kwargs)
|
|
add_attributes(
|
|
span, attrs, name, origin, *args, response=resp, **kwargs
|
|
)
|
|
return resp
|
|
else:
|
|
return fn(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
def async_trace(
|
|
*,
|
|
name: str,
|
|
origin: Optional[str] = None,
|
|
):
|
|
def decorator(fn):
|
|
@wraps(fn)
|
|
async def async_wrapper(*args, **kwargs):
|
|
hub_telemetry = HubTelemetry()
|
|
if hub_telemetry._enabled and hub_telemetry._tracer is not None:
|
|
with hub_telemetry._tracer.start_span(
|
|
name,
|
|
context=hub_telemetry.extract_current_context(),
|
|
set_status_on_exception=True,
|
|
) as span: # noqa
|
|
context = set_span_in_context(span)
|
|
hub_telemetry.inject_current_context(context=context)
|
|
|
|
nonlocal origin
|
|
origin = origin if origin is not None else name
|
|
add_attributes(span, {"async": True}, name, origin, *args, **kwargs)
|
|
return await fn(*args, **kwargs)
|
|
else:
|
|
return await fn(*args, **kwargs)
|
|
|
|
return async_wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
def _run_gen(fn, *args, **kwargs):
|
|
gen = fn(*args, **kwargs)
|
|
for item in gen:
|
|
yield item
|
|
|
|
|
|
def trace_stream(
|
|
*,
|
|
name: str,
|
|
origin: Optional[str] = None,
|
|
**attrs,
|
|
):
|
|
def decorator(fn):
|
|
@wraps(fn)
|
|
def wrapper(*args, **kwargs):
|
|
hub_telemetry = HubTelemetry()
|
|
if hub_telemetry._enabled and hub_telemetry._tracer is not None:
|
|
with hub_telemetry._tracer.start_span(
|
|
name,
|
|
context=hub_telemetry.extract_current_context(),
|
|
set_status_on_exception=True,
|
|
) as span: # noqa
|
|
context = set_span_in_context(span)
|
|
hub_telemetry.inject_current_context(context=context)
|
|
|
|
nonlocal origin
|
|
origin = origin if origin is not None else name
|
|
add_attributes(span, attrs, name, origin, *args, **kwargs)
|
|
return _run_gen(fn, *args, **kwargs)
|
|
else:
|
|
return fn(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
async def _run_async_gen(fn, *args, **kwargs) -> AsyncGenerator[Any, None]:
|
|
gen = fn(*args, **kwargs)
|
|
async for item in gen:
|
|
yield item
|
|
|
|
|
|
def async_trace_stream(
|
|
*,
|
|
name: str,
|
|
origin: Optional[str] = None,
|
|
**attrs,
|
|
):
|
|
def decorator(fn):
|
|
@wraps(fn)
|
|
def wrapper(*args, **kwargs):
|
|
hub_telemetry = HubTelemetry()
|
|
if hub_telemetry._enabled and hub_telemetry._tracer is not None:
|
|
with hub_telemetry._tracer.start_span(
|
|
name,
|
|
context=hub_telemetry.extract_current_context(),
|
|
set_status_on_exception=True,
|
|
) as span: # noqa
|
|
context = set_span_in_context(span)
|
|
hub_telemetry.inject_current_context(context=context)
|
|
|
|
nonlocal origin
|
|
origin = origin if origin is not None else name
|
|
add_attributes(span, attrs, name, origin, *args, **kwargs)
|
|
return fn(*args, **kwargs)
|
|
else:
|
|
return fn(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|