@azure-rest/ai-inference
Version:
Inference API for Azure-supported AI models
198 lines • 8.59 kB
JavaScript
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
import { isError } from "@azure/core-util";
/*
* Add event to span. Sample:
{
name: 'gen_ai.user.message',
attributes: {
'gen_ai.system': 'INFERENCE_GEN_AI_SYSTEM_NAME',
'gen_ai.event.content': `{"role":"user","content":"What's the weather like in Boston?"}`
},
time: [ 1725666879, 622695900 ],
droppedAttributesCount: 0
},
*/
/*
* Add event to span. Sample:
{
name: 'gen_ai.choice',
attributes: {
'gen_ai.system': 'INFERENCE_GEN_AI_SYSTEM_NAME',
'gen_ai.event.content': '{"finish_reason":"tool_calls","index":0,"message":{"content":""}}'
},
time: [ 1725666881, 780608000 ],
droppedAttributesCount: 0
}
*/
var TracingAttributesEnum;
(function (TracingAttributesEnum) {
TracingAttributesEnum["Operation_Name"] = "gen_ai.operation.name";
TracingAttributesEnum["Request_Model"] = "gen_ai.request.model";
TracingAttributesEnum["System"] = "gen_ai.system";
TracingAttributesEnum["Error_Type"] = "error.type";
TracingAttributesEnum["Server_Port"] = "server.port";
TracingAttributesEnum["Request_Frequency_Penalty"] = "gen_ai.request.frequency_penalty";
TracingAttributesEnum["Request_Max_Tokens"] = "gen_ai.request.max_tokens";
TracingAttributesEnum["Request_Presence_Penalty"] = "gen_ai.request.presence_penalty";
TracingAttributesEnum["Request_Stop_Sequences"] = "gen_ai.request.stop_sequences";
TracingAttributesEnum["Request_Temperature"] = "gen_ai.request.temperature";
TracingAttributesEnum["Request_Top_P"] = "gen_ai.request.top_p";
TracingAttributesEnum["Response_Finish_Reasons"] = "gen_ai.response.finish_reasons";
TracingAttributesEnum["Response_Id"] = "gen_ai.response.id";
TracingAttributesEnum["Response_Model"] = "gen_ai.response.model";
TracingAttributesEnum["Usage_Input_Tokens"] = "gen_ai.usage.input_tokens";
TracingAttributesEnum["Usage_Output_Tokens"] = "gen_ai.usage.output_tokens";
TracingAttributesEnum["Server_Address"] = "server.address";
})(TracingAttributesEnum || (TracingAttributesEnum = {}));
const INFERENCE_GEN_AI_SYSTEM_NAME = "az.ai.inference";
const isContentRecordingEnabled = () => envVarToBoolean("AZURE_TRACING_GEN_AI_CONTENT_RECORDING_ENABLED");
export function getRequestBody(request) {
return { body: JSON.parse(request.body) };
}
export function getSpanName(request) {
var _a;
const { body } = getRequestBody(request);
return `chat ${(_a = body === null || body === void 0 ? void 0 : body.model) !== null && _a !== void 0 ? _a : ""}`.trim();
}
export function onStartTracing(span, request, url) {
if (!span.isRecording()) {
return;
}
const urlObj = new URL(url);
const port = Number(urlObj.port) || (urlObj.protocol === "https:" ? undefined : 80);
if (port) {
span.setAttribute(TracingAttributesEnum.Server_Port, port);
}
span.setAttribute(TracingAttributesEnum.Server_Address, urlObj.hostname);
span.setAttribute(TracingAttributesEnum.Operation_Name, "chat");
span.setAttribute(TracingAttributesEnum.System, "az.ai.inference");
const { body } = getRequestBody(request);
if (!body)
return;
span.setAttribute(TracingAttributesEnum.Request_Model, body.model);
span.setAttribute(TracingAttributesEnum.Request_Frequency_Penalty, body.frequency_penalty);
span.setAttribute(TracingAttributesEnum.Request_Max_Tokens, body.max_tokens);
span.setAttribute(TracingAttributesEnum.Request_Presence_Penalty, body.presence_penalty);
span.setAttribute(TracingAttributesEnum.Request_Stop_Sequences, body.stop);
span.setAttribute(TracingAttributesEnum.Request_Temperature, body.temperature);
span.setAttribute(TracingAttributesEnum.Request_Top_P, body.top_p);
if (body.messages) {
addRequestChatMessageEvent(span, body.messages);
}
}
export function tryProcessResponse(span, response) {
var _a, _b, _c;
if (!span.isRecording()) {
return;
}
if (response === null || response === void 0 ? void 0 : response.bodyAsText) {
const body = JSON.parse(response.bodyAsText);
if ((_a = body.error) !== null && _a !== void 0 ? _a : body.message) {
span.setAttribute(TracingAttributesEnum.Error_Type, `${(_b = body.status) !== null && _b !== void 0 ? _b : body.statusCode}`);
span.setStatus({
status: "error",
error: (_c = body.error) !== null && _c !== void 0 ? _c : body.message, // message is not in the schema of the response, but it can present if there is crediential error
});
}
span.setAttribute(TracingAttributesEnum.Response_Id, body.id);
span.setAttribute(TracingAttributesEnum.Response_Model, body.model);
if (body.choices) {
span.setAttribute(TracingAttributesEnum.Response_Finish_Reasons, body.choices.map((choice) => choice.finish_reason).join(","));
}
if (body.usage) {
span.setAttribute(TracingAttributesEnum.Usage_Input_Tokens, body.usage.prompt_tokens);
span.setAttribute(TracingAttributesEnum.Usage_Output_Tokens, body.usage.completion_tokens);
}
addResponseChatMessageEvent(span, body);
}
}
export function tryProcessError(span, error) {
span.setStatus({
status: "error",
error: isError(error) ? error : undefined,
});
}
function addRequestChatMessageEvent(span, messages) {
messages.forEach((message) => {
var _a;
if (message.role) {
const content = {};
const chatMsg = message;
if (chatMsg.content) {
content.content = chatMsg.content;
}
if (!isContentRecordingEnabled()) {
content.content = "";
}
const assistantMsg = message;
if (assistantMsg.tool_calls) {
content.tool_calls = assistantMsg.tool_calls;
if (!isContentRecordingEnabled()) {
const toolCalls = JSON.parse(JSON.stringify(content.tool_calls));
toolCalls.forEach((toolCall) => {
if (toolCall.function.arguments) {
toolCall.function.arguments = "";
}
toolCall.function.name = "";
});
content.tool_calls = toolCalls;
}
}
const toolMsg = message;
if (toolMsg.tool_call_id) {
content.id = toolMsg.tool_call_id;
}
(_a = span.addEvent) === null || _a === void 0 ? void 0 : _a.call(span, `gen_ai.${message.role}.message`, {
attributes: {
"gen_ai.system": INFERENCE_GEN_AI_SYSTEM_NAME,
"gen_ai.event.content": JSON.stringify(content),
},
});
}
});
}
function addResponseChatMessageEvent(span, body) {
var _a;
if (!span.addEvent) {
return;
}
(_a = body === null || body === void 0 ? void 0 : body.choices) === null || _a === void 0 ? void 0 : _a.forEach((choice) => {
var _a;
let message = {};
if (choice.message.content) {
message.content = choice.message.content;
}
if (choice.message.tool_calls) {
message.toolCalls = choice.message.tool_calls;
}
if (!isContentRecordingEnabled()) {
message = JSON.parse(JSON.stringify(message));
message.content = "";
if (message.toolCalls) {
message.toolCalls.forEach((toolCall) => {
if (toolCall.function.arguments) {
toolCall.function.arguments = "";
}
toolCall.function.name = "";
});
}
}
const response = {
finish_reason: choice.finish_reason,
index: choice.index,
message,
};
const attributes = {
"gen_ai.system": INFERENCE_GEN_AI_SYSTEM_NAME,
"gen_ai.event.content": JSON.stringify(response),
};
(_a = span.addEvent) === null || _a === void 0 ? void 0 : _a.call(span, "gen_ai.choice", { attributes });
});
}
function envVarToBoolean(key) {
var _a;
const value = (_a = process.env[key]) !== null && _a !== void 0 ? _a : process.env[key.toLowerCase()];
return value !== "false" && value !== "0" && Boolean(value);
}
//# sourceMappingURL=tracingHelper.js.map