Add GUARDRAIL span kind to arize tracing

guardrails/telemetry/guard_tracing.py

@@ -10,6 +10,7 @@
     Union,
 )
 
+from openinference.semconv.trace import SpanAttributes
 from opentelemetry import context, trace
 from opentelemetry.trace import StatusCode, Tracer, Span, Link, get_tracer
 
@@ -153,6 +154,9 @@ def trace_stream_guard(
                     guard_span = new_span
                     add_guard_attributes(guard_span, history, res)
                     add_user_attributes(guard_span)
+                    new_span.set_attribute(
+                        SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
+                    )
                     yield res
         except StopIteration:
             next_exists = False
@@ -179,7 +183,9 @@ 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)
-
+            guard_span.set_attribute(
+                SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
+            )
             try:
                 result = _execute_fn(*args, **kwargs)
                 if isinstance(result, Iterator) and not isinstance(
@@ -218,6 +224,9 @@ async def trace_async_stream_guard(
 
                     add_guard_attributes(guard_span, history, res)
                     add_user_attributes(guard_span)
+                    guard_span.set_attribute(
+                        SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
+                    )
                     yield res
         except StopIteration:
             next_exists = False
@@ -259,7 +268,9 @@ 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)
-
+            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

@@ -8,6 +8,7 @@
     recursive_key_operation,
     redact,
 )
+from openinference.semconv.trace import SpanAttributes
 
 
 def trace_operation(
@@ -82,7 +83,7 @@ def trace_llm_call(
 
     if current_span is None:
         return
-
+    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

@@ -7,6 +7,7 @@
     Iterator,
     Optional,
 )
+from openinference.semconv.trace import SpanAttributes
 
 from opentelemetry import context, trace
 from opentelemetry.trace import StatusCode, Span
@@ -26,6 +27,7 @@
 )
 from guardrails.utils.safe_get import safe_get
 from guardrails.version import GUARDRAILS_VERSION
+from openinference.semconv.trace import SpanAttributes
 
 import sys
 
@@ -83,6 +85,9 @@ def trace_step_wrapper(*args, **kwargs) -> Iteration:
                 name="step",  # type: ignore
                 context=current_otel_context,  # type: ignore
             ) as step_span:
+                step_span.set_attribute(
+                    SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
+                )
                 try:
                     response = fn(*args, **kwargs)
                     add_step_attributes(step_span, response, *args, **kwargs)
@@ -111,6 +116,7 @@ def trace_stream_step_generator(
         name="step",  # type: ignore
         context=current_otel_context,  # type: ignore
     ) as step_span:
+        step_span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL")
         try:
             gen = fn(*args, **kwargs)
             next_exists = True
@@ -157,10 +163,14 @@ async def trace_async_step_wrapper(*args, **kwargs) -> Iteration:
                 name="step",  # type: ignore
                 context=current_otel_context,  # type: ignore
             ) as step_span:
+                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 +196,7 @@ async def trace_async_stream_step_generator(
         name="step",  # type: ignore
         context=current_otel_context,  # type: ignore
     ) as step_span:
+        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,7 @@
 
 from opentelemetry import context, trace
 from opentelemetry.trace import StatusCode, Tracer, Span
+from openinference.semconv.trace import SpanAttributes
 
 
 from guardrails.settings import settings
@@ -103,6 +104,10 @@ def trace_validator_wrapper(*args, **kwargs):
                     name=validator_span_name,  # type: ignore
                     context=current_otel_context,  # type: ignore
                 ) as validator_span:
+                    validator_span.set_attribute(
+                        SpanAttributes.OPENINFERENCE_SPAN_KIND, "GUARDRAIL"
+                    )
+
                     try:
                         resp = fn(*args, **kwargs)
                         add_user_attributes(validator_span)
@@ -167,6 +172,10 @@ async def trace_validator_wrapper(*args, **kwargs):
                     name=validator_span_name,  # type: ignore
                     context=current_otel_context,  # type: ignore
                 ) as validator_span:
+                    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 <[email protected]>"]
 license = "Apache License 2.0"