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)

const openai = createOpenAI({ apiKey: apiKey });
const wrappedOpenAI = wrapModelWithTinybird(
    openai('gpt-3.5-turbo'),
    process.env.NEXT_PUBLIC_TINYBIRD_API_URL!,
    process.env.TINYBIRD_TOKEN!,
    {
    event: 'search_filter',
    environment: process.env.NODE_ENV,
    project: 'ai-analytics',
    organization: 'your-org',
    }
);

import type { LanguageModelV1 } from '@ai-sdk/provider';

type TinybirdConfig = {
  event?: string;
  organization?: string;
  project?: string;
  environment?: string;
  user?: string;
  chatId?: string;
};

export function wrapModelWithTinybird(
  model: LanguageModelV1,
  tinybirdHost: string,
  tinybirdToken: string,
  config: TinybirdConfig = {}
) {
  const originalDoGenerate = model.doGenerate;
  const originalDoStream = model.doStream;

  const logToTinybird = async (
    messageId: string,
    startTime: Date,
    status: 'success' | 'error',
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    args: any[],
    result?: { text?: string; usage?: { promptTokens?: number; completionTokens?: number } },
    error?: Error
  ) => {
    const endTime = new Date();
    const duration = endTime.getTime() - startTime.getTime();

    const event = {
      start_time: startTime.toISOString(),
      end_time: endTime.toISOString(),
      message_id: messageId,
      model: model.modelId || 'unknown',
      provider: 'openai',
      duration,
      llm_api_duration_ms: duration,
      response: status === 'success' ? {
        id: messageId,
        object: 'chat.completion',
        usage: {
          prompt_tokens: result?.usage?.promptTokens || 0,
          completion_tokens: result?.usage?.completionTokens || 0,
          total_tokens: (result?.usage?.promptTokens || 0) + (result?.usage?.completionTokens || 0),
        },
        choices: [{ message: { content: result?.text ?? '' } }],
      } : undefined,
      messages: args[0]?.prompt ? [{ role: 'user', content: args[0].prompt }].map(m => ({
        role: String(m.role),
        content: String(m.content)
      })) : [],
      proxy_metadata: {
        organization: config.organization || '',
        project: config.project || '',
        environment: config.environment || '',
        chat_id: config.chatId || '',
      },
      user: config.user || 'unknown',
      standard_logging_object_status: status,
      standard_logging_object_response_time: duration,
      log_event_type: config.event || 'chat_completion',
      id: messageId,
      call_type: 'completion',
      cache_hit: false,
      ...(status === 'error' && {
        exception: error?.message || 'Unknown error',
        traceback: error?.stack || '',
      }),
    };

    // Send to Tinybird
    fetch(`${tinybirdHost}/v0/events?name=llm_events`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Authorization': `Bearer ${tinybirdToken}`,
      },
      body: JSON.stringify(event),
    }).catch(console.error);
  };

  model.doGenerate = async function (...args) {
    const startTime = new Date();
    const messageId = crypto.randomUUID();

    try {
      const result = await originalDoGenerate.apply(this, args);
      await logToTinybird(messageId, startTime, 'success', args, result);
      return result;
    } catch (error) {
      await logToTinybird(messageId, startTime, 'error', args, undefined, error as Error);
      throw error;
    }
  };

  model.doStream = async function (...args) {
    const startTime = new Date();
    const messageId = crypto.randomUUID();

    try {
      const result = await originalDoStream.apply(this, args);
      await logToTinybird(messageId, startTime, 'success', args, { text: '', usage: { promptTokens: 0, completionTokens: 0 } });
      return result;
    } catch (error) {
      await logToTinybird(messageId, startTime, 'error', args, undefined, error as Error);
      throw error;
    }
  };

  return model;
}