Add GUARDRAIL span kind to arize tracing

guardrails/telemetry/guard_tracing.py

@@ -10,6 +10,10 @@
     Union,
 )
 
+try:
+    from openinference.semconv.trace import SpanAttributes
+except ImportError:
+    SpanAttributes = None
 from opentelemetry import context, trace
 from opentelemetry.trace import StatusCode, Tracer, Span, Link, get_tracer
 
@@ -153,6 +157,10 @@ def trace_stream_guard(
                     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
@@ -179,7 +187,10 @@ def trace_guard_execution(
             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(
@@ -218,6 +229,10 @@ async def trace_async_stream_guard(
 
                     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
@@ -259,7 +274,10 @@ async def trace_async_guard_execution(
             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):

guardrails/telemetry/open_inference.py

@@ -9,6 +9,11 @@
     redact,
 )
 
+try:
+    from openinference.semconv.trace import SpanAttributes
+except ImportError:
+    SpanAttributes = None
+
 
 def trace_operation(
     *,
@@ -82,7 +87,8 @@ def trace_llm_call(
 
     if current_span is None:
         return
-
+    if SpanAttributes is not None:
+        current_span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL")
     ser_function_call = serialize(function_call)
     if ser_function_call:
         current_span.set_attribute("llm.function_call", ser_function_call)

guardrails/telemetry/runner_tracing.py

@@ -8,6 +8,11 @@
     Optional,
 )
 
+try:
+    from openinference.semconv.trace import SpanAttributes
+except ImportError:
+    SpanAttributes = None
+
 from opentelemetry import context, trace
 from opentelemetry.trace import StatusCode, Span
 
@@ -83,6 +88,10 @@ def trace_step_wrapper(*args, **kwargs) -> Iteration:
                 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)
@@ -111,6 +120,8 @@ def trace_stream_step_generator(
         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
@@ -157,10 +168,15 @@ async def trace_async_step_wrapper(*args, **kwargs) -> Iteration:
                 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))
@@ -186,6 +202,8 @@ async def trace_async_stream_step_generator(
         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

guardrails/telemetry/validator_tracing.py

@@ -9,6 +9,10 @@
 
 from opentelemetry import context, trace
 from opentelemetry.trace import StatusCode, Tracer, Span
+try:
+    from openinference.semconv.trace import SpanAttributes
+except ImportError:
+    SpanAttributes = None
 
 
 from guardrails.settings import settings
@@ -103,6 +107,11 @@ def trace_validator_wrapper(*args, **kwargs):
                     name=validator_span_name,  # type: ignore
                     context=current_otel_context,  # type: ignore
                 ) as validator_span:
+                    if SpanAttributes is not None:
+                        validator_span.set_attribute(
+                            SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
+                        )
+
                     try:
                         resp = fn(*args, **kwargs)
                         add_user_attributes(validator_span)
@@ -167,6 +176,11 @@ async def trace_validator_wrapper(*args, **kwargs):
                     name=validator_span_name,  # type: ignore
                     context=current_otel_context,  # type: ignore
                 ) as validator_span:
+                    if SpanAttributes is not None:
+                        validator_span.set_attribute(
+                            SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
+                        )  # see here for a list of span kinds: https://github.com/Arize-ai/openinference/blob/main/python/openinference-semantic-conventions/src/openinference/semconv/trace/__init__.py#L271
+
                     try:
                         resp = await fn(*args, **kwargs)
                         add_user_attributes(validator_span)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "guardrails-ai"
-version = "0.6.2"
+version = "0.6.3"
 description = "Adding guardrails to large language models."
 authors = ["Guardrails AI <contact@guardrailsai.com>"]
 license = "Apache License 2.0"