295 lines
10 KiB
Python
295 lines
10 KiB
Python
|
import inspect
|
||
|
|
from typing import (
|
||
|
|
Any,
|
||
|
|
AsyncIterator,
|
||
|
|
Awaitable,
|
||
|
|
Callable,
|
||
|
|
Coroutine,
|
||
|
|
Iterator,
|
||
|
|
Union,
|
||
|
|
)
|
||
|
|
|
||
|
|
try:
|
||
|
|
from openinference.semconv.trace import SpanAttributes # type: ignore
|
||
|
|
except ImportError:
|
||
|
|
SpanAttributes = None
|
||
|
|
from opentelemetry import context, trace
|
||
|
|
from opentelemetry.trace import StatusCode, Span, Link, get_tracer
|
||
|
|
|
||
|
|
from guardrails.settings import settings
|
||
|
|
from guardrails.classes.generic.stack import Stack
|
||
|
|
from guardrails.classes.history.call import Call
|
||
|
|
from guardrails.classes.output_type import OT
|
||
|
|
from guardrails.classes.validation_outcome import ValidationOutcome
|
||
|
|
from guardrails.telemetry.open_inference import trace_operation
|
||
|
|
from guardrails.telemetry.common import add_user_attributes
|
||
|
|
from guardrails.version import GUARDRAILS_VERSION
|
||
|
|
|
||
|
|
import sys
|
||
|
|
|
||
|
|
if sys.version_info.minor < 10:
|
||
|
|
from guardrails.utils.polyfills import anext
|
||
|
|
|
||
|
|
# from sentence_transformers import SentenceTransformer
|
||
|
|
# import numpy as np
|
||
|
|
# from numpy.linalg import norm
|
||
|
|
|
||
|
|
# model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
|
||
|
|
|
||
|
|
|
||
|
|
def add_guard_attributes(
|
||
|
|
guard_span: Span,
|
||
|
|
history: Stack[Call],
|
||
|
|
resp: ValidationOutcome,
|
||
|
|
):
|
||
|
|
messages = []
|
||
|
|
if history.last and history.last.iterations.last:
|
||
|
|
messages = history.last.iterations.last.inputs.messages or []
|
||
|
|
|
||
|
|
system_messages = [msg for msg in messages if msg["role"] == "system"]
|
||
|
|
system_message = system_messages[-1] if system_messages else {}
|
||
|
|
|
||
|
|
user_messages = [msg for msg in messages if msg["role"] == "user"]
|
||
|
|
user_message = user_messages[-1] if user_messages else {}
|
||
|
|
|
||
|
|
input_value = f"""
|
||
|
|
{system_message}
|
||
|
|
{user_message}
|
||
|
|
"""
|
||
|
|
trace_operation(
|
||
|
|
input_mime_type="text/plain",
|
||
|
|
input_value=input_value,
|
||
|
|
output_mime_type="text/plain",
|
||
|
|
output_value=resp.validated_output,
|
||
|
|
)
|
||
|
|
guard_span.set_attribute("type", "guardrails/guard")
|
||
|
|
guard_span.set_attribute("validation_passed", resp.validation_passed or False)
|
||
|
|
|
||
|
|
execution_id = history.last.id if history.last else None
|
||
|
|
if execution_id is not None:
|
||
|
|
guard_span.set_attribute("execution_id", execution_id)
|
||
|
|
|
||
|
|
token_consumption = history.last.tokens_consumed if history.last else None
|
||
|
|
if token_consumption is not None:
|
||
|
|
guard_span.set_attribute("token_consumption", token_consumption)
|
||
|
|
|
||
|
|
number_of_reasks = (
|
||
|
|
history.last.iterations.last.index
|
||
|
|
if history.last and history.last.iterations.last
|
||
|
|
else None
|
||
|
|
)
|
||
|
|
if number_of_reasks is not None:
|
||
|
|
guard_span.set_attribute("number_of_reasks", number_of_reasks)
|
||
|
|
|
||
|
|
number_of_llm_calls = number_of_reasks + 1 if number_of_reasks is not None else None
|
||
|
|
if number_of_llm_calls is not None:
|
||
|
|
guard_span.set_attribute("number_of_llm_calls", number_of_llm_calls)
|
||
|
|
|
||
|
|
# # FIXME: Find a lighter weight library to do this.
|
||
|
|
# raw_embed = model.encode(resp.raw_llm_output)
|
||
|
|
# validated_embed = model.encode(resp.validated_output)
|
||
|
|
# input_embed = model.encode(input_value)
|
||
|
|
|
||
|
|
# # define two arrays
|
||
|
|
# raw_embed_np = np.array(raw_embed)
|
||
|
|
# validated_embed_np = np.array(validated_embed)
|
||
|
|
# input_embed_np = np.array(input_embed)
|
||
|
|
|
||
|
|
# # compute cosine similarity
|
||
|
|
# raw_output_x_validated_output_cosine = (
|
||
|
|
# np.sum(raw_embed_np*validated_embed_np, axis=0)
|
||
|
|
# /
|
||
|
|
# (
|
||
|
|
# norm(raw_embed_np, axis=0)*norm(validated_embed_np, axis=0)
|
||
|
|
# )
|
||
|
|
# )
|
||
|
|
|
||
|
|
# input_x_validated_output_cosine = (
|
||
|
|
# np.sum(input_embed_np*validated_embed_np, axis=0)
|
||
|
|
# /
|
||
|
|
# (
|
||
|
|
# norm(input_embed_np, axis=0)*norm(validated_embed_np, axis=0)
|
||
|
|
# )
|
||
|
|
# )
|
||
|
|
|
||
|
|
# input_x_raw_output_cosine = (
|
||
|
|
# np.sum(input_embed_np*raw_embed_np, axis=0)
|
||
|
|
# /
|
||
|
|
# (
|
||
|
|
# norm(input_embed_np, axis=0)*norm(raw_embed_np, axis=0)
|
||
|
|
# )
|
||
|
|
# )
|
||
|
|
|
||
|
|
# guard_span.set_attribute(
|
||
|
|
# "raw_output_x_validated_output_cosine",
|
||
|
|
# float(str(raw_output_x_validated_output_cosine))
|
||
|
|
# )
|
||
|
|
# guard_span.set_attribute(
|
||
|
|
# "input_x_validated_output_cosine",
|
||
|
|
# float(str(input_x_validated_output_cosine))
|
||
|
|
# )
|
||
|
|
# guard_span.set_attribute(
|
||
|
|
# "input_x_raw_output_cosine",
|
||
|
|
# float(str(input_x_raw_output_cosine))
|
||
|
|
# )
|
||
|
|
|
||
|
|
|
||
|
|
def trace_stream_guard(
|
||
|
|
guard_span: Span,
|
||
|
|
result: Iterator[ValidationOutcome[OT]],
|
||
|
|
history: Stack[Call],
|
||
|
|
) -> Iterator[ValidationOutcome[OT]]:
|
||
|
|
next_exists = True
|
||
|
|
while next_exists:
|
||
|
|
try:
|
||
|
|
res = next(result) # type: ignore
|
||
|
|
# FIXME: This should only be called once;
|
||
|
|
# Accumulate the validated output and call at the end
|
||
|
|
if not guard_span.is_recording():
|
||
|
|
# Assuming you have a tracer instance
|
||
|
|
tracer = get_tracer(__name__)
|
||
|
|
# Create a new span and link it to the previous span
|
||
|
|
with tracer.start_as_current_span(
|
||
|
|
"stream_guard_span", # type: ignore
|
||
|
|
links=[Link(guard_span.get_span_context())],
|
||
|
|
) as new_span:
|
||
|
|
guard_span = new_span
|
||
|
|
add_guard_attributes(guard_span, history, res)
|
||
|
|
add_user_attributes(guard_span)
|
||
|
|
if SpanAttributes is not None:
|
||
|
|
new_span.set_attribute(
|
||
|
|
SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
|
||
|
|
)
|
||
|
|
yield res
|
||
|
|
except StopIteration:
|
||
|
|
next_exists = False
|
||
|
|
|
||
|
|
|
||
|
|
def trace_guard_execution(
|
||
|
|
guard_name: str,
|
||
|
|
history: Stack[Call],
|
||
|
|
_execute_fn: Callable[
|
||
|
|
..., Union[ValidationOutcome[OT], Iterator[ValidationOutcome[OT]]]
|
||
|
|
],
|
||
|
|
*args,
|
||
|
|
**kwargs,
|
||
|
|
) -> Union[ValidationOutcome[OT], Iterator[ValidationOutcome[OT]]]:
|
||
|
|
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="guard", # type: ignore
|
||
|
|
context=current_otel_context, # type: ignore
|
||
|
|
) as guard_span:
|
||
|
|
guard_span.set_attribute("guardrails.version", GUARDRAILS_VERSION)
|
||
|
|
guard_span.set_attribute("type", "guardrails/guard")
|
||
|
|
guard_span.set_attribute("guard.name", guard_name)
|
||
|
|
if SpanAttributes is not None:
|
||
|
|
guard_span.set_attribute(
|
||
|
|
SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
result = _execute_fn(*args, **kwargs)
|
||
|
|
if isinstance(result, Iterator) and not isinstance(
|
||
|
|
result, ValidationOutcome
|
||
|
|
):
|
||
|
|
return trace_stream_guard(guard_span, result, history)
|
||
|
|
|
||
|
|
add_guard_attributes(guard_span, history, result)
|
||
|
|
add_user_attributes(guard_span)
|
||
|
|
return result
|
||
|
|
except Exception as e:
|
||
|
|
guard_span.set_status(status=StatusCode.ERROR, description=str(e))
|
||
|
|
raise e
|
||
|
|
else:
|
||
|
|
return _execute_fn(*args, **kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
async def trace_async_stream_guard(
|
||
|
|
guard_span: Span,
|
||
|
|
result: AsyncIterator[ValidationOutcome[OT]],
|
||
|
|
history: Stack[Call],
|
||
|
|
) -> AsyncIterator[ValidationOutcome[OT]]:
|
||
|
|
next_exists = True
|
||
|
|
while next_exists:
|
||
|
|
try:
|
||
|
|
res = await anext(result) # type: ignore
|
||
|
|
if not guard_span.is_recording():
|
||
|
|
# Assuming you have a tracer instance
|
||
|
|
tracer = get_tracer(__name__)
|
||
|
|
# Create a new span and link it to the previous span
|
||
|
|
with tracer.start_as_current_span(
|
||
|
|
"async_stream_span", # type: ignore
|
||
|
|
links=[Link(guard_span.get_span_context())],
|
||
|
|
) as new_span:
|
||
|
|
guard_span = new_span
|
||
|
|
|
||
|
|
add_guard_attributes(guard_span, history, res)
|
||
|
|
add_user_attributes(guard_span)
|
||
|
|
if SpanAttributes is not None:
|
||
|
|
guard_span.set_attribute(
|
||
|
|
SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
|
||
|
|
)
|
||
|
|
yield res
|
||
|
|
except StopIteration:
|
||
|
|
next_exists = False
|
||
|
|
except StopAsyncIteration:
|
||
|
|
next_exists = False
|
||
|
|
|
||
|
|
|
||
|
|
async def trace_async_guard_execution(
|
||
|
|
guard_name: str,
|
||
|
|
history: Stack[Call],
|
||
|
|
_execute_fn: Callable[
|
||
|
|
...,
|
||
|
|
Coroutine[
|
||
|
|
Any,
|
||
|
|
Any,
|
||
|
|
Union[
|
||
|
|
ValidationOutcome[OT],
|
||
|
|
Awaitable[ValidationOutcome[OT]],
|
||
|
|
AsyncIterator[ValidationOutcome[OT]],
|
||
|
|
],
|
||
|
|
],
|
||
|
|
],
|
||
|
|
*args,
|
||
|
|
**kwargs,
|
||
|
|
) -> Union[
|
||
|
|
ValidationOutcome[OT],
|
||
|
|
Awaitable[ValidationOutcome[OT]],
|
||
|
|
AsyncIterator[ValidationOutcome[OT]],
|
||
|
|
]:
|
||
|
|
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="guard", # type: ignore
|
||
|
|
context=current_otel_context, # type: ignore
|
||
|
|
) as guard_span:
|
||
|
|
guard_span.set_attribute("guardrails.version", GUARDRAILS_VERSION)
|
||
|
|
guard_span.set_attribute("type", "guardrails/guard")
|
||
|
|
guard_span.set_attribute("guard.name", guard_name)
|
||
|
|
if SpanAttributes is not None:
|
||
|
|
guard_span.set_attribute(
|
||
|
|
SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
result = await _execute_fn(*args, **kwargs)
|
||
|
|
if isinstance(result, AsyncIterator):
|
||
|
|
return trace_async_stream_guard(guard_span, result, history)
|
||
|
|
|
||
|
|
res = result
|
||
|
|
if inspect.isawaitable(result):
|
||
|
|
res = await result
|
||
|
|
add_guard_attributes(guard_span, history, res) # type: ignore
|
||
|
|
add_user_attributes(guard_span)
|
||
|
|
return res
|
||
|
|
except Exception as e:
|
||
|
|
guard_span.set_status(status=StatusCode.ERROR, description=str(e))
|
||
|
|
add_user_attributes(guard_span)
|
||
|
|
raise e
|
||
|
|
else:
|
||
|
|
return await _execute_fn(*args, **kwargs)
|