diff --git a/python/openinference-instrumentation/examples/tracer.ipynb b/python/openinference-instrumentation/examples/tracer.ipynb index 0dbfca901..8b7f5d028 100644 --- a/python/openinference-instrumentation/examples/tracer.ipynb +++ b/python/openinference-instrumentation/examples/tracer.ipynb @@ -14,6 +14,7 @@ "from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor\n", "from opentelemetry.trace import Status, StatusCode, set_tracer_provider\n", "\n", + "from openinference.instrumentation import using_attributes\n", "from openinference.instrumentation.config import (\n", " OpenInferenceTracerProvider,\n", " chain,\n", @@ -21,6 +22,7 @@ " get_input_value_and_mime_type,\n", " get_openinference_span_kind,\n", " get_output_value_and_mime_type,\n", + " suppress_tracing,\n", ")\n", "from openinference.semconv.resource import ResourceAttributes" ] @@ -309,6 +311,82 @@ "\n", "decorated_chain_with_input_and_output_set_inside_the_wrapped_function(\"input\")" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Suppress Tracing" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [], + "source": [ + "with suppress_tracing():\n", + " with tracer.start_as_current_span(\n", + " \"THIS-SPAN-SHOULD-NOT-BE-TRACED\",\n", + " openinference_span_kind=\"chain\",\n", + " ) as span:\n", + " span.set_input(\"input\")\n", + " span.set_output(\"output\")\n", + " span.set_status(Status(StatusCode.OK))" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": {}, + "outputs": [], + "source": [ + "@chain\n", + "def decorated_chain_with_suppress_tracing(input: str) -> str:\n", + " return \"output\"\n", + "\n", + "\n", + "with suppress_tracing():\n", + " decorated_chain_with_suppress_tracing(\"input\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Context Attributes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with using_attributes(session_id=\"123\"):\n", + " with tracer.start_as_current_span(\n", + " \"chain-span-with-context-attributes\",\n", + " openinference_span_kind=\"chain\",\n", + " ) as span:\n", + " span.set_input(\"input\")\n", + " span.set_output(\"output\")\n", + " span.set_status(Status(StatusCode.OK))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@chain\n", + "def decorated_chain_with_context_attributes(input: str) -> str:\n", + " return \"output\"\n", + "\n", + "\n", + "with using_attributes(session_id=\"123\"):\n", + " decorated_chain_with_context_attributes(\"input\")" + ] } ], "metadata": { diff --git a/python/openinference-instrumentation/src/openinference/instrumentation/config.py b/python/openinference-instrumentation/src/openinference/instrumentation/config.py index fdf836d7a..997c47511 100644 --- a/python/openinference-instrumentation/src/openinference/instrumentation/config.py +++ b/python/openinference-instrumentation/src/openinference/instrumentation/config.py @@ -28,11 +28,13 @@ Context, attach, detach, + get_value, set_value, ) from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.id_generator import IdGenerator from opentelemetry.trace import ( + INVALID_SPAN, INVALID_SPAN_ID, INVALID_TRACE_ID, Link, @@ -60,6 +62,7 @@ SpanAttributes, ) +from .context_attributes import get_attributes_from_context from .helpers import safe_json_dumps from .logging import logger @@ -468,7 +471,10 @@ def sync_wrapper( ) as span: output = wrapped(*args, **kwargs) span.set_status(Status(StatusCode.OK)) - has_output = OUTPUT_MIME_TYPE in span.attributes or OUTPUT_VALUE in span.attributes + attributes = getattr( + span, "attributes", {} + ) # INVALID_SPAN does not have the attributes property + has_output = OUTPUT_MIME_TYPE in attributes or OUTPUT_VALUE in attributes if has_output: return output # don't overwrite if the output is set inside the wrapped function if isinstance(output, (str, int, float, bool)): @@ -506,7 +512,10 @@ async def async_wrapper( ) as span: output = await wrapped(*args, **kwargs) span.set_status(Status(StatusCode.OK)) - has_output = OUTPUT_MIME_TYPE in span.attributes or OUTPUT_VALUE in span.attributes + attributes = getattr( + span, "attributes", {} + ) # INVALID_SPAN does not have the attributes property + has_output = OUTPUT_MIME_TYPE in attributes or OUTPUT_VALUE in attributes if has_output: return output # don't overwrite if the output is set inside the wrapped function if isinstance(output, (str, int, float, bool)): @@ -660,6 +669,12 @@ def start_span( *, openinference_span_kind: Optional[OpenInferenceSpanKind] = None, ) -> OpenInferenceSpan: + span_wrapper_cls = OpenInferenceSpan + if openinference_span_kind is not None: + normalized_span_kind = _normalize_openinference_span_kind(openinference_span_kind) + span_wrapper_cls = _get_span_wrapper_cls(normalized_span_kind) + if get_value(_SUPPRESS_INSTRUMENTATION_KEY): + return span_wrapper_cls(INVALID_SPAN, self._self_config) tracer = cast(Tracer, self.__wrapped__) span = tracer.__class__.start_span( self, @@ -672,13 +687,10 @@ def start_span( record_exception=record_exception, set_status_on_exception=set_status_on_exception, ) - span_wrapper_cls = OpenInferenceSpan - if openinference_span_kind is not None: - normalized_span_kind = _normalize_openinference_span_kind(openinference_span_kind) - span_wrapper_cls = _get_span_wrapper_cls(normalized_span_kind) span = span_wrapper_cls(span, config=self._self_config) if attributes: span.set_attributes(dict(attributes)) + span.set_attributes(dict(get_attributes_from_context())) return span