Custom engines

import { AVA, registerEngine } from '@antv/ava';
import { SalesEngine } from './sales-engine';
 
registerEngine('sales', SalesEngine);
 
const ava = new AVA({
  llm: { model },
  engine: { type: 'sales' },
});

import { generateText, type LanguageModel } from 'ai';
import { executionResult } from '@antv/ava';
import type {
  AnalysisEngine,
  TranslationResult,
  DataModel,
  DataSourceConfig,
  ExecutionOptions,
  ExecutionResult,
  LLMConfig,
  Schema,
} from '@antv/ava';
 
type Sale = { region: string; revenue: number };
 
export class SalesEngine implements AnalysisEngine {
  readonly language = {
    name: 'Sales metrics JSON',
    fence: 'json',
    instructions:
      'Return exactly {"metric":"total_revenue"} or {"metric":"revenue_by_region"}. No other keys or operations are supported.',
  };
  private sales: Sale[] | null = null;
 
  private readonly model: LanguageModel;
 
  constructor(private readonly llm: LLMConfig) {
    if (typeof llm.model === 'string') {
      throw new Error('SalesEngine expects a configured AI SDK model object.');
    }
    this.model = llm.model;
  }
 
  async load(config: DataSourceConfig): Promise<void> {
    await this.dispose();
    if (config.type !== 'json' || !Array.isArray(config.options.data)) {
      throw new Error('SalesEngine requires an inline JSON array.');
    }
    this.sales = config.options.data.map((row) => {
      if (
        !row ||
        typeof row !== 'object' ||
        typeof row.region !== 'string' ||
        typeof row.revenue !== 'number' ||
        !Number.isFinite(row.revenue)
      ) {
        throw new Error('Each sale requires a region string and finite revenue number.');
      }
      return { region: row.region, revenue: row.revenue };
    });
  }
 
  async schema(): Promise<Schema> {
    if (!this.sales) throw new Error('Load a source first.');
    return {
      models: [
        {
          name: 'sales',
          tableReference: { table: 'sales' },
          columns: [
            { name: 'region', type: 'string' },
            { name: 'revenue', type: 'number' },
          ],
        },
      ],
    };
  }
 
  async getDSL(query: string, _dataModel: DataModel): Promise<TranslationResult> {
    if (!this.sales) throw new Error('Load a source first.');
    const { text, usage } = await generateText({
      model: this.model,
      maxRetries: this.llm.maxRetries ?? 3,
      system: `Convert a question into a sales metric. ${this.language.instructions}
Return raw JSON without Markdown. If neither metric answers the question, return {"metric":"unsupported"}.`,
      prompt: query,
    });
    return { dsl: text.trim(), usage };
  }
 
  async execute<T = Record<string, unknown>>(dsl: string, options?: ExecutionOptions): Promise<ExecutionResult<T>> {
    if (!this.sales) throw new Error('Load a source first.');
    const command = JSON.parse(dsl);
    if (
      !command ||
      Array.isArray(command) ||
      Object.keys(command).length !== 1 ||
      !['total_revenue', 'revenue_by_region'].includes(command.metric)
    ) {
      throw new Error('Expected one supported metric and no additional keys.');
    }
 
    let rows: Record<string, unknown>[];
    if (command.metric === 'total_revenue') {
      rows = [{ total_revenue: this.sales.reduce((sum, row) => sum + row.revenue, 0) }];
    } else {
      const totals = new Map<string, number>();
      for (const row of this.sales) {
        totals.set(row.region, (totals.get(row.region) ?? 0) + row.revenue);
      }
      rows = Array.from(totals, ([region, total_revenue]) => ({ region, total_revenue }));
    }
    if (rows.some((row) => !Number.isFinite(row.total_revenue))) {
      throw new Error('Revenue total exceeds the numeric range.');
    }
    const schema =
      command.metric === 'total_revenue'
        ? [{ name: 'total_revenue', type: 'number' }]
        : [
            { name: 'region', type: 'string' },
            { name: 'total_revenue', type: 'number' },
          ];
    return executionResult(rows as T[], schema, options);
  }
 
  async dispose(): Promise<void> {
    this.sales = null;
  }
}

import { createOpenAI } from '@ai-sdk/openai';
import { AVA, registerEngine } from '@antv/ava';
import { SalesEngine } from './sales-engine';
 
registerEngine('sales', SalesEngine);
 
const model = createOpenAI({
  apiKey: process.env.OPENAI_API_KEY,
  baseURL: process.env.OPENAI_BASE_URL,
}).chat(process.env.OPENAI_MODEL || 'gpt-4o-mini');
 
const ava = new AVA({
  llm: { model },
  engine: { type: 'sales' },
});
 
try {
  await ava.source({
    type: 'json',
    options: {
      data: [
        { region: 'East', revenue: 120 },
        { region: 'West', revenue: 90 },
        { region: 'East', revenue: 80 },
      ],
    },
  });
 
  const answer = await ava.analyze('What is total revenue by region?', {
    strategy: { type: 'direct', maxRetries: 0 },
    includeSummary: false,
  });
  console.log(answer.sql); // Expected command: {"metric":"revenue_by_region"}
  console.table(answer.data);
 
  // Execute a known command directly, without calling the model.
  const total = await ava.query('{"metric":"total_revenue"}');
  console.table(total.data); // [{ total_revenue: 290 }]
} finally {
  await ava.dispose();
}