Instrument LLM calls from Vercel AI SDK¶

SCHEMA >
    `model` LowCardinality(String) `json:$.model` DEFAULT 'unknown',
    `messages` Array(Map(String, String)) `json:$.messages[:]` DEFAULT [],
    `user` String `json:$.user` DEFAULT 'unknown',
    `start_time` DateTime `json:$.start_time` DEFAULT now(),
    `end_time` DateTime `json:$.end_time` DEFAULT now(),
    `id` String `json:$.id` DEFAULT '',
    `stream` Bool `json:$.stream` DEFAULT false,
    `call_type` LowCardinality(String) `json:$.call_type` DEFAULT 'unknown',
    `provider` LowCardinality(String) `json:$.provider` DEFAULT 'unknown',
    `api_key` String `json:$.api_key` DEFAULT '',
    `log_event_type` LowCardinality(String) `json:$.log_event_type` DEFAULT 'unknown',
    `llm_api_duration_ms` Float32 `json:$.llm_api_duration_ms` DEFAULT 0,
    `cache_hit` Bool `json:$.cache_hit` DEFAULT false,
    `response_status` LowCardinality(String) `json:$.standard_logging_object_status` DEFAULT 'unknown',
    `response_time` Float32 `json:$.standard_logging_object_response_time` DEFAULT 0,
    `proxy_metadata` String `json:$.proxy_metadata` DEFAULT '',
    `organization` String `json:$.proxy_metadata.organization` DEFAULT '',
    `environment` String `json:$.proxy_metadata.environment` DEFAULT '',
    `project` String `json:$.proxy_metadata.project` DEFAULT '',
    `chat_id` String `json:$.proxy_metadata.chat_id` DEFAULT '',
    `response` String `json:$.response` DEFAULT '',
    `response_id` String `json:$.response.id`,
    `response_object` String `json:$.response.object` DEFAULT 'unknown',
    `response_choices` Array(String) `json:$.response.choices[:]` DEFAULT [],
    `completion_tokens` UInt16 `json:$.response.usage.completion_tokens` DEFAULT 0,
    `prompt_tokens` UInt16 `json:$.response.usage.prompt_tokens` DEFAULT 0,
    `total_tokens` UInt16 `json:$.response.usage.total_tokens` DEFAULT 0,
    `cost` Float32 `json:$.cost` DEFAULT 0,
    `exception` String `json:$.exception` DEFAULT '',
    `traceback` String `json:$.traceback` DEFAULT '',
    `duration` Float32 `json:$.duration` DEFAULT 0


ENGINE MergeTree
ENGINE_SORTING_KEY start_time, organization, project, model
ENGINE_PARTITION_KEY toYYYYMM(start_time)

import type { LanguageModelMiddleware } from 'ai';

type TinybirdConfig = {
  host: string;
  token: string;
  datasource?: string;
  organization?: string;
  project?: string;
  environment?: string;
  user?: string;
  chatId?: string;
};

type Usage = {
  inputTokens: { total: number | undefined };
  outputTokens: { total: number | undefined };
};

export function tinybirdMiddleware(config: TinybirdConfig): LanguageModelMiddleware {
  const send = async (event: {
    model: { provider: string; modelId: string };
    startTime: Date;
    stream: boolean;
    usage?: Usage;
    text?: string;
    error?: unknown;
  }) => {
    const endTime = new Date();
    const duration = endTime.getTime() - event.startTime.getTime();
    const promptTokens = event.usage?.inputTokens.total ?? 0;
    const completionTokens = event.usage?.outputTokens.total ?? 0;
    const id = crypto.randomUUID();
    const failed = event.error !== undefined;
    const error = event.error instanceof Error ? event.error : undefined;

    await fetch(`${config.host}/v0/events?name=${config.datasource ?? 'llm_events'}`, {
      method: 'POST',
      headers: { Authorization: `Bearer ${config.token}` },
      body: JSON.stringify({
        id,
        message_id: id,
        start_time: event.startTime.toISOString(),
        end_time: endTime.toISOString(),
        model: event.model.modelId,
        provider: event.model.provider,
        stream: event.stream,
        user: config.user ?? 'unknown',
        call_type: 'completion',
        llm_api_duration_ms: duration,
        duration,
        standard_logging_object_status: failed ? 'failure' : 'success',
        standard_logging_object_response_time: duration,
        proxy_metadata: {
          organization: config.organization ?? '',
          project: config.project ?? '',
          environment: config.environment ?? '',
          chat_id: config.chatId ?? '',
        },
        response: {
          id,
          object: 'chat.completion',
          usage: {
            prompt_tokens: promptTokens,
            completion_tokens: completionTokens,
            total_tokens: promptTokens + completionTokens,
          },
          choices: event.text ? [event.text] : [],
        },
        exception: !failed ? '' : error ? error.message : typeof event.error === 'string' ? event.error : JSON.stringify(event.error),
        traceback: error?.stack ?? '',
      }),
    }).catch(console.error);
  };

  return {
    wrapGenerate: async ({ doGenerate, model }) => {
      const startTime = new Date();
      try {
        const result = await doGenerate();
        const text = result.content
          .map((part) => (part.type === 'text' ? part.text : ''))
          .join('');
        await send({ model, startTime, stream: false, usage: result.usage, text });
        return result;
      } catch (error) {
        await send({ model, startTime, stream: false, error });
        throw error;
      }
    },

    wrapStream: async ({ doStream, model }) => {
      const startTime = new Date();
      let text = '';
      try {
        const { stream, ...rest } = await doStream();
        // Log once the stream finishes, when usage is known.
        const logged = stream.pipeThrough(
          new TransformStream({
            transform(part, controller) {
              if (part.type === 'text-delta') text += part.delta;
              if (part.type === 'finish')
                void send({ model, startTime, stream: true, usage: part.usage, text });
              if (part.type === 'error')
                void send({ model, startTime, stream: true, error: part.error });
              controller.enqueue(part);
            },
          }),
        );
        return { stream: logged, ...rest };
      } catch (error) {
        await send({ model, startTime, stream: true, error });
        throw error;
      }
    },
  };
}

import { openai } from '@ai-sdk/openai';
import { generateText, wrapLanguageModel } from 'ai';
import { tinybirdMiddleware } from './lib/tinybird-middleware';

const model = wrapLanguageModel({
  model: openai('gpt-4o-mini'),
  middleware: tinybirdMiddleware({
    host: process.env.TINYBIRD_API_URL!,
    token: process.env.TINYBIRD_TOKEN!,
    organization: 'your-org',
    project: 'ai-analytics',
    environment: process.env.NODE_ENV,
  }),
});

const { text } = await generateText({ model, prompt: 'Summarize this week in one line' });