suppress instrumentation and context attrs

python/openinference-instrumentation/examples/tracer.ipynb

@@ -14,13 +14,15 @@
     "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",
     "    get_current_span,\n",
     "    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": {

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