removing exception when instrumenting multiple times, adding content … (Azure#38336)

* removing exception when instrumenting multiple times, adding content recording parameter

* fixing a typo

* Update version number and changelog

* Update model used in AOAI chat completions tests

* Update test recordings

* check tool related fix and readme code snippet update

* sphinx tool warning fix

---------

Co-authored-by: Marko Hietala <[email protected]>
Co-authored-by: Darren Cohen <[email protected]>

Author: 3 people

sdk/ai/azure-ai-inference/CHANGELOG.md

@@ -1,5 +1,14 @@
 # Release History
 
+## 1.0.0b6 (2024-11-15)
+
+### Features Added
+
+* OpenTelemetry tracing:
+  * Method `AIInferenceInstrumentor().instrument()` updated with an input argument `enable_content_recording`.
+  * Calling `AIInferenceInstrumentor().instrument()` twice no longer results in an exception.
+  * Added method `AIInferenceInstrumentor().is_content_recording_enabled()`
+
 ## 1.0.0b5 (2024-10-16)
 
 ### Features Added

sdk/ai/azure-ai-inference/README.md

@@ -301,10 +301,7 @@ response = client.complete(
                 "role": "assistant",
                 "content": "The main construction of the International Space Station (ISS) was completed between 1998 and 2011. During this period, more than 30 flights by US space shuttles and 40 by Russian rockets were conducted to transport components and modules to the station.",
             },
-            {
-                "role": "user",
-                "content": "And what was the estimated cost to build it?"
-            },
+            {"role": "user", "content": "And what was the estimated cost to build it?"},
         ]
     }
 )
@@ -578,6 +575,7 @@ Or configure it in the code with the following snippet:
 
 ```python
 from azure.core.settings import settings
+
 settings.tracing_implementation = "opentelemetry"
 ```
 
@@ -591,6 +589,7 @@ The final step is to enable Azure AI Inference instrumentation with the followin
 
 ```python
 from azure.ai.inference.tracing import AIInferenceInstrumentor
+
 # Instrument AI Inference API
 AIInferenceInstrumentor().instrument()
 ```
@@ -616,11 +615,13 @@ The `@tracer.start_as_current_span` decorator can be used to trace your own func
 
 ```python
 from opentelemetry.trace import get_tracer
+
 tracer = get_tracer(__name__)
 
+
 # The tracer.start_as_current_span decorator will trace the function call and enable adding additional attributes
 # to the span in the function implementation. Note that this will trace the function parameters and their values.
-@tracer.start_as_current_span("get_temperature") # type: ignore
+@tracer.start_as_current_span("get_temperature")  # type: ignore
 def get_temperature(city: str) -> str:
 
     # Adding attributes to the current span

sdk/ai/azure-ai-inference/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/ai/azure-ai-inference",
-  "Tag": "python/ai/azure-ai-inference_3934744053"
+  "Tag": "python/ai/azure-ai-inference_f4719b1cae"
 }

sdk/ai/azure-ai-inference/azure/ai/inference/_version.py

@@ -6,4 +6,4 @@
 # Changes may cause incorrect behavior and will be lost if the code is regenerated.
 # --------------------------------------------------------------------------
 
-VERSION = "1.0.0b5"
+VERSION = "1.0.0b6"

sdk/ai/azure-ai-inference/azure/ai/inference/tracing.py

@@ -63,15 +63,26 @@ def __init__(self):
         # and have a parameter that specifies the version to use.
         self._impl = _AIInferenceInstrumentorPreview()
 
-    def instrument(self) -> None:
+    def instrument(self, enable_content_recording: Optional[bool] = None) -> None:
         """
         Enable trace instrumentation for AI Inference.
 
-        Raises:
-            RuntimeError: If instrumentation is already enabled.
-
+        :param enable_content_recording: Whether content recording is enabled as part
+            of the traces or not. Content in this context refers to chat message content
+            and function call tool related function names, function parameter names and
+            values. True will enable content recording, False will disable it. If no value
+            s provided, then the value read from environment variable
+            AZURE_TRACING_GEN_AI_CONTENT_RECORDING_ENABLED is used. If the environment variable
+            is not found, then the value will default to False. Please note that successive calls
+            to instrument will always apply the content recording value provided with the most
+            recent call to instrument (including applying the environment variable if no value is
+            provided and defaulting to false if the environment variable is not found), even if
+            instrument was already previously called without uninstrument being called in between
+            the instrument calls.
+
+        :type enable_content_recording: bool, optional
         """
-        self._impl.instrument()
+        self._impl.instrument(enable_content_recording=enable_content_recording)
 
     def uninstrument(self) -> None:
         """
@@ -94,6 +105,15 @@ def is_instrumented(self) -> bool:
         """
         return self._impl.is_instrumented()
 
+    def is_content_recording_enabled(self) -> bool:
+        """
+        This function gets the content recording value.
+
+        :return: A bool value indicating whether content recording is enabled.
+        :rtype: bool
+        """
+        return self._impl.is_content_recording_enabled()
+
 
 class _AIInferenceInstrumentorPreview:
     """
@@ -108,37 +128,37 @@ def _str_to_bool(self, s):
             return False
         return str(s).lower() == "true"
 
-    def instrument(self):
+    def instrument(self, enable_content_recording: Optional[bool] = None):
         """
         Enable trace instrumentation for AI Inference.
 
-        Raises:
-            RuntimeError: If instrumentation is already enabled.
+        :param enable_content_recording: Whether content recording is enabled as part
+        of the traces or not. Content in this context refers to chat message content
+        and function call tool related function names, function parameter names and
+        values. True will enable content recording, False will disable it. If no value
+        is provided, then the value read from environment variable
+        AZURE_TRACING_GEN_AI_CONTENT_RECORDING_ENABLED is used. If the environment variable
+        is not found, then the value will default to False.
 
-        This method checks the environment variable
-        'AZURE_TRACING_GEN_AI_CONTENT_RECORDING_ENABLED' to determine
-        whether to enable content tracing.
+        :type enable_content_recording: bool, optional
         """
-        if self.is_instrumented():
-            raise RuntimeError("Already instrumented")
-
-        var_value = os.environ.get("AZURE_TRACING_GEN_AI_CONTENT_RECORDING_ENABLED")
-        enable_content_tracing = self._str_to_bool(var_value)
-        self._instrument_inference(enable_content_tracing)
+        if enable_content_recording is None:
+            var_value = os.environ.get("AZURE_TRACING_GEN_AI_CONTENT_RECORDING_ENABLED")
+            enable_content_recording = self._str_to_bool(var_value)
+        if not self.is_instrumented():
+            self._instrument_inference(enable_content_recording)
+        else:
+            self._set_content_recording_enabled(enable_content_recording=enable_content_recording)
 
     def uninstrument(self):
         """
         Disable trace instrumentation for AI Inference.
 
-        Raises:
-            RuntimeError: If instrumentation is not currently enabled.
-
         This method removes any active instrumentation, stopping the tracing
         of AI Inference.
         """
-        if not self.is_instrumented():
-            raise RuntimeError("Not instrumented")
-        self._uninstrument_inference()
+        if self.is_instrumented():
+            self._uninstrument_inference()
 
     def is_instrumented(self):
         """
@@ -149,6 +169,24 @@ def is_instrumented(self):
         """
         return self._is_instrumented()
 
+    def set_content_recording_enabled(self, enable_content_recording: bool = False) -> None:
+        """This function sets the content recording value.
+
+        :param enable_content_recording: Indicates whether tracing of message content should be enabled.
+                                    This also controls whether function call tool function names,
+                                    parameter names and parameter values are traced.
+        :type enable_content_recording: bool
+        """
+        self._set_content_recording_enabled(enable_content_recording=enable_content_recording)
+
+    def is_content_recording_enabled(self) -> bool:
+        """This function gets the content recording value.
+
+        :return: A bool value indicating whether content tracing is enabled.
+        :rtype bool
+        """
+        return self._is_content_recording_enabled()
+
     def _set_attributes(self, span: "AbstractSpan", *attrs: Tuple[str, Any]) -> None:
         for attr in attrs:
             key, value = attr
@@ -764,3 +802,22 @@ def _is_instrumented(self):
         :rtype: bool
         """
         return _inference_traces_enabled
+
+    def _set_content_recording_enabled(self, enable_content_recording: bool = False) -> None:
+        """This function sets the content recording value.
+
+        :param enable_content_recording: Indicates whether tracing of message content should be enabled.
+                                    This also controls whether function call tool function names,
+                                    parameter names and parameter values are traced.
+        :type enable_content_recording: bool
+        """
+        global _trace_inference_content  # pylint: disable=W0603
+        _trace_inference_content = enable_content_recording
+
+    def _is_content_recording_enabled(self) -> bool:
+        """This function gets the content recording value.
+
+        :return: A bool value indicating whether content tracing is enabled.
+        :rtype bool
+        """
+        return _trace_inference_content

sdk/ai/azure-ai-inference/tests/model_inference_test_base.py

@@ -49,13 +49,13 @@
 # hosted on Azure OpenAI (AOAI) endpoint.
 # TODO: When we have a MaaS model that supports chat completions with image input,
 # use that instead.
-# AZURE_OPENAI_CHAT_ENDPOINT=https://<endpont-name>.openai.azure.com/openai/deployments/gpt-4o-0806
+# AZURE_OPENAI_CHAT_ENDPOINT=https://<endpont-name>.openai.azure.com/openai/deployments/gpt-4o
 # AZURE_OPENAI_CHAT_KEY=<32-char-api-key>
 #
 ServicePreparerAOAIChatCompletions = functools.partial(
     EnvironmentVariableLoader,
     "azure_openai_chat",
-    azure_openai_chat_endpoint="https://your-deployment-name.openai.azure.com/openai/deployments/gpt-4o-0806",
+    azure_openai_chat_endpoint="https://your-deployment-name.openai.azure.com/openai/deployments/gpt-4o",
     azure_openai_chat_key="00000000000000000000000000000000",
 )
 

sdk/ai/azure-ai-inference/tests/test_model_inference_client.py

@@ -856,67 +856,133 @@ def test_instrumentation(self, **kwargs):
 
     @ServicePreparerChatCompletions()
     @recorded_by_proxy
-    def test_instrumenting_twice_causes_exception(self, **kwargs):
+    def test_instrumenting_twice_does_not_cause_exception(self, **kwargs):
         # Make sure code is not instrumented due to a previous test exception
         try:
             AIInferenceInstrumentor().uninstrument()
         except RuntimeError as e:
             pass
         client = self._create_chat_client(**kwargs)
         exception_caught = False
-        instrumented_once = False
         try:
             AIInferenceInstrumentor().instrument()
-            instrumented_once = True
             AIInferenceInstrumentor().instrument()
         except RuntimeError as e:
             exception_caught = True
             print(e)
         AIInferenceInstrumentor().uninstrument()
         client.close()
-        assert instrumented_once == True
-        assert exception_caught == True
+        assert exception_caught == False
+
+    @ServicePreparerChatCompletions()
+    @recorded_by_proxy
+    def test_uninstrumenting_uninstrumented_does_not_cause_exception(self, **kwargs):
+        # Make sure code is not instrumented due to a previous test exception
+        try:
+            AIInferenceInstrumentor().uninstrument()
+        except RuntimeError as e:
+            pass
+        client = self._create_chat_client(**kwargs)
+        exception_caught = False
+        try:
+            AIInferenceInstrumentor().uninstrument()
+        except RuntimeError as e:
+            exception_caught = True
+            print(e)
+        client.close()
+        assert exception_caught == False
 
     @ServicePreparerChatCompletions()
     @recorded_by_proxy
-    def test_uninstrumenting_uninstrumented_causes_exception(self, **kwargs):
+    def test_uninstrumenting_twice_does_not_cause_exception(self, **kwargs):
         # Make sure code is not instrumented due to a previous test exception
         try:
             AIInferenceInstrumentor().uninstrument()
         except RuntimeError as e:
             pass
         client = self._create_chat_client(**kwargs)
         exception_caught = False
+        uninstrumented_once = False
         try:
+            AIInferenceInstrumentor().instrument()
+            AIInferenceInstrumentor().uninstrument()
             AIInferenceInstrumentor().uninstrument()
         except RuntimeError as e:
             exception_caught = True
             print(e)
         client.close()
-        assert exception_caught == True
+        assert exception_caught == False
 
     @ServicePreparerChatCompletions()
     @recorded_by_proxy
-    def test_uninstrumenting_twice_causes_exception(self, **kwargs):
+    def test_is_content_recording_enabled(self, **kwargs):
         # Make sure code is not instrumented due to a previous test exception
         try:
             AIInferenceInstrumentor().uninstrument()
         except RuntimeError as e:
             pass
+        self.modify_env_var(CONTENT_TRACING_ENV_VARIABLE, "False")
         client = self._create_chat_client(**kwargs)
         exception_caught = False
         uninstrumented_once = False
         try:
+            # From environment variable instrumenting from uninstrumented
+            AIInferenceInstrumentor().instrument()
+            self.modify_env_var(CONTENT_TRACING_ENV_VARIABLE, "False")
             AIInferenceInstrumentor().instrument()
+            content_recording_enabled = AIInferenceInstrumentor().is_content_recording_enabled()
+            assert content_recording_enabled == False
             AIInferenceInstrumentor().uninstrument()
-            uninstrumented_once = True
+            self.modify_env_var(CONTENT_TRACING_ENV_VARIABLE, "True")
+            AIInferenceInstrumentor().instrument()
+            content_recording_enabled = AIInferenceInstrumentor().is_content_recording_enabled()
+            assert content_recording_enabled == True
             AIInferenceInstrumentor().uninstrument()
+            self.modify_env_var(CONTENT_TRACING_ENV_VARIABLE, "invalid")
+            AIInferenceInstrumentor().instrument()
+            content_recording_enabled = AIInferenceInstrumentor().is_content_recording_enabled()
+            assert content_recording_enabled == False
+
+            # From environment variable instrumenting from instrumented
+            self.modify_env_var(CONTENT_TRACING_ENV_VARIABLE, "True")
+            AIInferenceInstrumentor().instrument()
+            content_recording_enabled = AIInferenceInstrumentor().is_content_recording_enabled()
+            assert content_recording_enabled == True
+            self.modify_env_var(CONTENT_TRACING_ENV_VARIABLE, "True")
+            AIInferenceInstrumentor().instrument()
+            content_recording_enabled = AIInferenceInstrumentor().is_content_recording_enabled()
+            assert content_recording_enabled == True
+            self.modify_env_var(CONTENT_TRACING_ENV_VARIABLE, "invalid")
+            AIInferenceInstrumentor().instrument()
+            content_recording_enabled = AIInferenceInstrumentor().is_content_recording_enabled()
+            assert content_recording_enabled == False
+
+            # From parameter instrumenting from uninstrumented
+            AIInferenceInstrumentor().uninstrument()
+            self.modify_env_var(CONTENT_TRACING_ENV_VARIABLE, "True")
+            AIInferenceInstrumentor().instrument(enable_content_recording=False)
+            content_recording_enabled = AIInferenceInstrumentor().is_content_recording_enabled()
+            assert content_recording_enabled == False
+            AIInferenceInstrumentor().uninstrument()
+            self.modify_env_var(CONTENT_TRACING_ENV_VARIABLE, "False")
+            AIInferenceInstrumentor().instrument(enable_content_recording=True)
+            content_recording_enabled = AIInferenceInstrumentor().is_content_recording_enabled()
+            assert content_recording_enabled == True
+
+            # From parameter instrumenting from instrumented
+            self.modify_env_var(CONTENT_TRACING_ENV_VARIABLE, "True")
+            AIInferenceInstrumentor().instrument(enable_content_recording=False)
+            content_recording_enabled = AIInferenceInstrumentor().is_content_recording_enabled()
+            assert content_recording_enabled == False
+            self.modify_env_var(CONTENT_TRACING_ENV_VARIABLE, "False")
+            AIInferenceInstrumentor().instrument(enable_content_recording=True)
+            content_recording_enabled = AIInferenceInstrumentor().is_content_recording_enabled()
+            assert content_recording_enabled == True
         except RuntimeError as e:
             exception_caught = True
             print(e)
         client.close()
-        assert uninstrumented_once == True
-        assert exception_caught == True
+        assert exception_caught == False
 
     @ServicePreparerChatCompletions()
     @recorded_by_proxy