ai client support for tracing operations, inference instrumentation as…

.vscode/cspell.json

@@ -406,6 +406,9 @@
     "uamqp",
     "uksouth",
     "ukwest",
+    "uninstrument",
+    "uninstrumented",
+    "uninstrumenting",
     "unpad",
     "unpadder",
     "unpartial",

sdk/ai/azure-ai-client/azure/ai/client/_patch.py

@@ -15,7 +15,7 @@
 from ._serialization import Deserializer, Serializer
 from .operations import AgentsOperations, EndpointsOperations, EvaluationsOperations
 from ._client import AzureAIClient as ClientGenerated
-from .operations._patch import InferenceOperations
+from .operations._patch import InferenceOperations, TracingOperations
 
 
 class AzureAIClient(ClientGenerated):
@@ -150,6 +150,7 @@ def __init__(
         self.agents = AgentsOperations(self._client2, self._config2, self._serialize, self._deserialize)
         self.evaluations = EvaluationsOperations(self._client3, self._config3, self._serialize, self._deserialize)
         self.inference = InferenceOperations(self)
+        self.tracing = TracingOperations(self)
 
     def close(self) -> None:
         self._client1.close()

sdk/ai/azure-ai-client/azure/ai/client/operations/_patch.py

@@ -1128,10 +1128,29 @@ def upload_file(
         raise ValueError("Invalid parameters for upload_file. Please provide the necessary arguments.")
 
 
+class TracingOperations:
+
+    def __init__(self, outer_instance):
+        self.outer_instance = outer_instance
+
+    def create_inference_instrumentor(self, _version: "SemanticConventionsVersion" = None) -> "AIInferenceInstrumentor":
+        try:
+            from azure.ai.inference import AIInferenceInstrumentor, SemanticConventionsVersion
+        except ModuleNotFoundError as _:
+            raise ModuleNotFoundError(
+                "Azure AI Inference SDK is not installed. Please install it using 'pip install azure-ai-inference'"
+            )
+        if _version is None:
+            return AIInferenceInstrumentor()
+        else:
+            return AIInferenceInstrumentor(_version)
+
+
 __all__: List[str] = [
     "AgentsOperations",
     "EndpointsOperations",
     "InferenceOperations",
+    "TracingOperations"
 ]  # Add all objects you want publicly available to users at this package level
 
 

sdk/ai/azure-ai-client/samples/inference/sample_trace_chat_completions_client.py

@@ -0,0 +1,170 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+"""
+DESCRIPTION:
+    This sample demonstrates how to use Azure AI client to get Inference client and
+    instrument it with traces using Inference Instumentor that is also acquired
+    using the Azure AI Client.
+"""
+import os
+from opentelemetry import trace
+from azure.ai.client import AzureAIClient
+from azure.identity import DefaultAzureCredential
+from azure.ai.inference.models import SystemMessage, UserMessage, CompletionsFinishReason
+# opentelemetry-sdk is required for the opentelemetry.sdk imports.
+# You can install it with command "pip install opentelemetry-sdk".
+#from opentelemetry.sdk.trace import TracerProvider
+#from opentelemetry.sdk.trace.export import SimpleSpanProcessor, ConsoleSpanExporter
+
+ # [START trace_setting]
+from azure.core.settings import settings
+settings.tracing_implementation = "opentelemetry"
+# [END trace_setting]
+
+# Setup tracing to console
+# Requires opentelemetry-sdk
+#exporter = ConsoleSpanExporter()
+#trace.set_tracer_provider(TracerProvider())
+#tracer = trace.get_tracer(__name__)
+#trace.get_tracer_provider().add_span_processor(SimpleSpanProcessor(exporter))
+
+ # [START trace_function]
+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
+def get_temperature(city: str) -> str:
+
+    # Adding attributes to the current span
+    span = trace.get_current_span()
+    span.set_attribute("requested_city", city)
+
+    if city == "Seattle":
+        return "75"
+    elif city == "New York City":
+        return "80"
+    else:
+        return "Unavailable"
+ # [END trace_function]
+
+
+def get_weather(city: str) -> str:
+    if city == "Seattle":
+        return "Nice weather"
+    elif city == "New York City":
+        return "Good weather"
+    else:
+        return "Unavailable"
+
+
+def chat_completion_with_function_call(client):
+    import json
+    from azure.ai.inference.models import ToolMessage, AssistantMessage, ChatCompletionsToolCall, ChatCompletionsToolDefinition, FunctionDefinition
+
+    weather_description = ChatCompletionsToolDefinition(
+        function=FunctionDefinition(
+            name="get_weather",
+            description="Returns description of the weather in the specified city",
+            parameters={
+                "type": "object",
+                "properties": {
+                    "city": {
+                        "type": "string",
+                        "description": "The name of the city for which weather info is requested",
+                    },
+                },
+                "required": ["city"],
+            },
+        )
+    )
+
+    temperature_in_city = ChatCompletionsToolDefinition(
+        function=FunctionDefinition(
+            name="get_temperature",
+            description="Returns the current temperature for the specified city",
+            parameters={
+                "type": "object",
+                "properties": {
+                    "city": {
+                        "type": "string",
+                        "description": "The name of the city for which temperature info is requested",
+                    },
+                },
+                "required": ["city"],
+            },
+        )
+    )
+
+    messages=[
+        SystemMessage(content="You are a helpful assistant."),
+        UserMessage(content="What is the weather and temperature in Seattle?"),
+    ]
+
+    response = client.complete(messages=messages, tools=[weather_description, temperature_in_city])
+
+    if response.choices[0].finish_reason == CompletionsFinishReason.TOOL_CALLS:
+        # Append the previous model response to the chat history
+        messages.append(AssistantMessage(tool_calls=response.choices[0].message.tool_calls))
+        # The tool should be of type function call.
+        if response.choices[0].message.tool_calls is not None and len(response.choices[0].message.tool_calls) > 0:
+            for tool_call in response.choices[0].message.tool_calls:
+                if type(tool_call) is ChatCompletionsToolCall:
+                    function_args = json.loads(tool_call.function.arguments.replace("'", '"'))
+                    print(f"Calling function `{tool_call.function.name}` with arguments {function_args}")
+                    callable_func = globals()[tool_call.function.name]
+                    function_response = callable_func(**function_args)
+                    print(f"Function response = {function_response}")
+                    # Provide the tool response to the model, by appending it to the chat history
+                    messages.append(ToolMessage(tool_call_id=tool_call.id, content=function_response))
+                    # With the additional tools information on hand, get another response from the model
+            response = client.complete(messages=messages, tools=[weather_description, temperature_in_city])
+
+    print(f"Model response = {response.choices[0].message.content}")
+
+
+def main():
+    # Create an Azure AI Client from a connection string, copied from your AI Studio project.
+    # It should have the format "<Endpoint>;<AzureSubscriptionId>;<ResourceGroup>;<WorkspaceName>"
+    ai_client = AzureAIClient.from_connection_string(
+        credential=DefaultAzureCredential(),
+        connection=os.environ["AI_CLIENT_CONNECTION_STRING"],
+    )
+
+    # Or, you can create the Azure AI Client by giving all required parameters directly
+    # ai_client = AzureAIClient(
+    #     credential=DefaultAzureCredential(),
+    #     endpoint=os.environ["AI_CLIENT_ENDPOINT"],
+    #     subscription_id=os.environ["AI_CLIENT_SUBSCRIPTION_ID"],
+    #     resource_group_name=os.environ["AI_CLIENT_RESOURCE_GROUP_NAME"],
+    #     workspace_name=os.environ["AI_CLIENT_WORKSPACE_NAME"],
+    # )
+
+    # Get an authenticated OpenAI client for your default Azure OpenAI connection:
+    client = ai_client.inference.get_azure_openai_client()
+
+    # [START instrument_inferencing]
+    # Speficy a semantic conventions version used for the instrumentation if needed.
+    #instrumentor = ai_client.tracing.create_inference_instrumentor(SemanticConventionsVersion.PREVIEW)
+
+    # If you do not specify semantic conventions version, then the default version supported by the current
+    # Inference SDK will be returned.
+    instrumentor = ai_client.tracing.create_inference_instrumentor()
+
+    # Instrument AI Inference API
+    instrumentor.instrument()
+    # [END instrument_inferencing]
+
+    chat_completion_with_function_call(client)
+
+    # [START uninstrument_inferencing]
+    instrumentor.uninstrument()
+    # [END uninstrument_inferencing]
+
+
+if __name__ == "__main__":
+    main()

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

@@ -57,6 +57,14 @@ To update an existing installation of the package, use:
 pip install --upgrade azure-ai-inference
 ```
 
+If you want to install Azure AI Inferencing package with support for OpenTelemetry based tracing, use the following command:
+
+```bash
+pip install azure-ai-inference[trace]
+```
+
+
+
 ## Key concepts
 
 ### Create and authenticate a client directly, using API key or GitHub token
@@ -530,6 +538,87 @@ For more information, see [Configure logging in the Azure libraries for Python](
 
 To report issues with the client library, or request additional features, please open a GitHub issue [here](https://github.com/Azure/azure-sdk-for-python/issues)
 
+## Tracing
+
+The Azure AI Inferencing API Tracing library provides tracing for Azure AI Inference client library for Python. Refer to Installation chapter above for installation instructions.
+
+### Setup
+
+The environment variable AZURE_TRACING_GEN_AI_CONTENT_RECORDING_ENABLED controls whether the actual message contents will be recorded in the traces or not. By default, the message contents are not recorded as part of the trace. When message content recording is disabled any function call tool related function names, function parameter names and function parameter values are also not recorded in the trace. Set the value of the environment variable to "true" (case insensitive) for the message contents to be recorded as part of the trace. Any other value will cause the message contents not to be recorded.
+
+You also need to configure the tracing implementation in your code by setting `AZURE_SDK_TRACING_IMPLEMENTATION` to `opentelemetry` or configuring it in the code with the following snippet:
+
+<!-- SNIPPET:sample_chat_completions_with_tracing.trace_setting -->
+
+```python
+from azure.core.settings import settings
+settings.tracing_implementation = "opentelemetry"
+```
+
+<!-- END SNIPPET -->
+
+Please refer to [azure-core-tracing-documentation](https://learn.microsoft.com/python/api/overview/azure/core-tracing-opentelemetry-readme) for more information.
+
+### Exporting Traces with OpenTelemetry
+
+Azure AI Inference is instrumented with OpenTelemetry. In order to enable tracing you need to configure OpenTelemetry to export traces to your observability backend. 
+Refer to [Azure SDK tracing in Python](https://learn.microsoft.com/python/api/overview/azure/core-tracing-opentelemetry-readme?view=azure-python-preview) for more details.
+
+Refer to [Azure Monitor OpenTelemetry documentation](https://learn.microsoft.com/azure/azure-monitor/app/opentelemetry-enable?tabs=python) for the details on how to send Azure AI Inference traces to Azure Monitor and create Azure Monitor resource.
+
+### Instrumentation
+
+Use the AIInferenceInstrumentor to instrument the Azure AI Inferencing API for LLM tracing, this will cause the LLM traces to be emitted from Azure AI Inferencing API.
+
+<!-- SNIPPET:sample_chat_completions_with_tracing.instrument_inferencing -->
+
+```python
+from azure.core.tracing.ai.inference import AIInferenceInstrumentor
+# Instrument AI Inference API
+AIInferenceInstrumentor().instrument()
+```
+
+<!-- END SNIPPET -->
+
+
+It is also possible to uninstrument the Azure AI Inferencing API by using the uninstrument call. After this call, the traces will no longer be emitted by the Azure AI Inferencing API until instrument is called again.
+
+<!-- SNIPPET:sample_chat_completions_with_tracing.uninstrument_inferencing -->
+
+```python
+AIInferenceInstrumentor().uninstrument()
+```
+
+<!-- END SNIPPET -->
+
+### Tracing Your Own Functions
+The @tracer.start_as_current_span decorator can be used to trace your own functions. This will trace the function parameters and their values. You can also add further attributes to the span in the function implementation as demonstrated below. Note that you will have to setup the tracer in your code before using the decorator. More information is available [here](https://opentelemetry.io/docs/languages/python/).
+
+<!-- SNIPPET:sample_chat_completions_with_tracing.trace_function -->
+
+```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
+def get_temperature(city: str) -> str:
+
+    # Adding attributes to the current span
+    span = trace.get_current_span()
+    span.set_attribute("requested_city", city)
+
+    if city == "Seattle":
+        return "75"
+    elif city == "New York City":
+        return "80"
+    else:
+        return "Unavailable"
+```
+
+<!-- END SNIPPET -->
+
 ## Next steps
 
 * Have a look at the [Samples](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/ai/azure-ai-inference/samples) folder, containing fully runnable Python code for doing inference using synchronous and asynchronous clients.

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

@@ -9,6 +9,8 @@
 from ._patch import ChatCompletionsClient
 from ._patch import EmbeddingsClient
 from ._patch import ImageEmbeddingsClient
+from ._patch import AIInferenceInstrumentor
+from ._patch import SemanticConventionsVersion
 from ._version import VERSION
 
 __version__ = VERSION
@@ -22,6 +24,8 @@
     "ChatCompletionsClient",
     "EmbeddingsClient",
     "ImageEmbeddingsClient",
+    "SemanticConventionsVersion",
+    "AIInferenceInstrumentor",
 ]