1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184
|
# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
"""
DESCRIPTION:
This sample demonstrates how to use tracing with the Inference client library.
Azure AI Inference is instrumented with OpenTelemetry. In order to enable tracing
you need to configure OpenTelemetry to export traces to your observability backend.
This sample shows how to print traces to stdout.
This sample assumes the AI model is hosted on a Serverless API or
Managed Compute endpoint. For GitHub Models or Azure OpenAI endpoints,
the client constructor needs to be modified. See package documentation:
https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/ai/azure-ai-inference/README.md#key-concepts
USAGE:
python sample_chat_completions_with_tracing.py
Set these two or three environment variables before running the sample:
1) AZURE_AI_CHAT_ENDPOINT - Your endpoint URL, in the form
https://<your-deployment-name>.<your-azure-region>.models.ai.azure.com
where `your-deployment-name` is your unique AI Model deployment name, and
`your-azure-region` is the Azure region where your model is deployed.
2) AZURE_AI_CHAT_KEY - Your model key. Keep it secret.
3) AZURE_TRACING_GEN_AI_CONTENT_RECORDING_ENABLED - Optional. Set to 'true'
for detailed traces, including chat request and response messages.
"""
import os
from opentelemetry import trace
# Install opentelemetry with command "pip install opentelemetry-sdk".
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, ConsoleSpanExporter
from azure.ai.inference import ChatCompletionsClient
from azure.ai.inference.models import SystemMessage, UserMessage, CompletionsFinishReason
from azure.core.credentials import AzureKeyCredential
# [START trace_setting]
from azure.core.settings import settings
settings.tracing_implementation = "opentelemetry"
# [END trace_setting]
# Setup tracing to console
# Requires opentelemetry-sdk
span_exporter = ConsoleSpanExporter()
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
trace.set_tracer_provider(tracer_provider)
# [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(key, endpoint):
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"],
},
)
)
client = ChatCompletionsClient(endpoint=endpoint, credential=AzureKeyCredential(key))
messages = [
SystemMessage("You are a helpful assistant."),
UserMessage("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(function_response, tool_call_id=tool_call.id))
# 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():
# [START instrument_inferencing]
from azure.ai.inference.tracing import AIInferenceInstrumentor
# Instrument AI Inference API
AIInferenceInstrumentor().instrument()
# [END instrument_inferencing]
try:
endpoint = os.environ["AZURE_AI_CHAT_ENDPOINT"]
key = os.environ["AZURE_AI_CHAT_KEY"]
except KeyError:
print("Missing environment variable 'AZURE_AI_CHAT_ENDPOINT' or 'AZURE_AI_CHAT_KEY'")
print("Set them before running this sample.")
exit()
chat_completion_with_function_call(key, endpoint)
# [START uninstrument_inferencing]
AIInferenceInstrumentor().uninstrument()
# [END uninstrument_inferencing]
if __name__ == "__main__":
main()
|