import { createOpenAI } from '@ai-sdk/openai';
import { axiomAIMiddleware } from 'axiom/ai';
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/ai';
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,
      },
    })
  }
)

import { withSpan } from 'axiom/ai';
import { Scorer } from 'axiom/ai/evals/scorers';
import { onlineEval } from 'axiom/ai/evals/online';

const formatScorer = Scorer('format-check', ({ output }: { output: string }) => {
  return output.length > 0;
});

await withSpan({ capability: 'qa', step: 'answer' }, async () => {
  const response = await generateText({ model, messages });
  void onlineEval(
    { capability: 'qa', step: 'answer' },
    { output: response.text, scorers: [formatScorer] }
  );
  return response.text;
});