import { createOpenAI } from '@ai-sdk/openai';
import { axiomAIMiddleware } from 'axiom';
import { wrapLanguageModel } from 'ai';

const openai = createOpenAI({
  apiKey: process.env.OPENAI_API_KEY!,
  compatibility: 'strict',
});

const model = openai('gpt-4o-mini');

export const gpt4oMini = wrapLanguageModel({
  model,
  middleware: [axiomAIMiddleware({ model })],
});

import { tool } from 'ai';
import { wrapTool } from 'axiom';
import { z } from 'zod';

const getWeather = tool({
  description: 'Get current weather for a city',
  parameters: z.object({
    city: z.string().describe('The city name'),
    country: z.string().describe('The country code'),
  }),
  execute: async ({ city, country }) => {
    // Your tool implementation
    return {
      city,
      country,
      temperature: 22,
      condition: 'sunny',
    };
  },
});

// Wrap the tool for observability
const wrappedWeatherTool = wrapTool('weatherTool', weatherTool);

const result = await withSpan(
  { capability: 'weather_bot', step: 'get_weather' },
  (span) => {
    return generateText({
      model: gpt4oMini,
      messages: [{ role: 'user', content: 'What is the weather in London?' }],
      tools: {
        getWeather: wrappedWeatherTool,
      },
    })
  }
)