221 lines
7.7 KiB
Python
221 lines
7.7 KiB
Python
import json
|
|
from pydantic import BaseModel
|
|
from typing import Any, Callable, Dict, Optional, Union, List
|
|
from opentelemetry.baggage import get_baggage
|
|
from opentelemetry import context
|
|
from opentelemetry.context import Context
|
|
from opentelemetry.trace import Span
|
|
|
|
from guardrails.logger import logger
|
|
|
|
|
|
def get_current_context() -> Union[Context, None]:
|
|
otel_current_context = (
|
|
context.get_current()
|
|
if context is not None and hasattr(context, "get_current")
|
|
else None
|
|
)
|
|
return otel_current_context
|
|
|
|
|
|
def get_span(span: Optional[Span] = None) -> Optional[Span]:
|
|
if span is not None and hasattr(span, "add_event"):
|
|
return span
|
|
try:
|
|
from opentelemetry import trace
|
|
|
|
current_context = get_current_context()
|
|
current_span = trace.get_current_span(current_context)
|
|
return current_span
|
|
except Exception as e:
|
|
logger.error(e)
|
|
return None
|
|
|
|
|
|
def serialize(val: Any) -> Optional[str]:
|
|
try:
|
|
if val is None:
|
|
return None
|
|
if hasattr(val, "to_dict"):
|
|
return json.dumps(val.to_dict())
|
|
elif isinstance(val, BaseModel):
|
|
return val.model_dump_json(exclude_none=True, by_alias=True)
|
|
elif hasattr(val, "__dict__"):
|
|
return json.dumps(val.__dict__)
|
|
elif isinstance(val, dict) or isinstance(val, list):
|
|
return json.dumps(val)
|
|
return str(val)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def to_dict(val: Any) -> Dict:
|
|
try:
|
|
if val is None:
|
|
return {}
|
|
elif isinstance(val, dict):
|
|
return val
|
|
elif hasattr(val, "to_dict"):
|
|
return val.to_dict()
|
|
elif hasattr(val, "__dict__"):
|
|
return val.__dict__
|
|
else:
|
|
return dict(val)
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def wrap_with_otel_context(
|
|
outer_scope_otel_context: Context, func: Callable[..., Any]
|
|
) -> Callable[..., Any]:
|
|
"""This function is designed to ensure that a given OpenTelemetry context
|
|
is applied when executing a specified function. It is particularly useful
|
|
for preserving the trace context when a guardrails is executed in a
|
|
different execution flow or when integrating with other frameworks.
|
|
|
|
Args:
|
|
outer_scope_otel_context (Context): The OpenTelemetry context to apply
|
|
when executing the function.
|
|
func (Callable[..., Any]): The function to be executed within
|
|
the given OpenTelemetry context.
|
|
|
|
Returns:
|
|
Callable[..., Any]: A wrapped version of 'func' that, when called,
|
|
executes with 'outer_scope_otel_context' applied.
|
|
"""
|
|
|
|
def wrapped_func(*args: Any, **kwargs: Any) -> Any:
|
|
# Attach the specified OpenTelemetry context before executing 'func'
|
|
token = context.attach(outer_scope_otel_context)
|
|
try:
|
|
# Execute 'func' within the attached context
|
|
return func(*args, **kwargs)
|
|
finally:
|
|
# Ensure the context is detached after execution
|
|
# to maintain correct context management
|
|
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
|
|
|
|
|
|
def redact(value: str) -> str:
|
|
"""Redacts all but the last four characters of the given string.
|
|
|
|
Args:
|
|
value (str): The string to be redacted.
|
|
|
|
Returns:
|
|
str: The redacted string with all but the last four characters
|
|
replaced by asterisks.
|
|
"""
|
|
redaction_length = len(value) - 4
|
|
stars = "*" * redaction_length
|
|
return f"{stars}{value[-4:]}"
|
|
|
|
|
|
def ismatchingkey(
|
|
target_key: str,
|
|
keys_to_match: tuple[str, ...] = ("key", "token", "password"),
|
|
) -> bool:
|
|
"""Check if the target key contains any of the specified keys to match.
|
|
|
|
Args:
|
|
target_key (str): The key to be checked.
|
|
keys_to_match (tuple[str, ...], optional): A tuple of keys to match
|
|
against the target key. Defaults to ("key", "token").
|
|
|
|
Returns:
|
|
bool: True if any of the keys to match are found in the target key,
|
|
False otherwise.
|
|
"""
|
|
for k in keys_to_match:
|
|
if k in target_key:
|
|
return True
|
|
return False
|
|
|
|
|
|
def can_convert_to_dict(s: str) -> bool:
|
|
"""Check if a string can be converted to a dictionary.
|
|
|
|
This function attempts to load the input string as JSON. If successful,
|
|
it returns True, indicating that the string can be converted to a dictionary.
|
|
Otherwise, it catches ValueError and TypeError exceptions and returns False.
|
|
|
|
Args:
|
|
s (str): The input string to be checked.
|
|
|
|
Returns:
|
|
bool: True if the string can be converted to a dictionary, False otherwise.
|
|
"""
|
|
try:
|
|
json.loads(s)
|
|
return True
|
|
except (ValueError, TypeError):
|
|
return False
|
|
|
|
|
|
def recursive_key_operation(
|
|
data: Optional[Union[Dict[str, Any], List[Any], str]],
|
|
operation: Callable[[str], str],
|
|
keys_to_match: List[str] = ["key", "token", "password"],
|
|
) -> Optional[Union[Dict[str, Any], List[Any], str]]:
|
|
"""Recursively traverses a dictionary, list, or JSON string and applies a
|
|
specified operation to the values of keys that match any in the
|
|
`keys_to_match` list. This function is useful for masking sensitive data
|
|
(e.g., keys, tokens, passwords) in nested structures.
|
|
|
|
Args:
|
|
data (Optional[Union[Dict[str, Any], List[Any], str]]): The input data
|
|
to traverse. This can bea dictionary, list, or JSON string. If a
|
|
JSON string is provided, it will be parsed into a dictionary before
|
|
processing.
|
|
|
|
operation (Callable[[str], str]): A function that takes a string value
|
|
and returns a modified string. This operation is applied to the values
|
|
of keys that match any in `keys_to_match`.
|
|
keys_to_match (List[str]): A list of keys to search for in the data. If
|
|
a key matche any in this list, the corresponding value will be processed
|
|
by the `operation`. Defaults to ["key", "token", "password"].
|
|
|
|
Returns:
|
|
Optional[Union[Dict[str, Any], List[Any], str]]: The modified data structure
|
|
with the operation applied to the values of matched keys. The return type
|
|
matches the input type (dict, list, or str).
|
|
"""
|
|
if isinstance(data, str) and can_convert_to_dict(data):
|
|
data_dict = json.loads(data)
|
|
data = str(recursive_key_operation(data_dict, operation, keys_to_match))
|
|
elif isinstance(data, dict):
|
|
for key, value in data.items():
|
|
if ismatchingkey(key, tuple(keys_to_match)) and isinstance(value, str):
|
|
# Apply the operation to the value of the matched key
|
|
data[key] = operation(value)
|
|
else:
|
|
# Recursively process nested dictionaries or lists
|
|
data[key] = recursive_key_operation(value, operation, keys_to_match)
|
|
elif isinstance(data, list):
|
|
for i in range(len(data)):
|
|
data[i] = recursive_key_operation(data[i], operation, keys_to_match)
|
|
|
|
return data
|