ChatCompletion icon
Switch icon
ChatCompletion icon
ChatCompletion icon
ChatCompletion icon

Route AI Requests to the Optimal Model Based on Query Classification

Build an AI model router on Kestra. Classify prompts with OpenAI, then route coding, reasoning, search, and general queries to the best-fit LLM.

Categories
AICore

Build an AI model router that classifies each incoming prompt and dispatches it to the most suitable large language model, so you stop paying frontier prices for simple questions and stop sending hard problems to weak models. A lightweight OpenAI classifier reads the request, tags it as coding, reasoning, search, or general, and a single routing decision picks the right provider for the job. This solves the core problem of multi-provider AI: choosing the right model per request, with full observability over every classification and response.

How it works

  1. The classify_request task (io.kestra.plugin.openai.ChatCompletion, model gpt-4o-mini) classifies the user prompt into one of four categories using a constrained jsonResponseSchema and a small maxTokens budget, keeping classification fast and cheap.
  2. The route_to_model task (io.kestra.plugin.core.flow.Switch) reads the category with a jq expression over the classifier output and selects a branch.
  3. Coding prompts go to coding_model (io.kestra.plugin.anthropic.ChatCompletion, Claude Opus). Reasoning prompts go to reasoning_model (io.kestra.plugin.gemini.ChatCompletion, Gemini 2.5 Flash). Search prompts go to search_model (io.kestra.plugin.perplexity.ChatCompletion, Sonar) for real-time web retrieval.
  4. Anything else falls through to the Switch defaults, hitting general_model (io.kestra.plugin.openai.ChatCompletion, gpt-4o-mini) for cost-efficient general answers.

What you get

  • A working, auditable routing layer across four AI providers in one declarative flow.
  • Cheap pre-classification with structured JSON output instead of brittle string parsing.
  • Per-request model selection that matches each task to a specialized model.
  • Full execution history showing what was classified, where it routed, and what came back.

Who it's for

  • AI platform and chatbot teams serving diverse query types.
  • Cost-conscious teams avoiding frontier models for trivial prompts.
  • Engineers building observable multi-provider LLM pipelines.

Why orchestrate this with Kestra

A provider SDK can call one model, but it cannot coordinate classification, branching, retries, and lineage across several vendors. Kestra adds event and schedule triggers, automatic retries on transient API errors, and full execution lineage so every routing decision is logged and replayable. The entire pipeline is declarative YAML you can version and review, filling the gap that no single LLM provider's own client can cover: vendor-agnostic routing with auditability.

Prerequisites

  • A Kestra instance (>= 0.16).
  • Active accounts with OpenAI, Anthropic, Google Gemini, and Perplexity.

Secrets

  • OPENAI_API_KEY: used by the classifier and the general fallback model.
  • ANTHROPIC_API_KEY: used by the coding branch.
  • GEMINI_API_KEY: used by the reasoning branch.
  • PERPLEXITY_API_KEY: used by the search branch.

Quick start

  1. Add OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY, and PERPLEXITY_API_KEY as Kestra secrets.
  2. Create the flow from this blueprint.
  3. Run it with the default question input, or supply your own prompt.
  4. Inspect the execution to see the classification and the branch that handled it.

How to extend

  • Add new categories and Switch cases for more specialized models.
  • Weight routing by cost or latency targets per category.
  • A/B test two models on the same category and compare outputs.
  • Log classification and response quality to a database for continuous tuning.

Links

See How

New to Kestra?

Use blueprints to kickstart your first workflows.