axiom/dist/index.js

Axiom AI SDK provides - an API to wrap your AI calls with observability instrumentation. - offline evals

import {
  Attr,
  RedactionPolicy,
  axiomAIMiddleware,
  axiomAIMiddlewareV1,
  axiomAIMiddlewareV2,
  classifyToolError,
  createAppScope,
  createStartActiveSpan,
  getGlobalTracer,
  getRedactionPolicy,
  getTracer,
  handleMaybeRedactedAttribute,
  initAxiomAI,
  resetAxiomAI,
  setAxiomBaseAttributes,
  setScopeAttributes,
  withSpan,
  wrapAISDKModel
} from "./chunk-PZCJLQFO.js";
import "./chunk-X2LH7XLM.js";
import {
  init_esm_shims
} from "./chunk-4VNFFUM5.js";

// src/index.ts
init_esm_shims();

// src/otel/wrapTool.ts
init_esm_shims();
import "@opentelemetry/api";

// src/util/typedEntries.ts
init_esm_shims();
function typedEntries(obj) {
  return Object.entries(obj);
}

// src/otel/wrapTool.ts
function wrapTool(toolName, tool) {
  if (!tool || typeof tool !== "object") {
    console.error("Invalid tool provided to wrapTool, returning unwrapped tool");
    return tool;
  }
  if (!("execute" in tool) || typeof tool.execute !== "function") {
    console.error(
      "Cannot wrap a tool that does not have an execute method, returning unwrapped tool"
    );
    return tool;
  }
  const originalExecute = tool.execute;
  return {
    ...tool,
    execute: async (...executeArgs) => {
      const [args, opts] = executeArgs;
      const tracer = getTracer();
      const startActiveSpan = createStartActiveSpan(tracer);
      const spanName = `${Attr.GenAI.Operation.Name_Values.ExecuteTool} ${toolName}`;
      return startActiveSpan(spanName, null, async (span) => {
        const redactionPolicy = getRedactionPolicy();
        setAxiomBaseAttributes(span);
        setScopeAttributes(span);
        const toolCallId = opts && typeof opts === "object" ? opts.toolCallId || opts.toolCall?.toolCallId : void 0;
        if (toolCallId) {
          span.setAttribute(Attr.GenAI.Tool.CallID, toolCallId);
        }
        span.setAttribute(Attr.GenAI.Operation.Name, Attr.GenAI.Operation.Name_Values.ExecuteTool);
        span.setAttribute(Attr.GenAI.Tool.Name, toolName);
        const type = args && typeof args === "object" && "type" in args ? args.type : "function";
        span.setAttribute(Attr.GenAI.Tool.Type, type);
        if (tool.description) {
          span.setAttribute(Attr.GenAI.Tool.Description, tool.description);
        }
        if (redactionPolicy.mirrorToolPayloadOnToolSpan) {
          try {
            handleMaybeRedactedAttribute(
              span,
              Attr.GenAI.Tool.Arguments,
              JSON.stringify(args),
              redactionPolicy.captureMessageContent
            );
          } catch (_error) {
            span.setAttribute(Attr.GenAI.Tool.Arguments, "[Unable to serialize arguments]");
          }
        }
        try {
          const result = await originalExecute(args, opts);
          if (redactionPolicy.mirrorToolPayloadOnToolSpan) {
            try {
              handleMaybeRedactedAttribute(
                span,
                Attr.GenAI.Tool.Message,
                JSON.stringify(result),
                redactionPolicy.captureMessageContent
              );
            } catch (_error) {
              span.setAttribute(Attr.GenAI.Tool.Message, "[Unable to serialize result]");
            }
          }
          return result;
        } catch (err) {
          classifyToolError(err, span);
          throw err;
        }
      });
    }
  };
}
function wrapTools(tools) {
  if (!tools || typeof tools !== "object") {
    console.error("Invalid tools object provided to wrapTools");
    return tools;
  }
  const wrappedTools = {};
  for (const [toolName, tool] of typedEntries(tools)) {
    wrappedTools[toolName] = wrapTool(toolName, tool);
  }
  return wrappedTools;
}

// src/template.ts
init_esm_shims();
import { Type as TypeBox } from "@sinclair/typebox";
var createTemplateType = (schema) => schema;
var createString = (...args) => createTemplateType(TypeBox.String(...args));
var createNumber = (...args) => createTemplateType(TypeBox.Number(...args));
var createInteger = (...args) => createTemplateType(TypeBox.Integer(...args));
var createBoolean = (...args) => createTemplateType(TypeBox.Boolean(...args));
var createNull = (...args) => createTemplateType(TypeBox.Null(...args));
var createLiteral = (value, ...args) => createTemplateType(TypeBox.Literal(value, ...args));
var createUndefined = (...args) => createTemplateType(TypeBox.Undefined(...args));
var createArray = (items, ...args) => createTemplateType(TypeBox.Array(items, ...args));
var createObject = (properties, ...args) => createTemplateType(TypeBox.Object(properties, ...args));
var createRecord = (value, ...args) => createTemplateType(TypeBox.Record(TypeBox.String(), value, ...args));
var createTuple = (types, ...args) => {
  return createTemplateType(TypeBox.Tuple(types, ...args));
};
var createOptional = (schema, ...args) => createTemplateType(TypeBox.Optional(schema, true, ...args));
var createUnion = (schemas, ...args) => createTemplateType(TypeBox.Union(schemas, ...args));
var Type = {
  // Primitives
  String: createString,
  Number: createNumber,
  Integer: createInteger,
  Boolean: createBoolean,
  Null: createNull,
  Undefined: createUndefined,
  // Literals
  Literal: createLiteral,
  // Collections - only accept Template types
  Array: createArray,
  Object: createObject,
  Record: createRecord,
  Tuple: createTuple,
  // Modifiers - only accept Template types
  Optional: createOptional,
  Union: createUnion
};

// src/prompt/index.ts
init_esm_shims();
var getParser = async () => {
  const handlebars = await import("./handlebars-AIP3BZD5.js").then((m) => m.handlebarsParse);
  return handlebars;
};
var parse = async (prompt, {
  context
}) => {
  const messagesPromises = prompt.messages.map(async (message) => {
    const parser = await getParser();
    return {
      ...message,
      content: await parser(message.content, { context })
    };
  });
  const parsedMessages = await Promise.all(messagesPromises);
  const promptMetadata = {
    id: prompt.promptId,
    name: prompt.name,
    slug: prompt.slug,
    version: prompt.version
  };
  if (parsedMessages.length > 0) {
    const lastMessage = parsedMessages[parsedMessages.length - 1];
    lastMessage.providerOptions = {
      ...lastMessage.providerOptions,
      _axiomMeta: promptMetadata
    };
    lastMessage.providerMetadata = {
      ...lastMessage.providerMetadata,
      _axiomMeta: promptMetadata
    };
  }
  const messages = new Proxy(parsedMessages, {
    get(target, prop, receiver) {
      if (prop === "_axiomMeta") {
        return promptMetadata;
      }
      return Reflect.get(target, prop, receiver);
    }
  });
  return {
    ...prompt,
    messages
  };
};
export {
  RedactionPolicy,
  axiomAIMiddleware,
  axiomAIMiddlewareV1,
  axiomAIMiddlewareV2,
  createAppScope,
  Type as experimental_Template,
  Type as experimental_Type,
  parse as experimental_parse,
  getGlobalTracer,
  initAxiomAI,
  resetAxiomAI,
  withSpan,
  wrapAISDKModel,
  wrapTool,
  wrapTools
};
//# sourceMappingURL=index.js.map