Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions guardrails/telemetry/common.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import json
from typing import Any, Callable, Dict, Optional, Union
from opentelemetry.baggage import get_baggage
from opentelemetry import context
from opentelemetry.context import Context
from opentelemetry.trace import Tracer, Span
Expand Down Expand Up @@ -103,3 +104,23 @@ def wrapped_func(*args: Any, **kwargs: Any) -> Any:
context.detach(token)

return wrapped_func


def add_user_attributes(span: Span):
try:
client_ip = get_baggage("client.ip") or "unknown"
user_agent = get_baggage("http.user_agent") or "unknown"
referrer = get_baggage("http.referrer") or "unknown"
user_id = get_baggage("user.id") or "unknown"
organization = get_baggage("organization") or "unknown"
app = get_baggage("app") or "unknown"

span.set_attribute("client.ip", str(client_ip))
span.set_attribute("http.user_agent", str(user_agent))
span.set_attribute("http.referrer", str(referrer))
span.set_attribute("user.id", str(user_id))
span.set_attribute("organization", str(organization))
span.set_attribute("app", str(app))
except Exception as e:
logger.warning("Error loading baggage user information", e)
pass
6 changes: 6 additions & 0 deletions guardrails/telemetry/guard_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
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


Expand Down Expand Up @@ -141,6 +142,7 @@ def trace_stream_guard(
# FIXME: This should only be called once;
# Accumulate the validated output and call at the end
add_guard_attributes(guard_span, history, res)
add_user_attributes(guard_span)
yield res
except StopIteration:
next_exists = False
Expand Down Expand Up @@ -175,6 +177,7 @@ def trace_guard_execution(
):
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))
Expand All @@ -193,6 +196,7 @@ async def trace_async_stream_guard(
try:
res = await anext(result) # type: ignore
add_guard_attributes(guard_span, history, res)
add_user_attributes(guard_span)
yield res
except StopIteration:
next_exists = False
Expand Down Expand Up @@ -244,9 +248,11 @@ async def trace_async_guard_execution(
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)
3 changes: 1 addition & 2 deletions guardrails/telemetry/open_inference.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
from typing import Any, Dict, List, Optional

from guardrails.telemetry.common import get_span, to_dict
from guardrails.utils.serialization_utils import serialize
from guardrails.telemetry.common import get_span, to_dict, serialize


def trace_operation(
Expand Down
8 changes: 6 additions & 2 deletions guardrails/telemetry/runner_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,8 @@
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 get_tracer
from guardrails.telemetry.common import get_tracer, add_user_attributes, serialize
from guardrails.utils.safe_get import safe_get
from guardrails.utils.serialization_utils import serialize
from guardrails.version import GUARDRAILS_VERSION


Expand Down Expand Up @@ -73,10 +72,12 @@ def trace_step_wrapper(*args, **kwargs) -> Iteration:
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)
Expand Down Expand Up @@ -112,6 +113,7 @@ def trace_stream_step_generator(
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

Expand Down Expand Up @@ -145,10 +147,12 @@ async def trace_async_step_wrapper(*args, **kwargs) -> Iteration:
) as step_span:
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

Expand Down
7 changes: 5 additions & 2 deletions guardrails/telemetry/validator_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,10 @@

from guardrails.settings import settings
from guardrails.classes.validation.validation_result import ValidationResult
from guardrails.telemetry.common import get_tracer
from guardrails.telemetry.common import get_tracer, 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.utils.serialization_utils import serialize
from guardrails.version import GUARDRAILS_VERSION


Expand Down Expand Up @@ -106,6 +105,7 @@ def trace_validator_wrapper(*args, **kwargs):
) as validator_span:
try:
resp = fn(*args, **kwargs)
add_user_attributes(validator_span)
add_validator_attributes(
*args,
validator_span=validator_span,
Expand All @@ -122,6 +122,7 @@ def trace_validator_wrapper(*args, **kwargs):
validator_span.set_status(
status=StatusCode.ERROR, description=str(e)
)
add_user_attributes(validator_span)
add_validator_attributes(
*args,
validator_span=validator_span,
Expand Down Expand Up @@ -168,6 +169,7 @@ async def trace_validator_wrapper(*args, **kwargs):
) as validator_span:
try:
resp = await fn(*args, **kwargs)
add_user_attributes(validator_span)
add_validator_attributes(
*args,
validator_span=validator_span,
Expand All @@ -184,6 +186,7 @@ async def trace_validator_wrapper(*args, **kwargs):
validator_span.set_status(
status=StatusCode.ERROR, description=str(e)
)
add_user_attributes(validator_span)
add_validator_attributes(
*args,
validator_span=validator_span,
Expand Down
Loading