axiom
Version:
Axiom AI SDK provides - an API to wrap your AI calls with observability instrumentation. - offline evals - online evals
1,578 lines (1,564 loc) • 99.3 kB
JavaScript
import {
WITHSPAN_BAGGAGE_KEY,
WITHSPAN_REDACTION_POLICY_KEY,
getGlobalTracer,
getRedactionPolicy,
handleMaybeRedactedAttribute,
package_default
} from "./chunk-PU64TWX4.js";
import {
isValidName
} from "./chunk-MM5FFQJT.js";
import {
Attr
} from "./chunk-4TKUTT24.js";
import {
__publicField
} from "./chunk-KEXKKQVW.js";
// src/otel/withSpan.ts
import {
context as context2,
propagation as propagation2,
trace as trace2,
SpanStatusCode as SpanStatusCode3
} from "@opentelemetry/api";
// src/otel/utils/wrapperUtils.ts
import {
trace,
context,
propagation,
SpanStatusCode as SpanStatusCode2
} from "@opentelemetry/api";
// src/schema.ts
var SCHEMA_VERSION = "0.0.2";
var SCHEMA_BASE_URL = "https://axiom.co/ai/schemas/";
var SCHEMA_URL = `${SCHEMA_BASE_URL}${SCHEMA_VERSION}`;
// src/otel/startActiveSpan.ts
import { SpanStatusCode } from "@opentelemetry/api";
var createStartActiveSpan = (tracer) => async (name, options, fn, callbacks) => {
return tracer.startActiveSpan(name, { ...options ?? {} }, async (span) => {
try {
const result = await fn(span);
callbacks?.onSuccess?.(span);
return result;
} catch (error) {
callbacks?.onError?.(error, span);
if (error instanceof Error) {
span.recordException(error);
span.setStatus({
code: SpanStatusCode.ERROR,
message: error.message
});
}
throw error;
} finally {
callbacks?.onFinally?.(span);
if (!callbacks?.manualEnd) {
span.end();
}
}
});
};
// src/otel/utils/wrapperUtils.ts
function ensureNumber(value) {
const v = typeof value === "number" ? value : typeof value === "string" ? Number(value) : void 0;
if (Number.isNaN(v)) return void 0;
if (v === Infinity || v === -Infinity) return void 0;
return v;
}
function classifyError(err) {
if (err == null) return void 0;
if (err instanceof Error) {
const name = err.name.toLowerCase();
if (name.includes("timeout")) return "timeout";
if (name.includes("abort")) return "timeout";
if (name.includes("network") || name.includes("fetch")) return "network";
if (name.includes("validation")) return "validation";
if (name.includes("auth")) return "authentication";
if (name.includes("parse") || name.includes("json")) return "parsing";
if (name.includes("permission") || name.includes("forbidden")) return "authorization";
if (name.includes("rate") && name.includes("limit")) return "rate_limit";
if (name.includes("quota") || name.includes("limit")) return "quota_exceeded";
return void 0;
}
return void 0;
}
function classifyToolError(err, span) {
if (err instanceof Error) {
span.recordException(err);
} else {
span.recordException({
message: String(err),
name: "UnknownError"
});
}
span.setStatus({
code: SpanStatusCode2.ERROR,
message: err instanceof Error ? err.message : String(err)
});
let errorType = "unknown";
let statusCode;
if (err && typeof err === "object") {
const errObj = err;
const name = errObj.name?.toLowerCase() || "";
const message = errObj.message?.toLowerCase() || "";
if (name.includes("timeout") || name.includes("abort") || message.includes("timeout")) {
errorType = "timeout";
} else if (name.includes("validation") || errObj.code === "VALIDATION_ERROR" || message.includes("validation")) {
errorType = "validation";
} else if (name.includes("fetch") || name.includes("network") || message.includes("network") || message.includes("fetch failed")) {
errorType = "network";
statusCode = errObj.status || errObj.code;
} else if (name.includes("auth") || message.includes("auth") || message.includes("unauthorized")) {
errorType = "authentication";
} else if (name.includes("permission") || name.includes("forbidden") || message.includes("forbidden")) {
errorType = "authorization";
} else if (name.includes("rate") && (name.includes("limit") || message.includes("rate limit"))) {
errorType = "rate_limit";
} else if (name.includes("quota") || message.includes("quota") || message.includes("limit exceeded")) {
errorType = "quota_exceeded";
} else if (name.includes("parse") || name.includes("json") || message.includes("json") || message.includes("parse")) {
errorType = "parsing";
}
}
span.setAttribute(Attr.Error.Type, errorType);
if (err instanceof Error && err.message) {
span.setAttribute(Attr.Error.Message, err.message);
}
if (statusCode !== void 0) {
span.setAttribute(Attr.HTTP.Response.StatusCode, statusCode);
}
}
function isNoOpTracerProvider() {
const provider = trace.getTracerProvider();
if (provider.constructor.name === "NoopTracerProvider") {
return true;
}
if (typeof provider.getTracer !== "function") {
return true;
}
return false;
}
function getTracer() {
const tracer = getGlobalTracer();
if (isNoOpTracerProvider()) {
const isDebug = process.env.AXIOM_DEBUG === "true";
if (!isDebug) {
console.warn(
"[AxiomAI] No TracerProvider registered - spans will be no-op. Make sure to call initAxiomAI() after your OpenTelemetry SDK has started (sdk.start())."
);
}
}
return tracer;
}
function createGenAISpanName(operation, suffix) {
return suffix ? `${operation} ${suffix}` : operation;
}
function setScopeAttributes(span) {
const bag = propagation.getActiveBaggage();
if (bag) {
const capability = bag.getEntry("capability")?.value;
if (capability) {
span.setAttribute(Attr.GenAI.Capability.Name, capability);
}
const step = bag.getEntry("step")?.value;
if (step) {
span.setAttribute(Attr.GenAI.Step.Name, step);
}
const conversationId = bag.getEntry("conversationId")?.value;
if (conversationId) {
span.setAttribute(Attr.GenAI.Conversation.ID, conversationId);
}
}
}
function setAxiomBaseAttributes(span) {
span.setAttributes({
[Attr.Axiom.GenAI.SchemaURL]: SCHEMA_URL,
[Attr.Axiom.GenAI.SDK.Name]: package_default.name,
[Attr.Axiom.GenAI.SDK.Version]: package_default.version
});
}
function setBaseAttributes(span, provider, modelId) {
span.setAttributes({
[Attr.GenAI.Operation.Name]: Attr.GenAI.Operation.Name_Values.Chat,
[Attr.GenAI.Request.Model]: modelId
});
const systemValue = mapVercelSDKProviderToOTelProvider(provider);
if (systemValue) {
span.setAttribute(Attr.GenAI.Provider.Name, systemValue);
}
setAxiomBaseAttributes(span);
}
function setRequestParameterAttributes(span, params) {
const {
maxTokens,
frequencyPenalty,
presencePenalty,
temperature,
topP,
topK,
seed,
stopSequences
} = params;
if (maxTokens !== void 0) {
span.setAttribute(Attr.GenAI.Request.MaxTokens, maxTokens);
}
if (frequencyPenalty !== void 0) {
span.setAttribute(Attr.GenAI.Request.FrequencyPenalty, frequencyPenalty);
}
if (presencePenalty !== void 0) {
span.setAttribute(Attr.GenAI.Request.PresencePenalty, presencePenalty);
}
if (temperature !== void 0) {
span.setAttribute(Attr.GenAI.Request.Temperature, temperature);
}
if (topP !== void 0) {
span.setAttribute(Attr.GenAI.Request.TopP, topP);
}
if (topK !== void 0) {
span.setAttribute(Attr.GenAI.Request.TopK, topK);
}
if (seed !== void 0) {
span.setAttribute(Attr.GenAI.Request.Seed, seed);
}
if (stopSequences && stopSequences.length > 0) {
span.setAttribute(Attr.GenAI.Request.StopSequences, JSON.stringify(stopSequences));
}
}
function createStreamChildSpan(parentSpan, operationName) {
const tracer = getTracer();
const ctx = context.active();
const spanContext = trace.setSpan(ctx, parentSpan);
const childSpan = tracer.startSpan(operationName, void 0, spanContext);
return childSpan;
}
function recordSpanError(span, err) {
if (err instanceof Error) {
span.recordException(err);
} else {
span.recordException({
message: String(err),
name: "UnknownError"
});
}
span.setStatus({
code: SpanStatusCode2.ERROR,
message: err instanceof Error ? err.message : String(err)
});
const errorType = classifyError(err);
span.setAttribute(Attr.Error.Type, errorType ?? "unknown");
if (err instanceof Error && err.message) {
span.setAttribute(Attr.Error.Message, err.message);
}
if (err && typeof err === "object" && "status" in err) {
span.setAttribute(Attr.HTTP.Response.StatusCode, err.status);
}
}
var spanContextStore = /* @__PURE__ */ new WeakMap();
var spanContextFactories = {
v1: () => ({ version: "v1", originalPrompt: [], rawCall: void 0 }),
v2: () => ({ version: "v2", originalPrompt: [], originalV2Prompt: void 0 }),
v3: () => ({ version: "v3", originalPrompt: [], originalV3Prompt: void 0 })
};
async function withSpanHandling(modelId, operation, options) {
const bag = propagation.getActiveBaggage();
const isWithinWithSpan = bag?.getEntry(WITHSPAN_BAGGAGE_KEY)?.value === "true";
const createSpanContext = spanContextFactories[options?.version ?? "v1"];
const name = createGenAISpanName(Attr.GenAI.Operation.Name_Values.Chat, modelId);
if (isWithinWithSpan) {
const activeSpan = trace.getActiveSpan();
if (!activeSpan) {
throw new Error("Expected active span when within withSpan");
}
activeSpan.updateName(name);
let spanContext = spanContextStore.get(activeSpan);
if (!spanContext) {
spanContext = createSpanContext();
spanContextStore.set(activeSpan, spanContext);
}
const lease = {
owned: false,
end: () => {
}
// No-op: we don't own this span
};
try {
return await operation(activeSpan, spanContext, lease);
} catch (err) {
recordSpanError(activeSpan, err);
throw err;
}
} else {
const tracer = getTracer();
const startActiveSpan = createStartActiveSpan(tracer);
return startActiveSpan(
name,
null,
async (span) => {
const lease = {
owned: true,
end: () => span.end()
};
return await operation(span, createSpanContext(), lease);
},
{
manualEnd: options?.streaming ?? false,
onError: (err, span) => {
const errorType = classifyError(err);
span.setAttribute(Attr.Error.Type, errorType ?? "unknown");
if (err instanceof Error && err.message) {
span.setAttribute(Attr.Error.Message, err.message);
}
if (err && typeof err === "object" && "status" in err) {
span.setAttribute(Attr.HTTP.Response.StatusCode, err.status);
}
}
}
);
}
}
function determineOutputTypeV1(options) {
if (options.responseFormat?.type) {
switch (options.responseFormat.type) {
case "json":
return Attr.GenAI.Output.Type_Values.Json;
case "text":
return Attr.GenAI.Output.Type_Values.Text;
}
}
if (options.mode?.type === "object-json" || options.mode?.type === "object-tool") {
return Attr.GenAI.Output.Type_Values.Json;
}
if (options.mode?.type === "regular") {
return Attr.GenAI.Output.Type_Values.Text;
}
return void 0;
}
function determineOutputTypeV2(options) {
if (options.responseFormat?.type) {
switch (options.responseFormat.type) {
case "json":
return Attr.GenAI.Output.Type_Values.Json;
case "text":
return Attr.GenAI.Output.Type_Values.Text;
}
}
return void 0;
}
function determineOutputTypeV3(options) {
if (options.responseFormat?.type) {
switch (options.responseFormat.type) {
case "json":
return Attr.GenAI.Output.Type_Values.Json;
case "text":
return Attr.GenAI.Output.Type_Values.Text;
}
}
return void 0;
}
function mapVercelSDKProviderToOTelProvider(vercelSDKProvider) {
if (vercelSDKProvider === "openai-compatible") {
return void 0;
}
switch (vercelSDKProvider) {
case "amazon-bedrock":
return Attr.GenAI.Provider.Name_Values.AWSBedrock;
case "anthropic":
case "anthropic.messages":
return Attr.GenAI.Provider.Name_Values.Anthropic;
case "assemblyai":
case "assemblyai.transcription":
return Attr.GenAI.Provider.Name_Values.AssemblyAI;
case "deepgram":
case "deepgram.transcription":
return Attr.GenAI.Provider.Name_Values.Deepgram;
case "gateway":
return Attr.GenAI.Provider.Name_Values.Vercel;
case "gladia":
case "gladia.transcription":
return Attr.GenAI.Provider.Name_Values.Gladia;
case "google":
case "google.generative-ai":
return Attr.GenAI.Provider.Name_Values.GCPGemini;
case "groq":
return Attr.GenAI.Provider.Name_Values.Groq;
case "mistral":
return Attr.GenAI.Provider.Name_Values.MistralAI;
case "openai":
return Attr.GenAI.Provider.Name_Values.OpenAI;
case "perplexity":
return Attr.GenAI.Provider.Name_Values.Perplexity;
case "replicate":
return Attr.GenAI.Provider.Name_Values.Replicate;
case "revai":
case "revai.transcription":
return Attr.GenAI.Provider.Name_Values.RevAI;
case "togetherai":
return Attr.GenAI.Provider.Name_Values.TogetherAI;
case "xai":
return Attr.GenAI.Provider.Name_Values.XAI;
// startswith + fall through
default: {
if (vercelSDKProvider.startsWith("azure.")) {
return Attr.GenAI.Provider.Name_Values.AzureAIOpenAI;
}
if (vercelSDKProvider.startsWith("cerebras.")) {
return Attr.GenAI.Provider.Name_Values.Cerebras;
}
if (vercelSDKProvider.startsWith("cohere.")) {
return Attr.GenAI.Provider.Name_Values.Cohere;
}
if (vercelSDKProvider.startsWith("deepinfra.")) {
return Attr.GenAI.Provider.Name_Values.DeepInfra;
}
if (vercelSDKProvider.startsWith("deepseek.")) {
return Attr.GenAI.Provider.Name_Values.Deepseek;
}
if (vercelSDKProvider.startsWith("elevenlabs.")) {
return Attr.GenAI.Provider.Name_Values.ElevenLabs;
}
if (vercelSDKProvider.startsWith("fal.")) {
return Attr.GenAI.Provider.Name_Values.Fal;
}
if (vercelSDKProvider.startsWith("fireworks.")) {
return Attr.GenAI.Provider.Name_Values.Fireworks;
}
if (vercelSDKProvider.startsWith("google.vertex.")) {
return Attr.GenAI.Provider.Name_Values.GCPVertexAI;
}
if (vercelSDKProvider.startsWith("groq.")) {
return Attr.GenAI.Provider.Name_Values.Groq;
}
if (vercelSDKProvider.startsWith("hume.")) {
return Attr.GenAI.Provider.Name_Values.Hume;
}
if (vercelSDKProvider.startsWith("lmnt.")) {
return Attr.GenAI.Provider.Name_Values.Lmnt;
}
if (vercelSDKProvider.startsWith("luma.")) {
return Attr.GenAI.Provider.Name_Values.Luma;
}
if (vercelSDKProvider.startsWith("mistral.")) {
return Attr.GenAI.Provider.Name_Values.MistralAI;
}
if (vercelSDKProvider.startsWith("openai.")) {
return Attr.GenAI.Provider.Name_Values.OpenAI;
}
if (vercelSDKProvider.startsWith("vercel.")) {
return Attr.GenAI.Provider.Name_Values.Vercel;
}
if (vercelSDKProvider.startsWith("vertex.anthropic.")) {
return Attr.GenAI.Provider.Name_Values.GCPVertexAI;
}
if (vercelSDKProvider.startsWith("xai.")) {
return Attr.GenAI.Provider.Name_Values.XAI;
}
const s = vercelSDKProvider.split(".");
if (s.length === 2) {
return s[0];
}
return void 0;
}
}
}
// src/otel/withSpan.ts
function withSpan(meta, fn, opts) {
const tracer = opts?.tracer ?? getTracer();
const span = tracer.startSpan("chat");
const spanContext = trace2.setSpan(context2.active(), span);
return context2.with(spanContext, async () => {
const capabilityValidation = isValidName(meta.capability);
if (!capabilityValidation.valid) {
console.warn(`[AxiomAI] Invalid capability name: ${capabilityValidation.error}. `);
}
const stepValidation = isValidName(meta.step);
if (!stepValidation.valid) {
console.warn(`[AxiomAI] Invalid step name: ${stepValidation.error}. `);
}
if (!span.isRecording()) {
const provider = trace2.getTracerProvider();
const providerIsNoOp = provider.constructor.name === "NoopTracerProvider";
if (providerIsNoOp) {
const isDebug = process.env.AXIOM_DEBUG === "true";
if (!isDebug) {
console.warn(
"[AxiomAI] No TracerProvider registered - spans are no-op. Make sure to call initAxiomAI() after your OpenTelemetry SDK has started."
);
}
}
}
const bag = propagation2.createBaggage({
capability: { value: meta.capability },
step: { value: meta.step },
// TODO: maybe we can just check the active span name instead?
[WITHSPAN_BAGGAGE_KEY]: { value: "true" },
// Mark that we're inside withSpan
// Store serialized redaction policy if provided
...opts?.redactionPolicy && {
[WITHSPAN_REDACTION_POLICY_KEY]: { value: JSON.stringify(opts.redactionPolicy) }
},
// Store conversation ID if provided
...meta.conversationId && {
conversationId: { value: meta.conversationId }
}
});
const ctx = propagation2.setBaggage(context2.active(), bag);
let spanEnded = false;
const safeEndSpan = () => {
if (!spanEnded) {
spanEnded = true;
span.end();
}
};
const timeoutMs = opts?.timeoutMs ?? 6e5;
const timeoutId = setTimeout(() => {
safeEndSpan();
}, timeoutMs);
try {
const result = await context2.with(ctx, () => fn(span));
if (result instanceof Response && result.body) {
if (result.body.locked) {
console.warn("[AxiomAI] Response body is already locked, cannot instrument stream");
clearTimeout(timeoutId);
safeEndSpan();
return result;
}
const originalReader = result.body.getReader();
const wrappedStream = new ReadableStream({
async pull(controller) {
try {
const { value, done } = await context2.with(ctx, () => originalReader.read());
if (done) {
originalReader.releaseLock?.();
clearTimeout(timeoutId);
span.setStatus({ code: SpanStatusCode3.OK });
safeEndSpan();
controller.close();
} else {
controller.enqueue(value);
}
} catch (err) {
originalReader.releaseLock?.();
clearTimeout(timeoutId);
span.recordException(err);
span.setStatus({
code: SpanStatusCode3.ERROR,
message: err instanceof Error ? err.message : String(err)
});
safeEndSpan();
controller.error(err);
}
},
async cancel(reason) {
try {
originalReader.releaseLock?.();
clearTimeout(timeoutId);
if (reason instanceof Error) {
span.recordException(reason);
} else if (reason) {
span.recordException({ message: String(reason), name: "CancelError" });
}
span.setStatus({
code: SpanStatusCode3.ERROR,
message: reason instanceof Error ? reason.message : String(reason)
});
safeEndSpan();
await originalReader.cancel(reason);
} catch (_err) {
}
}
});
return new Response(wrappedStream, {
status: result.status,
statusText: result.statusText,
headers: result.headers
});
}
if (result && typeof result === "object" && "textStream" in result) {
console.warn(
"[AxiomAI] Detected streaming object with textStream. For proper span lifecycle, call .toUIMessageStreamResponse() or similar inside withSpan, not after."
);
clearTimeout(timeoutId);
safeEndSpan();
return result;
}
clearTimeout(timeoutId);
span.setStatus({ code: SpanStatusCode3.OK });
safeEndSpan();
return result;
} catch (err) {
clearTimeout(timeoutId);
span.recordException(err);
span.setStatus({
code: SpanStatusCode3.ERROR,
message: err instanceof Error ? err.message : String(err)
});
safeEndSpan();
throw err;
}
});
}
// src/otel/middleware.ts
import "@opentelemetry/api";
// src/otel/utils/contentSanitizer.ts
import { createHash } from "crypto";
function extractImageMetadata(url) {
if (url.startsWith("data:")) {
const [header, base64Data] = url.split(",");
const formatMatch = header.match(/data:image\/(\w+)/);
const format = formatMatch?.[1];
const sizeBytes = base64Data ? Math.floor(base64Data.length * 3 / 4) : 0;
const hash = base64Data ? createHash("sha256").update(base64Data).digest("hex").slice(0, 16) : "unknown";
return {
format,
size_bytes: sizeBytes,
hash,
is_data_url: true
};
} else {
const hash = createHash("sha256").update(url).digest("hex").slice(0, 16);
return {
hash,
is_data_url: false
};
}
}
function sanitizeImageUrl(url, detail) {
const metadata = extractImageMetadata(url);
if (metadata.is_data_url) {
const formatPart = metadata.format ? `:${metadata.format}` : "";
const sizePart = metadata.size_bytes ? `:${metadata.size_bytes}b` : "";
return {
url: `[IMAGE${formatPart}${sizePart}:${metadata.hash}]`,
detail,
...metadata
};
} else {
return {
url,
detail,
...metadata
};
}
}
function sanitizeMultimodalContent(content) {
if (Array.isArray(content)) {
return content.map((part) => {
if (part && typeof part === "object" && "type" in part && part.type === "image_url") {
const imagePart = part;
if (imagePart.image_url?.url) {
return {
...part,
image_url: sanitizeImageUrl(imagePart.image_url.url, imagePart.image_url.detail)
};
}
}
return part;
});
}
return content;
}
// src/otel/completionUtils.ts
function createSimpleCompletion({ text }) {
const assistantMessage = {
role: "assistant",
content: text ?? ""
};
return [assistantMessage];
}
// src/util/promptUtils.ts
function appendToolCalls(prompt, toolCalls, toolResults, assistantContent) {
const updatedPrompt = [...prompt];
updatedPrompt.push({
role: "assistant",
content: assistantContent || null,
tool_calls: toolCalls.map((toolCall) => ({
id: toolCall.toolCallId,
function: {
name: toolCall.toolName,
arguments: typeof toolCall.args === "string" ? toolCall.args : JSON.stringify(toolCall.args)
},
type: "function"
}))
});
for (const toolCall of toolCalls) {
const realToolResult = toolResults.get(toolCall.toolName);
if (realToolResult) {
updatedPrompt.push({
role: "tool",
tool_call_id: toolCall.toolCallId,
content: JSON.stringify(realToolResult)
});
}
}
return updatedPrompt;
}
function extractToolResultsFromRawPrompt(rawPrompt) {
const toolResultsMap = /* @__PURE__ */ new Map();
if (!Array.isArray(rawPrompt)) {
return toolResultsMap;
}
for (const message of rawPrompt) {
if (message?.role === "user" && Array.isArray(message.parts)) {
for (const part of message.parts) {
if (part?.functionResponse) {
const functionResponse = part.functionResponse;
if (functionResponse.name && functionResponse.response) {
toolResultsMap.set(
functionResponse.name,
functionResponse.response.content || functionResponse.response
);
}
}
}
}
if (message?.role === "tool" && message?.tool_call_id && message?.content) {
}
}
return toolResultsMap;
}
function extractToolResultsFromPromptV2(prompt) {
const idToName = /* @__PURE__ */ new Map();
const results = /* @__PURE__ */ new Map();
for (const message of prompt) {
if (message.role === "assistant" && Array.isArray(message.content)) {
for (const part of message.content) {
if (part.type === "tool-call") {
idToName.set(part.toolCallId, part.toolName);
}
}
}
}
for (const message of prompt) {
if (message.role === "tool" && Array.isArray(message.content)) {
for (const part of message.content) {
if (part.toolCallId && part.output !== void 0) {
const toolName = idToName.get(part.toolCallId);
if (toolName) {
results.set(toolName, part.output);
}
}
}
}
}
return results;
}
function extractToolResultsFromPromptV3(prompt) {
const idToName = /* @__PURE__ */ new Map();
const results = /* @__PURE__ */ new Map();
for (const message of prompt) {
if (message.role === "assistant" && Array.isArray(message.content)) {
for (const part of message.content) {
if (part.type === "tool-call") {
idToName.set(part.toolCallId, part.toolName);
}
}
}
}
for (const message of prompt) {
if (message.role === "tool" && Array.isArray(message.content)) {
for (const part of message.content) {
if (part.type === "tool-result" && part.toolCallId && part.output !== void 0) {
const toolName = idToName.get(part.toolCallId);
if (toolName) {
results.set(toolName, part.output);
}
}
}
}
}
return results;
}
// src/otel/utils/normalized.ts
function normalizeV1ToolCall(toolCall) {
return {
toolCallId: toolCall.toolCallId,
toolName: toolCall.toolName,
args: typeof toolCall.args === "string" ? toolCall.args : JSON.stringify(toolCall.args),
toolCallType: "function"
};
}
function normalizeV2ToolCall(toolCall) {
return {
toolCallId: toolCall.toolCallId,
toolName: toolCall.toolName,
args: typeof toolCall.input === "string" ? toolCall.input.replace(/:\s+/g, ":") : JSON.stringify(toolCall.input),
toolCallType: "function"
};
}
function normalizeV1ToolCalls(toolCalls) {
return toolCalls.map(normalizeV1ToolCall);
}
function normalizeV2ToolCalls(toolCalls) {
return toolCalls.map(normalizeV2ToolCall);
}
function normalizeV3ToolCall(toolCall) {
return {
toolCallId: toolCall.toolCallId,
toolName: toolCall.toolName,
args: toolCall.input,
// V3 input is always a string
toolCallType: "function"
};
}
function normalizeV3ToolCalls(toolCalls) {
return toolCalls.map(normalizeV3ToolCall);
}
function promptV1ToOpenAI(prompt) {
const results = [];
for (const message of prompt) {
switch (message.role) {
case "system":
results.push({
role: "system",
content: message.content
});
break;
case "assistant":
const textPart = message.content.find((part) => part.type === "text");
const toolCallParts = message.content.filter(
(part) => part.type === "tool-call"
);
results.push({
role: "assistant",
content: textPart?.text || null,
...toolCallParts.length > 0 ? {
tool_calls: toolCallParts.map((part) => ({
id: part.toolCallId,
function: {
name: part.toolName,
arguments: JSON.stringify(part.args)
},
type: "function"
}))
} : {}
});
break;
case "user":
results.push({
role: "user",
content: message.content.map((part) => {
switch (part.type) {
case "text":
return {
type: "text",
text: part.text
};
case "image":
return {
type: "image_url",
image_url: {
url: part.image.toString()
}
};
default:
return {
type: "text",
text: `[${part.type}]` + (typeof part === "object" && part !== null ? JSON.stringify(part) : String(part))
};
}
})
});
break;
case "tool":
for (const part of message.content) {
results.push({
role: "tool",
tool_call_id: part.toolCallId,
content: JSON.stringify(part.result)
});
}
break;
}
}
return results;
}
function promptV2ToOpenAI(prompt) {
const results = [];
for (const message of prompt) {
switch (message.role) {
case "system":
results.push({
role: "system",
content: message.content
});
break;
case "assistant":
const textContent = message.content.find(
(part) => part.type === "text"
);
const toolCalls = message.content.filter(
(part) => part.type === "tool-call"
);
results.push({
role: "assistant",
content: textContent?.text || null,
...toolCalls.length > 0 ? {
tool_calls: toolCalls.map((part) => ({
id: part.toolCallId,
function: {
name: part.toolName,
arguments: typeof part.input === "string" ? part.input : JSON.stringify(part.input)
},
type: "function"
}))
} : {}
});
break;
case "user":
results.push({
role: "user",
content: message.content.map((part) => {
switch (part.type) {
case "text":
return {
type: "text",
text: part.text
};
case "image":
return {
type: "image_url",
image_url: {
url: part.image.toString()
}
};
default:
return part;
}
})
});
break;
case "tool":
for (const part of message.content) {
results.push({
role: "tool",
tool_call_id: part.toolCallId,
content: formatV2ToolCallOutput(part.output)
});
}
break;
}
}
return results;
}
function formatV2ToolCallOutput(output) {
switch (output.type) {
case "text":
return output.value;
case "json":
return typeof output.value === "string" ? output.value : JSON.stringify(output.value);
case "error-text":
return output.value;
case "error-json":
return typeof output.value === "string" ? output.value : JSON.stringify(output.value);
case "content":
return JSON.stringify(output.value);
}
}
function promptV3ToOpenAI(prompt) {
const results = [];
for (const message of prompt) {
switch (message.role) {
case "system":
results.push({
role: "system",
content: message.content
});
break;
case "assistant":
const textContent = message.content.find(
(part) => part.type === "text"
);
const toolCalls = message.content.filter(
(part) => part.type === "tool-call"
);
results.push({
role: "assistant",
content: textContent?.text || null,
...toolCalls.length > 0 ? {
tool_calls: toolCalls.map((part) => ({
id: part.toolCallId,
function: {
name: part.toolName,
// V3 input can be unknown (usually object), convert to string
arguments: typeof part.input === "string" ? part.input : JSON.stringify(part.input)
},
type: "function"
}))
} : {}
});
break;
case "user":
results.push({
role: "user",
content: message.content.map((part) => {
switch (part.type) {
case "text":
return {
type: "text",
text: part.text
};
case "file":
return {
type: "text",
text: `[file: ${part.mediaType}]`
};
default:
return part;
}
})
});
break;
case "tool":
for (const part of message.content) {
if (part.type === "tool-result") {
results.push({
role: "tool",
tool_call_id: part.toolCallId,
content: formatV3ToolCallOutput(part.output)
});
}
}
break;
}
}
return results;
}
function formatV3ToolCallOutput(output) {
switch (output.type) {
case "text":
return output.value;
case "json":
return typeof output.value === "string" ? output.value : JSON.stringify(output.value);
case "error-text":
return output.value;
case "error-json":
return typeof output.value === "string" ? output.value : JSON.stringify(output.value);
case "content":
return JSON.stringify(output.value);
case "execution-denied":
return output.reason ?? "Execution denied";
}
}
// src/util/currentUnixTime.ts
function currentUnixTime() {
return Date.now() / 1e3;
}
// src/otel/streaming/aggregators.ts
var ToolCallAggregator = class {
constructor() {
__publicField(this, "calls", {});
}
handleChunk(chunk) {
switch (chunk.type) {
case "tool-call":
this.calls[chunk.toolCallId] = {
toolCallType: chunk.toolCallType,
toolCallId: chunk.toolCallId,
toolName: chunk.toolName,
args: chunk.args
};
break;
case "tool-call-delta":
if (!this.calls[chunk.toolCallId]) {
this.calls[chunk.toolCallId] = {
toolCallType: chunk.toolCallType,
toolCallId: chunk.toolCallId,
toolName: chunk.toolName,
args: ""
};
}
this.calls[chunk.toolCallId].args += chunk.argsTextDelta;
break;
}
}
get result() {
return Object.values(this.calls);
}
};
var TextAggregator = class {
constructor() {
__publicField(this, "content", "");
}
feed(chunk) {
if (chunk.type === "text-delta") {
this.content += chunk.textDelta;
}
}
get text() {
return this.content || void 0;
}
};
var StreamStats = class {
constructor() {
__publicField(this, "startTime");
__publicField(this, "timeToFirstToken");
__publicField(this, "_usage");
__publicField(this, "_finishReason");
__publicField(this, "_responseId");
__publicField(this, "_responseModelId");
this.startTime = currentUnixTime();
}
feed(chunk) {
if (this.timeToFirstToken === void 0) {
this.timeToFirstToken = currentUnixTime() - this.startTime;
}
switch (chunk.type) {
case "response-metadata":
if (chunk.id) {
this._responseId = chunk.id;
}
if (chunk.modelId) {
this._responseModelId = chunk.modelId;
}
break;
case "finish":
this._usage = chunk.usage;
this._finishReason = chunk.finishReason;
break;
}
}
get result() {
return {
response: this._responseId || this._responseModelId ? {
id: this._responseId,
modelId: this._responseModelId
} : void 0,
finishReason: this._finishReason,
usage: this._usage
};
}
get firstTokenTime() {
return this.timeToFirstToken;
}
};
var ToolCallAggregatorV2 = class {
constructor() {
__publicField(this, "calls", {});
}
handleChunk(chunk) {
if (chunk.type === "tool-call") {
this.calls[chunk.toolCallId] = chunk;
}
}
get result() {
return Object.values(this.calls);
}
};
var TextAggregatorV2 = class {
constructor() {
__publicField(this, "content", "");
}
feed(chunk) {
switch (chunk.type) {
case "text-start":
this.content = "";
break;
case "text-delta":
this.content += chunk.delta;
break;
case "text-end":
break;
}
}
get text() {
return this.content || void 0;
}
};
var StreamStatsV2 = class {
constructor() {
__publicField(this, "startTime");
__publicField(this, "timeToFirstToken");
__publicField(this, "_usage");
__publicField(this, "_finishReason");
__publicField(this, "_responseMetadata");
this.startTime = currentUnixTime();
}
feed(chunk) {
if (this.timeToFirstToken === void 0) {
this.timeToFirstToken = currentUnixTime() - this.startTime;
}
switch (chunk.type) {
case "response-metadata":
this._responseMetadata = {
id: chunk.id,
modelId: chunk.modelId,
timestamp: chunk.timestamp
};
break;
case "finish":
this._usage = chunk.usage;
this._finishReason = chunk.finishReason;
break;
}
}
get result() {
return {
response: this._responseMetadata,
finishReason: this._finishReason,
usage: this._usage
};
}
get firstTokenTime() {
return this.timeToFirstToken;
}
};
var ToolCallAggregatorV3 = class {
constructor() {
__publicField(this, "calls", {});
__publicField(this, "pendingInputs", {});
}
handleChunk(chunk) {
switch (chunk.type) {
case "tool-call":
this.calls[chunk.toolCallId] = chunk;
break;
case "tool-input-start":
this.pendingInputs[chunk.id] = { toolName: chunk.toolName, input: "" };
break;
case "tool-input-delta":
if (this.pendingInputs[chunk.id]) {
this.pendingInputs[chunk.id].input += chunk.delta;
}
break;
case "tool-input-end":
break;
}
}
get result() {
return Object.values(this.calls);
}
};
var TextAggregatorV3 = class {
constructor() {
__publicField(this, "content", "");
}
feed(chunk) {
switch (chunk.type) {
case "text-start":
this.content = "";
break;
case "text-delta":
this.content += chunk.delta;
break;
case "text-end":
break;
}
}
get text() {
return this.content || void 0;
}
};
var StreamStatsV3 = class {
constructor() {
__publicField(this, "startTime");
__publicField(this, "timeToFirstToken");
__publicField(this, "_usage");
__publicField(this, "_finishReason");
__publicField(this, "_responseMetadata");
this.startTime = currentUnixTime();
}
feed(chunk) {
if (this.timeToFirstToken === void 0) {
this.timeToFirstToken = currentUnixTime() - this.startTime;
}
switch (chunk.type) {
case "response-metadata":
this._responseMetadata = {
id: chunk.id,
modelId: chunk.modelId,
timestamp: chunk.timestamp
};
break;
case "finish":
this._usage = chunk.usage;
this._finishReason = chunk.finishReason;
break;
}
}
get result() {
return {
response: this._responseMetadata,
finishReason: this._finishReason,
usage: this._usage
};
}
get firstTokenTime() {
return this.timeToFirstToken;
}
};
// src/otel/middleware.ts
var appendPromptMetadataToSpan = (span, messages) => {
const lastMessage = messages?.[messages.length - 1];
let axiomMeta;
if ("providerMetadata" in lastMessage) {
axiomMeta = lastMessage?.providerMetadata?._axiomMeta;
} else if ("providerOptions" in lastMessage) {
axiomMeta = lastMessage?.providerOptions?._axiomMeta;
}
if (axiomMeta) {
if (axiomMeta.id) span.setAttribute(Attr.GenAI.PromptMetadata.ID, axiomMeta.id);
if (axiomMeta.name) span.setAttribute(Attr.GenAI.PromptMetadata.Name, axiomMeta.name);
if (axiomMeta.slug) span.setAttribute(Attr.GenAI.PromptMetadata.Slug, axiomMeta.slug);
if (axiomMeta.version) span.setAttribute(Attr.GenAI.PromptMetadata.Version, axiomMeta.version);
}
};
function setStreamStepAttributes(span, values) {
const inputTokens = ensureNumber(values.inputTokens);
if (inputTokens !== void 0) {
span.setAttribute(Attr.GenAI.Usage.InputTokens, inputTokens);
}
const outputTokens = ensureNumber(values.outputTokens);
if (outputTokens !== void 0) {
span.setAttribute(Attr.GenAI.Usage.OutputTokens, outputTokens);
}
if (values.finishReason) {
span.setAttribute(Attr.GenAI.Response.FinishReasons, JSON.stringify([values.finishReason]));
}
}
function axiomAIMiddlewareV1() {
return {
wrapGenerate: async ({ doGenerate, params, model }) => {
return withSpanHandling(
model.modelId,
async (span, commonContext, _lease) => {
const context3 = commonContext;
appendPromptMetadataToSpan(span, params.prompt);
setScopeAttributes(span);
setPreCallAttributesV1(span, params, context3, model);
const res = await doGenerate();
context3.rawCall = res.rawCall;
await setPostCallAttributesV1(span, res, context3, model);
return res;
},
{ version: "v1" }
);
},
wrapStream: async ({ doStream, params, model }) => {
return withSpanHandling(
model.modelId,
async (span, commonContext, lease) => {
const context3 = commonContext;
appendPromptMetadataToSpan(span, params.prompt);
setScopeAttributes(span);
setPreCallAttributesV1(span, params, context3, model);
const { stream, ...head } = await doStream();
const childSpan = createStreamChildSpan(span, `chat ${model.modelId} stream`);
const stats = new StreamStats();
const toolAggregator = new ToolCallAggregator();
const textAggregator = new TextAggregator();
return {
...head,
stream: stream.pipeThrough(
new TransformStream({
transform(chunk, controller) {
try {
stats.feed(chunk);
toolAggregator.handleChunk(chunk);
textAggregator.feed(chunk);
controller.enqueue(chunk);
} catch (err) {
classifyToolError(err, childSpan);
childSpan.end();
if (lease.owned) lease.end();
controller.error(err);
}
},
async flush(controller) {
try {
const statsResult = stats.result;
await setPostCallAttributesV1(
span,
{
...head,
...statsResult,
toolCalls: toolAggregator.result.length > 0 ? toolAggregator.result : void 0,
text: textAggregator.text
},
context3,
model
);
setStreamStepAttributes(childSpan, {
inputTokens: statsResult.usage?.promptTokens,
outputTokens: statsResult.usage?.completionTokens,
finishReason: statsResult.finishReason
});
childSpan.end();
if (lease.owned) lease.end();
controller.terminate();
} catch (err) {
classifyToolError(err, childSpan);
childSpan.end();
if (lease.owned) lease.end();
controller.error(err);
}
}
})
)
};
},
{ streaming: true, version: "v1" }
// Don't auto-end span, we'll end it when stream completes
);
}
};
}
function axiomAIMiddleware(config) {
if (config.model.specificationVersion === "v1") {
return axiomAIMiddlewareV1();
} else if (config.model.specificationVersion === "v2") {
return axiomAIMiddlewareV2();
} else if (config.model.specificationVersion === "v3") {
return axiomAIMiddlewareV3();
} else {
console.warn(
// @ts-expect-error - not allowed at type level, but users can still do it...
`Unsupported model specification version: ${JSON.stringify(config.model.specificationVersion)}. Creating no-op middleware instead.`
);
return {};
}
}
function axiomAIMiddlewareV2() {
return {
wrapGenerate: async ({ doGenerate, params, model }) => {
return withSpanHandling(
model.modelId,
async (span, commonContext, _lease) => {
const context3 = commonContext;
appendPromptMetadataToSpan(span, params.prompt);
setScopeAttributes(span);
setPreCallAttributesV2(span, params, context3, model);
const res = await doGenerate();
await setPostCallAttributesV2(span, res, context3, model);
return res;
},
{ version: "v2" }
);
},
wrapStream: async ({ doStream, params, model }) => {
return withSpanHandling(
model.modelId,
async (span, commonContext, lease) => {
const context3 = commonContext;
appendPromptMetadataToSpan(span, params.prompt);
setScopeAttributes(span);
setPreCallAttributesV2(span, params, context3, model);
const ret = await doStream();
const childSpan = createStreamChildSpan(span, `chat ${model.modelId} stream`);
const stats = new StreamStatsV2();
const toolAggregator = new ToolCallAggregatorV2();
const textAggregator = new TextAggregatorV2();
return {
...ret,
stream: ret.stream.pipeThrough(
new TransformStream({
transform(chunk, controller) {
try {
stats.feed(chunk);
toolAggregator.handleChunk(chunk);
textAggregator.feed(chunk);
controller.enqueue(chunk);
} catch (err) {
classifyToolError(err, childSpan);
childSpan.end();
if (lease.owned) lease.end();
controller.error(err);
}
},
async flush(controller) {
try {
const statsResult = stats.result;
const streamResult = {
...statsResult,
content: [
...textAggregator.text ? [{ type: "text", text: textAggregator.text }] : [],
...toolAggregator.result
]
};
await setPostCallAttributesV2(span, streamResult, context3, model);
setStreamStepAttributes(childSpan, {
inputTokens: statsResult.usage?.inputTokens,
outputTokens: statsResult.usage?.outputTokens,
finishReason: statsResult.finishReason
});
childSpan.end();
if (lease.owned) lease.end();
controller.terminate();
} catch (err) {
classifyToolError(err, childSpan);
childSpan.end();
if (lease.owned) lease.end();
controller.error(err);
}
}
})
)
};
},
{ streaming: true, version: "v2" }
// Don't auto-end span, we'll end it when stream completes
);
}
};
}
function setPreCallAttributesV1(span, options, context3, model) {
const redactionPolicy = getRedactionPolicy();
const {
prompt,
maxTokens,
frequencyPenalty,
presencePenalty,
temperature,
topP,
topK,
seed,
stopSequences,
responseFormat,
mode
} = options;
const processedPrompt = promptV1ToOpenAI(prompt);
context3.originalPrompt = processedPrompt;
handleMaybeRedactedAttribute(
span,
Attr.GenAI.Input.Messages,
JSON.stringify(sanitizeMultimodalContent(processedPrompt)),
redactionPolicy.captureMessageContent
);
setBaseAttributes(span, model.provider, model.modelId);
const outputType = determineOutputTypeV1({ responseFormat, mode });
if (outputType) {
span.setAttribute(Attr.GenAI.Output.Type, outputType);
}
setRequestParameterAttributes(span, {
maxTokens,
frequencyPenalty,
presencePenalty,
temperature,
topP,
topK,
seed,
stopSequences
});
}
async function setPostCallAttributesV1(span, result, context3, _model) {
const redactionPolicy = getRedactionPolicy();
if (result.toolCalls && result.toolCalls.length > 0) {
const originalPrompt = context3.originalPrompt || [];
const normalizedToolCalls = normalizeV1ToolCalls(result.toolCalls);
const toolResultsMap = context3.rawCall?.rawPrompt ? extractToolResultsFromRawPrompt(context3.rawCall.rawPrompt) : /* @__PURE__ */ new Map();
const updatedPrompt = appendToolCalls(
originalPrompt,
normalizedToolCalls,
toolResultsMap,
result.text
);
handleMaybeRedactedAttribute(
span,
Attr.GenAI.Input.Messages,
JSON.stringify(sanitizeMultimodalContent(updatedPrompt)),
redactionPolicy.captureMessageContent
);
}
if (result.text) {
const completion = createSimpleCompletion({
text: resu