@axiomhq/ai/dist/index.d.ts

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

import { Tracer, Span } from '@opentelemetry/api';
import { LanguageModelV1 } from '@ai-sdk/providerv1';
import { LanguageModelV2, LanguageModelV2Middleware } from '@ai-sdk/providerv2';
import { LanguageModelV1Middleware } from 'aiv4';
import * as _sinclair_typebox from '@sinclair/typebox';
import { TSchema as TSchema$1, Static, Type as Type$1 } from '@sinclair/typebox';
export { A as experimental_AxiomReporter } from './reporter-B9Ly79ok.js';
import 'vitest';
import 'vitest/node.js';

/**
 * Register this in your `instrumentation.ts` to set up @axiomhq/ai.
 * This function stores the tracer's scope information globally to enable Context Propagation
 * and Instrumentation Scope. The tracer will be available across all execution contexts including Next.js.
 *
 * This function is idempotent - calling it multiple times with the same scope has no additional cost.
 *
 * @param config
 * @param config.tracer - The tracer that you are using in your application.
 */
declare function initAxiomAI(config: {
    tracer: Tracer;
}): void;
/**
 * Get a tracer using the globally stored scope information
 * Fall back to package.json defaults if not set
 */
declare function getGlobalTracer(): Tracer;
/**
 * Reset AxiomAI configuration (useful for testing)
 */
declare function resetAxiomAI(): void;

/**
 * Wraps an AI SDK model to provide OpenTelemetry instrumentation.
 *
 * Supports both AI SDK v4 (LanguageModelV1) and v5 (LanguageModelV2) models.
 *
 * @param model - Language model implementing LanguageModelV1 or LanguageModelV2 interface
 * @returns Wrapped model with identical interface but added instrumentation
 */
declare function wrapAISDKModel<T extends LanguageModelV1 | LanguageModelV2>(model: T): T;

type WithSpanMeta = {
    capability: string;
    step: string;
};
/**
 * Wrap this around Vercel AI SDK functions such as `generateText` to wrap them in a span.
 */
declare function withSpan<Return>(meta: WithSpanMeta, fn: (span: Span) => Promise<Return>, opts?: {
    tracer?: Tracer;
}): Promise<Return>;

interface ToolLike {
    execute?: (...args: any[]) => any;
    description?: string;
    [key: string]: any;
}
/**
 * Wraps a tool to create child spans when the tool's execute method is called.
 *
 * @param toolName The name of the tool (key from the tools object) - span name will be `execute_tool <toolName>`
 * @param tool The tool to wrap
 * @returns The same tool but with a wrapped execute method that creates spans
 */
declare function wrapTool<T extends ToolLike>(toolName: string, tool: T): T;
/**
 * Wraps multiple tools to create child spans when their execute methods are called.
 *
 * @param tools An object containing tools to wrap
 * @returns The same object with all tools wrapped
 */
declare function wrapTools<T extends Record<string, ToolLike>>(tools: T): T;

interface AxiomTelemetryConfig {
}
/**
 * Creates Axiom telemetry middleware for LanguageModelV1
 */
declare function axiomAIMiddlewareV1(): LanguageModelV1Middleware;
/**
 * Creates unified Axiom telemetry middleware that works with both V1 and V2 models
 */
declare function axiomAIMiddleware(config: {
    model: LanguageModelV1;
}): LanguageModelV1Middleware;
declare function axiomAIMiddleware(config: {
    model: LanguageModelV2;
}): LanguageModelV2Middleware;
/**
 * Creates Axiom telemetry middleware for LanguageModelV2
 */
declare function axiomAIMiddlewareV2(): LanguageModelV2Middleware;

/**
 * Template-friendly TypeBox types for use with Handlebars.
 *
 * These types are designed to be easily serializable and work well in template contexts
 * where data needs to be rendered as strings or used in conditionals and loops.
 *
 * All composite types (Array, Object, etc.) only accept other Template types.
 */
declare const SchemaBrand: unique symbol;
type TSchema<T extends TSchema$1 = TSchema$1> = T & {
    [SchemaBrand]: true;
};
/**
 * Utility type to infer TypeScript types from {@link TSchema}.
 */
type InferSchema<T extends TSchema> = Static<T>;
/**
 * Utility type to infer context types from {@link Prompt} arguments.
 *
 * Used with {@link parse} function to ensure type safety when providing context values.
 */
type InferContext<T extends Record<string, TSchema>> = Static<ReturnType<typeof createObject<T>>>;
type OmitFirst<T extends any[]> = T extends [any, ...infer R] ? R : never;
declare const createObject: <T extends Record<string, TSchema>>(properties: T, ...args: OmitFirst<Parameters<typeof Type$1.Object>>) => TSchema<_sinclair_typebox.TObject<T>>;
declare const Type: {
    readonly String: (options?: _sinclair_typebox.StringOptions | undefined) => TSchema<_sinclair_typebox.TString>;
    readonly Number: (options?: _sinclair_typebox.NumberOptions | undefined) => TSchema<_sinclair_typebox.TNumber>;
    readonly Integer: (options?: _sinclair_typebox.IntegerOptions | undefined) => TSchema<_sinclair_typebox.TInteger>;
    readonly Boolean: (options?: _sinclair_typebox.SchemaOptions | undefined) => TSchema<_sinclair_typebox.TBoolean>;
    readonly Null: (options?: _sinclair_typebox.SchemaOptions | undefined) => TSchema<_sinclair_typebox.TNull>;
    readonly Undefined: (options?: _sinclair_typebox.SchemaOptions | undefined) => TSchema<_sinclair_typebox.TUndefined>;
    readonly Literal: <T extends string | number | boolean>(value: T, options?: _sinclair_typebox.SchemaOptions | undefined) => TSchema<_sinclair_typebox.TLiteral<T>>;
    readonly Array: <T extends TSchema>(items: T, options?: _sinclair_typebox.ArrayOptions | undefined) => TSchema<_sinclair_typebox.TArray<T>>;
    readonly Object: <T extends Record<string, TSchema>>(properties: T, options?: _sinclair_typebox.ObjectOptions | undefined) => TSchema<_sinclair_typebox.TObject<T>>;
    readonly Record: <V extends TSchema>(value: V, options?: _sinclair_typebox.ObjectOptions | undefined) => TSchema<_sinclair_typebox.TRecord<_sinclair_typebox.TString, V>>;
    readonly Tuple: <T extends TSchema[]>(types: [...T], options?: _sinclair_typebox.SchemaOptions | undefined) => TSchema<_sinclair_typebox.TTuple<T>>;
    readonly Optional: <T extends TSchema>(schema: T) => TSchema<T extends _sinclair_typebox.TOptional<infer S extends TSchema$1> ? _sinclair_typebox.TOptional<S> : _sinclair_typebox.Ensure<_sinclair_typebox.TOptional<T>>>;
    readonly Union: <T extends TSchema[]>(schemas: [...T], options?: _sinclair_typebox.SchemaOptions | undefined) => TSchema<_sinclair_typebox.Union<T>>;
};

/**
 * Configuration options for language model generation.
 *
 * @experimental This API is experimental and may change in future versions.
 */
type Options = {
    /** Maximum number of tokens to generate */
    maxOutputTokens?: number;
    /** Controls randomness in generation (0.0 to 2.0) */
    temperature?: number;
    /** Controls nucleus sampling for token selection */
    topP?: number;
    /** Controls top-k sampling for token selection */
    topK?: number;
    /** Penalty for repeating content (presence) */
    presencePenalty?: number;
    /** Penalty for repeating tokens (frequency) */
    frequencyPenalty?: number;
    /** Sequences that will stop generation */
    stopSequences?: string[];
    /** Seed for deterministic generation */
    seed?: number;
    /** Maximum number of retry attempts */
    maxRetries?: number;
};
/**
 * Complete prompt definition with all metadata and versioning information.
 *
 * Extended version of {@link PromptInput} with additional versioning and identification.
 * Used with {@link parse} to process templates and generate {@link ParsedMessagesArray}.
 *
 * @experimental This API is experimental and may change in future versions.
 */
type Prompt = {
    /** Human-readable name for the prompt */
    name: string;
    /** Immutable user-defined identifier for the prompt */
    slug: string;
    /** Array of messages that make up the conversation */
    messages: {
        role: 'system' | 'user' | 'assistant' | 'tool';
        content: string;
    }[];
    /** The language model to use for this prompt */
    model: string;
    /** Optional generation parameters */
    options?: Options;
    /** {@link TSchema} format arguments for API communication */
    arguments: Record<string, TSchema>;
    /** Optional description of the prompt's purpose */
    description?: string;
    /** Version identifier for the prompt */
    version: string;
    /** Optional unique identifier for the prompt */
    promptId?: string;
};

/**
 * Metadata structure for Axiom prompt tracking and instrumentation.
 *
 * @experimental This API is experimental and may change in future versions.
 */
interface AxiomPromptMetadata {
    /** Unique identifier for the prompt */
    id?: string;
    /** Human-readable name for the prompt */
    name?: string;
    /** Inmutable user-defined identifier for the prompt */
    slug?: string;
    /** Version identifier for the prompt */
    version?: string;
}
type PromptMessage = Prompt['messages'][number];
/**
 * A message that has been processed and attached metadata.
 *
 * Extends the base message type with {@link AxiomPromptMetadata} for instrumentation tracking.
 *
 * @experimental This API is experimental and may change in future versions.
 */
interface ParsedMessage extends PromptMessage {
    /** Provider options that may include Axiom metadata for propagation */
    providerOptions?: {
        /** Internal {@link AxiomPromptMetadata} for prompt tracking */
        _axiomMeta?: AxiomPromptMetadata;
        [key: string]: any;
    };
    /** Provider metadata that may include Axiom metadata */
    providerMetadata?: {
        /** Internal {@link AxiomPromptMetadata} for prompt tracking */
        _axiomMeta?: AxiomPromptMetadata;
        [key: string]: any;
    };
}
type ExtendMessage<T> = T extends object ? T & ParsedMessage : T;
type ExtendedMessages<T extends readonly unknown[]> = {
    [K in keyof T]: ExtendMessage<T[K]>;
};
/**
 * Array of parsed messages with attached Axiom metadata for prompt tracking.
 *
 * Returned by {@link parse} function and contains {@link ParsedMessage} items with
 * accessible {@link AxiomPromptMetadata} via proxy.
 *
 * @experimental This API is experimental and may change in future versions.
 */
type ParsedMessagesArray<T extends readonly PromptMessage[] = readonly PromptMessage[]> = ExtendedMessages<T> & {
    /** {@link AxiomPromptMetadata} accessible via proxy for prompt tracking */
    _axiomMeta: AxiomPromptMetadata;
};

/**
 * Parses a prompt template by replacing variables with provided context values.
 *
 * This function processes Handlebars templates in prompt messages and attaches metadata
 * for instrumentation and tracking.
 *
 * @experimental This API is experimental and may change in future versions.
 *
 * @param prompt - The {@link Prompt} template to parse
 * @param options - Parsing options
 * @param options.context - Context values to substitute into the template
 * @returns Promise that resolves to the parsed prompt with processed messages;
 */
declare const parse: <TPrompt extends Prompt, TMessages extends TPrompt["messages"] = TPrompt["messages"]>(prompt: TPrompt & {
    messages: TMessages;
}, { context, }: {
    context: InferContext<TPrompt["arguments"]>;
}) => Promise<Omit<TPrompt, "messages"> & {
    messages: ParsedMessagesArray<TMessages>;
}>;

export { type AxiomTelemetryConfig, axiomAIMiddleware, axiomAIMiddlewareV1, axiomAIMiddlewareV2, type AxiomPromptMetadata as experimental_AxiomPromptMetadata, type InferContext as experimental_InferContext, type InferSchema as experimental_InferSchema, type ParsedMessage as experimental_ParsedMessage, type ParsedMessagesArray as experimental_ParsedMessagesArray, type Prompt as experimental_Prompt, type TSchema as experimental_TSchema, Type as experimental_Template, Type as experimental_Type, parse as experimental_parse, getGlobalTracer, initAxiomAI, resetAxiomAI, withSpan, wrapAISDKModel, wrapTool, wrapTools };