338 lines
12 KiB
Python
338 lines
12 KiB
Python
|
import json
|
||
|
|
from functools import wraps
|
||
|
|
from typing import (
|
||
|
|
AsyncIterator,
|
||
|
|
Awaitable,
|
||
|
|
Callable,
|
||
|
|
Iterator,
|
||
|
|
Optional,
|
||
|
|
)
|
||
|
|
|
||
|
|
try:
|
||
|
|
from openinference.semconv.trace import SpanAttributes # type: ignore
|
||
|
|
except ImportError:
|
||
|
|
SpanAttributes = None
|
||
|
|
|
||
|
|
from opentelemetry import context, trace
|
||
|
|
from opentelemetry.trace import StatusCode, Span
|
||
|
|
|
||
|
|
from guardrails.classes.history.iteration import Iteration
|
||
|
|
from guardrails.classes.llm.llm_response import LLMResponse
|
||
|
|
from guardrails.settings import settings
|
||
|
|
from guardrails.classes.output_type import OT
|
||
|
|
from guardrails.classes.validation_outcome import ValidationOutcome
|
||
|
|
from guardrails.stores.context import get_guard_name
|
||
|
|
from guardrails.telemetry.common import (
|
||
|
|
add_user_attributes,
|
||
|
|
serialize,
|
||
|
|
recursive_key_operation,
|
||
|
|
redact,
|
||
|
|
)
|
||
|
|
from guardrails.utils.safe_get import safe_get
|
||
|
|
from guardrails.version import GUARDRAILS_VERSION
|
||
|
|
|
||
|
|
import sys
|
||
|
|
|
||
|
|
if sys.version_info.minor < 10:
|
||
|
|
from guardrails.utils.polyfills import anext
|
||
|
|
|
||
|
|
#########################################
|
||
|
|
### START Runner.step Instrumentation ###
|
||
|
|
#########################################
|
||
|
|
|
||
|
|
|
||
|
|
# TODO: Track input arguments and outputs explicitly as named attributes
|
||
|
|
def add_step_attributes(
|
||
|
|
step_span: Span, response: Optional[Iteration], *args, **kwargs
|
||
|
|
):
|
||
|
|
step_number = safe_get(args, 1, kwargs.get("index", 0))
|
||
|
|
guard_name = get_guard_name()
|
||
|
|
|
||
|
|
step_span.set_attribute("guardrails.version", GUARDRAILS_VERSION)
|
||
|
|
step_span.set_attribute("type", "guardrails/guard/step")
|
||
|
|
step_span.set_attribute("guard.name", guard_name)
|
||
|
|
step_span.set_attribute("step.index", step_number)
|
||
|
|
|
||
|
|
ser_args = [serialize(arg) for arg in args]
|
||
|
|
ser_kwargs = {k: serialize(v) for k, v in kwargs.items()}
|
||
|
|
|
||
|
|
inputs = {
|
||
|
|
"args": [sarg for sarg in ser_args if sarg is not None],
|
||
|
|
"kwargs": {k: v for k, v in ser_kwargs.items() if v is not None},
|
||
|
|
}
|
||
|
|
for k in inputs:
|
||
|
|
inputs[k] = recursive_key_operation(inputs[k], redact)
|
||
|
|
|
||
|
|
step_span.set_attribute("input.mime_type", "application/json")
|
||
|
|
step_span.set_attribute("input.value", json.dumps(inputs))
|
||
|
|
|
||
|
|
ser_output = serialize(response)
|
||
|
|
if ser_output:
|
||
|
|
step_span.set_attribute("output.mime_type", "application/json")
|
||
|
|
step_span.set_attribute(
|
||
|
|
"output.value",
|
||
|
|
(json.dumps(ser_output) if isinstance(ser_output, dict) else ser_output),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def trace_step(fn: Callable[..., Iteration]):
|
||
|
|
@wraps(fn)
|
||
|
|
def trace_step_wrapper(*args, **kwargs) -> Iteration:
|
||
|
|
if not settings.disable_tracing:
|
||
|
|
current_otel_context = context.get_current()
|
||
|
|
tracer = trace.get_tracer("guardrails-ai", GUARDRAILS_VERSION)
|
||
|
|
|
||
|
|
with tracer.start_as_current_span(
|
||
|
|
name="step", # type: ignore
|
||
|
|
context=current_otel_context, # type: ignore
|
||
|
|
) as step_span:
|
||
|
|
if SpanAttributes is not None:
|
||
|
|
step_span.set_attribute(
|
||
|
|
SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
response = fn(*args, **kwargs)
|
||
|
|
add_step_attributes(step_span, response, *args, **kwargs)
|
||
|
|
add_user_attributes(step_span)
|
||
|
|
return response
|
||
|
|
except Exception as e:
|
||
|
|
step_span.set_status(status=StatusCode.ERROR, description=str(e))
|
||
|
|
add_step_attributes(step_span, None, *args, **kwargs)
|
||
|
|
add_user_attributes(step_span)
|
||
|
|
raise e
|
||
|
|
else:
|
||
|
|
return fn(*args, **kwargs)
|
||
|
|
|
||
|
|
return trace_step_wrapper
|
||
|
|
|
||
|
|
|
||
|
|
def trace_stream_step_generator(
|
||
|
|
fn: Callable[..., Iterator[ValidationOutcome[OT]]], *args, **kwargs
|
||
|
|
) -> Iterator[ValidationOutcome[OT]]:
|
||
|
|
current_otel_context = context.get_current()
|
||
|
|
tracer = trace.get_tracer("guardrails-ai", GUARDRAILS_VERSION)
|
||
|
|
|
||
|
|
exception = None
|
||
|
|
with tracer.start_as_current_span(
|
||
|
|
name="step", # type: ignore
|
||
|
|
context=current_otel_context, # type: ignore
|
||
|
|
) as step_span:
|
||
|
|
if SpanAttributes is not None:
|
||
|
|
step_span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL")
|
||
|
|
try:
|
||
|
|
gen = fn(*args, **kwargs)
|
||
|
|
next_exists = True
|
||
|
|
while next_exists:
|
||
|
|
try:
|
||
|
|
res = next(gen)
|
||
|
|
yield res
|
||
|
|
except StopIteration:
|
||
|
|
next_exists = False
|
||
|
|
except Exception as e:
|
||
|
|
step_span.set_status(status=StatusCode.ERROR, description=str(e))
|
||
|
|
exception = e
|
||
|
|
finally:
|
||
|
|
call = safe_get(args, 8, kwargs.get("call_log", None))
|
||
|
|
iteration = call.iterations.last if call else None
|
||
|
|
add_step_attributes(step_span, iteration, *args, **kwargs)
|
||
|
|
add_user_attributes(step_span)
|
||
|
|
if exception:
|
||
|
|
raise exception
|
||
|
|
|
||
|
|
|
||
|
|
def trace_stream_step(
|
||
|
|
fn: Callable[..., Iterator[ValidationOutcome[OT]]],
|
||
|
|
) -> Callable[..., Iterator[ValidationOutcome[OT]]]:
|
||
|
|
@wraps(fn)
|
||
|
|
def trace_stream_step_wrapper(*args, **kwargs) -> Iterator[ValidationOutcome[OT]]:
|
||
|
|
if not settings.disable_tracing:
|
||
|
|
return trace_stream_step_generator(fn, *args, **kwargs)
|
||
|
|
else:
|
||
|
|
return fn(*args, **kwargs)
|
||
|
|
|
||
|
|
return trace_stream_step_wrapper
|
||
|
|
|
||
|
|
|
||
|
|
def trace_async_step(fn: Callable[..., Awaitable[Iteration]]):
|
||
|
|
@wraps(fn)
|
||
|
|
async def trace_async_step_wrapper(*args, **kwargs) -> Iteration:
|
||
|
|
if not settings.disable_tracing:
|
||
|
|
current_otel_context = context.get_current()
|
||
|
|
tracer = trace.get_tracer("guardrails-ai", GUARDRAILS_VERSION)
|
||
|
|
|
||
|
|
with tracer.start_as_current_span(
|
||
|
|
name="step", # type: ignore
|
||
|
|
context=current_otel_context, # type: ignore
|
||
|
|
) as step_span:
|
||
|
|
if SpanAttributes is not None:
|
||
|
|
step_span.set_attribute(
|
||
|
|
SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
response = await fn(*args, **kwargs)
|
||
|
|
add_user_attributes(step_span)
|
||
|
|
add_step_attributes(step_span, response, *args, **kwargs)
|
||
|
|
|
||
|
|
return response
|
||
|
|
except Exception as e:
|
||
|
|
step_span.set_status(status=StatusCode.ERROR, description=str(e))
|
||
|
|
add_user_attributes(step_span)
|
||
|
|
add_step_attributes(step_span, None, *args, **kwargs)
|
||
|
|
raise e
|
||
|
|
|
||
|
|
else:
|
||
|
|
return await fn(*args, **kwargs)
|
||
|
|
|
||
|
|
return trace_async_step_wrapper
|
||
|
|
|
||
|
|
|
||
|
|
async def trace_async_stream_step_generator(
|
||
|
|
fn: Callable[..., AsyncIterator[ValidationOutcome[OT]]], *args, **kwargs
|
||
|
|
) -> AsyncIterator[ValidationOutcome[OT]]:
|
||
|
|
current_otel_context = context.get_current()
|
||
|
|
tracer = trace.get_tracer("guardrails-ai", GUARDRAILS_VERSION)
|
||
|
|
|
||
|
|
exception = None
|
||
|
|
with tracer.start_as_current_span(
|
||
|
|
name="step", # type: ignore
|
||
|
|
context=current_otel_context, # type: ignore
|
||
|
|
) as step_span:
|
||
|
|
if SpanAttributes is not None:
|
||
|
|
step_span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL")
|
||
|
|
try:
|
||
|
|
gen = fn(*args, **kwargs)
|
||
|
|
next_exists = True
|
||
|
|
while next_exists:
|
||
|
|
try:
|
||
|
|
res = await anext(gen) # type: ignore
|
||
|
|
yield res
|
||
|
|
except StopIteration:
|
||
|
|
next_exists = False
|
||
|
|
except StopAsyncIteration:
|
||
|
|
next_exists = False
|
||
|
|
except Exception as e:
|
||
|
|
step_span.set_status(status=StatusCode.ERROR, description=str(e))
|
||
|
|
exception = e
|
||
|
|
finally:
|
||
|
|
call = safe_get(args, 3, kwargs.get("call_log", None))
|
||
|
|
iteration = call.iterations.last if call else None
|
||
|
|
add_step_attributes(step_span, iteration, *args, **kwargs)
|
||
|
|
if exception:
|
||
|
|
raise exception
|
||
|
|
|
||
|
|
|
||
|
|
def trace_async_stream_step(
|
||
|
|
fn: Callable[..., AsyncIterator[ValidationOutcome[OT]]],
|
||
|
|
):
|
||
|
|
@wraps(fn)
|
||
|
|
async def trace_async_stream_step_wrapper(
|
||
|
|
*args, **kwargs
|
||
|
|
) -> AsyncIterator[ValidationOutcome[OT]]:
|
||
|
|
if not settings.disable_tracing:
|
||
|
|
return trace_async_stream_step_generator(fn, *args, **kwargs)
|
||
|
|
else:
|
||
|
|
return fn(*args, **kwargs)
|
||
|
|
|
||
|
|
return trace_async_stream_step_wrapper
|
||
|
|
|
||
|
|
|
||
|
|
#########################################
|
||
|
|
#### END Runner.step Instrumentation ####
|
||
|
|
#########################################
|
||
|
|
|
||
|
|
|
||
|
|
#########################################
|
||
|
|
### START Runner.call Instrumentation ###
|
||
|
|
#########################################
|
||
|
|
|
||
|
|
|
||
|
|
# TODO: Track input arguments and outputs explicitly as named attributes
|
||
|
|
def add_call_attributes(
|
||
|
|
call_span: Span, response: Optional[LLMResponse], *args, **kwargs
|
||
|
|
):
|
||
|
|
guard_name = get_guard_name()
|
||
|
|
|
||
|
|
call_span.set_attribute("guardrails.version", GUARDRAILS_VERSION)
|
||
|
|
call_span.set_attribute("type", "guardrails/guard/step/call")
|
||
|
|
call_span.set_attribute("guard.name", guard_name)
|
||
|
|
|
||
|
|
ser_args = [serialize(arg) for arg in args]
|
||
|
|
ser_kwargs = {k: serialize(v) for k, v in kwargs.items()}
|
||
|
|
inputs = {
|
||
|
|
"args": [sarg for sarg in ser_args if sarg is not None],
|
||
|
|
"kwargs": {k: v for k, v in ser_kwargs.items() if v is not None},
|
||
|
|
}
|
||
|
|
for k in inputs:
|
||
|
|
inputs[k] = recursive_key_operation(inputs[k], redact)
|
||
|
|
call_span.set_attribute("input.mime_type", "application/json")
|
||
|
|
call_span.set_attribute("input.value", json.dumps(inputs))
|
||
|
|
|
||
|
|
ser_output = serialize(response)
|
||
|
|
if ser_output:
|
||
|
|
call_span.set_attribute("output.mime_type", "application/json")
|
||
|
|
call_span.set_attribute(
|
||
|
|
"output.value",
|
||
|
|
(json.dumps(ser_output) if isinstance(ser_output, dict) else ser_output),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def trace_call(fn: Callable[..., LLMResponse]):
|
||
|
|
@wraps(fn)
|
||
|
|
def trace_call_wrapper(*args, **kwargs):
|
||
|
|
if not settings.disable_tracing:
|
||
|
|
current_otel_context = context.get_current()
|
||
|
|
tracer = trace.get_tracer("guardrails-ai", GUARDRAILS_VERSION)
|
||
|
|
|
||
|
|
with tracer.start_as_current_span(
|
||
|
|
name="call", # type: ignore
|
||
|
|
context=current_otel_context, # type: ignore
|
||
|
|
) as call_span:
|
||
|
|
try:
|
||
|
|
response = fn(*args, **kwargs)
|
||
|
|
if isinstance(response, LLMResponse) and (
|
||
|
|
response.async_stream_output or response.stream_output
|
||
|
|
):
|
||
|
|
# TODO: Iterate, add a call attr each time
|
||
|
|
return response
|
||
|
|
add_call_attributes(call_span, response, *args, **kwargs)
|
||
|
|
return response
|
||
|
|
except Exception as e:
|
||
|
|
call_span.set_status(status=StatusCode.ERROR, description=str(e))
|
||
|
|
add_call_attributes(call_span, None, *args, **kwargs)
|
||
|
|
raise e
|
||
|
|
else:
|
||
|
|
return fn(*args, **kwargs)
|
||
|
|
|
||
|
|
return trace_call_wrapper
|
||
|
|
|
||
|
|
|
||
|
|
def trace_async_call(fn: Callable[..., Awaitable[LLMResponse]]):
|
||
|
|
@wraps(fn)
|
||
|
|
async def trace_async_call_wrapper(*args, **kwargs):
|
||
|
|
if not settings.disable_tracing:
|
||
|
|
current_otel_context = context.get_current()
|
||
|
|
tracer = trace.get_tracer("guardrails-ai", GUARDRAILS_VERSION)
|
||
|
|
|
||
|
|
with tracer.start_as_current_span(
|
||
|
|
name="call", # type: ignore
|
||
|
|
context=current_otel_context, # type: ignore
|
||
|
|
) as call_span:
|
||
|
|
try:
|
||
|
|
response = await fn(*args, **kwargs)
|
||
|
|
add_call_attributes(call_span, response, *args, **kwargs)
|
||
|
|
return response
|
||
|
|
except Exception as e:
|
||
|
|
call_span.set_status(status=StatusCode.ERROR, description=str(e))
|
||
|
|
add_call_attributes(call_span, None, *args, **kwargs)
|
||
|
|
raise e
|
||
|
|
|
||
|
|
else:
|
||
|
|
return await fn(*args, **kwargs)
|
||
|
|
|
||
|
|
return trace_async_call_wrapper
|
||
|
|
|
||
|
|
|
||
|
|
#########################################
|
||
|
|
#### END Runner.call Instrumentation ####
|
||
|
|
#########################################
|