Observe

Learn how to observe your deployed AI capabilities in production using Axiom’s AI SDK to capture telemetry.

The Observe stage is about understanding how your deployed generative AI capabilities perform in the real world. After creating and evaluating a capability, observing its production behavior is crucial for identifying unexpected issues, tracking costs, and gathering the data needed for future improvements.

Capturing telemetry with the @axiomhq/ai SDK

The foundation of the Observe stage is Axiom’s SDK, which integrates with your app to capture detailed OpenTelemetry traces for every AI interaction.

The initial release of @axiomhq/ai is focused on providing deep integration with TypeScript applications, particularly those using Vercel’s AI SDK to interact with frontier models.

Instrumenting AI SDK calls

The easiest way to get started is by wrapping your existing AI model client. The @axiomhq/ai package provides helper functions for popular libraries like Vercel’s AI SDK.

The wrapAISDKModel function takes an existing AI model object and returns an instrumented version that will automatically generate trace data for every call.

OpenAI

// src/shared/openai.ts

import { createOpenAI } from '@ai-sdk/openai';
import { wrapAISDKModel } from '@axiomhq/ai';

const openaiProvider = createOpenAI({
  apiKey: process.env.OPENAI_API_KEY,
});

// Wrap the model to enable automatic tracing
export const gpt4o = wrapAISDKModel(openaiProvider('gpt-4o'));
export const gpt4oMini = wrapAISDKModel(openaiProvider('gpt-4o-mini'));

Adding context with withSpan

While wrapAISDKModel handles the automatic instrumentation, the withSpan function allows you to add crucial business context to your traces. It creates a parent span around your LLM call and attaches metadata about the capability and step being executed.

// src/app/page.tsx

import { withSpan } from '@axiomhq/ai';
import { generateText } from 'ai';
import { gpt4o } from '@/shared/openai';

export default async function Page() {
  const userId = 123;

  // Use withSpan to define the capability and step
  const res = await withSpan({ capability: 'get_capital', step: 'generate_answer' }, (span) => {
    // You have access to the OTel span to add custom attributes
    span.setAttribute('user_id', userId);

    return generateText({
      model: gpt4o, // Use the wrapped model
      messages: [
        {
          role: 'user',
          content: 'What is the capital of Spain?',
        },
      ],
    });
  });

  return <p>{res.text}</p>;
}

Instrumenting tool calls with wrapTool

For many AI capabilities, the LLM call is only part of the story. If your capability uses tools to interact with external data or services, observing the performance and outcome of those tools is critical. The Axiom AI SDK provides the wrapTool and wrapTools functions to automatically instrument your Vercel AI SDK tool definitions.

The wrapTool helper takes your tool’s name and its definition and returns an instrumented version. This wrapper creates a dedicated child span for every tool execution, capturing its arguments, output, and any errors.

// src/app/generate-text/page.tsx
import { tool } from 'ai';
import { z } from 'zod';
import { wrapTool } from '@axiomhq/ai';
import { generateText } from 'ai';
import { gpt4o } from '@/shared/openai';

// In your generateText call, provide wrapped tools
const { text, toolResults } = await generateText({
  model: gpt4o,
  messages: [
    { role: 'system', content: 'You are a helpful assistant.' },
    { role: 'user', content: 'How do I get from Paris to Berlin?' },
  ],
  tools: {
    // Wrap each tool with its name
    findDirections: wrapTool(
      'findDirections', // The name of the tool
      tool({
        description: 'Find directions to a location',
        inputSchema: z.object({
          from: z.string(),
          to: z.string(),
        }),
        execute: async (params) => {
          // Your tool logic here...
          return { directions: `To get from ${params.from} to ${params.to}, use a teleporter.` };
        },
      })
    )
  }
});

Complete instrumentation example

Full end-to-end code example

Here’s how all three instrumentation functions work together in a single, real-world example:

// src/app/page.tsx

import { withSpan, wrapAISDKModel, wrapTool } from '@axiomhq/ai';
import { generateText, tool } from 'ai';
import { createOpenAI } from '@ai-sdk/openai';
import { z } from 'zod';

// 1. Create and wrap the AI model client
const openaiProvider = createOpenAI({
  apiKey: process.env.OPENAI_API_KEY,
});
const gpt4o = wrapAISDKModel(openaiProvider('gpt-4o'));

// 2. Define and wrap your tool(s)
const findDirectionsTool = wrapTool(
  'findDirections', // The tool name must be passed to the wrapper
  tool({
    description: 'Find directions to a location',
    inputSchema: z.object({ from: z.string(), to: z.string() }),
    execute: async ({ from, to }) => ({
      directions: `To get from ${from} to ${to}, use a teleporter.`,
    }),
  })
);

// 3. In your application logic, use `withSpan` to add context
//    and call the AI model with your wrapped tools.
export default async function Page() {
  const userId = 123;

  const { text } = await withSpan({ capability: 'get_directions', step: 'generate_ai_response' }, async (span) => {
    // You have access to the OTel span to add custom attributes
    span.setAttribute('user_id', userId);

    return generateText({
      model: gpt4o, // Use the wrapped model
      messages: [
        { role: 'system', content: 'You are a helpful assistant.' },
        { role: 'user', content: 'How do I get from Paris to Berlin?' },
      ],
      tools: {
        findDirections: findDirectionsTool, // Use the wrapped tool
      },
    });
  });

  return <p>{text}</p>;
}

This demonstrates the three key steps to rich observability:

1. wrapAISDKModel: Automatically captures telemetry for the LLM provider call
2. wrapTool: Instruments the tool execution with detailed spans
3. withSpan: Creates a parent span that ties everything together under a business capability

Setting up instrumentation

The Axiom AI SDK is built on the OpenTelemetry standard. To send traces, you need to configure a Node.js or edge-compatible tracer that exports data to Axiom.

Configuring the tracer

You must configure an OTLP trace exporter pointing to your Axiom instance. This is typically done in a dedicated instrumentation file that’s loaded before your application starts.

// src/instrumentation.ts

import 'dotenv/config'; // Make sure to load environment variables
import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-http';
import { resourceFromAttributes } from '@opentelemetry/resources';
import { NodeSDK } from '@opentelemetry/sdk-node';
import { SimpleSpanProcessor } from '@opentelemetry/sdk-trace-node';
import { ATTR_SERVICE_NAME } from '@opentelemetry/semantic-conventions';
import { trace } from "@opentelemetry/api";
import { initAxiomAI } from '@axiomhq/ai';

const tracer = trace.getTracer("my-tracer");

// Configure the NodeSDK to export traces to your Axiom dataset
const sdk = new NodeSDK({
  resource: resourceFromAttributes({
    [ATTR_SERVICE_NAME]: 'my-ai-app', // Replace with your service name
  }),
  spanProcessor: new SimpleSpanProcessor(
    new OTLPTraceExporter({
      url: `https://api.axiom.co/v1/traces`,
      headers: {
        Authorization: `Bearer ${process.env.AXIOM_TOKEN}`,
        'X-Axiom-Dataset': process.env.AXIOM_DATASET!,
      },
    }),
  ),
});

// Start the SDK
sdk.start();

// Initialize the Axiom AI SDK with the configured tracer
initAxiomAI({ tracer });

Set your Axiom credentials as environment variables AXIOM_TOKEN and AXIOM_DATASET.

Understanding your AI telemetry

Once instrumented, every LLM call will send a detailed span to your Axiom dataset. These spans are enriched with standardized gen_ai.* attributes that make your AI interactions easy to query and analyze.

Key attributes include:

• gen_ai.capability.name: The high-level capability name you defined in withSpan.
• gen_ai.step.name: The specific step within the capability.
• gen_ai.request.model: The model requested for the completion.
• gen_ai.response.model: The model that actually fulfilled the request.
• gen_ai.usage.input_tokens: The number of tokens in the prompt.
• gen_ai.usage.output_tokens: The number of tokens in the generated response.
• gen_ai.prompt: The full, rendered prompt or message history sent to the model (as a JSON string).
• gen_ai.completion: The full response from the model, including tool calls (as a JSON string).
• gen_ai.response.finish_reasons: The reason the model stopped generating tokens (e.g., stop, tool-calls).
• gen_ai.tool.name: The name of the executed tool.
• gen_ai.tool.arguments: The arguments passed to the tool (as a JSON string).
• gen_ai.tool.message: The result returned by the tool (as a JSON string).

Visualizing traces in the console

Visualizing and making sense of this telemetry data is a core part of the Axiom Console experience.

• Coming soon
  A dedicated AI Trace Waterfall view will visualize single and multi-step LLM workflows, with clear input/output inspection at each stage.
• Coming soon
  A pre-built Gen AI OTel Dashboard will automatically appear for any dataset receiving AI telemetry. It will feature elements for tracking cost per invocation, time-to-first-token, call counts by model, and error rates.

What’s next?

Now that you are capturing and analyzing production telemetry, the next step is to use these insights to improve your capability.

Learn more in the Iterate page.